def to_wavefunction_properties(
job_output: pb.JobOutput,
atomic_input: AtomicInput) -> WavefunctionProperties:
"""Extract WavefunctionProperties from JobOutput protobuf message"""
jo_dict = MessageToDict(job_output, preserving_proto_field_name=True)
return WavefunctionProperties(
basis=BasisSet(
name=atomic_input.model.basis,
center_data={}, # TODO: need to fill out
atom_map=[], # TODO: need to fill out
),
restricted=atomic_input.keywords.get("restricted", True),
scf_eigenvalues_a=jo_dict.get("orba_energies"),
scf_occupations_a=jo_dict.get("orba_occupations"),
scf_eigenvalues_b=jo_dict.get("orbb_energies", []),
scf_occupations_b=jo_dict.get("orbb_occupations", []),
)
"""
Tests the DQM compute dispatch module
"""
import numpy as np
import pytest
from qcelemental.models import AtomicInput, BasisSet
from qcelemental.tests.test_model_results import center_data
import qcengine as qcng
from qcengine.testing import has_program, using
qcsk_bs = BasisSet(name="custom_basis", center_data=center_data, atom_map=["bs_sto3g_h", "bs_sto3g_h"])
_canonical_methods = [
("dftd3", {"method": "b3lyp-d3"}, {}),
("qcore", {"method": "pbe", "basis": "6-31G"}, {}),
("molpro", {"method": "hf", "basis": "6-31G"}, {}),
("mopac", {"method": "PM6"}, {}),
("mp2d", {"method": "MP2-DMP2"}, {}),
("nwchem", {"method": "hf", "basis": "6-31G"}, {}),
("openmm", {"method": "openff-1.0.0", "basis": "smirnoff"}, {}),
("psi4", {"method": "hf", "basis": "6-31G"}, {}),
("qchem", {"method": "hf", "basis": "6-31G"}, {}),
("rdkit", {"method": "UFF"}, {}),
("terachem_pbs", {"method": "b3lyp", "basis": "6-31G"}, {}),
("torchani", {"method": "ANI1x"}, {}),
("turbomole", {"method": "pbe", "basis": "6-31G"}, {}),
("xtb", {"method": "GFN2-xTB"}, {}),
("adcc", {"method": "adc2", "basis": "6-31G"}, {"n_triplets": 3}),
("gcp", {"method": "hf3c"}, {}),
("mrchem", {"method": "blyp"}, {"world_prec": 1.0e-3}),
def parse_output(self, output: Dict[str, Any],
input_model: "AtomicInput") -> "AtomicResult":
wavefunction_map = {
"orbitals_alpha": "scf_orbitals_a",
"orbitals_beta": "scf_orbitals_b",
"density_alpha": "scf_density_a",
"density_beta": "scf_density_b",
"fock_alpha": "scf_fock_a",
"fock_beta": "scf_fock_b",
"eigenvalues_alpha": "scf_eigenvalues_a",
"eigenvalues_beta": "scf_eigenvalues_b",
"occupations_alpha": "scf_occupations_a",
"occupations_beta": "scf_occupations_b",
}
output_data = input_model.dict()
output_data["return_result"] = output[input_model.driver.value]
# Always build a wavefunction, it will be stripped
obas = output["wavefunction"]["ao_basis"]
for k, center in obas["center_data"].items():
# Convert basis set, cannot handle arrays
for shell in center["electron_shells"]:
shell.pop("normalized_primitives", None)
for el_k in ["coefficients", "exponents", "angular_momentum"]:
shell[el_k] = shell[el_k].tolist()
if center["ecp_potentials"] is not None:
for shell in center["ecp_potentials"]:
shell.pop("ecp_potentials", None)
for ecp_k in [
"angular_momentum", "r_exponents",
"gaussian_exponents", "coefficients"
]:
shell[ecp_k] = shell[ecp_k].tolist()
basis_set = BasisSet(name=str(input_model.model.basis),
center_data=obas["center_data"],
atom_map=obas["atom_map"])
wavefunction = {"basis": basis_set}
for key, qcschema_key in wavefunction_map.items():
qcore_data = output["wavefunction"].get(key, None)
if qcore_data is None:
continue
if ("density" in key) or ("fock" in key):
qcore_data = reorder_row_and_column_ao_indices(
qcore_data, basis_set, self._qcore_to_cca_ao_order)
# Handles orbitals and 1D
elif "orbitals" in key:
qcore_data = reorder_column_ao_indices(
qcore_data, basis_set, self._qcore_to_cca_ao_order)
elif "eigenvalues" in key:
qcore_data = reorder_column_ao_indices(
qcore_data.reshape(1, -1), basis_set,
self._qcore_to_cca_ao_order).ravel()
elif "occupations" in key:
tmp = np.zeros(basis_set.nbf)
tmp[:qcore_data.shape[0]] = qcore_data
qcore_data = reorder_column_ao_indices(
tmp.reshape(1, -1), basis_set,
self._qcore_to_cca_ao_order).ravel()
else:
raise KeyError("Wavefunction conversion key not understood")
wavefunction[qcschema_key] = qcore_data
wavefunction["restricted"] = True
if "scf_eigenvalues_b" in wavefunction:
wavefunction["restricted"] = False
output_data["wavefunction"] = wavefunction
# Handle remaining top level keys
properties = {
"calcinfo_nbasis": basis_set.nbf,
"calcinfo_nmo": basis_set.nbf,
"calcinfo_nalpha": np.sum(wavefunction["scf_occupations_a"] > 0),
"calcinfo_natom": input_model.molecule.symbols.shape[0],
"return_energy": output["energy"],
}
if wavefunction["restricted"]:
properties["calcinfo_nbeta"] = properties["calcinfo_nalpha"]
else:
properties["calcinfo_nbeta"] = np.sum(
wavefunction["scf_occupations_b"] > 0)
output_data["properties"] = properties
output_data["schema_name"] = "qcschema_output"
output_data["success"] = True
return AtomicResult(**output_data)
def parse_output(self, outfiles: Dict[str, str],
input_model: "AtomicInput") -> "AtomicResult":
scf_map = {"energy": "scf_total_energy", "n_iter": "scf_iterations"}
dft_map = scf_map.copy()
hf_map = scf_map.copy()
xtb_map = scf_map.copy()
energy_command_map = {"dft": dft_map, "hf": hf_map, "xtb": xtb_map}
extras_map = {"converged": "scf_converged"}
wavefunction_map = {
"restricted": {
"orbitals": "scf_orbitals_a",
"density": "scf_density_a",
"fock": "scf_fock_a",
"eigenvalues": "scf_eigenvalues_a",
"occupations": "scf_occupations_a",
},
"unrestricted": {
"orbitals_alpha": "scf_orbitals_a",
"orbitals_beta": "scf_orbitals_b",
"density_alpha": "scf_density_a",
"density_beta": "scf_density_b",
"fock_alpha": "scf_fock_a",
"fock_beta": "scf_fock_b",
"eigenvalues_alpha": "scf_eigenvalues_a",
"eigenvalues_beta": "scf_eigenvalues_b",
"occupations_alpha": "scf_occupations_a",
"occupations_beta": "scf_occupations_b",
},
}
# Determine the energy_command
energy_command = self.determine_energy_command(
input_model.model.method)
gradient_map = {"gradient": "gradient"}
gradient_map.update({"energy": "scf_total_energy"})
# TODO Uncomment once entos adds scf_map to gradient json results
# gradient_map.update(energy_command_map[energy_command])
hessian_map = {"hessian": "hessian"}
hessian_map.update(energy_command_map[energy_command])
# Determine whether to use the energy map or the gradient map
if input_model.driver == "energy":
entos_map = energy_command_map[energy_command]
elif input_model.driver == "gradient":
entos_map = gradient_map
elif input_model.driver == "hessian":
entos_map = hessian_map
else:
raise NotImplementedError(
f"Driver {input_model.driver} not implemented for entos.")
# Parse the results.json output from entos
properties = {}
load_results = json.loads(outfiles["results.json"])
entos_results = load_results["json_results"]
for key in entos_map.keys():
if key in entos_results:
properties[entos_map[key]] = entos_results[key]
# Parse calcinfo_* properties from the results.json
if "ao_basis" in entos_results.keys():
properties["calcinfo_nbasis"] = entos_results["ao_basis"][
"__Basis"]["n_functions"]
if "structure" in entos_results.keys():
properties["calcinfo_natom"] = len(
entos_results["structure"]["__Atoms"]["atoms"])
# Parse wavefunction quantities from entos_results
wavefunction = {}
if input_model.protocols.wavefunction != "none":
# First parse basis set information
if "ao_basis" in entos_results.keys():
atom_map = [
item[0]
for item in entos_results["structure"]["__Atoms"]["atoms"]
]
# Each item in electron_shells is a dictionary containing info for one basis function
electron_shells_by_center = {}
for basis_item in entos_results["ao_basis"]["__Basis"][
"electron_shells"]:
center_index = basis_item["center_index"]
electron_shell_info = {
"angular_momentum": [basis_item["angular_momentum"]],
"harmonic_type":
basis_item["function_type"].split("_")[-1],
"exponents": basis_item["exponents"],
"coefficients": basis_item["coefficients"],
}
if center_index not in electron_shells_by_center:
electron_shells_by_center[center_index] = [
electron_shell_info
]
else:
electron_shells_by_center[center_index].append(
electron_shell_info)
# Construct center_data from electron_shells_by_center
# Note: Duplicate atoms will over write each other
center_data = {}
for i in range(len(electron_shells_by_center)):
basis_center_info = {
"electron_shells": electron_shells_by_center[i]
}
center_data[atom_map[i]] = basis_center_info
# Construct BasisSet
basis_info = {
"name": input_model.model.basis,
# "description": "", # None provided by entos
"center_data": center_data,
"atom_map": atom_map,
"nbf": entos_results["ao_basis"]["__Basis"]["n_functions"],
}
basis_set = BasisSet(**basis_info)
wavefunction["basis"] = basis_set
else:
raise KeyError(
f"Basis set information not found so wavefunction protocol {input_model.protocols.wavefunction} is not available."
)
# Now parse wavefunction information
n_channels = entos_results["n_channels"]
if n_channels == 1:
wavefunction["restricted"] = True
for key in wavefunction_map["restricted"].keys():
if key in entos_results:
if "orbitals" in key:
orbitals_transposed = reorder_column_ao_indices(
np.array(entos_results[key]), basis_set,
self._entos_to_cca_ao_order)
wavefunction[wavefunction_map["restricted"][
key]] = orbitals_transposed.transpose()
elif "density" in key or "fock" in key:
wavefunction[wavefunction_map["restricted"][
key]] = reorder_row_and_column_ao_indices(
entos_results[key], basis_set,
self._entos_to_cca_ao_order)
else:
wavefunction[wavefunction_map["restricted"]
[key]] = entos_results[key]
# TODO Add a test in QCEngineRecords
elif n_channels == 2:
wavefunction["restricted"] = False
for key in wavefunction_map["unrestricted"].keys():
if key in entos_results:
if "orbitals" in key:
orbitals_transposed = reorder_column_ao_indices(
np.array(entos_results[key]), basis_set,
self._entos_to_cca_ao_order)
wavefunction[wavefunction_map["restricted"][
key]] = orbitals_transposed.transpose()
elif "density" in key or "fock" in key:
wavefunction[wavefunction_map["restricted"][
key]] = reorder_row_and_column_ao_indices(
entos_results[key], basis_set,
self._entos_to_cca_ao_order)
else:
wavefunction[wavefunction_map["restricted"]
[key]] = entos_results[key]
# Parse results for the extras_map from results.json
extras = {}
for key in extras_map.keys():
if key in entos_results:
extras[extras_map[key]] = entos_results[key]
# Initialize output_data by copying over input_model.dict()
output_data = input_model.dict()
# Determine the correct return_result
if input_model.driver == "energy":
if "scf_total_energy" in properties:
output_data["return_result"] = properties["scf_total_energy"]
else:
raise KeyError(
f"Could not find {input_model.model} total energy")
elif input_model.driver == "gradient" or input_model.driver == "hessian":
if input_model.driver in properties:
output_data["return_result"] = properties.pop(
input_model.driver)
else:
raise KeyError(f"{input_model.driver} not found.")
else:
raise NotImplementedError(
f"Driver {input_model.driver} not implemented for entos.")
output_data["properties"] = properties
if input_model.protocols.wavefunction != "none":
output_data["wavefunction"] = wavefunction
output_data["extras"].update(extras)
output_data["schema_name"] = "qcschema_output"
output_data["success"] = True
return AtomicResult(**output_data)