mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/calculators/ORCA5.py

@@ -0,0 +1,6 @@
+from pysisyphus.calculators.ORCA import ORCA
+
+
+class ORCA5(ORCA):
+
+    conf_key = "orca5"

pysisyphus/calculators/OpenMM.py

@@ -0,0 +1,88 @@
+import numpy as np
+
+try:
+    from simtk.openmm import app
+    from simtk import openmm as mm
+    from simtk import unit
+except ModuleNotFoundError:
+    print("OpenMM is not installed!")
+
+from pysisyphus.constants import BOHR2ANG
+from pysisyphus.calculators.Calculator import Calculator
+
+
+class OpenMM(Calculator):
+    def __init__(self, topology, params, **kwargs):
+        super().__init__(**kwargs)
+
+        self.topology = topology
+        self.params = params
+        self.system = self.topology.createSystem(
+            self.params, nonbondedMethod=app.NoCutoff
+        )
+
+        dummy_int = mm.VerletIntegrator(0.001)
+        self.simulation = app.Simulation(self.topology, self.system, dummy_int)
+        self.context = self.simulation.context
+
+        self.en2au = 1 / unit.hartree / unit.AVOGADRO_CONSTANT_NA
+        self.frc2au = 1 / (unit.hartree / unit.bohr) / unit.AVOGADRO_CONSTANT_NA
+
+    def get_charges(self):
+        (nb_force,) = [
+            force
+            for force in self.system.getForces()
+            if force.__class__.__name__ == "NonbondedForce"
+        ]
+        charges = list()
+        for i in range(nb_force.getNumParticles()):
+            charge, *_ = nb_force.getParticleParameters(i)
+            charges.append(charge.value_in_unit(unit.elementary_charge))
+        charges = np.array(charges)
+        return charges
+
+    def report_charges(self):
+        """Total charge of the system.
+
+        Adapted from modeller.py _addIons() in OpenMM.
+        """
+
+        charges = self.get_charges()
+        total_charge = charges.sum()
+        charges_str = np.array2string(charges, precision=4)
+        self.log(f"OpenMM charges: {charges_str}")
+        self.log(f"OpenMM system total charge: {total_charge:.6f} e")
+
+    def get_dipole_moment(self, coords3d, reference=None, masses=None):
+        def get_com(coords3d):
+            total_mass = sum(masses)
+            com = (1 / total_mass * coords3d * masses[:, None]).sum(axis=0)
+            return com
+
+        charges = self.get_charges()
+        ref_dict = {
+            None: lambda c3d: c3d,
+            "centroid": lambda c3d: c3d.mean(axis=1)[None, :],
+            "com": lambda c3d: c3d - get_com(c3d)[None, :],
+        }
+        coords3d = ref_dict[reference](coords3d)
+        # Dipole moment vector in bohr * elementary_charge
+        dip_moms = (coords3d * charges[:, None]).sum(axis=0)
+        tot_dip_mom = np.linalg.norm(dip_moms)
+        return dip_moms, tot_dip_mom
+
+    def get_energy(self, atoms, coords, **prepare_kwargs):
+        results = self.get_forces(atoms, coords, **prepare_kwargs)
+        return {"energy": results["energy"]}
+
+    def get_forces(self, atoms, coords, **prepare_kwargs):
+        positions = coords.reshape(-1, 3) * BOHR2ANG / 10  # 2d array in nm
+        self.context.setPositions(positions)
+        state = self.context.getState(getEnergy=True, getForces=True)
+        energy = state.getPotentialEnergy() * self.en2au
+        forces = state.getForces(asNumpy=True)
+        forces = (forces * self.frc2au).reshape(-1)
+
+        energy = float(energy)
+        forces = np.array(forces)
+        return {"energy": energy, "forces": forces}

pysisyphus/calculators/OpenMolcas.py

