From 2fe27e1642a0a31faf749a28e0c659bea306184d Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 20:20:24 +0300 Subject: [PATCH 01/15] Rewrite kinetic_density einsum to escape intermediet representation with shape [grid, orbitals, orbitals]. --- grad_dft/molecule.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/grad_dft/molecule.py b/grad_dft/molecule.py index 580430b..ea1f060 100644 --- a/grad_dft/molecule.py +++ b/grad_dft/molecule.py @@ -498,8 +498,8 @@ def kinetic_density( Array The kinetic energy density. Shape: (n_spin, n_grid_points) """ - - return 0.5 * jnp.einsum("...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision) + temp = jnp.einsum("sab,raj->srbj", rdm1, grad_ao, precision=precision) + return 0.5 * jnp.einsum("srbj,rbj->rs", temp, grad_ao, precision=precision) @jaxtyped @typechecked From c1c233717313a82bcd1376899d9ea42a516332be Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 20:33:31 +0300 Subject: [PATCH 02/15] Add tests on equivalence. --- tests/unit/test_kinetic_density_eq.py | 44 +++++++++++++++++++++++++++ 1 file changed, 44 insertions(+) create mode 100644 tests/unit/test_kinetic_density_eq.py diff --git a/tests/unit/test_kinetic_density_eq.py b/tests/unit/test_kinetic_density_eq.py new file mode 100644 index 0000000..e708ade --- /dev/null +++ b/tests/unit/test_kinetic_density_eq.py @@ -0,0 +1,44 @@ +from functools import partial +import pytest + +from grad_dft.molecule import kinetic_density +import jax.numpy as jnp +import jax +from jax.lax import Precision + + +@partial(jax.jit, static_argnames="precision") +def kinetic_density_old( + rdm1: jax.Array, + grad_ao: jax.Array, + precision: Precision = Precision.HIGHEST +) -> jax.Array: + return 0.5 * jnp.einsum( + "...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision + ) + + +@pytest.mark.parametrize("spin", [1, 2]) +@pytest.mark.parametrize("orbitals", [10, 25]) +@pytest.mark.parametrize("grid", [30, 50]) +@pytest.mark.parametrize("precision", [ + Precision.HIGHEST, + Precision.HIGH, + Precision.DEFAULT, +]) +def test_kinetic_density_equivalence(spin, orbitals, grid, precision): + rng = jax.random.PRNGKey(0) + key1, key2 = jax.random.split(rng) + + rdm1_real = jax.random.normal(key1, (spin, orbitals, orbitals)) + rdm1 = 0.5 * (rdm1_real + rdm1_real.transpose(0, 2, 1)) + grad_ao = jax.random.normal(key2, (grid, orbitals, 3)) + result_old = kinetic_density_old(rdm1, grad_ao, precision) + result_new = kinetic_density(rdm1, grad_ao, precision) + + jnp.testing.assert_allclose( + result_old, result_new, + rtol=1e-6, atol=1e-6, + err_msg=f"Results differ for spin={spin}, orbitals={orbitals}, " + f"grid={grid}, precision={precision}" + ) \ No newline at end of file From 674e97a3b3c3f17a1afa21332f35ced0b37ed3ac Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 20:40:39 +0300 Subject: [PATCH 03/15] Fix deps. --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index a8f10c0..f9f63ce 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,4 +9,4 @@ typeguard==2.13.3 typing_extensions>=4.8.0 jaxtyping pytest>=7.4.3 -chex>=0.1.91 +chex>=0.1.90 From 222061339dd6106187435a950bebfac226edb43a Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 20:55:04 +0300 Subject: [PATCH 04/15] Fix tests. --- grad_dft/train.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/grad_dft/train.py b/grad_dft/train.py index e81680f..6fbcc2a 100644 --- a/grad_dft/train.py +++ b/grad_dft/train.py @@ -420,7 +420,7 @@ def dm21_grad_regularization(molecule: Molecule, F: Float[Array, "spin ao ao"]) prefactors = numerator / safe_denominator - dE = jnp.clip(0.5 * jnp.sum(prefactors * factors), a_min=-10, a_max=10) + dE = jnp.clip(0.5 * jnp.sum(prefactors * factors), min=-10, max=10) return dE**2 @@ -579,10 +579,10 @@ def sq_electron_err_int( Scalar: the value epsilon described above ---------- """ - pred_density = jnp.clip(pred_density, a_min=clip_cte) - truth_density = jnp.clip(truth_density, a_min=clip_cte) - diff_up = jnp.clip(jnp.clip(pred_density[:, 0] - truth_density[:, 0], a_min=clip_cte) ** 2, a_min=clip_cte) - diff_dn = jnp.clip(jnp.clip(pred_density[:, 1] - truth_density[:, 1], a_min=clip_cte) ** 2, a_min=clip_cte) + pred_density = jnp.clip(pred_density, min=clip_cte) + truth_density = jnp.clip(truth_density, min=clip_cte) + diff_up = jnp.clip(jnp.clip(pred_density[:, 0] - truth_density[:, 0], min=clip_cte) ** 2, min=clip_cte) + diff_dn = jnp.clip(jnp.clip(pred_density[:, 1] - truth_density[:, 1], min=clip_cte) ** 2, min=clip_cte) err_int = jnp.sum(diff_up * atoms.grid.weights) + jnp.sum(diff_dn * atoms.grid.weights) return err_int From 4fb5aa9768c54713bebba5131f0565249dc327fc Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 21:00:10 +0300 Subject: [PATCH 05/15] Add kinetic density test to CI. --- .github/workflows/install_and_test.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml index 37c0c22..ca8bfeb 100644 --- a/.github/workflows/install_and_test.yml +++ b/.github/workflows/install_and_test.yml @@ -32,6 +32,7 @@ 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 From 516885a7af6a96596d616c43f98de7d2282594fa Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 21:13:56 +0300 Subject: [PATCH 06/15] Fix assert. --- tests/unit/test_kinetic_density_eq.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_kinetic_density_eq.py b/tests/unit/test_kinetic_density_eq.py index e708ade..0f2c8e4 100644 --- a/tests/unit/test_kinetic_density_eq.py +++ b/tests/unit/test_kinetic_density_eq.py @@ -5,6 +5,7 @@ import jax.numpy as jnp import jax from jax.lax import Precision +import numpy as np @partial(jax.jit, static_argnames="precision") @@ -16,7 +17,7 @@ def kinetic_density_old( return 0.5 * jnp.einsum( "...ab,raj,rbj->r...", rdm1, grad_ao, grad_ao, precision=precision ) - + @pytest.mark.parametrize("spin", [1, 2]) @pytest.mark.parametrize("orbitals", [10, 25]) @@ -36,7 +37,7 @@ def test_kinetic_density_equivalence(spin, orbitals, grid, precision): result_old = kinetic_density_old(rdm1, grad_ao, precision) result_new = kinetic_density(rdm1, grad_ao, precision) - jnp.testing.assert_allclose( + np.testing.assert_allclose( result_old, result_new, rtol=1e-6, atol=1e-6, err_msg=f"Results differ for spin={spin}, orbitals={orbitals}, " From a0eec6b16d9567a5d025584ad7cf51143f7277c9 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 21:33:32 +0300 Subject: [PATCH 07/15] Increase tolerance. --- tests/unit/test_kinetic_density_eq.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_kinetic_density_eq.py b/tests/unit/test_kinetic_density_eq.py index 0f2c8e4..f5672dd 100644 --- a/tests/unit/test_kinetic_density_eq.py +++ b/tests/unit/test_kinetic_density_eq.py @@ -39,7 +39,7 @@ def test_kinetic_density_equivalence(spin, orbitals, grid, precision): np.testing.assert_allclose( result_old, result_new, - rtol=1e-6, atol=1e-6, + rtol=1e-5, atol=1e-5, err_msg=f"Results differ for spin={spin}, orbitals={orbitals}, " f"grid={grid}, precision={precision}" ) \ No newline at end of file From e2919e5c7cb1613a1934b03a97ff476691f53478 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sat, 25 Apr 2026 21:45:05 +0300 Subject: [PATCH 08/15] Use integers. --- tests/unit/test_kinetic_density_eq.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_kinetic_density_eq.py b/tests/unit/test_kinetic_density_eq.py index f5672dd..6e416a1 100644 --- a/tests/unit/test_kinetic_density_eq.py +++ b/tests/unit/test_kinetic_density_eq.py @@ -31,9 +31,9 @@ def test_kinetic_density_equivalence(spin, orbitals, grid, precision): rng = jax.random.PRNGKey(0) key1, key2 = jax.random.split(rng) - rdm1_real = jax.random.normal(key1, (spin, orbitals, orbitals)) + rdm1_real = jax.random.randint(key1, (spin, orbitals, orbitals), minval=-80, maxval=80).astype(jnp.float64) rdm1 = 0.5 * (rdm1_real + rdm1_real.transpose(0, 2, 1)) - grad_ao = jax.random.normal(key2, (grid, orbitals, 3)) + grad_ao = jax.random.randint(key2, (grid, orbitals, 3), minval=-80, maxval=80).astype(jnp.float64) result_old = kinetic_density_old(rdm1, grad_ao, precision) result_new = kinetic_density(rdm1, grad_ao, precision) From c3ae42fce1ee8f29c90d3f01559fa33d8b055005 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Thu, 2 Jul 2026 16:33:32 +0300 Subject: [PATCH 09/15] Modify dependencies usage. --- .../density_functional_approximation_dm21/neural_numint.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py index 7c62fa2..18940a2 100644 --- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py +++ b/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.load(self._model_path) 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") @@ -351,7 +351,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.load(spec) with tf.Session() as session: session.run(tf.global_variables_initializer()) functional_and_derivatives.export(export_path, session) From 07493566f4be8d4a62b0924c1ab38fd3938efc63 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Thu, 2 Jul 2026 16:58:39 +0300 Subject: [PATCH 10/15] Try another way. --- .../density_functional_approximation_dm21/neural_numint.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py index 18940a2..5461b7a 100644 --- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py +++ b/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.load(self._model_path) + self._functional = hub.KerasLayer(self._model_path) 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") @@ -351,7 +351,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.load(spec) + functional_and_derivatives = hub.KerasLayer(spec) with tf.Session() as session: session.run(tf.global_variables_initializer()) functional_and_derivatives.export(export_path, session) From 39f2f22804a6e7b18743752f159739ab75781f27 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Thu, 2 Jul 2026 17:17:30 +0300 Subject: [PATCH 11/15] Fix issue with as_dict. --- .../density_functional_approximation_dm21/neural_numint.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py index 5461b7a..9673d3f 100644 --- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py +++ b/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.KerasLayer(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) @@ -351,7 +351,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.KerasLayer(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) From c18ebb7eb21c73f89a659133ace41c585adfcebb Mon Sep 17 00:00:00 2001 From: Anya497 Date: Thu, 2 Jul 2026 17:26:08 +0300 Subject: [PATCH 12/15] Try remove add_signature calling. --- .../density_functional_approximation_dm21/neural_numint.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py index 9673d3f..089457c 100644 --- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py +++ b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py @@ -318,11 +318,6 @@ def _build_graph(self, batch_dim: Optional[int] = None): "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, From b5a10e2744686986b39cd005f703688d09014c64 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sun, 5 Jul 2026 11:49:11 +0300 Subject: [PATCH 13/15] Run CI on python 3.12 --- .github/workflows/install_and_test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml index ca8bfeb..de10def 100644 --- a/.github/workflows/install_and_test.yml +++ b/.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.12"] os: [ubuntu-latest] steps: - uses: actions/checkout@v3 From f42a89c91f5f1bca397ad3df72502fe6d6d17db4 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sun, 5 Jul 2026 12:03:32 +0300 Subject: [PATCH 14/15] Fix warnings and error with clip. --- .../neural_numint.py | 8 ---- grad_dft/functional.py | 30 +++++++------ grad_dft/molecule.py | 42 +++++++------------ grad_dft/popular_functionals.py | 17 ++++---- grad_dft/solid.py | 33 +++++---------- 5 files changed, 47 insertions(+), 83 deletions(-) diff --git a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py index 089457c..811b92a 100644 --- a/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py +++ b/grad_dft/external/density_functional_approximation_dm21/density_functional_approximation_dm21/neural_numint.py @@ -311,14 +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), - } - def export_functional_and_derivatives( self, export_path: str, diff --git a/grad_dft/functional.py b/grad_dft/functional.py index 6e5ffef..e55a00d 100644 --- a/grad_dft/functional.py +++ b/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)) diff --git a/grad_dft/molecule.py b/grad_dft/molecule.py index ea1f060..a0ef2f1 100644 --- a/grad_dft/molecule.py +++ b/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"], @@ -501,9 +496,8 @@ def kinetic_density( temp = jnp.einsum("sab,raj->srbj", rdm1, grad_ao, precision=precision) return 0.5 * jnp.einsum("srbj,rbj->rs", temp, 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"], diff --git a/grad_dft/popular_functionals.py b/grad_dft/popular_functionals.py index 69d4c03..e3e7340 100644 --- a/grad_dft/popular_functionals.py +++ b/grad_dft/popular_functionals.py @@ -38,7 +38,7 @@ def lsda_x_e(rho: Float[Array, "grid spin"], clip_cte) -> Float[Array, "grid"]: ------- Float[Array, "grid"] """ - rho = jnp.clip(rho, a_min=clip_cte) + rho = jnp.clip(rho, clip_cte) lda_es = ( -3.0 / 4.0 @@ -69,14 +69,14 @@ def b88_x_e(rho: Float[Array, "grid spin"], grad_rho: Float[Array, "grid spin di beta = 0.0042 - rho = jnp.clip(rho, a_min=clip_cte) + rho = jnp.clip(rho, clip_cte) # LDA preprocessing data: Note that we duplicate the density to sum and divide in the last eq. - log_rho = jnp.log2(jnp.clip(rho, a_min=clip_cte)) + log_rho = jnp.log2(jnp.clip(rho, clip_cte)) grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1) - 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 # GGA preprocessing data log_x_sigma = log_grad_rho_norm - 4 / 3.0 * log_rho @@ -124,7 +124,7 @@ def pw92_c_e(rho: Float[Array, "grid spin"], clip_cte: float = 1e-30) -> Float[A 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)) @@ -161,7 +161,7 @@ def vwn_c_e(rho: Float[Array, "grid spin"], clip_cte: float = 1e-30) -> Float[Ar x0 = jnp.array([[-0.10498, -0.325]]) rho = jnp.where(rho > clip_cte, rho, 0.0) - 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)) # 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 log_x = log_rs / 2 @@ -232,7 +232,7 @@ def lyp_c_e(rho: Float[Array, "grid spin"], grad_rho: Float[Array, "grid spin 3" d = 0.349 CF = (3 / 10) * (3 * jnp.pi**2) ** (2 / 3) - rho = jnp.clip(rho, a_min=clip_cte) + rho = jnp.clip(rho, clip_cte) grad_rho_norm_sq = jnp.sum(grad_rho**2, axis=-1) @@ -325,8 +325,7 @@ def b3lyp_exhf_densities(molecule: Molecule, clip_cte: float = 1e-30, *_, **__) return jnp.stack((lda_e, b88_e, vwn_e, lyp_e), axis=1) -@jaxtyped -@typechecked +@jaxtyped(typechecker=typechecked) def b3lyp_combine(features: Float[Array, "grid densities"], ehf: Float[Array, "omega spin grid"]) -> Float[Array, "grid densities+1"]: r""" Auxiliary function to combine the non Hartree-Fock features of B3LYP functional diff --git a/grad_dft/solid.py b/grad_dft/solid.py index f6a3ea2..238f6bb 100644 --- a/grad_dft/solid.py +++ b/grad_dft/solid.py @@ -266,9 +266,8 @@ def get_mo_grads(self, *args, **kwargs): return orbital_grad(self.mo_coeff, self.mo_occ, self.fock, *args, **kwargs) -@jaxtyped -@typechecked @partial(jit, static_argnames=["precision"]) +@jaxtyped(typechecker=typechecked) def one_body_energy( rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"], h1e: Complex[Array, "n_kpt n_orbitals n_orbitals"], @@ -296,9 +295,8 @@ def one_body_energy( h1e_energy = jnp.einsum("k,kij,kji->", weights, rdm1, h1e, precision=precision) return h1e_energy.real -@jaxtyped -@typechecked @partial(jit, static_argnames=["precision"]) +@jaxtyped(typechecker=typechecked) def coulomb_potential( rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"], rep_tensor: Complex[Array, "n_kpt n_kpt n_orbitals n_orbitals n_orbitals n_orbitals"], @@ -327,9 +325,8 @@ def coulomb_potential( return v_k -@jaxtyped -@typechecked @partial(jit, static_argnames=["precision"]) +@jaxtyped(typechecker=typechecked) def coulomb_energy( rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"], rep_tensor: Complex[Array, "n_kpt n_kpt n_orbitals n_orbitals n_orbitals n_orbitals"], @@ -361,9 +358,8 @@ def coulomb_energy( coulomb_energy = jnp.einsum("k,kij,kji->", weights, rdm1, v_k)/2.0 return coulomb_energy.real -@jaxtyped -@typechecked @partial(jit, static_argnames=["precision"]) +@jaxtyped(typechecker=typechecked) def non_xc( rdm1: Complex[Array, "n_kpt n_orbitals n_orbitals"], h1e: Complex[Array, "n_kpt n_orbitals n_orbitals"], @@ -407,9 +403,8 @@ def non_xc( return nuclear_repulsion + kinetic_and_external + coulomb -@jaxtyped -@typechecked @partial(jit, static_argnames=["precision"]) +@jaxtyped(typechecker=typechecked) def make_rdm1( mo_coeff: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], mo_occ: Float[Array, "n_spin n_kpt n_orbitals"], @@ -433,8 +428,7 @@ def make_rdm1( return jnp.einsum("skij,skj,sklj -> skil", mo_coeff, mo_occ, mo_coeff.conj(), precision=precision) -@jaxtyped -@typechecked +@jaxtyped(typechecker=typechecked) def get_occ( mo_energies: Float[Array, "n_spin n_kpt n_orbitals"], nelecs: Int[Array, "spin"], @@ -485,9 +479,8 @@ def assign_values(i, mo_occ): 1BZ. """ -@jaxtyped -@typechecked @partial(jit, static_argnames="precision") +@jaxtyped(typechecker=typechecked) def density(rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], ao: Complex[Array, "n_kpt n_flat_grid n_orbitals"], weights: Float[Array, "n_kpts_or_n_ir_kpts"], @@ -517,9 +510,8 @@ def density(rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], den = jnp.einsum("k,skab,kra,krb->rs", weights, rdm1, ao, ao, precision=precision).real return den -@jaxtyped -@typechecked @partial(jit, static_argnames="precision") +@jaxtyped(typechecker=typechecked) def grad_density( rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], ao: Complex[Array, "n_kpt n_flat_grid n_orbitals"], @@ -554,9 +546,8 @@ def grad_density( return 2 * jnp.einsum("k,...kab,kra,krbj->r...j", weights, rdm1, ao, grad_ao, precision=precision).real -@jaxtyped -@typechecked @partial(jit, static_argnames="precision") +@jaxtyped(typechecker=typechecked) def lapl_density( rdm1: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], ao: Complex[Array, "n_kpt n_flat_grid n_orbitals"], @@ -594,9 +585,8 @@ def lapl_density( "k,...kab,kraj,krbj->r...", weights, rdm1, grad_ao, grad_ao, precision=precision ) + 2 * jnp.einsum("k,...kab,kra,krbi->r...", weights, rdm1, ao, grad_2_ao, precision=precision)).real -@jaxtyped -@typechecked @partial(jit, static_argnames="precision") +@jaxtyped(typechecker=typechecked) def kinetic_density( rdm1 : Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], grad_ao: Complex[Array, "n_kpt n_flat_grid n_orbitals 3"], @@ -628,9 +618,8 @@ def kinetic_density( return 0.5 * jnp.einsum("k,...kab,kraj,krbj->r...", weights, rdm1, grad_ao, grad_ao, precision=precision).real -@jaxtyped -@typechecked @partial(jit, static_argnames="precision") +@jaxtyped(typechecker=typechecked) def orbital_grad( mo_coeff: Complex[Array, "n_spin n_kpt n_orbitals n_orbitals"], mo_occ: Float[Array, "n_spin n_kpt n_orbitals"], From a7cbdbcda8f63ef8395f1b3315432ebbafd3b2c5 Mon Sep 17 00:00:00 2001 From: Anya497 Date: Sun, 5 Jul 2026 13:19:04 +0300 Subject: [PATCH 15/15] Add debug step. --- .github/workflows/install_and_test.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/install_and_test.yml b/.github/workflows/install_and_test.yml index de10def..590b6ee 100644 --- a/.github/workflows/install_and_test.yml +++ b/.github/workflows/install_and_test.yml @@ -35,6 +35,10 @@ jobs: pytest -v tests/unit/test_kinetic_density_eq.py pytest -v tests/unit/test_eigenproblem.py pytest -v tests/unit/test_loss.py + - name: Debug + run: | + pip list + python --version - name: Run integration tests run: | pytest -v tests/integration/molecules/test_non_xc_energy.py