diff --git a/iodata/formats/json.py b/iodata/formats/json.py index 18ea7026a..5e8d7090c 100644 --- a/iodata/formats/json.py +++ b/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 diff --git a/iodata/test/data/LiCl_STO4G_basis.json b/iodata/test/data/LiCl_STO4G_basis.json new file mode 100644 index 000000000..57d8bd527 --- /dev/null +++ b/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"] +} \ No newline at end of file diff --git a/iodata/test/test_json.py b/iodata/test/test_json.py index a2d243f1b..45397fedf 100644 --- a/iodata/test/test_json.py +++ b/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