diff --git a/qcmanybody/__init__.py b/qcmanybody/__init__.py index c637536..bfb3f82 100644 --- a/qcmanybody/__init__.py +++ b/qcmanybody/__init__.py @@ -2,6 +2,6 @@ from .manybody import ManyBodyCalculator from .models import BsseEnum -from .utils import labeler, delabeler +from .utils import labeler, delabeler, resize_gradient, resize_hessian __version__ = version("qcmanybody") diff --git a/qcmanybody/manybody.py b/qcmanybody/manybody.py index 9953207..a627b36 100644 --- a/qcmanybody/manybody.py +++ b/qcmanybody/manybody.py @@ -2,6 +2,7 @@ import logging import math +from collections import defaultdict from typing import Set, Dict, Tuple, Union, Literal, Mapping, Any, Sequence import numpy as np @@ -19,8 +20,8 @@ find_shape, shaped_zero, all_same_shape, - expand_hessian, - expand_gradient, + resize_hessian, + resize_gradient, ) logger = logging.getLogger(__name__) @@ -126,11 +127,11 @@ def compute_map(self) -> Dict[str, Dict[str, Dict[int, Set[FragBasIndex]]]]: return self.mc_compute_dict - def expand_gradient(self, grad: np.ndarray, bas: Tuple[int, ...]) -> np.ndarray: - return expand_gradient(grad, bas, self.fragment_size_dict, self.fragment_slice_dict) + def resize_gradient(self, grad: np.ndarray, bas: Tuple[int, ...], *, reverse: bool = False) -> np.ndarray: + return resize_gradient(grad, bas, self.fragment_size_dict, self.fragment_slice_dict, reverse=reverse) - def expand_hessian(self, hess: np.ndarray, bas: Tuple[int, ...]) -> np.ndarray: - return expand_hessian(hess, bas, self.fragment_size_dict, self.fragment_slice_dict) + def resize_hessian(self, hess: np.ndarray, bas: Tuple[int, ...], *, reverse: bool = False) -> np.ndarray: + return resize_hessian(hess, bas, self.fragment_size_dict, self.fragment_slice_dict, reverse=reverse) def iterate_molecules(self) -> Tuple[str, str, Molecule]: """Iterate over all the molecules needed for the computation. @@ -206,7 +207,7 @@ def _assemble_nbody_components( first_key = next(iter(component_results.keys())) property_shape = find_shape(component_results[first_key]) - # Final dictionaries + # Accumulation dictionaries # * {bsse_type}_by_level is filled by sum_cluster_data to contain for NOCP # & CP the summed total energies (or other property) of each nb-body. That is: # * NOCP: {1: 1b@1b, 2: 2b@2b, ..., max_nbody: max_nbody-b@max_nbody-b} and @@ -218,6 +219,17 @@ def _assemble_nbody_components( nocp_by_level = {n: shaped_zero(property_shape) for n in range(1, nbodies[-1] + 1)} vmfc_by_level = {n: shaped_zero(property_shape) for n in range(1, nbodies[-1] + 1)} + # * {bsse_type}_body_dict is usually filled with total energies (or other property). + # Multiple model chemistry levels may be involved. + # Generally, all consecutive keys between 1 and max_nbody will be present in the body_dict, + # but if supersystem_ie_only=T, only 1b and nfr-b are present, or if "supersystem" in levels, ??? + # * TOT: {1: 1b@1b, 2: 2b tot prop with bsse_type treatment, ..., max_nbody: max_nbody-b tot prop with bsse_type treatment} + # If 1b@1b (monomers in monomer basis) aren't available, which can happen when return_total_data=F + # and 1b@1b aren't otherwise needed, body_dict contains interaction energies (or other property). + # * IE: {1: shaped_zero, 2: 2b interaction prop using bsse_type, ..., max_nbody: max_nbody-b interaction prop using bsse_type} + # For both TOT and IE cases, body_dict values are cummulative, not additive. For TOT, total, + # interaction, and contribution data in ManyBodyResultProperties can be computed in + # collect_vars. For IE, interaction and contribution data can be computed. cp_body_dict = {n: shaped_zero(property_shape) for n in range(1, nbodies[-1] + 1)} nocp_body_dict = {n: shaped_zero(property_shape) for n in range(1, nbodies[-1] + 1)} vmfc_body_dict = {n: shaped_zero(property_shape) for n in range(1, nbodies[-1] + 1)} @@ -321,6 +333,7 @@ def _analyze( # results per model chemistry mc_results = {} + species_results = {} # sort by nbody level, ignore supersystem sorted_nbodies = [(k, v) for k, v in self.nbodies_per_mc_level.items() if v != ["supersystem"]] @@ -435,7 +448,8 @@ def analyze( """ # All properties that were passed to us - available_properties = set() + # * seed with "energy" so free/no-op jobs can process + available_properties = set(["energy"]) for property_data in component_results.values(): available_properties.update(property_data.keys()) @@ -455,23 +469,25 @@ def analyze( component_results_inv["energy"] = {'["dummy", [1000], [1000]]': 0.0} # Actually analyze + is_embedded = bool(self.embedding_charges) + component_properties = defaultdict(dict) all_results = {} + nbody_dict = {} +# all_results["energy_body_dict"] = {"cp": {1: 0.0}} for property_label, property_results in component_results_inv.items(): # Expand gradient and hessian if property_label == "gradient": - property_results = {k: self.expand_gradient(v, delabeler(k)[2]) for k, v in property_results.items()} + property_results = {k: self.resize_gradient(v, delabeler(k)[2]) for k, v in property_results.items()} if property_label == "hessian": - property_results = {k: self.expand_hessian(v, delabeler(k)[2]) for k, v in property_results.items()} + property_results = {k: self.resize_hessian(v, delabeler(k)[2]) for k, v in property_results.items()} r = self._analyze(property_label, property_results) + for k, v in property_results.items(): + component_properties[k]["calcinfo_natom"] = len(self.molecule.symbols) + component_properties[k][f"return_{property_label}"] = v all_results.update(r) - # Analyze the total results - nbody_dict = {} - - is_embedded = bool(self.embedding_charges) - for bt in self.bsse_type: print_nbody_energy( all_results["energy_body_dict"][bt], @@ -480,12 +496,22 @@ def analyze( is_embedded, ) - if not self.has_supersystem: # skipped levels? - nbody_dict.update( - collect_vars(bt.upper(), all_results["energy_body_dict"][bt], self.max_nbody, is_embedded, self.supersystem_ie_only) - ) + for property_label in available_properties: + for bt in self.bsse_type: + if not self.has_supersystem: # skipped levels? + nbody_dict.update( + collect_vars( + bt.upper(), + property_label.upper(), + all_results[f"{property_label}_body_dict"][bt], + self.max_nbody, + is_embedded, + self.supersystem_ie_only, + ) + ) all_results["results"] = nbody_dict + all_results["component_properties"] = component_properties # Make dictionary with "1cp", "2cp", etc ebd = all_results["energy_body_dict"] diff --git a/qcmanybody/models/__init__.py b/qcmanybody/models/__init__.py index 6547d24..beb208b 100644 --- a/qcmanybody/models/__init__.py +++ b/qcmanybody/models/__init__.py @@ -1 +1,8 @@ -from .manybody_v1 import BsseEnum, FragBasIndex +from .manybody_pydv1 import ( + BsseEnum, + FragBasIndex, + ManyBodyInput, + ManyBodyKeywords, + ManyBodyResult, + ManyBodyResultProperties, +) diff --git a/qcmanybody/models/manybody_v1.py b/qcmanybody/models/manybody_pydv1.py similarity index 98% rename from qcmanybody/models/manybody_v1.py rename to qcmanybody/models/manybody_pydv1.py index e07f17c..1b9f00f 100644 --- a/qcmanybody/models/manybody_v1.py +++ b/qcmanybody/models/manybody_pydv1.py @@ -13,7 +13,7 @@ #from .basemodels import ExtendedConfigDict, ProtoModel from qcelemental.models.common_models import Model from qcelemental.models.molecule import Molecule -from qcelemental.models.results import AtomicResultProtocols +from qcelemental.models.results import AtomicResultProperties, AtomicResultProtocols from qcelemental.models import DriverEnum, ProtoModel, Provenance @@ -516,6 +516,11 @@ class ManyBodyResult(SuccessfulResultBase): "all results regardless of if they failed or succeeded by checking `result.success`.", ) properties: ManyBodyResultProperties = Field(..., description=str(ManyBodyResultProperties.__doc__)) + component_properties: Dict[str, AtomicResultProperties] = Field( + ..., + description="The key results for each subsystem species computed. Keys contain modelchem, real and ghost information (e.g., `'[\"(auto)\", [2], [1, 2, 3]]'`). Values are total e/g/H/property results. Array values, if present, are sized and shaped for the full supersystem.", + + ) return_result: Union[float, Array[float], Dict[str, Any]] = Field( ..., description="The primary return specified by the :attr:`~qcelemental.models.AtomicInput.driver` field. Scalar if energy; array if gradient or hessian; dictionary with property keys if properties.", diff --git a/qcmanybody/models/qcng_computer.py b/qcmanybody/models/qcng_computer.py index 069bace..6376e48 100644 --- a/qcmanybody/models/qcng_computer.py +++ b/qcmanybody/models/qcng_computer.py @@ -21,7 +21,7 @@ from qcelemental.models import FailedOperation, Molecule, DriverEnum, ProtoModel, AtomicResult, AtomicInput import qcengine as qcng -from .manybody_v1 import BsseEnum, ManyBodyKeywords, ManyBodyInput, ManyBodyResult, ManyBodyResultProperties +from .manybody_pydv1 import BsseEnum, ManyBodyKeywords, ManyBodyInput, ManyBodyResult, ManyBodyResultProperties from qcmanybody import ManyBodyCalculator from qcmanybody.utils import delabeler, provenance_stamp @@ -741,6 +741,7 @@ def get_results(self, client: Optional["qcportal.FractalClient"] = None, externa ret_ptype = ret_energy if self.driver == "energy" else external_results.pop(f"ret_{self.driver.name}") ret_gradient = external_results.pop("ret_gradient", None) nbody_number = external_results.pop("nbody_number") + component_properties = external_results.pop("component_properties") # load QCVariables qcvars = { @@ -845,10 +846,10 @@ def get_results(self, client: Optional["qcportal.FractalClient"] = None, externa #'molecule': self.molecule, # v2: 'properties': {**atprop.model_dump(), **properties}, 'properties': {**atprop.dict(), **properties}, + 'component_properties': component_properties, 'provenance': provenance_stamp(__name__), 'extras': { 'qcvars': qcvars, -# 'component_results': component_results, }, 'return_result': ret_ptype, 'success': True, diff --git a/qcmanybody/models/test_mbe_he4_multilevel.py b/qcmanybody/models/test_mbe_he4_multilevel.py index 6078111..8e514c7 100644 --- a/qcmanybody/models/test_mbe_he4_multilevel.py +++ b/qcmanybody/models/test_mbe_he4_multilevel.py @@ -8,7 +8,7 @@ # v2: from qcelemental.models.procedures_manybody import AtomicSpecification, ManyBodyKeywords, ManyBodyInput from qcelemental.testing import compare_values -from qcmanybody.models.manybody_v1 import AtomicSpecification, ManyBodyKeywords, ManyBodyInput +from qcmanybody.models.manybody_pydv1 import AtomicSpecification, ManyBodyKeywords, ManyBodyInput from qcmanybody.models.qcng_computer import ManyBodyComputerQCNG, qcvars_to_manybodyproperties import qcengine as qcng diff --git a/qcmanybody/models/test_mbe_he4_singlelevel.py b/qcmanybody/models/test_mbe_he4_singlelevel.py index 52fc2d5..b2f86ec 100644 --- a/qcmanybody/models/test_mbe_he4_singlelevel.py +++ b/qcmanybody/models/test_mbe_he4_singlelevel.py @@ -7,7 +7,7 @@ # v2: from qcelemental.models.procedures_manybody import AtomicSpecification, ManyBodyKeywords, ManyBodyInput from qcelemental.testing import compare_values -from qcmanybody.models.manybody_v1 import AtomicSpecification, ManyBodyKeywords, ManyBodyInput +from qcmanybody.models.manybody_pydv1 import AtomicSpecification, ManyBodyKeywords, ManyBodyInput from qcmanybody.models.qcng_computer import ManyBodyComputerQCNG, qcvars_to_manybodyproperties import qcengine as qcng @@ -18,6 +18,12 @@ def skprop(qcvar): return qcvars_to_manybodyproperties[qcvar] +he4_refs_species = { + "NO": ('[\"(auto)\", [1, 2, 4], [1, 2, 4]]', -8.644153798224503), + "CP": ('[\"(auto)\", [1, 2, 4], [1, 2, 3, 4]]', -8.64425850181438), + "VM": ('[\"(auto)\", [1, 4], [1, 2, 4]]', -5.765439283351071), +} + he4_refs_conv = { "CP-CORRECTED TOTAL ENERGY THROUGH 1-BODY": -11.530668717083888, "CP-CORRECTED TOTAL ENERGY THROUGH 2-BODY": -11.522467757090013, @@ -490,6 +496,10 @@ def test_nbody_he4_single(program, basis, keywords, mbe_keywords, anskey, bodyke assert qcv not in ret.extras["qcvars"]["nbody"], f"[y] {qcv=} wrongly present" assert getattr(ret.properties, skp) is None + if "3b" in kwdsln and not progln.endswith("df"): + k, v = he4_refs_species[anskey[:2]] + assert compare_values(v, ret.component_properties[k].return_energy, atol=atol, label=f"[h] species {k}") + for qcv, ref in { "CURRENT ENERGY": ans, }.items(): diff --git a/qcmanybody/models/test_mbe_keywords.py b/qcmanybody/models/test_mbe_keywords.py index 71ed8d1..05c8df2 100644 --- a/qcmanybody/models/test_mbe_keywords.py +++ b/qcmanybody/models/test_mbe_keywords.py @@ -9,7 +9,7 @@ from pydantic import ValidationError from qcelemental.models import DriverEnum, Molecule -from qcmanybody.models.manybody_v1 import BsseEnum, ManyBodyInput +from qcmanybody.models.manybody_pydv1 import BsseEnum, ManyBodyInput import qcengine as qcng diff --git a/qcmanybody/models/test_utils.py b/qcmanybody/models/test_utils.py new file mode 100644 index 0000000..3745e37 --- /dev/null +++ b/qcmanybody/models/test_utils.py @@ -0,0 +1,154 @@ +import pytest +import numpy as np + +from qcelemental.testing import compare_values + +import qcmanybody as qcmb + +@pytest.fixture(scope="function") +def mbe_data(): + henehh = Molecule(symbols=["He", "Ne", "H", "H"], fragments=[[0], [1], [2, 3]], geometry=[0, 0, 0, 2, 0, 0, 0, 1, 0, 0, -1, 0]) + return henehh + +f3grads = { + "full": np.array([ + [ 0.598726, 0. , 0. ], + [-0.960726, 0. , 0. ], + [ 0.181 , -0.858448, 0. ], + [ 0.181 , 0.858448, 0. ], + ]), + "b13": np.array([ + [ 0.598726, 0. , 0. ], +#[-0.960726, 0. , 0. ], + [ 0.181 , -0.858448, 0. ], + [ 0.181 , 0.858448, 0. ], + ]), + "b13_": np.array([ + [ 0.598726, 0. , 0. ], + [ 0. , 0. , 0. ], + [ 0.181 , -0.858448, 0. ], + [ 0.181 , 0.858448, 0. ], + ]), + "b3": np.array([ + #[ 0.598726, 0. , 0. ], + #[-0.960726, 0. , 0. ], + [ 0.181 , -0.858448, 0. ], + [ 0.181 , 0.858448, 0. ], + ]), + "b3_": np.array([ + [ 0. , 0. , 0. ], + [ 0. , 0. , 0. ], + [ 0.181 , -0.858448, 0. ], + [ 0.181 , 0.858448, 0. ], + ]), + } + +f3hesses = { + "full": np.array([ + [ 0.035028, 0. , 0. , -1.398945, 0. , -0. , 0.681958, 0.052467, -0. , 0.681958, -0.052467, 0. ], + [ 0. , 3.339155, 0. , 0. , 0.120345, 0. , -0.179019, -1.72975 , -0. , 0.179019, -1.72975 , -0. ], + [ 0. , 0. , -1.232377, 0. , 0. , 0.299363, 0. , -0. , 0.466507, -0. , -0. , 0.466507], + [-1.398945, 0. , 0. , 2.147922, 0. , 0. , -0.374488, 0.060761, 0. , -0.374488, -0.060761, -0. ], + [ 0. , 0.120345, 0. , 0. , -0.252472, 0. , 0.227891, 0.066064, -0. , -0.227891, 0.066064, -0. ], + [-0. , 0. , 0.299363, 0. , 0. , -0.480363, 0. , -0. , 0.0905 , 0. , -0. , 0.0905 ], + [ 0.681958, -0.179019, 0. , -0.374488, 0.227891, 0. , -0.355068, -0.08105 , -0. , 0.047598, 0.032178, -0. ], + [ 0.052467, -1.72975 , -0. , 0.060761, 0.066064, -0. , -0.08105 , 2.276729, 0. , -0.032178, -0.613043, -0. ], + [-0. , -0. , 0.466507, 0. , -0. , 0.0905 , -0. , 0. , -0.707728, 0. , -0. , 0.150721], + [ 0.681958, 0.179019, -0. , -0.374488, -0.227891, 0. , 0.047598, -0.032178, 0. , -0.355068, 0.08105 , 0. ], + [-0.052467, -1.72975 , -0. , -0.060761, 0.066064, -0. , 0.032178, -0.613043, -0. , 0.08105 , 2.276729, 0. ], + [ 0. , -0. , 0.466507, -0. , -0. , 0.0905 , -0. , -0. , 0.150721, 0. , 0. , -0.707728], + ]), + "b13": np.array([ + [ 0.035028, 0. , 0. , 0.681958, 0.052467, -0. , 0.681958, -0.052467, 0. ], + [ 0. , 3.339155, 0. , -0.179019, -1.72975 , -0. , 0.179019, -1.72975 , -0. ], + [ 0. , 0. , -1.232377, 0. , -0. , 0.466507, -0. , -0. , 0.466507], +#[-1.398945, 0. , 0. , -0.374488, 0.060761, 0. , -0.374488, -0.060761, -0. ], +#[ 0. , 0.120345, 0. , 0.227891, 0.066064, -0. , -0.227891, 0.066064, -0. ], +#[-0. , 0. , 0.299363, 0. , -0. , 0.0905 , 0. , -0. , 0.0905 ], + [ 0.681958, -0.179019, 0. , -0.355068, -0.08105 , -0. , 0.047598, 0.032178, -0. ], + [ 0.052467, -1.72975 , -0. , -0.08105 , 2.276729, 0. , -0.032178, -0.613043, -0. ], + [-0. , -0. , 0.466507, -0. , 0. , -0.707728, 0. , -0. , 0.150721], + [ 0.681958, 0.179019, -0. , 0.047598, -0.032178, 0. , -0.355068, 0.08105 , 0. ], + [-0.052467, -1.72975 , -0. , 0.032178, -0.613043, -0. , 0.08105 , 2.276729, 0. ], + [ 0. , -0. , 0.466507, -0. , -0. , 0.150721, 0. , 0. , -0.707728], + ]), + "b13_": np.array([ + [ 0.035028, 0. , 0. , 0. , 0. , -0. , 0.681958, 0.052467, -0. , 0.681958, -0.052467, 0. ], + [ 0. , 3.339155, 0. , 0. , 0. , 0. , -0.179019, -1.72975 , -0. , 0.179019, -1.72975 , -0. ], + [ 0. , 0. , -1.232377, 0. , 0. , 0. , 0. , -0. , 0.466507, -0. , -0. , 0.466507], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0.681958, -0.179019, 0. , 0. , 0. , 0. , -0.355068, -0.08105 , -0. , 0.047598, 0.032178, -0. ], + [ 0.052467, -1.72975 , -0. , 0. , 0. , -0. , -0.08105 , 2.276729, 0. , -0.032178, -0.613043, -0. ], + [-0. , -0. , 0.466507, 0. , -0. , 0. , -0. , 0. , -0.707728, 0. , -0. , 0.150721], + [ 0.681958, 0.179019, -0. , -0. , 0. , 0. , 0.047598, -0.032178, 0. , -0.355068, 0.08105 , 0. ], + [-0.052467, -1.72975 , -0. , -0. , 0. , -0. , 0.032178, -0.613043, -0. , 0.08105 , 2.276729, 0. ], + [ 0. , -0. , 0.466507, -0. , -0. , 0. , -0. , -0. , 0.150721, 0. , 0. , -0.707728], + + ]), + "b3": np.array([ +#[ 0.681958, 0.052467, -0. , 0.681958, -0.052467, 0. ], +#[ -0.179019, -1.72975 , -0. , 0.179019, -1.72975 , -0. ], +#[ 0. , -0. , 0.466507, -0. , -0. , 0.466507], +#[ -0.374488, 0.060761, 0. , -0.374488, -0.060761, -0. ], +#[ 0.227891, 0.066064, -0. , -0.227891, 0.066064, -0. ], +#[ 0. , -0. , 0.0905 , 0. , -0. , 0.0905 ], + [ -0.355068, -0.08105 , -0. , 0.047598, 0.032178, -0. ], + [ -0.08105 , 2.276729, 0. , -0.032178, -0.613043, -0. ], + [ -0. , 0. , -0.707728, 0. , -0. , 0.150721], + [ 0.047598, -0.032178, 0. , -0.355068, 0.08105 , 0. ], + [ 0.032178, -0.613043, -0. , 0.08105 , 2.276729, 0. ], + [ -0. , -0. , 0.150721, 0. , 0. , -0.707728], + + ]), + "b3_": np.array([ + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. , 0. ], + [ 0. , 0. , 0. , 0. , 0. , 0. , -0.355068, -0.08105 , -0. , 0.047598, 0.032178, -0. ], + [ 0. , -0. , -0. , 0. , 0. , -0. , -0.08105 , 2.276729, 0. , -0.032178, -0.613043, -0. ], + [-0. , -0. , 0. , 0. , -0. , 0. , -0. , 0. , -0.707728, 0. , -0. , 0.150721], + [ 0. , 0. , -0. , 0. , -0. , 0. , 0.047598, -0.032178, 0. , -0.355068, 0.08105 , 0. ], + [-0. , -0. , -0. , 0. , 0. , -0. , 0.032178, -0.613043, -0. , 0.08105 , 2.276729, 0. ], + [ 0. , -0. , 0. , -0. , -0. , 0. , -0. , -0. , 0.150721, 0. , 0. , -0.707728], + ]), +} + + +@pytest.mark.parametrize("gin,bas,reverse,gans", [ + pytest.param(f3grads["full"], [1, 2, 3], False, f3grads["full"]), # idempotent + pytest.param(f3grads["full"], [1, 2, 3], True, f3grads["full"]), # idempotent + pytest.param(f3grads["b13"], [1, 3], False, f3grads["b13_"]), + pytest.param(f3grads["b13_"], [1, 3], True, f3grads["b13"]), + pytest.param(f3grads["full"], [1, 3], True, f3grads["b13"]), + pytest.param(f3grads["b3"], [3], False, f3grads["b3_"]), + pytest.param(f3grads["b3_"], [3], True, f3grads["b3"]), + pytest.param(f3grads["full"], [3], True, f3grads["b3"]), + pytest.param(f3grads["full"], [], False, np.zeros((4, 3))), # zero + pytest.param(f3grads["full"], [], True, np.zeros((0, 3))), # collapse +]) +def test_resize_gradient(gin, bas, reverse, gans): + gout = qcmb.resize_gradient(gin, bas, {1: 1, 2: 1, 3: 2}, {1: slice(0, 1), 2: slice(1, 2), 3: slice(2, 4)}, reverse=reverse) + assert compare_values(gans, gout, atol=1e-5, label="resize_gradient") + + +@pytest.mark.parametrize("hin,bas,reverse,hans", [ + pytest.param(f3hesses["full"], [1, 2, 3], False, f3hesses["full"]), # idempotent + pytest.param(f3hesses["full"], [1, 2, 3], True, f3hesses["full"]), # idempotent + pytest.param(f3hesses["b13"], [1, 3], False, f3hesses["b13_"]), + pytest.param(f3hesses["b13_"], [1, 3], True, f3hesses["b13"]), + pytest.param(f3hesses["full"], [1, 3], True, f3hesses["b13"]), + pytest.param(f3hesses["b3"], [3], False, f3hesses["b3_"]), + pytest.param(f3hesses["b3_"], [3], True, f3hesses["b3"]), + pytest.param(f3hesses["full"], [3], True, f3hesses["b3"]), + pytest.param(f3hesses["full"], [], False, np.zeros((12, 12))), # zero + pytest.param(f3hesses["full"], [], True, np.zeros((0, 0))), # collapse +]) +def test_resize_hessian(hin, bas, reverse, hans): + hout = qcmb.resize_hessian(hin, bas, {1: 1, 2: 1, 3: 2}, {1: slice(0, 1), 2: slice(1, 2), 3: slice(2, 4)}, reverse=reverse) + assert compare_values(hans, hout, atol=1e-5, label="resize_hessian") + diff --git a/qcmanybody/utils.py b/qcmanybody/utils.py index aabd6b0..0e39858 100644 --- a/qcmanybody/utils.py +++ b/qcmanybody/utils.py @@ -1,7 +1,7 @@ from __future__ import annotations import json -from typing import Tuple, Dict, Union, Iterable +from typing import Dict, Mapping, Tuple, Union, Iterable import numpy as np from qcelemental import constants @@ -44,32 +44,94 @@ def all_same_shape(it: Iterable[Union[float, np.ndarray]]) -> bool: raise TypeError(f"Expected float or np.ndarray, got {type(first)}") -def expand_gradient( - grad: np.ndarray, bas: Tuple[int, ...], fragment_size_dict: Dict[int, int], fragment_slice_dict: Dict[int, slice] +def resize_gradient( + grad: np.ndarray, + bas: Tuple[int, ...], + fragment_size_dict: Dict[int, int], + fragment_slice_dict: Dict[int, slice], + *, + reverse: bool = False, ) -> np.ndarray: - """ - Expands a gradient calculated for a cluster to the full system - """ + """Pads or extracts a gradient array between subsystem and full supersystem sizes. + + Parameters + ---------- + grad + Gradient matrix of natural size for *bas*, (3 * , 3). + If `reverse=True`, gradient matrix of supersystem size, (3 * , 3). + bas + 1-indexed fragments active in *grad*. + If `reverse=True`, 1-indexed fragments to be extracted from *grad*. + fragment_size_dict + Dictionary containing the number of atoms of each 1-indexed fragment. + For He--HOOH--Me cluster, `{1: 1, 2: 4, 3: 5}`. + fragment_slice_dict + Dictionary containing slices that index the gradient matrix for each of the 1-indexed fragments. + For He--HOOH--Me cluster, `{1: slice(0, 1), 2: slice(1, 5), 3: slice(5, 10)}`. + reverse + If True, contract *grad* from supersystem size and shape that which is natural for *bas*. + + Returns + ------- + Gradient array padded with zeros to supersystem size, (3 * , 3). + If reverse=True, gradient array extracted to natural size, (3 * , 3). - nat = sum(fragment_size_dict.values()) + """ + if reverse: + nat = sum(fragment_size_dict[ifr] for ifr in bas) + else: + nat = sum(fragment_size_dict.values()) ret = np.zeros((nat, 3)) + start = 0 for ifr in bas: end = start + fragment_size_dict[ifr] - ret[fragment_slice_dict[ifr]] = grad[start:end] + if reverse: + ret[start:end] = grad[fragment_slice_dict[ifr]] + else: + ret[fragment_slice_dict[ifr]] = grad[start:end] start += fragment_size_dict[ifr] return ret -def expand_hessian( - hess: np.ndarray, bas: Tuple[int, ...], fragment_size_dict: Dict[int, int], fragment_slice_dict: Dict[int, slice] +def resize_hessian( + hess: np.ndarray, + bas: Tuple[int, ...], + fragment_size_dict: Dict[int, int], + fragment_slice_dict: Dict[int, slice], + *, + reverse: bool = False, ) -> np.ndarray: - """ - Expands a hessian calculated for a cluster to the full system - """ + """Pads or extracts a Hessian array between subsystem and full supersystem sizes. + + Parameters + ---------- + grad + Hessian matrix of natural size for *bas*, (3 * , 3 * ). + If `reverse=True`, Hessian matrix of supersystem size, (3 * , 3 * ). + bas + 1-indexed fragments active in *hess*. + If `reverse=True`, 1-indexed fragments to be extracted from *hess*. + fragment_size_dict + Dictionary containing the number of atoms of each 1-indexed fragment. + For He--HOOH--Me cluster, `{1: 1, 2: 4, 3: 5}`. + fragment_slice_dict + Dictionary containing slices that index the gradient matrix for each of the 1-indexed fragments. + For He--HOOH--Me cluster, `{1: slice(0, 1), 2: slice(1, 5), 3: slice(5, 10)}`. + reverse + If True, contract *hess* from supersystem size and shape that which is natural for *bas*. + + Returns + ------- + Hessian array padded with zeros to supersystem size, (3 * , 3 * ). + If reverse=True, Hessian array extracted to natural size, (3 * , 3 * ). - nat = sum(fragment_size_dict.values()) + """ + if reverse: + nat = sum(fragment_size_dict[ifr] for ifr in bas) + else: + nat = sum(fragment_size_dict.values()) ret = np.zeros((nat * 3, nat * 3)) # Build up start and end slices @@ -85,12 +147,33 @@ def expand_hessian( for abs_sl1, rel_sl1 in zip(abs_slices, rel_slices): for abs_sl2, rel_sl2 in zip(abs_slices, rel_slices): - ret[abs_sl1, abs_sl2] = hess[rel_sl1, rel_sl2] + if reverse: + ret[rel_sl1, rel_sl2] = hess[abs_sl1, abs_sl2] + else: + ret[abs_sl1, abs_sl2] = hess[rel_sl1, rel_sl2] return ret def labeler(mc_level_lbl: str, frag: Tuple[int, ...], bas: Tuple[int, ...]) -> str: + """Form label from model chemistry id and fragment and basis indices. + + Parameters + ---------- + mc_level_lbl + Key identifying the model chemistry. May be `"(auto)"`. Often the + ManyBodyInput.specification.specification keys. + frag + List of 1-indexed fragments active in the supersystem. + bas + List of 1-indexed fragments with active basis sets in the supersystem. + All those in *frag* plus any ghost. + + Returns + ------- + str + JSON string from inputs: `labeler("mp2",(1), (1, 2))` returns `'["mp2", 1, [1, 2]]'`. + """ return json.dumps([str(mc_level_lbl), frag, bas]) @@ -146,9 +229,38 @@ def print_nbody_energy( print(info) -def collect_vars(bsse, body_dict, max_nbody: int, embedding: bool = False, supersystem_ie_only: bool = False): +def collect_vars( + bsse: str, + prop: str, + body_dict: Mapping[int, Union[float, np.ndarray]], + max_nbody: int, + embedding: bool = False, + supersystem_ie_only: bool = False, + ) -> Dict: + """From *body_dict*, construct QCVariables. + + Parameters + ---------- + bsse + Uppercase label for a single many-body treatment, generally a value of BsseEnum. + prop + Uppercase label for a single property, generally a value of DriverEnum. + body_dict + Dictionary of minimal per-body info already specialized for property *prop* and treatment *bsse*. May contain either total data or interaction data, both cummulative not additive, from 1-body to max_nbody-body (see also *supersystem_ie_only*). Interaction data signaled by zero float or array for 1-body. May contain multiple model chemistry levels. + max_nbody + _description_ + embedding, optional + Is charge embedding enabled, by default False? + supersystem_ie_only, optional + Is data available in *body_dict* only for 1-body (possibly zero) and nfr-body levels? By default False: data is available for consecutive levels, up to max_nbody-body. + + Returns + ------- + _description_. Empty return if *embedding* enabled. + """ previous_e = body_dict[1] - tot_e = previous_e != 0.0 + property_shape = find_shape(previous_e) + tot_e = bool(np.count_nonzero(previous_e)) nbody_range = list(body_dict) nbody_range.sort() res = {} @@ -156,26 +268,26 @@ def collect_vars(bsse, body_dict, max_nbody: int, embedding: bool = False, super return res if tot_e: - res[f"{bsse}-CORRECTED TOTAL ENERGY"] = body_dict[max_nbody] - res[f"{bsse}-CORRECTED INTERACTION ENERGY"] = body_dict[max_nbody] - body_dict[1] - res[f"{bsse}-CORRECTED INTERACTION ENERGY THROUGH 1-BODY"] = 0.0 + res[f"{bsse}-CORRECTED TOTAL {prop}"] = body_dict[max_nbody] + res[f"{bsse}-CORRECTED INTERACTION {prop}"] = body_dict[max_nbody] - body_dict[1] + res[f"{bsse}-CORRECTED INTERACTION {prop} THROUGH 1-BODY"] = shaped_zero(property_shape) if supersystem_ie_only: nfr = nbody_range[-1] for nb in [nfr]: - res[f"{bsse}-CORRECTED INTERACTION ENERGY THROUGH {nb}-BODY"] = body_dict[nb] - body_dict[1] + res[f"{bsse}-CORRECTED INTERACTION {prop} THROUGH {nb}-BODY"] = body_dict[nb] - body_dict[1] if nb == 2: - res[f"{bsse}-CORRECTED {nb}-BODY CONTRIBUTION TO ENERGY"] = body_dict[nb] - body_dict[nb - 1] + res[f"{bsse}-CORRECTED {nb}-BODY CONTRIBUTION TO {prop}"] = body_dict[nb] - body_dict[nb - 1] if tot_e: for nb in [1, nfr]: - res[f"{bsse}-CORRECTED TOTAL ENERGY THROUGH {nb}-BODY"] = body_dict[nb] + res[f"{bsse}-CORRECTED TOTAL {prop} THROUGH {nb}-BODY"] = body_dict[nb] else: for nb in range(2, max(nbody_range) + 1): - res[f"{bsse}-CORRECTED INTERACTION ENERGY THROUGH {nb}-BODY"] = body_dict[nb] - body_dict[1] - res[f"{bsse}-CORRECTED {nb}-BODY CONTRIBUTION TO ENERGY"] = body_dict[nb] - body_dict[nb - 1] + res[f"{bsse}-CORRECTED INTERACTION {prop} THROUGH {nb}-BODY"] = body_dict[nb] - body_dict[1] + res[f"{bsse}-CORRECTED {nb}-BODY CONTRIBUTION TO {prop}"] = body_dict[nb] - body_dict[nb - 1] if tot_e: for nb in nbody_range: - res[f"{bsse}-CORRECTED TOTAL ENERGY THROUGH {nb}-BODY"] = body_dict[nb] + res[f"{bsse}-CORRECTED TOTAL {prop} THROUGH {nb}-BODY"] = body_dict[nb] return res