Add QCSchema Basis

iodata/formats/json.py

@@ -571,6 +571,7 @@
 
 import numpy as np
 
+from ..basis import CCA_CONVENTIONS, Shell, MolecularBasis
 from ..docstrings import document_dump_one, document_load_one
 from ..iodata import IOData
 from ..periodic import num2sym, sym2num
@@ -976,55 +977,132 @@ def _find_passthrough_dict(result: dict, keys: set) -> dict:
     return passthrough_dict
 
 
-def _load_qcschema_basis(_result: dict, _lit: LineIterator) -> dict:
+def _load_qcschema_basis(result: dict, lit: LineIterator) -> dict:
     """Load qcschema_basis properties.
 
     Parameters
     ----------
-    _result
+    result
         The JSON dict loaded from file.
-    _lit
+    lit
         The line iterator holding the file data.
 
     Returns
     -------
     basis_dict
-        ...
+        Dictionary containing ``obasis``, ``obasis_name`` & ``extra`` keys.
 
+    """
+    extra_dict = dict()
+    basis_dict = _parse_basis_keys(result, lit)
+    extra_dict["basis"] = basis_dict["extra"]
 
-    Raises
-    ------
-    NotImplementedError
-        QCSchema Basis schema is not yet implemented in IOData.
+    basis_dict["extra"] = extra_dict
+    basis_dict["extra"]["schema_name"] = "qcschema_basis"
 
-    """
-    # basis_dict = {}
-    # return basis_dict
-    raise NotImplementedError("qcschema_basis is not yet implemented in IOData.")
+    return basis_dict
 
 
-def _parse_basis_keys(_basis: dict, _lit: LineIterator) -> dict:
+def _parse_basis_keys(basis: dict, lit: LineIterator) -> dict:
     """Parse basis keys for a QCSchema input, output, or basis file.
 
     Parameters
     ----------
-    _basis
+    basis
         The basis dictionary from a QCSchema basis file or QCSchema input or output 'method' key.
-    _lit
+    lit
         The line iterator holding the file data.
 
     Returns
     -------
     basis_dict
