add charge embedding (#15)

* works

* fix arg

* fix test

qcmanybody/manybody.py

@@ -35,10 +35,9 @@ def __init__(
         levels: Mapping[Union[int, Literal["supersystem"]], str],
         return_total_data: bool,
         supersystem_ie_only: bool,
+        embedding_charges: Mapping[int, Sequence[float]],
     ):
-        # TODO
-        self.embedding_charges = {}
-
+        self.embedding_charges = embedding_charges
         self.molecule = molecule
         self.bsse_type = [BsseEnum(x) for x in bsse_type]
         self.return_total_data = return_total_data
@@ -197,13 +196,13 @@ def iterate_molecules(self) -> Iterable[Tuple[str, str, Molecule]]:
                     mol = self.molecule.get_fragment(real_atoms_0, ghost_atoms_0, orient=False, group_fragments=False)
                     mol = mol.copy(update={"fix_com": True, "fix_orientation": True})
 
-                    # if self.embedding_charges:
-                    #    embedding_frags = list(set(range(1, self.nfragments + 1)) - set(pair[1]))
-                    #    charges = []
-                    #    for frag in embedding_frags:
-                    #        positions = self.molecule.extract_subsets(frag).geometry().np.tolist()
-                    #        charges.extend([[chg, i] for i, chg in zip(positions, self.embedding_charges[frag])])
-                    #    data['keywords']['function_kwargs'].update({'external_potentials': charges})
+                    if self.embedding_charges:
+                        embedding_frags = list(set(range(1, self.nfragments + 1)) - set(basis_atoms))
+                        charges = []
+                        for ifr in embedding_frags:
+                            positions = self.molecule.get_fragment(ifr-1).geometry.tolist()
+                            charges.extend([[chg, i] for i, chg in zip(positions, self.embedding_charges[ifr])])
+                        mol.extras["embedding_charges"] = charges
 
                     done_molecules.add(label)
                     yield mc, label, mol

qcmanybody/models/manybody_pydv1.py

@@ -99,6 +99,8 @@ class ManyBodyKeywords(ProtoModel):
         {},
         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?
+            # TODO embedding charges irrelevant to CP (basis sets always present)?
         json_schema_extra={
             "shape": ["nfr", "<varies: nat in ifr>"],
         },

qcmanybody/models/qcng_computer.py

@@ -212,10 +212,12 @@ def set_bsse_type(cls, v: Any) -> List[BsseEnum]:
     @classmethod
     # v2: def set_embedding_charges(cls, v: Any, info: FieldValidationInfo) -> Dict[int, List[float]]:
     def set_embedding_charges(cls, v, values): # -> Dict[int, List[float]]:
+        print(f"hit embedding_charges validator with {v}", end="")
         # v2: if len(v) != info.data["nfragments"]:
-        if len(v) != values["nfragments"]:
-            raise ValueError("embedding_charges dict should have entries for each 1-indexed fragment.")
+        if len(v) != len(values["molecule"].fragments):
+            raise ValueError(f"embedding_charges dict should have entries for each 1-indexed fragment ({len(values['molecule'].fragments)}).")
 
+        print(f" ... setting embedding_charges={v}")
         return v
 
     # v2: @field_validator("return_total_data")
@@ -397,6 +399,7 @@ def from_qcschema_ben(cls, input_model: ManyBodyInput, build_tasks: bool = True)
             specifications[mtd] = {}
             specifications[mtd]["program"] = spec.pop("program")
             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)
 
@@ -406,6 +409,7 @@ def from_qcschema_ben(cls, input_model: ManyBodyInput, build_tasks: bool = True)
             comp_levels,
             computer_model.return_total_data,
             computer_model.supersystem_ie_only,
+            computer_model.embedding_charges,
         )
 
         if not build_tasks:
@@ -416,6 +420,15 @@ def from_qcschema_ben(cls, input_model: ManyBodyInput, build_tasks: bool = True)
         for chem, label, imol in calculator_cls.iterate_molecules():
             inp = AtomicInput(molecule=imol, **specifications[chem]["specification"])
 
+            if imol.extras.get("embedding_charges"):  # or test on self.embedding_charges ?
+                if specifications[chem]["program"] == "psi4":
+                    charges = imol.extras["embedding_charges"]
+                    fkw = inp.keywords.get("function_kwargs", {})
+                    fkw.update({"external_potentials": charges})
+                    inp.keywords["function_kwargs"] = fkw
+                else:
+                    raise RuntimeError(f"Don't know how to handle external charges in {specifications[chem]['program']}")
+
             _, real, bas = delabeler(label)
             result = qcng.compute(inp, specifications[chem]["program"])