component data (#9)

* tidying

* more misc

* poor fix of tests

* fix 1b

* rename stuff

* test component

* test resize fns

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")

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"]

qcmanybody/models/__init__.py

@@ -1 +1,8 @@
-from .manybody_v1 import BsseEnum, FragBasIndex
+from .manybody_pydv1 import (
+    BsseEnum,
+    FragBasIndex,
+    ManyBodyInput,
+    ManyBodyKeywords,
+    ManyBodyResult,
+    ManyBodyResultProperties,
+)

qcmanybody/models/manybody_v1.py → 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.",

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,

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

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():

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