Merge pull request #371 from tovrstra/attrs-basis

Convert Shell attributes to arrays

iodata/basis.py

@@ -32,7 +32,7 @@
 import numpy as np
 from numpy.typing import NDArray
 
-from .attrutils import validate_shape
+from .attrutils import convert_array_to, validate_shape
 
 __all__ = (
     "angmom_sti",
@@ -114,37 +114,49 @@ class Shell:
     icenter: int = attrs.field()
     """An integer index specifying the row in the atcoords array of IOData object."""
 
-    angmoms: list[int] = attrs.field(validator=validate_shape(("coeffs", 1)))
+    angmoms: NDArray[int] = attrs.field(
+        converter=convert_array_to(int),
+        validator=validate_shape(("coeffs", 1)),
+    )
     """An integer array of angular momentum quantum numbers, non-negative, with shape (ncon,).
 
-    In the case of ordinary (not generalized) contractions, this list contains one element.
+    In the case of ordinary (not generalized) contractions, this array contains one element.
 
-    For generalized contractions, this list contains multiple elements.
+    For generalized contractions, this array contains multiple elements.
     The same angular momentum may or may not appear multiple times.
 
     The most common form of generalized contraction is the SP shell,
     e.g. as found in the Pople basis sets (6-31G and others),
-    in which case this list is ``[0, 1]``.
+    in which case this array is ``[0, 1]``.
 
     Other forms of generalized contractions exist,
     but only some quantum chemistry codes have efficient implementations for them.
     For example, the ANO-RCC basis for carbon has 8 S-type basis functions,
     all different linear combinations of the same 14 Gaussian primitives.
-    In this case, this list is ``[0, 0, 0, 0, 0, 0, 0, 0]``.
+    In this case, this array is ``[0, 0, 0, 0, 0, 0, 0, 0]``.
     """
 
-    kinds: list[str] = attrs.field(validator=validate_shape(("coeffs", 1)))
+    kinds: NDArray[str] = attrs.field(
+        converter=convert_array_to(str),
+        validator=validate_shape(("coeffs", 1)),
+    )
     """
-    List of strings describing the kind of contractions:
+    Array of strings describing the kind of contractions:
     ``'c'`` for Cartesian and ``'p'`` for pure.
     Pure functions are only allowed for ``angmom > 1``.
     The length equals the number of contractions (``ncon = len(kinds)``).
     """
 
-    exponents: NDArray[float] = attrs.field(validator=validate_shape(("coeffs", 0)))
+    exponents: NDArray[float] = attrs.field(
+        converter=convert_array_to(float),
+        validator=validate_shape(("coeffs", 0)),
+    )
     """The array containing the exponents of the primitives, with shape (nexp,)."""
 
-    coeffs: NDArray[float] = attrs.field(validator=validate_shape(("exponents", 0), ("kinds", 0)))
+    coeffs: NDArray[float] = attrs.field(
+        converter=convert_array_to(float),
+        validator=validate_shape(("exponents", 0), ("kinds", 0)),
+    )
     """
     The array containing the coefficients of the normalized primitives in each contraction;
     shape = (nexp, ncon).

iodata/formats/fchk.py

@@ -667,7 +667,7 @@ def dump_one(f: TextIO, data: IOData):
                 shell_types.append(shell.angmoms[0])
             elif shell.ncon == 1 and shell.kinds == ["p"]:
                 shell_types.append(-1 * shell.angmoms[0])
-            elif shell.ncon == 2 and shell.angmoms == [0, 1]:
+            elif shell.ncon == 2 and (shell.angmoms == [0, 1]).all():
                 shell_types.append(-1)
             else:
                 raise DumpError("Cannot identify type of shell!", f)

iodata/test/common.py

@@ -112,7 +112,7 @@ def compare_mols(mol1, mol2, atol=1.0e-8, rtol=0.0):
         for shell1, shell2 in zip(mol1.obasis.shells, mol2.obasis.shells):
             assert shell1.icenter == shell2.icenter
             assert_equal(shell1.angmoms, shell2.angmoms)
-            assert shell1.kinds == shell2.kinds
+            assert_equal(shell1.kinds, shell2.kinds)
             assert_allclose(shell1.exponents, shell2.exponents, atol=atol, rtol=rtol)
             assert_allclose(shell1.coeffs, shell2.coeffs, atol=atol, rtol=rtol)
         assert mol1.obasis.primitive_normalization == mol2.obasis.primitive_normalization

iodata/test/test_cp2klog.py

@@ -41,11 +41,11 @@ def test_atom_si_uks():
     assert_allclose(mol.mo.energiesb, [-0.334567, -0.092237, -0.092237, -0.092237], atol=1.0e-4)
     assert_allclose(mol.energy, -3.761587698067, atol=1.0e-10)
     assert len(mol.obasis.shells) == 3
-    assert mol.obasis.shells[0].kinds == ["c", "c"]
+    assert_equal(mol.obasis.shells[0].kinds, ["c", "c"])
     assert_equal(mol.obasis.shells[1].angmoms, [1, 1])
-    assert mol.obasis.shells[1].kinds == ["c", "c"]
+    assert_equal(mol.obasis.shells[1].kinds, ["c", "c"])
     assert_equal(mol.obasis.shells[2].angmoms, [2])
-    assert mol.obasis.shells[2].kinds == ["p"]
+    assert_equal(mol.obasis.shells[2].kinds, ["p"])
     # check mo normalization
     olp = compute_overlap(mol.obasis, mol.atcoords)
     check_orthonormal(mol.mo.coeffsa, olp)
@@ -63,11 +63,11 @@ def test_atom_o_rks():
     assert_allclose(mol.energy, -15.464982778766, atol=1.0e-10)
     assert_equal(mol.obasis.shells[0].angmoms, [0, 0])
     assert len(mol.obasis.shells) == 3
-    assert mol.obasis.shells[0].kinds == ["c", "c"]
+    assert_equal(mol.obasis.shells[0].kinds, ["c", "c"])
     assert_equal(mol.obasis.shells[1].angmoms, [1, 1])
-    assert mol.obasis.shells[1].kinds == ["c", "c"]
+    assert_equal(mol.obasis.shells[1].kinds, ["c", "c"])
     assert_equal(mol.obasis.shells[2].angmoms, [2])
-    assert mol.obasis.shells[2].kinds == ["p"]
+    assert_equal(mol.obasis.shells[2].kinds, ["p"])
     # check mo normalization
     olp = compute_overlap(mol.obasis, mol.atcoords)
     check_orthonormal(mol.mo.coeffs, olp)

iodata/test/test_fchk.py

@@ -65,17 +65,17 @@ def test_load_fchk_hf_sto3g_num():
     assert len(mol.obasis.shells) == 3
     shell0 = mol.obasis.shells[0]
     assert shell0.icenter == 0
-    assert shell0.angmoms == [0]
-    assert shell0.kinds == ["c"]
+    assert_equal(shell0.angmoms, [0])
+    assert_equal(shell0.kinds, ["c"])
     assert_allclose(shell0.exponents, np.array([1.66679134e02, 3.03608123e01, 8.21682067e00]))
     assert_allclose(shell0.coeffs, np.array([[1.54328967e-01], [5.35328142e-01], [4.44634542e-01]]))
     assert shell0.nexp == 3
     assert shell0.ncon == 1
     assert shell0.nbasis == 1
     shell1 = mol.obasis.shells[1]
     assert shell1.icenter == 0
-    assert shell1.angmoms == [0, 1]
-    assert shell1.kinds == ["c", "c"]
+    assert_equal(shell1.angmoms, [0, 1])
+    assert_equal(shell1.kinds, ["c", "c"])
     assert_allclose(shell1.exponents, np.array([6.46480325e00, 1.50228124e00, 4.88588486e-01]))
     assert_allclose(
         shell1.coeffs,

iodata/test/test_molden.py

@@ -612,8 +612,8 @@ def test_mixed_pure_cartesian(tmpdir):
         atcoords=[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]],
         obasis=MolecularBasis(
             [
-                Shell(0, [2], ["c"], np.array([1.0]), np.array([[1.0]])),
-                Shell(0, [2], ["p"], np.array([1.0]), np.array([[1.0]])),
+                Shell(0, [2], ["c"], [1.0], [[1.0]]),
+                Shell(0, [2], ["p"], [1.0], [[1.0]]),
             ],
             HORTON2_CONVENTIONS,
             "L2",

iodata/test/test_overlap.py

@@ -57,9 +57,9 @@ def test_factorial2_float_array_argument():
 
 def test_normalization_basics_segmented():
     for angmom in range(7):
-        shells = [Shell(0, [angmom], ["c"], np.array([0.23]), np.array([[1.0]]))]
+        shells = [Shell(0, [angmom], ["c"], [0.23], [[1.0]])]
         if angmom >= 2:
-            shells.append(Shell(0, [angmom], ["p"], np.array([0.23]), np.array([[1.0]])))
+            shells.append(Shell(0, [angmom], ["p"], [0.23], [[1.0]]))
         obasis = MolecularBasis(shells, OVERLAP_CONVENTIONS, "L2")
         atcoords = np.zeros((1, 3))
         overlap = compute_overlap(obasis, atcoords)
@@ -68,7 +68,7 @@ def test_normalization_basics_segmented():
 
 def test_normalization_basics_generalized():
     for angmom in range(2, 7):
-        shells = [Shell(0, [angmom] * 2, ["c", "p"], np.array([0.23]), np.array([[1.0, 1.0]]))]
+        shells = [Shell(0, [angmom] * 2, ["c", "p"], [0.23], [[1.0, 1.0]])]
         obasis = MolecularBasis(shells, OVERLAP_CONVENTIONS, "L2")
         atcoords = np.zeros((1, 3))
         overlap = compute_overlap(obasis, atcoords)