-        Dictionary containing ...
-
-    Raises
-    ------
-    NotImplementedError
-        QCSchema Basis schema is not yet implemented in IOData.
+        Dictionary containing ``obasis``, ``obasis_name``, & ``extra`` keys.
 
     """
-    raise NotImplementedError("qcschema_basis is not yet implemented in IOData.")
+    # Check for required properties:
+    # NOTE: description is optional in QCElemental, required in v1.dev
+    basis_keys = {
+        "schema_name",
+        "schema_version",
+        "name",
+        "center_data",
+        "atom_map",
+    }
+    for key in basis_keys:
+        if key not in basis:
+            raise FileFormatWarning(
+                "{}: QCSchema `qcschema_basis` requires '{}' key".format(lit.filename, key)
+            )
+
+    basis_dict = dict()
+    extra_dict = dict()
+    extra_dict["schema_name"] = basis["schema_name"]
+    extra_dict["schema_version"] = basis["schema_version"]
+    basis_dict["obasis_name"] = basis["name"]
+    # Load basis data
+    center_data = basis["center_data"]
+    atom_map = basis["atom_map"]
+    center_shells = dict()
+    # Center_data is composed of basis_center, each has req'd electron_shells, ecp_electrons,
+    # and optional ecp_potentials
+    for center in center_data:
+        # Initiate lists for building basis
+        center_shells[center] = list()
+        # QCElemental example omits ecp_electrons for cases with default value (0)
+        if "electron_shells" not in center_data[center]:
+            raise FileFormatError(
+                "{}: Basis center {} requires `electron_shells` key".format(lit.filename, center)
+            )
+        if "ecp_electrons" in center_data[center]:
+            ecp_electrons = center["ecp_electrons"]
+        else:
+            ecp_electrons = 0
+        shells = center_data[center]["electron_shells"]
+        for shell in shells:
+            # electron_shell requires angular_momentum, harmonic_type, exponents, coefficients
+            for key in ["angular_momentum", "harmonic_type", "exponents", "coefficients"]:
+                if key not in shell:
+                    raise FileFormatError(
+                        "{}: Basis center {} contains a shell missing '{}' key".format(
+                            lit.filename, center, key
+                        )
+                    )
+            # Load shell data
+            if shell["harmonic_type"] not in {"cartesian", "spherical"}:
+                raise FileFormatError(
+                    "{}: `harmonic_type` must be `cartesian` or `spherical`".format(
+                        lit.filename,
+                    )
+                )
+            angmoms = shell["angular_momentum"]
+            exps = np.array(shell["exponents"])
+            coeffs = np.array([[x for x in segment] for segment in shell["coefficients"]])
+            coeffs = coeffs.T
+            kinds = [shell["harmonic_type"] for _ in range(len(angmoms))]
+
+            # Gather shell components
+            center_shells[center].append(
+                {"angmoms": angmoms, "kinds": kinds, "exponents": exps, "coeffs": coeffs}
+            )
+
+    # Build obasis shells using the atom_map
+    # Each atom in atom_map corresponds to a key in center_shells
+    obasis_shells = list()
+    for i, atom in enumerate(atom_map):
+        for shell in center_shells[atom]:
+            # Unpack angmoms, kinds, exponents, coeffs into obasis
+            obasis_shells.append(Shell(icenter=i, **shell))
+
+    # These are assumed within QCSchema
+    conventions = CCA_CONVENTIONS
+    prim_norm = "L2"
+
+    basis_dict["obasis"] = MolecularBasis(
+        shells=obasis_shells, conventions=conventions, primitive_normalization=prim_norm
+    )
+    basis_dict["extra"] = extra_dict
+
+    return basis_dict
 
 
 def _load_qcschema_input(result: dict, lit: LineIterator) -> dict:
@@ -1464,8 +1542,7 @@ def dump_one(f: TextIO, data: IOData):
     if schema_name == "qcschema_molecule":
         return_dict = _dump_qcschema_molecule(data)
     elif schema_name == "qcschema_basis":
-        raise NotImplementedError("{} not yet implemented in IOData.".format(schema_name))
-        # return_dict = _dump_qcschema_basis(data)
+        return_dict = _dump_qcschema_basis(data)
     elif schema_name == "qcschema_input":
         return_dict = _dump_qcschema_input(data)
     elif schema_name == "qcschema_output":
@@ -1742,3 +1819,47 @@ def _dump_qcschema_output(data: IOData) -> dict:
             output_dict[k] = data.extra["input"]["unparsed"][k]
 
     return output_dict
+
+
+def _dump_qcschema_basis(data: IOData) -> dict:
+    """Dump relevant attributes from IOData to `qcschema_basis`.
+
+    Using this function requires one entry in the `extra` dict: 'schema_name' = 'qcschema_basis'.
+
+    Parameters
+    ----------
+    data
+        The IOData instance to dump to file.
+
+    Returns
+    -------
+    basis_dict
+        The dict that will produce the QCSchema JSON file.
+
+    """
+    basis_dict = {"schema_name": "qcschema_basis", "schema_version": 1}
+    if not data.obasis_name:
+        raise FileFormatError("qcschema_basis requires `obasis_name`")
+    if not data.obasis:
+        raise FileFormatError("qcschema_basis requires `obasis`")
+
+    basis_dict["name"] = data.obasis_name
+    basis_dict["center_data"] = dict()
+    for shell in data.obasis.shells:
+        if len(set(shell.kinds)) > 1:
+            raise FileFormatError("qcschema_basis does not support mixed kinds in one shell."
+                                  "To support this functionality consider the BasisSetExchange's"
+                                  "similar JSON schema.")
+        if shell.icenter not in basis_dict["center_data"]:
+            basis_dict["center_data"][str(shell.icenter)] = {"electron_shells": list()}
+        basis_dict["center_data"][str(shell.icenter)]["electron_shells"].append(
+            {
+                "angular_momentum": shell.angmoms,
+                "exponents": shell.exponents.tolist(),
+                "coefficients": shell.coeffs.tolist(),
+                "harmonic_type": shell.kinds[0]
+            }
+        )
+        basis_dict["atom_map"] = list(basis_dict["center_data"].keys())
+
+    return basis_dict

iodata/test/data/LiCl_STO4G_basis.json

@@ -0,0 +1,137 @@
+{
+  "schema_name": "qcschema_basis",
+  "schema_version": 1,
+  "name": "STO-4G",
+  "center_data": {
+    "0": {
+      "electron_shells": [
+        {
+          "angular_momentum": [
+            0
+          ],
+          "exponents": [
+            0.3774960873E+02,
+            0.6907713307E+01,
+            0.1919038397E+01,
+            0.6369115922E+00
+          ],
+          "coefficients": [
+            [
+              0.5675242080E-01,
+              0.2601413550E+00,
+              0.5328461143E+00,
+              0.2916254405E+00
+            ]
+          ],
+          "harmonic_type": "cartesian"
+        },
+        {
+          "angular_momentum": [
+            0,
+            1
+          ],
+          "exponents": [
+            0.1487042352E+01,
+            0.3219127620E+00,
+            0.1046660300E+00,
+            0.4019868296E-01
+          ],
+          "coefficients": [
+            [
+              -0.6220714565E-01,
+              0.2976804596E-04,
+              0.5588549221E+00,
+              0.4977673218E+00
+            ],
+            [
+              0.4368434884E-01,
+              0.2863793984E+00,
+              0.5835753141E+00,
+              0.2463134378E+00
+            ]
+          ],
+          "harmonic_type": "cartesian"
+        }
+      ]
+    },
+    "1": {
+      "electron_shells": [
+        {
+          "angular_momentum": [
+            0
+          ],
+          "exponents": [
+            0.1408260576E+04,
+            0.2576943351E+03,
+            0.7159030805E+02,
+            0.2376017966E+02
+          ],
+          "coefficients": [
+            [
+              0.5675242080E-01,
+              0.2601413550E+00,
+              0.5328461143E+00,
+              0.2916254405E+00
+            ]
+          ],
+          "harmonic_type": "cartesian"
+        },
+        {
+          "angular_momentum": [
+            0,
+            1
+          ],
+          "exponents": [
+            0.9105253261E+02,
+            0.1971091961E+02,
+            0.6408766434E+01,
+            0.2461390482E+01
+          ],
+          "coefficients": [
+            [
+              -0.6220714565E-01,
+              0.2976804596E-04,
+              0.5588549221E+00,
+              0.4977673218E+00
+            ],
+            [
+              0.4368434884E-01,
+              0.2863793984E+00,
+              0.5835753141E+00,
+              0.2463134378E+00
+            ]
+          ],
+          "harmonic_type": "cartesian"
+        },
+        {
+          "angular_momentum": [
+            0,
+            1
+          ],
+          "exponents": [
+            0.4064728946E+01,
+            0.1117815984E+01,
+            0.4399626124E+00,
+            0.1958035957E+00
+          ],
+          "coefficients": [
+            [
+              -0.8529019644E-01,
+              -0.2132074034E+00,
+              0.5920843928E+00,
+              0.6115584746E+00
+            ],
+            [
+              -0.2504945181E-01,
+              0.1686604461E+00,
+              0.6409553151E+00,
+              0.2779508957E+00
+            ]
+          ],
+          "harmonic_type": "cartesian"
+        }
+      ]
+    }
+  },
+  "atom_map": ["0", "1"]
+}

iodata/test/test_json.py

@@ -25,6 +25,7 @@
 import pytest
 
 from ..api import dump_one, load_one
+from ..basis import CCA_CONVENTIONS
 from ..utils import FileFormatError, FileFormatWarning
 
 
@@ -405,3 +406,22 @@ def test_inout_qcschema_output(tmpdir, filename):
     if "provenance" in mol2["molecule"]:
         del mol2["molecule"]["provenance"]
     assert mol1 == mol2
+
+
+# BASIS_FILES: {filename}
+BASIS_FILES = [
+    ("LiCl_STO4G_basis.json"),
+]
+
+
+@pytest.mark.parametrize("filename", BASIS_FILES)
+def test_qcschema_basis(filename):
+    """Test qcschema_basis with quick check."""
+    with path("iodata.test.data", filename) as qcschema_basis:
+        basis = load_one(str(qcschema_basis))
+
+    print(basis.obasis)
+    assert basis.obasis_name == "STO-4G"
+    assert basis.obasis.conventions == CCA_CONVENTIONS
+    assert basis.obasis.primitive_normalization == "L2"
+    assert len(basis.obasis.shells) == 5