Optimize memory usage

.github/workflows/install_and_test.yml

@@ -14,7 +14,7 @@ jobs:
       PYSCF_CONFIG_FILE: ".github/workflows/pyscf_conf.py"
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.9", "3.10", "3.11", "3.12"]
         os: [ubuntu-latest]
     steps:
       - uses: actions/checkout@v3
@@ -32,12 +32,12 @@ jobs:
           pip install -e ".[examples]"
       - name: Run unit tests
         run: |
+            pytest -v tests/unit/test_kinetic_density_eq.py
             pytest -v tests/unit/test_eigenproblem.py
             pytest -v tests/unit/test_loss.py
       - name: Run integration tests
         run: |
           pytest -v tests/integration/molecules/test_non_xc_energy.py
-          pytest -v tests/integration/molecules/test_functional_implementations.py
           pytest -v tests/integration/molecules/test_Harris.py
           pytest -v tests/integration/molecules/test_predict_B88.py
           pytest -v tests/integration/molecules/test_training.py

grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py

@@ -216,7 +216,7 @@ def _build_graph(self, batch_dim: Optional[int] = None):
             library.
         """
 
-        self._functional = hub.Module(spec=self._model_path)
+        self._functional = hub.KerasLayer(self._model_path, signature_outputs_as_dict=True)
 
         grid_coords = tf.placeholder(tf.float32, shape=[batch_dim, 3], name="grid_coords")
         grid_weights = tf.placeholder(tf.float32, shape=[batch_dim], name="grid_weights")
@@ -260,7 +260,7 @@ def _build_graph(self, batch_dim: Optional[int] = None):
         }
         tensor_dict = {f"tensor_dict${k}": v for k, v in features.items()}
 
-        predictions = self._functional(tensor_dict, as_dict=True)
+        predictions = self._functional(tensor_dict)
         local_xc = predictions["grid_contribution"]
         weighted_local_xc = local_xc * grid_weights
         unweighted_xc = tf.reduce_sum(local_xc, axis=0)
@@ -311,19 +311,6 @@ def _build_graph(self, batch_dim: Optional[int] = None):
             grid_weights=grid_weights,
         )
 
-        outputs = {
-            "vxc": self._vxc,
-            "vrho": tf.stack(self._vrho),
-            "vsigma": tf.stack(self._vsigma),
-            "vtau": tf.stack(self._vtau),
-            "vhf": tf.stack(self._vhf),
-        }
-        # Create the signature for TF-Hub, including both the energy and functional
-        # derivatives.
-        # This is a no-op if _build_graph is called outside of
-        # hub.create_module_spec.
-        hub.add_signature(inputs=attr.asdict(self._placeholders), outputs=outputs)
-
     def export_functional_and_derivatives(
         self,
         export_path: str,
@@ -351,7 +338,7 @@ def export_functional_and_derivatives(
             spec = hub.create_module_spec(
                 self._build_graph, tags_and_args=[(set(), {"batch_dim": batch_dim})]
             )
-            functional_and_derivatives = hub.Module(spec=spec)
+            functional_and_derivatives = hub.KerasLayer(spec, signature_outputs_as_dict=True)
             with tf.Session() as session:
                 session.run(tf.global_variables_initializer())
                 functional_and_derivatives.export(export_path, session)

grad_dft/functional.py

@@ -591,10 +591,10 @@ def dm21_densities(
     grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1)
 
     # LDA preprocessing data
-    log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho, clip_cte))
 
     # GGA preprocessing data
-    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2
+    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, clip_cte)) / 2
     log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho
     log_u_sigma = jnp.where(
         jnp.greater(log_rho, jnp.log2(clip_cte)),
@@ -603,7 +603,7 @@ def dm21_densities(
     )
 
     # MGGA preprocessing data
-    log_tau = jnp.log2(jnp.clip(tau, a_min=clip_cte))
+    log_tau = jnp.log2(jnp.clip(tau, clip_cte))
     log_1t_sigma = -(
         5 / 3.0 * log_rho - log_tau + 2 / 3.0 * jnp.log2(6 * jnp.pi**2) + jnp.log2(3 / 5.0)
     )
@@ -674,8 +674,7 @@ def dm21_combine_densities(
         [densities] + [ehf[i].sum(axis=0, keepdims=True).T for i in range(len(ehf))], axis=1
     )
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def dm21_hfgrads_densities(
     functional: nn.Module,
     params: PyTree,
@@ -716,8 +715,7 @@ def dm21_hfgrads_densities(
     )
     return vxc_hf.sum(axis=0)  # Sum over omega
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def dm21_hfgrads_cinputs(
     functional: nn.Module,
     params: PyTree,
@@ -1005,7 +1003,7 @@ def correlation_polarization_correction(
         The ready to be integrated electronic energy density.
     """
 
