From d94c88c0dc3b7613152d50eefd0fd71bb9a04979 Mon Sep 17 00:00:00 2001
From: Mihai Polceanu <mihaip@graphcore.ai>
Date: Wed, 8 Nov 2023 22:49:35 +0000
Subject: [PATCH] added option to store AO grid in fp16 to reduce memory

---
 pyscf_ipu/nanoDFT/nanoDFT.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/pyscf_ipu/nanoDFT/nanoDFT.py b/pyscf_ipu/nanoDFT/nanoDFT.py
index 44518ef..d67c4c2 100644
--- a/pyscf_ipu/nanoDFT/nanoDFT.py
+++ b/pyscf_ipu/nanoDFT/nanoDFT.py
@@ -269,6 +269,10 @@ def init_dft_tensors_cpu(mol, opts, DIIS_iters=9):
     else:
         grid_coords = grids.coords
     density_matrix  = pyscf.scf.hf.init_guess_by_minao(mol)         # (N,N)=(66,66) for C6H6.
+    
+    if opts.fp16_grid:
+        # reduce grid_AO memory by half
+        grid_AO = grid_AO.astype(np.float16)
 
     # TODO(): Add integral math formulas for kinetic/nuclear/O/ERI.
     kinetic         = mol.intor_symmetric('int1e_kin')              # (N,N)
@@ -619,7 +623,8 @@ def nanoDFT_options(
         vis_num_error: bool = False,
         molecule_name: str = None,
         screen_tol: float = 1e-9,
-        fast_shells: bool = False
+        fast_shells: bool = False,
+        fp16_grid: bool = False
 ):
     """
     nanoDFT