fix test_numeric: trim arrays to its real data

Author: enjyashraf18

atomdb/species.py

@@ -179,6 +179,7 @@ class DensitySpline:
     def __init__(self, x, y, log=False):
         r"""Initialize the CubicSpline instance."""
         self._log = log
+        x, y = self._trim_padded_data(x, y)
         self._obj = CubicSpline(
             x,
             # Clip y values to >= ε^2 if using log because they have to be above 0;
@@ -189,6 +190,13 @@ def __init__(self, x, y, log=False):
             extrapolate=True,
         )
 
+    def _trim_padded_data(self, data_x, data_y, tol=1e-10):
+        """Trim padded zeros from the end of arrays."""
+        trimmed_x = trim_padded_array(data_x, tol)
+        effective_length = len(trimmed_x)
+        trimmed_y = data_y[:effective_length]
+        return trimmed_x, trimmed_y
+
     def __call__(self, x, deriv=0):
         r"""
         Compute the interpolation at some x-values.
@@ -1002,9 +1010,33 @@ def get_species_data(DATASETS_H5FILE, folder_path, elem, DATASET_PROPERTY_CONFIG
     for prop in ("atmass", "cov_radius", "vdw_radius", "at_radius", "polarizability", "dispersion"):
         fields[prop] = get_scalar_data(prop, fields["atnum"], fields["nelec"])
 
+    if "rs" in fields:
+        rs_trimmed_array = trim_padded_array(fields["rs"])
+        effective_length = len(rs_trimmed_array)
+        for config in DATASET_PROPERTY_CONFIGS:
+            if "Carray_property" in config:
+                # Trim only if spins == "no" or spins is not specified
+                if config.get("spins") == "no" or ("spins" not in config):
+                    fields[config["Carray_property"]] = fields[config["Carray_property"]][
+                        :effective_length
+                    ]
     return fields
 
 
+def trim_padded_array(data, tol=1e-10):
+    """Trim padded zeros from the end of arrays."""
+    # Find non-zero elements
+    non_zero_mask = np.abs(data) > tol
+    non_zero_indices = np.where(non_zero_mask)[0]
+
+    if len(non_zero_indices) == 0:
+        return data[:0]
+
+    # Get the last non-zero index
+    last_non_zero_index = non_zero_indices[-1]
+    return data[: (last_non_zero_index + 1)]
+
+
 def raw_datafile(
     suffix,
     elem,

atomdb/test/test_numeric.py

@@ -45,8 +45,8 @@ def test_numerical_hf_data_h():
     sp = load("H", 0, 2, dataset="numeric", datapath=TEST_DATAPATH)
 
     # check shape radial grid and total density arrays
-    assert sp._data.rs.shape == (1000,)
-    assert sp._data.dens_tot.shape == sp._data.rs.shape
+    npoints = len(sp._data.rs)
+    assert sp._data.dens_tot.shape == (npoints,)
 
     # check radial grid and total density arrays values
     assert all(sp._data.rs >= 0.0)
@@ -81,8 +81,8 @@ def test_numerical_hf_data_h_anion():
     assert_almost_equal(sp.energy, -0.487929734301232, decimal=10)
 
     # check shape radial grid and total density arrays
-    assert sp._data.rs.shape == (1000,)
-    assert sp._data.dens_tot.shape == sp._data.rs.shape
+    npoints = len(sp._data.rs)
+    assert sp._data.dens_tot.shape == (npoints,)
 
     # reference radial values sample and corresponding indices
     ref_rs = np.array(
@@ -135,7 +135,7 @@ def test_numerical_hf_energy_especies(atom, mult, energy):
 
 
 @pytest.mark.parametrize(
-    "atom, mult, npoints, nelec", [("Be", 1, 1000, 4.0), ("Cl", 2, 1000, 17.0), ("Ne", 1, 1000, 10.0)]
+    "atom, mult, npoints, nelec", [("Be", 1, 146, 4.0), ("Cl", 2, 164, 17.0), ("Ne", 1, 151, 10.0)]
 )
 def test_numerical_hf_atomic_density(atom, mult, npoints, nelec):
     # load atomic and density data
@@ -230,6 +230,7 @@ def test_numerical_hf_density_laplacian(atom, charge, mult):
     ref_lapl = np.load(f"{TEST_DATAPATH}/numeric/db/{fname}")
 
     # check interpolated Laplacian of density values against reference values
-    assert np.allclose(laplacian_dens, ref_lapl, atol=1e-10)
+    meaningful_length = len(ref_lapl)
+    assert np.allclose(laplacian_dens[:meaningful_length], ref_lapl, atol=1e-10)
     # for r=0, the Laplacian function in not well defined and is set to zero
     assert np.allclose(laplacian_dens[0], [0.0], atol=1e-10)