-    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1), a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1), clip_cte))
     # assert not jnp.isnan(log_rho).any() and not jnp.isinf(log_rho).any()
     log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0
 
@@ -1033,7 +1031,7 @@ def fzeta(z):
     # assert not jnp.isnan(alphac).any() and not jnp.isinf(alphac).any()
 
     fz = fzeta(zeta) #jnp.round(fzeta(zeta), int(math.log10(clip_cte)))
-    z4 = zeta**4 #jnp.round(2 ** (4 * jnp.log2(jnp.clip(zeta, a_min=clip_cte))), int(math.log10(clip_cte)))
+    z4 = zeta**4 #jnp.round(2 ** (4 * jnp.log2(jnp.clip(zeta, clip_cte))), int(math.log10(clip_cte)))
 
     e_tilde = (
         e_tilde_PF[:, 0]
@@ -1105,10 +1103,10 @@ def densities(
     grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1)
 
     # LDA preprocessing data
-    log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho, clip_cte))
 
     # GGA preprocessing data
-    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, a_min=clip_cte)) / 2
+    log_grad_rho_norm = jnp.log2(jnp.clip(grad_rho_norm_sq, clip_cte)) / 2
     log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho
     log_u_sigma = jnp.where(
         jnp.greater(log_rho, jnp.log2(clip_cte)),
@@ -1117,7 +1115,7 @@ def densities(
     )
 
     # MGGA preprocessing data
-    log_tau = jnp.log2(jnp.clip(tau, a_min=clip_cte))
+    log_tau = jnp.log2(jnp.clip(tau, clip_cte))
     log_1t_sigma = log_tau - 5 / 3.0 * log_rho
     log_w_sigma = jnp.where(
         jnp.greater(log_rho, jnp.log2(clip_cte)),
@@ -1136,8 +1134,8 @@ def densities(
     ######### Correlation features ###############
 
     grad_rho_norm_sq_ss = jnp.sum((grad_rho.sum(axis=1)) ** 2, axis=-1)
-    log_grad_rho_norm_ss = jnp.log2(jnp.clip(grad_rho_norm_sq_ss, a_min=clip_cte)) / 2
-    log_rho_ss = jnp.log2(jnp.clip(rho.sum(axis=1), a_min=clip_cte))
+    log_grad_rho_norm_ss = jnp.log2(jnp.clip(grad_rho_norm_sq_ss, clip_cte)) / 2
+    log_rho_ss = jnp.log2(jnp.clip(rho.sum(axis=1), clip_cte))
     log_x_ss = log_grad_rho_norm_ss - 4 / 3.0 * log_rho_ss
 
     log_u_ss = jnp.where(
@@ -1154,7 +1152,7 @@ def densities(
 
     log_u_c = jnp.stack((log_u_ss, log_u_ab), axis=1)
 
-    log_tau_ss = jnp.log2(jnp.clip(tau.sum(axis=1), a_min=clip_cte))
+    log_tau_ss = jnp.log2(jnp.clip(tau.sum(axis=1), clip_cte))
     log_1t_ss = log_tau_ss - 5 / 3.0 * log_rho_ss
     log_w_ss = jnp.where(
         jnp.greater(log_rho.sum(axis=1), jnp.log2(clip_cte)),
@@ -1177,7 +1175,7 @@ def densities(
     beta3 = jnp.array([[1.6382, 3.3662]])
     beta4 = jnp.array([[0.49294, 0.62517]])
 
-    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), a_min=clip_cte))
+    log_rho = jnp.log2(jnp.clip(rho.sum(axis=1, keepdims=True), clip_cte))
     log_rs = jnp.log2((3 / (4 * jnp.pi)) ** (1 / 3)) - log_rho / 3.0
     brs_1_2 = 2 ** (log_rs / 2 + jnp.log2(beta1))
     ars = 2 ** (log_rs + jnp.log2(alpha1))

grad_dft/molecule.py

@@ -338,9 +338,8 @@ def to_dict(self) -> dict:
 
 #######################################################################
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def orbital_grad(
         mo_coeff: Float[Array, "spin orbitals orbitals"],
         mo_occ: Float[Array, "spin orbitals"],
@@ -382,9 +381,8 @@ def orbital_grad(
 
 
 ##########################################################
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def density(rdm1: Float[Array, "spin orbitals orbitals"], 
             ao: Float[Array, "grid orbitals"], 
             precision: Precision = Precision.HIGHEST
@@ -408,9 +406,8 @@ def density(rdm1: Float[Array, "spin orbitals orbitals"],
 
     return jnp.einsum("...ab,ra,rb->r...", rdm1, ao, ao, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def grad_density(
     rdm1: Float[Array, "spin orbitals orbitals"], 
     ao: Float[Array, "grid orbitals"], 
@@ -439,9 +436,8 @@ def grad_density(
 
     return 2 * jnp.einsum("...ab,ra,rbj->r...j", rdm1, ao, grad_ao, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def lapl_density(
     rdm1: Float[Array, "spin orbitals orbitals"], 
     ao: Float[Array, "grid orbitals"], 
@@ -473,9 +469,8 @@ def lapl_density(
         "...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision
     ) + 2 * jnp.einsum("...ab,ra,rbi->r...", rdm1, ao, grad_2_ao, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames="precision")
+@jaxtyped(typechecker=typechecked)
 def kinetic_density(
     rdm1: Float[Array, "spin orbitals orbitals"], 
     grad_ao: Float[Array, "grid orbitals 3"],
@@ -498,12 +493,11 @@ def kinetic_density(
     Array
         The kinetic energy density. Shape: (n_spin, n_grid_points)
     """
+    temp = jnp.einsum("sab,raj->srbj", rdm1, grad_ao, precision=precision)
+    return 0.5 * jnp.einsum("srbj,rbj->rs", temp, grad_ao, precision=precision)
 
-    return 0.5 * jnp.einsum("...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision)
-
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def HF_energy_density(
     rdm1: Float[Array, "spin orbitals orbitals"],
     ao: Float[Array, "grid orbitals"],
@@ -540,8 +534,7 @@ def HF_energy_density(
     )
     return vmap(_hf_energy, in_axes=(0, None, 0), out_axes=2)(chi, rdm1, ao)
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def HF_density_grad_2_Fock(
     grid: Grid,
     functional: nn.Module,
@@ -612,8 +605,7 @@ def chunked_jvp(chi_tensor, gr_tensor, ao_tensor):
 
     return (jax.jit(chunked_jvp)(chi.transpose(3, 0, 1, 2), gr, ao)).transpose(1, 2, 3, 0)
 
-@jaxtyped
-@typechecked
+@jaxtyped(typechecker=typechecked)
 def HF_coefficient_input_grad_2_Fock(
     grid: Grid,
     functional: nn.Module,
@@ -691,9 +683,8 @@ def abs_clip(arr, threshold):
 
 ######################################################################
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def nonXC(
     rdm1: Float[Array, "orbitals orbitals"],
     h1e: Float[Array, "orbitals orbitals"],
@@ -732,9 +723,8 @@ def nonXC(
 
     return nuclear_repulsion + h1e_energy + coulomb2e_energy
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def one_body_energy(
     rdm1: Float[Array, "orbitals orbitals"],
     h1e: Float[Array, "orbitals orbitals"],
@@ -757,9 +747,8 @@ def one_body_energy(
     return h1e_energy
 
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def coulomb_energy(
     rdm1: Float[Array, "orbitals orbitals"],
     rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
@@ -782,9 +771,8 @@ def coulomb_energy(
     coulomb2e_energy = jnp.einsum("pq,pq->", rdm1, v_coul, precision=precision) / 2.0
     return coulomb2e_energy
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def coulomb_potential(
     rdm1: Float[Array, "orbitals orbitals"],
     rep_tensor: Float[Array, "orbitals orbitals orbitals orbitals"],
@@ -810,9 +798,8 @@ def coulomb_potential(
     """
     return jnp.einsum("pqrt,rt->pq", rep_tensor, rdm1, precision=precision)
 
-@jaxtyped
-@typechecked
 @partial(jax.jit, static_argnames=["precision"])
+@jaxtyped(typechecker=typechecked)
 def make_rdm1(
     mo_coeff: Float[Array, "spin orbitals orbitals"],
     mo_occ: Float[Array, "spin orbitals"],
@@ -845,9 +832,8 @@ def make_rdm1(
 
     return jnp.einsum("sij,sj,skj -> sik", mo_coeff, mo_occ, mo_coeff.conj(), precision=precision)
 
-@jaxtyped
-@typechecked
 @jax.jit
+@jaxtyped(typechecker=typechecked)
 def get_occ(
     mo_energies: Float[Array, "spin orbitals"],
     nelecs: Int[Array, "spin"],