@@ -0,0 +1,281 @@
+import os
+from pathlib import Path
+import re
+import warnings
+
+import numpy as np
+
+from pysisyphus.calculators.Calculator import Calculator
+from pysisyphus.constants import BOHR2ANG
+from pysisyphus.xyzloader import make_xyz_str
+
+
+class OpenMolcas(Calculator):
+
+    conf_key = "openmolcas"
+    _set_plans = (
+        ("rasorb", "inporb"),
+        "jobiph",
+    )
+
+    def __init__(
+        self,
+        basis,
+        inporb,
+        rasscf=None,
+        gateway=None,
+        mcpdft=None,
+        track=True,
+        **kwargs,
+    ):
+        super(OpenMolcas, self).__init__(**kwargs)
+
+        assert self.pal == 1, (
+            "RI SA-CASSCF analytical gradients do not work correctly in "
+            "parallel (yet). Consider using pal=1 instead of the current "
+            f"pal={self.pal}!"
+        )
+        env = os.environ
+        if not ("MOLCAS" in env):
+            warnings.warn(
+                "$MOLCAS environment variable is not set! Couldn't find it in the "
+                "environment!"
+            )
+
+        self.basis = basis
+        inporb = Path(inporb).absolute()
+        self.inporb = inporb
+        if rasscf is None:
+            rasscf = {}
+        self.rasscf = rasscf
+        if gateway is None:
+            gateway = {}
+        self.gateway = gateway
+        if mcpdft is None:
+            mcpdft = {}
+        self.mcpdft = mcpdft
+        self.track = track
+
+        self.to_keep = ("RasOrb", "out", "in", "JobIph", "rasscf.molden")
+        self.jobiph = ""
+
+        self.inp_fn = "openmolcas.in"
+        self.out_fn = "openmolcas.out"
+        self.float_regex = r"([\d\.\-E]+)"
+
+        self.openmolcas_input = """
+        >> copy {inporb}  $Project.RasOrb
+        &gateway
+         coord
+          {xyz_str}
+         basis
+          {basis}
+         group
+          nosym
+         {gateway_kwargs}
+
+        &seward
+         doanalytical
+
+        &rasscf
+         charge
+          {charge}
+         spin
+          {mult}
+         fileorb
+          $Project.RasOrb
+         {rasscf_kwargs}
+
+        {mcpdft}
+
+        >> copy $Project.JobIph $CurrDir/$Project.JobIph
+
+        {rassi}
+
+        {alaska}
+        """
+
+        self.parser_funcs = {
+            "energy": self.parse_energy,
+            "grad": self.parse_gradient,
+        }
+
+        self.base_cmd = self.get_cmd()
+
+    def reattach(self, last_calc_cycle):
+        self.inporb = self.make_fn("RasOrb", last_calc_cycle)
+        self.jobiph = self.make_fn("JobIph", last_calc_cycle)
+        self.log(f"restarted. using {self.inporb}, {self.jobiph}")
+
+    def build_str_from_dict(self, dct):
+        strs = list()
+        for key, val in dct.items():
+            if val is None:
+                val = ""
+            strs.append(f"{key}\n{val}")
+        return "\n".join(strs)
+
+    def build_gateway_str(self):
+        return self.build_str_from_dict(self.gateway)
+
+    def build_rasscf_str(self):
+        return self.build_str_from_dict(self.rasscf)
+
+    def build_rassi_str(self):
+        # In the first iteration self.jobiph isn't set yet.
+        if (not self.track) or (self.calc_counter == 0):
+            return ""
+        else:
+            # JOB001 corresponds to the current iteration,
+            # JOB002 to the previous iteration.
+            return f"""
+            >> copy $Project.JobIph JOB001
+            >> copy {self.jobiph} JOB002
+            &rassi
+             track
+            """
+
+    def build_mcpdft_str(self):
+        if not self.mcpdft:
+            return ""
+
+        mcpdft_kwargs = self.build_str_from_dict(self.mcpdft)
+        return f"&mcpdft\n{mcpdft_kwargs}"
+
+    def get_pal_env(self):
+        env_copy = os.environ.copy()
+        env_copy["MOLCAS_NPROCS"] = str(self.pal)
+
+        return env_copy
+
+    def prepare_coords(self, atoms, coords):
+        coords = coords * BOHR2ANG
+        return make_xyz_str(atoms, coords.reshape((-1, 3)))
+
+    def prepare_input(self, atoms, coords, calc_type):
+        self.log(f"using inporb: {self.inporb}")
+        xyz_str = self.prepare_coords(atoms, coords)
+        alaska_str = "&alaska\npnew" if calc_type == "grad" else ""
+        inp = self.openmolcas_input.format(
+            inporb=self.inporb,
+            xyz_str=xyz_str,
+            basis=self.basis,
+            charge=self.charge,
+            mult=self.mult,
+            gateway_kwargs=self.build_gateway_str(),
+            rasscf_kwargs=self.build_rasscf_str(),
+            mcpdft=self.build_mcpdft_str(),
+            rassi=self.build_rassi_str(),
+            alaska=alaska_str,
+        )
+        return inp
+
+    def run_calculation(self, atoms, coords, calc_type="energy"):
+        inp = self.prepare_input(atoms, coords, calc_type)
+        add_args = ("-clean", "-oe", self.out_fn)
+        env_copy = self.get_pal_env()
+        env_copy["MOLCAS_PROJECT"] = f"{self.name}_{self.calc_counter}"
+        kwargs = {
+            "calc": calc_type,
+            "add_args": add_args,
+            "env": env_copy,
+        }
+        results = self.run(inp, **kwargs)
+        return results
+
+    def get_energy(self, atoms, coords):
+        return self.run_calculation(atoms, coords, "energy")
+
+    def get_forces(self, atoms, coords):
+        return self.run_calculation(atoms, coords, "grad")
+
+    def get_root(self):
+        ras = self.rasscf
+        return int(ras.get("mdrlxroot", ras.get("rlxroot", 1)))
+
+    def parse_energies(self, text):
+        if self.mcpdft:
+            root_re = r"PDFT Root\s*\d+\s*Total energy:\s*" + self.float_regex
+            matches = re.findall(root_re, text)
+            root_energies = np.array(matches, dtype=float)
+            root = self.get_root()
+            energy = root_energies[root - 1]
+        else:
+            # Energy of root for which gradient was computed
+            energy_regex = r"RASSCF state energy =\s*" + self.float_regex
+            energy = float(re.search(energy_regex, text).groups()[0])
+            # All state average energies
+            root_re = "RASSCF root number.+Total energy.+?" + self.float_regex
+            matches = re.findall(root_re, text)
+            root_energies = np.array(matches, dtype=float)
+        return energy, root_energies
+
+    def parse_energy(self, path):
+        with open(path / self.out_fn) as handle:
+            text = handle.read()
+        energy, _ = self.parse_energies(text)
+        return {
+            "energy": energy,
+        }
+
+    def parse_gradient(self, path):
+        results = {}
+        gradient_fn = os.path.join(path, self.out_fn)
+        with open(gradient_fn) as handle:
+            text = handle.read()
+
+        # Search for the block containing the gradient table
+        regex = r"Molecular gradients(.+?)--- Stop Module:\s*alaska"
+        floats = [self.float_regex for i in range(3)]
+        line_regex = r"([A-Z\d]+)\s*" + r"\s*".join(floats)
+
+        mobj = re.search(regex, text, re.DOTALL)
+        gradient = list()
+        for line in mobj.groups()[0].split("\n"):
+            # Now look for the lines containing the gradient
+            mobj = re.match(line_regex, line.strip())
+            if not mobj:
+                continue
+            # Discard first column (atom+number)
+            gradient.append(mobj.groups()[1:])
+        gradient = np.array(gradient, dtype=float).flatten()
+
+        if self.track and self.calc_counter > 0:
+            self.parse_rassi_track(path)
+
+        energy, _ = self.parse_energies(text)
+
+        results["energy"] = energy
+        # results["sa_energies"] = sa_energies
+        results["forces"] = -gradient
+
+        return results
+
+    def parse_rassi_track(self, path):
+        gradient_fn = path / self.out_fn
+        with open(gradient_fn) as handle:
+            text = handle.read()
+        track_re = (
+            r"Initial root:\s*(\d+)\s*Overlaps with current "
+            r"states:(.+)New root:\s*(\d+)"
+        )
+        # overlap_re = "OVERLAP MATRIX FOR THE ORIGINAL STATES:(.+?)##"
+        mobj = re.search(track_re, text, re.DOTALL)
+
+        initial_root, overlaps, new_root = mobj.groups()
+        overlaps = np.array(overlaps.strip().split(), dtype=float).reshape(-1, 2)
+        # Filters for overlaps > 10% (0.1**2 ~ 0.31622)
+        thresh = 0.1
+        inds = np.where(np.abs(overlaps[:, 1]) > thresh**0.5)
+        ov_perc_str = ", ".join([f"{nr:.0f}: {ov**2:.2%}" for nr, ov in overlaps[inds]])
+        self.log(
+            f"Overlaps between previous root {initial_root} and "
+            f"new roots bigger {thresh:.0%}:  {ov_perc_str}. Will "
+            f"use root {new_root} for the following gradient calculation."
+        )
+        if new_root != initial_root:
+            self.log("Found a root flip!")
+        self.mdrlxroot = new_root
+
+    def __str__(self):
+        return "OpenMolcas calculator"