Add pre-fitted atomic Gaussian parameter loader and data

MANIFEST.in

@@ -1,4 +1,4 @@
 include README.md
 include LICENSE
 recursive-include grid/src/data *.npz
-global-exclude *.pyc __pycache__ .DS_Store
+global-exclude *.pyc __pycache__ .DS_Store

src/grid/__init__.py

@@ -35,3 +35,5 @@
 from grid.periodicgrid import *
 from grid.rtransform import *
 from grid.ngrid import *
+from grid.coulomb import *
+

src/grid/coulomb.py

@@ -26,10 +26,20 @@
 
 from __future__ import annotations
 
+import json
+from importlib.resources import files
 import numpy as np
 from scipy.special import erf
 
-__all__ = ["coulomb_gaussian_p", "coulomb_gaussian_s", "coulomb_potential"]
+from grid.utils import num2sym, sym2num
+
+
+__all__ = [
+    "coulomb_gaussian_p",
+    "coulomb_gaussian_s",
+    "coulomb_potential",
+    "load_atomic_gaussian_params",
+]
 
 # Distance threshold below which the r->0 analytical limit is used
 # instead of erf(x)/x to avoid division by zero at atomic nuclei.
@@ -238,3 +248,70 @@ def coulomb_potential(
             V += c * coulomb_gaussian_p(r, alpha, normalized=normalized)
 
     return V
+
+
+_ATOMIC_GAUSS_PARAMS_CACHE: dict | None = None
+
+
+def load_atomic_gaussian_params(element: str | int) -> tuple[np.ndarray, np.ndarray]:
+    """Load pre-fitted s-type Gaussian parameters for a given element.
+
+    Parameters
+    ----------
+    element : str or int
+        Element symbol (case-insensitive, e.g., "H", "C") or atomic number (e.g., 1, 6).
+
+    Returns
+    -------
+    coeffs_s : np.ndarray
+        Coefficients of the normalized s-type Gaussians.
+    alphas_s : np.ndarray
+        Exponents of the s-type Gaussians.
+
+    Raises
+    ------
+    ValueError
+        If pre-fitted parameters are not available for the element.
+    TypeError
+        If the element parameter has an invalid type.
+    """
+    if isinstance(element, str):
+        json_symbol = element.strip().title()
+        if json_symbol not in sym2num:
+            raise ValueError(
+                f"Unknown element symbol: '{element}'. "
+                f"Use a valid symbol such as 'H', 'C', 'N', 'O', 'Cl'."
+            )
+
+    elif isinstance(element, (int, np.integer)):
+        atnum = int(element)
+        json_symbol = num2sym.get(atnum)
+        if json_symbol is None:
+            raise ValueError(
+                f"Unknown atomic number: {atnum}. "
+                f"Must be between 1 and 118."
+            )
+    else:
+        raise TypeError(f"element must be a string or integer; got {type(element)}")
+
+    global _ATOMIC_GAUSS_PARAMS_CACHE
+    if _ATOMIC_GAUSS_PARAMS_CACHE is None:
+        try:
+            path = files("grid.data").joinpath("atomic_gauss_params.json")
+            with path.open("r", encoding="utf-8") as f:
+                _ATOMIC_GAUSS_PARAMS_CACHE = json.load(f)
+        except Exception as e:
+            raise ValueError(
+                f"Could not load pre-fitted Gaussian parameters from data: {e}"
+            ) from e
+
+    try:
+        data = _ATOMIC_GAUSS_PARAMS_CACHE[json_symbol]
+    except KeyError as e:
+        raise ValueError(
+            f"Pre-fitted Gaussian parameters are not available for element: {json_symbol}"
+        ) from e
+    coeffs_s = np.asarray(data["coeffs_s"], dtype=float)
+    alphas_s = np.asarray(data["alphas_s"], dtype=float)
+
+    return coeffs_s, alphas_s

src/grid/data/atomic_gauss_params.json

@@ -0,0 +1,194 @@
+{
+  "H": {
+    "coeffs_s": [
+      2.00748526502357e-08,
+      3.314086085258354e-07,
+      6.414728908403625e-07,
+      2.910120019422806e-06,
+      1.169435093185597e-05,
+      4.1850680933381154e-05,
+      0.00016631725407837694,
+      0.0006132764002116911,
+      0.0023124173488736103,
+      0.008312797303960385,
+      0.028567320400773935,
+      0.08757483702030483,
+      0.21788912723883622,
+      0.3566371897730248,
+      0.2598257779286926,
+      0.03697311221759376,
+      0.0012855905887847435
+    ],
+    "alphas_s": [
+      13739.0854166734,
+      1829.8696247655,
+      934.48913472921,
+      477.230689612032,
+      243.715119463182,
+      124.461944187287,
+      63.5609952513411,
+      32.4597220758638,
+      16.5767315800501,
+      8.46550778330153,
+      4.32321786011102,
+      2.20780762884063,
+      1.12749685157938,
+      0.575797063890464,
+      0.294051604951678,
+      0.150168091845475,
+      0.07668876968794858
+    ]
+  },
+  "C": {
+    "coeffs_s": [
+      5.894723103300583e-06,
+      2.0635070043091174e-05,
+      0.00018715066301433259,
+      0.00036029859600310104,
+      0.001988795020955914,
+      0.0060704473135859685,
+      0.0245314953150802,
+      0.07338393107647316,
+      0.23203826850409184,
+      0.4702171969739957,
+      0.7968878189081399,
+      0.2697305408775165,
+      0.24553524504245708,
+      3.5153129058837447,
+      0.27863778573369447
+    ],
+    "alphas_s": [
+      26903.1860742976,
+      13739.0854166734,
+      3583.15866840946,
+      1829.8696247655,
+      934.489134729211,
+      477.230689612032,
+      243.715119463182,
+      124.461944187287,
+      63.5609952513412,
+      32.4597220758638,
+      16.5767315800501,
+      8.46550778330153,
+      1.12749685157938,
+      0.575797063890465,
+      0.0766887696879486
+    ]
+  },
+  "N": {
+    "coeffs_s": [
+      8.197232108001016e-07,
+      9.997368964801968e-05,
+      0.00019434755305028855,
+      0.0009258494859595727,
+      0.0032365377469431036,
+      0.012145317606417624,
+      0.04122202653070691,
+      0.13128460644333742,
+      0.3400527411876631,
+      0.6475026339991764,
+      0.6675682196705341,
+      2.993596943657288,
+      1.0821003670657379,
+      1.12031828185087,
+      0.0030724785923507457
+    ],
+    "alphas_s": [
+      201995.358823884,
+      7016.361094383,
+      3583.15866840946,
+      1829.8696247655,
+      934.489134729211,
+      477.230689612032,
+      243.715119463182,
+      124.461944187287,
+      63.5609952513412,
+      32.4597220758638,
+      16.5767315800501,
+      1.12749685157938,
+      0.575797063890465,
+      0.294051604951678,
+      0.150168091845475
+    ]
+  },
+  "O": {
+    "coeffs_s": [
+      5.009541670630619e-06,
+      1.7553971507812904e-05,
+      0.00015918253361386784,
+      0.0003080424448463953,
+      0.001694236412541336,
+      0.005216001888759684,
+      0.02115298710403933,
+      0.06453933153335781,
+      0.20850144407008608,
+      0.4454735348098692,
+      0.786412336281571,
+      0.3153416859276515,
+      0.8815501738100202,
+      4.238787122136956,
+      0.261215531613026,
+      1.3436858672806462
+    ],
+    "alphas_s": [
+      52680.4659115021,
+      26903.1860742975,
+      7016.36109438299,
+      3583.15866840945,
+      1829.8696247655,
+      934.48913472921,
+      477.230689612032,
+      243.715119463182,
+      124.461944187287,
+      63.5609952513411,
+      32.4597220758638,
+      16.5767315800501,
+      2.20780762884063,
+      1.12749685157938,
+      0.294051604951678,
+      0.150168091845475
+    ]
+  },
+  "Cl": {
+    "coeffs_s": [
+      6.664702268984005e-08,
+      8.917187590865874e-07,
+      4.64609925513115e-06,
+      2.592823074586147e-06,
+      6.471232236659015e-05,
+      0.00010097326217239535,
+      0.0007372237647551631,
+      0.001892144459242878,
+      0.009297899882357784,
+      0.026647792671694457,
+      0.10329057994254036,
+      0.26096677847075644,
+      0.6120984545408865,
+      0.6900141284561521,
+      1.3652433276964213,
+      6.231393701987117,
+      3.71510313011454,
+      5.198668137511898
+    ],
+    "alphas_s": [
+      2969784.65079438,
+      395537.152566831,
+      201995.358823884,
+      103156.238855453,
+      52680.4659115022,
+      26903.1860742976,
+      13739.0854166734,
+      7016.361094383,
+      3583.15866840946,
+      1829.8696247655,
+      934.489134729211,
+      477.230689612032,
+      243.715119463182,
+      124.461944187287,
+      16.5767315800501,
+      8.46550778330153,
+      0.575797063890465,
+      0.294051604951678
+    ]
+  }
+}

src/grid/tests/test_coulomb.py

@@ -20,10 +20,16 @@
 """Tests for analytical Coulomb potentials of Gaussians."""
 
 import numpy as np
+import pytest
 from numpy.testing import assert_allclose
 from scipy.special import erf
 
-from grid.coulomb import coulomb_gaussian_p, coulomb_gaussian_s, coulomb_potential
+from grid.coulomb import (
+    coulomb_gaussian_p,
+    coulomb_gaussian_s,
+    coulomb_potential,
+    load_atomic_gaussian_params,
+)
 
 
 def test_coulomb_gaussian_s():
@@ -98,3 +104,38 @@ def test_coulomb_potential():
         v_expected += c * val
 
     assert_allclose(v_s, v_expected)
+
+
+def test_load_atomic_gaussian_params():
+
+    # Helper func to check array shapes and non-negativity
+    def _check_params(coeffs, alphas):
+        assert isinstance(coeffs, np.ndarray)
+        assert isinstance(alphas, np.ndarray)
+        assert len(coeffs) == len(alphas) > 0
+        assert np.all(coeffs >= 0)
+        assert np.all(alphas > 0)
+
+    # Test valid elements using symbols (case-insensitive and with spaces)
+    for symbol in ("H", " C ", "cl", "N", "O"):
+        _check_params(*load_atomic_gaussian_params(symbol))
+
+    # Test valid elements using atomic numbers
+    for atnum in (1, 6, 7, 8, 17):
+        _check_params(*load_atomic_gaussian_params(atnum))
+
+    # Symbol and atomic number must resolve to identical arrays
+    c1, a1 = load_atomic_gaussian_params("Cl")
+    c2, a2 = load_atomic_gaussian_params(17)
+    assert_allclose(c1, c2)
+    assert_allclose(a1, a2)
+
+    # Error cases
+    with pytest.raises(ValueError, match="Pre-fitted Gaussian parameters are not available"):
+        load_atomic_gaussian_params("Xe")
+
+    with pytest.raises(ValueError, match="Pre-fitted Gaussian parameters are not available"):
+        load_atomic_gaussian_params(2)
+
+    with pytest.raises(TypeError, match="element must be a string or integer"):
+        load_atomic_gaussian_params(1.0)