component_results, protocols, and GeneralizedOptimizationInput/Result (…

…#25)

* much

* asdf

* genopt

* extras

qcmanybody/models/generalized_optimization.py

@@ -0,0 +1,29 @@
+from typing import Annotated, List, Literal, Union
+
+try:
+    from pydantic.v1 import Field
+except ImportError:
+    from pydantic import Field
+
+from qcelemental.models import AtomicResult, OptimizationInput, OptimizationResult
+from qcelemental.models.procedures import QCInputSpecification
+
+from .manybody_input_pydv1 import ManyBodySpecification
+from .manybody_output_pydv1 import ManyBodyResult
+
+
+# note that qcel AtomicResult.schema_name needs editing
+ResultTrajectories = Annotated[Union[AtomicResult, ManyBodyResult], Field(discriminator='schema_name')]
+
+class GeneralizedOptimizationInput(OptimizationInput):
+    schema_name: Literal["qcschema_generalizedoptimizationinput"] = "qcschema_generalizedoptimizationinput"
+    schema_version: int = 1
+    input_specification: Union[QCInputSpecification, ManyBodySpecification] = Field(..., description="ordinary or mbe grad spec")
+
+
+class GeneralizedOptimizationResult(OptimizationResult):
+    schema_name: Literal["qcschema_generalizedoptimizationresult"] = "qcschema_generalizedoptimizationresult"
+    trajectory: List[ResultTrajectories] = Field(
+        ..., description="A list of ordered Result objects for each step in the optimization."
+    )
+    input_specification: Union[QCInputSpecification, ManyBodySpecification] = Field(..., description="ordinary or mbe grad spec")

qcmanybody/models/manybody_input_pydv1.py

@@ -60,14 +60,11 @@ class SuccessfulResultBase(ResultBase):
 
 # ====  Protocols  ==============================================================
 
-class ManyBodyProtocolEnum(str, Enum):
-    """
-    Which atomic evaluations to keep in a many body evaluation.
-    """
+class ComponentResultsProtocolEnum(str, Enum):
+    r"""Which component results to preserve in a many body result; usually AtomicResults."""
 
     all = "all"
-    all_real = "all_real"
-    largest_body = "largest_body"
+    # max_nbody = "max_nbody"
     none = "none"
 
 
@@ -76,9 +73,9 @@ class ManyBodyProtocols(ProtoModel):
     Protocols regarding the manipulation of a ManyBody output data.
     """
 
-    atomics: ManyBodyProtocolEnum = Field(
-        ManyBodyProtocolEnum.all,
-        description=str(ManyBodyProtocolEnum.__doc__),
+    component_results: ComponentResultsProtocolEnum = Field(
+        ComponentResultsProtocolEnum.none,
+        description=str(ComponentResultsProtocolEnum.__doc__)
     )
 
     # v2: model_config = ExtendedConfigDict(force_skip_defaults=True)
@@ -95,6 +92,7 @@ class BsseEnum(str, Enum):
     cp = "cp"      # Boys-Bernardi counterpoise correction; site-site functional counterpoise (SSFC)
     vmfc = "vmfc"  # Valiron-Mayer function counterpoise
     ssfc = "cp"
+    mbe = "nocp"
 
     def formal(self):
         return {
@@ -123,8 +121,8 @@ class ManyBodyKeywords(ProtoModel):
         description="Requested BSSE treatments. First in list determines which interaction or total "
             "energy/gradient/Hessian returned.",
     )
-    embedding_charges: Dict[int, List[float]] = Field(
-        {},
+    embedding_charges: Optional[Dict[int, List[float]]] = Field(
+        None,
         description="Atom-centered point charges to be used on molecule fragments whose basis sets are not included in "
             "the computation. Keys: 1-based index of fragment. Values: list of atom charges for that fragment.",
             # TODO embedding charges should sum to fragment charge, right? enforce?
@@ -200,6 +198,7 @@ class ManyBodySpecification(ProtoModel):
     #provenance: Provenance = Field(Provenance(**provenance_stamp(__name__)), description=Provenance.__doc__)
     keywords: ManyBodyKeywords = Field(..., description=ManyBodyKeywords.__doc__)
     #program: str = Field(..., description="The program for which the Specification is intended.")
+    protocols: ManyBodyProtocols = Field(ManyBodyProtocols(), description=str(ManyBodyProtocols.__doc__))
     driver: DriverEnum = Field(
         ...,
         description="The computation driver; i.e., energy, gradient, hessian.",
@@ -235,4 +234,7 @@ class ManyBodyInput(ProtoModel):
         ...,
         description="Target molecule for many-body expansion (MBE) or interaction energy (IE) analysis.",
     )
-    #protocols
+    extras: Dict[str, Any] = Field(
+        {},
+        description="Additional information to bundle with the computation. Use for schema development and scratch space.",
+    )

qcmanybody/models/manybody_output_pydv1.py

@@ -9,7 +9,7 @@
     from pydantic import create_model, Field, validator
 
 from qcelemental.models.types import Array
-from qcelemental.models.results import AtomicResultProperties #, AtomicResultProtocols
+from qcelemental.models.results import AtomicResult, AtomicResultProperties
 from qcelemental.models import ProtoModel, Provenance
 
 from .manybody_input_pydv1 import ManyBodyInput, SuccessfulResultBase
@@ -320,6 +320,7 @@ class ManyBodyResult(SuccessfulResultBase):
         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.",
 
     )
+    component_results: Dict[str, AtomicResult] = Field({}, description="Detailed results")
     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.",
@@ -330,3 +331,13 @@ class ManyBodyResult(SuccessfulResultBase):
     )
     stderr: Optional[str] = Field(None, description="The standard error of the program execution.")
     success: Literal[True] = Field(True, description="Always `True` for a successful result")
+
+    @validator("component_results")
+    def _component_results(cls, value, values):
+        crp = values["input_data"].specification.protocols.component_results
+        if crp == "all":
+            return value
+        elif crp == "none":
+            return {}
+        else:
+            raise ValueError(f"Protocol `component_resutls:{crp}` is not understood")

qcmanybody/qcng_computer.py

@@ -154,8 +154,8 @@ class ManyBodyComputerQCNG(BaseComputerQCNG):
             "through finite difference energies or MBE energy through composite method), this field refers to the "
             "*target* derivative, not any *means* specification.",
     )
-    embedding_charges: Dict[int, List[float]] = Field(
-        {},
+    embedding_charges: Optional[Dict[int, List[float]]] = Field(
+        None,
         description="Atom-centered point charges to be used to speed up nbody-level convergence. Charges are placed on "
         "molecule fragments whose basis sets are not included in the computation. (An implication is that charges aren't "
         "invoked for bsse_type=cp.) Keys: 1-based index of fragment. Values: list of atom charges for that fragment.",
@@ -224,7 +224,7 @@ def set_embedding_charges(cls, v, values): # -> Dict[int, List[float]]:
         print(f"hit embedding_charges validator with {v}", end="")
         nfr = len(values["molecule"].fragments)
         # v2: if len(v) != info.data["nfragments"]:
-        if len(v) != nfr:
+        if v and len(v) != nfr:
             raise ValueError(f"embedding_charges dict should have entries for each 1-indexed fragment ({nfr})")
 
         print(f" ... setting embedding_charges={v}")
@@ -401,7 +401,6 @@ def from_manybodyinput(cls, input_model: ManyBodyInput, build_tasks: bool = True
             specifications[mtd]["specification"] = spec
             specifications[mtd]["specification"]["driver"] = computer_model.driver  # overrides atomic driver with mb driver
             specifications[mtd]["specification"].pop("schema_name", None)
-            specifications[mtd]["specification"].pop("protocols", None)
 
         computer_model.qcmb_calculator = ManyBodyCalculator(
             computer_model.molecule,
@@ -415,10 +414,12 @@ def from_manybodyinput(cls, input_model: ManyBodyInput, build_tasks: bool = True
         if not build_tasks:
             return computer_model
 
+        component_properties = {}
         component_results = {}
 
         for chem, label, imol in computer_model.qcmb_calculator.iterate_molecules():
             inp = AtomicInput(molecule=imol, **specifications[chem]["specification"])
+            # inp = AtomicInput(molecule=imol, **specifications[chem]["specification"], extras={"psiapi": True})  # faster for p4
 
             if imol.extras.get("embedding_charges"):  # or test on self.embedding_charges ?
                 if specifications[chem]["program"] == "psi4":
@@ -431,6 +432,7 @@ def from_manybodyinput(cls, input_model: ManyBodyInput, build_tasks: bool = True
 
             _, real, bas = delabeler(label)
             result = qcng.compute(inp, specifications[chem]["program"])
+            component_results[label] = result
 
             if not result.success:
                 print(result.error.error_message)
@@ -439,25 +441,25 @@ def from_manybodyinput(cls, input_model: ManyBodyInput, build_tasks: bool = True
             # pull out stuff
             props = {"energy", "gradient", "hessian"}
 
-            component_results[label] = {}
+            component_properties[label] = {}
 
             for p in props:
                 if hasattr(result.properties, f"return_{p}"):
                     v = getattr(result.properties, f"return_{p}")
                     # print(f"  {label} {p}: {v}")
                     if v is not None:
-                        component_results[label][p] = v
+                        component_properties[label][p] = v
 
-        print("\n<<<  (ZZ 2) QCEngine harness ManyBodyComputerQCNG.from_qcschema_ben component_results  >>>")
-        pprint.pprint(component_results, width=200)
+        print("\n<<<  (ZZ 2) QCEngine harness ManyBodyComputerQCNG.from_qcschema_ben component_properties  >>>")
+        pprint.pprint(component_properties, width=200)
 
         print("start to analyze")
-        analyze_back = computer_model.qcmb_calculator.analyze(component_results)
-        analyze_back["nbody_number"] = len(component_results)
+        analyze_back = computer_model.qcmb_calculator.analyze(component_properties)
+        analyze_back["nbody_number"] = len(component_properties)
         print("\n<<<  (ZZ 3) QCEngine harness ManyBodyComputerQCNG.from_qcschema_ben analyze_back  >>>")
         pprint.pprint(analyze_back, width=200)
 
-        return computer_model.get_results(external_results=analyze_back)
+        return computer_model.get_results(external_results=analyze_back, component_results=component_results)
 
     def plan(self):
         # uncalled function
@@ -471,7 +473,7 @@ def compute(self, client: Optional["qcportal.FractalClient"] = None) -> None:
         for t in self.task_list.values():
             t.compute(client=client)
 
-    def get_results(self, external_results: Dict, client: Optional["qcportal.FractalClient"] = None) -> ManyBodyResult:
+    def get_results(self, external_results: Dict, component_results: Dict, client: Optional["qcportal.FractalClient"] = None) -> ManyBodyResult:
         """Return results as ManyBody-flavored QCSchema."""
 
         ret_energy = external_results.pop("ret_energy")
@@ -540,7 +542,7 @@ def get_results(self, external_results: Dict, client: Optional["qcportal.Fractal
                 qcvars[qcv] = val
 
         # v2: component_results = self.model_dump()['task_list']  # TODO when/where include the indiv outputs
-        component_results = self.dict()['task_list']  # TODO when/where include the indiv outputs
+        #?component_results = self.dict()['task_list']  # TODO when/where include the indiv outputs
 #        for k, val in component_results.items():
 #            val['molecule'] = val['molecule'].to_schema(dtype=2)
 
@@ -554,6 +556,7 @@ def get_results(self, external_results: Dict, client: Optional["qcportal.Fractal
                 # v2: 'properties': {**atprop.model_dump(), **properties},
                 'properties': {**atprop.dict(), **properties},
                 'component_properties': component_properties,
+                "component_results": component_results,
                 'provenance': provenance_stamp(__name__),
                 'extras': {
                     'qcvars': qcvars,

qcmanybody/tests/test_mbe_he4_singlelevel.py

@@ -621,8 +621,8 @@ def test_nbody_he4_single(program, basis, keywords, mbe_keywords, anskey, bodyke
     #   addressing 4-body formulas
     # e, wfn = energy('MP2/aug-cc-pVDZ', molecule=he_tetramer, ...)
 
-    atomic_spec = AtomicSpecification(model={"method": "mp2", "basis": basis}, program=program, driver="energy", keywords=keywords)
-    mbe_model = ManyBodyInput(specification={"specification": atomic_spec, "keywords": mbe_keywords, "driver": "energy"}, molecule=he_tetramer)
+    atomic_spec = AtomicSpecification(model={"method": "mp2", "basis": basis}, program=program, driver="energy", keywords=keywords, protocols={"stdout": False})
+    mbe_model = ManyBodyInput(specification={"specification": atomic_spec, "keywords": mbe_keywords, "driver": "energy", "protocols": {"component_results": "all"}}, molecule=he_tetramer)
 
     # qcng: ret = qcng.compute_procedure(mbe_model, "manybody", raise_error=True)
     ret = ManyBodyComputerQCNG.from_manybodyinput(mbe_model)
@@ -669,6 +669,10 @@ def test_nbody_he4_single(program, basis, keywords, mbe_keywords, anskey, bodyke
     assert compare_values(ans, ret.return_result, atol=atol, label=f"[g] ret")
 
     assert ret.properties.calcinfo_nmbe == ref_nmbe, f"[i] {ret.properties.calcinfo_nmbe=} != {ref_nmbe}"
+    assert len(ret.component_results) == ref_nmbe, f"[k] {len(ret.component_results)=} != {ref_nmbe}; mbe protocol did not take"
+    if ref_nmbe > 0:
+        an_atres = next(iter(ret.component_results.values()))
+        assert an_atres.stdout is None, f"[l] atomic protocol did not take"
 
     if outstrs and progln != "psi4_df":
         for stdoutkey in outstrs:

qcmanybody/tests/test_mbe_keywords.py

@@ -214,6 +214,8 @@ def test_mbe_level_fails(mbe_data, kws, errmsg):
     pytest.param({"bsse_type": ["ssFC", "nocp"]}, [BsseEnum.cp, BsseEnum.nocp]),
     pytest.param({"bsse_type": ["ssfc", "cp"]}, [BsseEnum.cp]),
     pytest.param({"bsse_type": ["ssfc", "vmfc", "nocp", "cp"]}, [BsseEnum.cp, BsseEnum.vmfc, BsseEnum.nocp]),
+    pytest.param({"bsse_type": "mbE"}, [BsseEnum.nocp]),
+    pytest.param({"bsse_type": ["ssfc", "cp"]}, [BsseEnum.cp]),
     pytest.param({"bsse_type": "mycp"}, "error"),
     pytest.param({"bsse_type": ["CP", "mycp"]}, "error"),
 ])