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

pysisyphus/calculators/Psi4.py

@@ -0,0 +1,218 @@
+import re
+import textwrap
+
+import numpy as np
+
+from pysisyphus.calculators.Calculator import Calculator
+
+
+class Psi4(Calculator):
+
+    conf_key = "psi4"
+
+    def __init__(
+        self,
+        method,
+        basis,
+        to_set=None,
+        to_import=None,
+        pcm="iefpcm",
+        solvent=None,
+        write_fchk=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.method = method
+        self.basis = basis
+        self.to_set = {} if to_set is None else dict(to_set)
+        self.to_import = [] if to_import is None else list(to_import)
+        self.pcm = pcm
+        self.solvent = solvent
+        self.write_fchk = write_fchk
+
+        self.inp_fn = "psi4.inp"
+        self.out_fn = "psi4.out"
+        self.to_keep = ("inp", "psi4.out", "grad.npy", "hessian.npy")
+
+        self.parser_funcs = {
+            "energy": self.parse_energy,
+            "grad": self.parse_grad,
+            "hessian": self.parse_hessian,
+        }
+
+        self.base_cmd = self.get_cmd()
+
+        self.inp = textwrap.dedent(
+            """
+        {to_import}
+        molecule mol{{
+          {xyz}
+          {charge} {mult}
+        no_reorient
+        no_com  
+        symmetry c1
+        }}
+
+        set_num_threads({pal})
+        memory {mem} MB
+
+        {basis}
+        {to_set}
+        {pcm}
+
+        {method}
+
+        {fchk}
+        """
+        )
+
+    def get_fchk_str(self):
+        fchk_str = ""
+        if self.write_fchk:
+            fchk_fn = self.make_fn("wfn.fchk")
+            fchk_str = (
+                "fchk_writer = psi4.FCHKWriter(wfn)\n" f"fchk_writer.write('{fchk_fn}')"
+            )
+        return fchk_str
+
+    def prepare_input(self, atoms, coords, calc_type):
+        xyz = self.prepare_coords(atoms, coords)
+
+        calc_types = {
+            "energy": "E, wfn = energy('{}', return_wfn=True)",
+            # Right now we don't need the wavefunction
+            # "grad": "G, wfn = gradient('{}', return_wfn=True)\n" \
+            # "G_arr = np.array(G)\n" \
+            # "np.save('grad', G_arr)",
+            # "hessian": "H, wfn = hessian('{}', return_wfn=True)\n" \
+            # "H_arr = np.array(H)\n" \
+            # "np.save('hessian', H_arr)",
+            "grad": "G, wfn = gradient('{}', return_wfn=True)\n"
+            "G_arr = np.array(G)\n"
+            "np.save('grad', G_arr)",
+            "hessian": "H, wfn = hessian('{}', return_wfn=True)\n"
+            "H_arr = np.array(H)\n"
+            "np.save('hessian', H_arr)",
+        }
+        method = calc_types[calc_type].format(self.method)
+        wfn_path = self.make_fn("wfn.npy")
+        method += f"\nWavefunction.to_file(wfn, '{wfn_path}')"
+        method += "\nprint('PARSE ENERGY:', wfn.energy())"
+        set_strs = [f"set {key} {value}" for key, value in self.to_set.items()]
+        set_strs = "\n".join(set_strs)
+        import_strs = [f"import {mod}" for mod in self.to_import]
+        import_strs = "\n".join(import_strs)
+
+        # Basis section
+        basis = self.basis
+        # Construct more complex basis input
+        if isinstance(basis, dict):
+            # Check if a global basis is given for all atoms. This must come
+            # first, otherwise Psi4 throws an error.
+            basis_lines = [
+                "basis {",
+            ]
+            try:
+                basis_lines.append(f"assign {basis['assign']}")
+            except KeyError:
+                pass
+            # Add remaining lines
+            basis_lines.extend(
+                [
+                    f"assign {atms} {bas}"
+                    for atms, bas in basis.items()
+                    if atms != "assign"
+                ]
+            )
+            basis_lines.append("}")
+
+            basis = "\n".join(basis_lines)
+        # Use set when self.basis is a string
+        else:
+            basis = f"set basis {basis}"
+
+        # PCM section
+        pcm = ""
+        if self.solvent:
+            pcm = textwrap.dedent(
+                f"""
+            set pcm true
+
+            pcm = {{
+                Medium {{
+                    SolverType = {self.pcm}
+                    Solvent = {self.solvent}
+                }}
+
+                Cavity {{
+                    Type = GePol
+                }}
+            }}
+            """
+            )
+
+        inp = self.inp.format(
+            xyz=xyz,
+            charge=self.charge,
+            mult=self.mult,
+            basis=basis,
+            to_import=import_strs,
+            to_set=set_strs,
+            pcm=pcm,
+            method=method,
+            pal=self.pal,
+            mem=self.mem,
+            fchk=self.get_fchk_str(),
+        )
+        # inp = "\n".join([line.strip() for line in inp.split("\n")])
+        return inp
+
+    def get_energy(self, atoms, coords, **prepare_kwargs):
+        calc_type = "energy"
+        inp = self.prepare_input(atoms, coords, calc_type)
+        results = self.run(inp, calc="energy")
+        return results
+
+    def get_forces(self, atoms, coords, **prepare_kwargs):
+        calc_type = "grad"
+        inp = self.prepare_input(atoms, coords, calc_type)
+        results = self.run(inp, calc="grad")
+        return results
+
+    def get_hessian(self, atoms, coords, **prepare_kwargs):
+        calc_type = "hessian"
+        inp = self.prepare_input(atoms, coords, calc_type)
+        results = self.run(inp, calc="hessian")
+        return results
+
+    def run_calculation(self, atoms, coords, **prepare_kwargs):
+        return self.get_energy(atoms, coords)
+
+    def parse_energy(self, path):
+        with open(path / "psi4.out") as handle:
+            text = handle.read()
+        en_regex = re.compile(r"PARSE ENERGY: ([\d\-\.]+)")
+        mobj = en_regex.search(text)
+        result = {"energy": float(mobj[1])}
+        return result
+
+    def parse_grad(self, path):
+        gradient = np.load(path / "grad.npy")
+        forces = -gradient.flatten()
+        result = {
+            "forces": forces,
+        }
+        result.update(self.parse_energy(path))
+        return result
+
+    def parse_hessian(self, path):
+        hessian = np.load(path / "hessian.npy")
+        result = {
+            "hessian": hessian,
+        }
+        result.update(self.parse_energy(path))
+        return result
+
+    def __str__(self):
+        return f"Psi4({self.name})"

pysisyphus/calculators/PyPsi4.py

@@ -0,0 +1,37 @@
+try:
+    import psi4
+except ImportError:
+    pass
+import numpy as np
+
+from pysisyphus.calculators.Calculator import Calculator
+from pysisyphus.constants import BOHR2ANG
+from pysisyphus.xyzloader import make_xyz_str
+
+
+class PyPsi4(Calculator):
+    def __init__(self, method, basis, **kwargs):
+        super().__init__(**kwargs)
+
+        self.method = method
+        self.basis = basis
+
+        self.meth_bas = f"{self.method}/{self.basis}"
+
+        psi4.core.be_quiet()
+        psi4.set_memory(f"{self.pal*self.mem} MB")
+
+    def get_forces(self, atoms, coords):
+        xyz = make_xyz_str(atoms, coords.reshape(-1, 3) * BOHR2ANG)
+        mol = psi4.geometry(xyz)
+        energy, wfn = psi4.energy(self.meth_bas, return_wfn=True)
+        gradient = psi4.gradient(self.meth_bas, molecule=mol, ref_wfn=wfn)
+
+        results = {
+            "energy": energy,
+            "forces": -np.asarray(gradient).flatten(),
+        }
+        return results
+
+    def __str__(self):
+        return f"PyPsi4({self.name})"

pysisyphus/calculators/PySCF.py

@@ -0,0 +1,341 @@
+import os
+import shutil
+
+import numpy as np
+import pyscf
+from pyscf import grad, gto, lib, hessian, qmmm, tddft
+
+from pysisyphus.calculators.OverlapCalculator import OverlapCalculator
+from pysisyphus.helpers import geom_loader
+
+
+class PySCF(OverlapCalculator):
+    conf_key = "pyscf"
+    drivers = {
+        # key: (method, unrestricted?)
+        ("dft", False): pyscf.dft.RKS,
+        ("dft", True): pyscf.dft.UKS,
+        ("scf", False): pyscf.scf.RHF,
+        ("scf", True): pyscf.scf.UHF,
+        ("mp2", False): pyscf.mp.MP2,
+        ("mp2", True): pyscf.mp.UMP2,
+    }
+    multisteps = {
+        "scf": ("scf",),
+        "dft": ("dft",),
+        "mp2": ("scf", "mp2"),
+        "tddft": ("dft", "tddft"),
+        "tda": ("dft", "tda"),
+    }
+    pruning_method = {
+        "nwchem": pyscf.dft.gen_grid.nwchem_prune,
+        "sg1": pyscf.dft.gen_grid.sg1_prune,
+        "treutler": pyscf.dft.gen_grid.treutler_prune,
+        "none": None,
+    }
+
+    def __init__(
+        self,
+        basis,
+        xc=None,
+        method="scf",
+        root=None,
+        nstates=None,
+        auxbasis=None,
+        keep_chk=True,
+        verbose=0,
+        unrestricted=None,
+        grid_level=3,
+        pruning="nwchem",
+        use_gpu=False,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.basis = basis
+        self.xc = xc
+        self.method = method.lower()
+        if self.method in ("tda", "tddft") and self.xc is None:
+            self.multisteps[self.method] = ("scf", self.method)
+        if self.xc and self.method != "tddft":
+            self.method = "dft"
+        self.root = root
+        self.nstates = nstates
+        if self.method == "tddft":
+            assert self.nstates, "nstates must be set with method='tddft'!"
+        if self.track:
+            assert self.root <= self.nstates, (
+                "'root' must be smaller " "than 'nstates'!"
+            )
+        self.auxbasis = auxbasis
+        self.keep_chk = keep_chk
+        self.verbose = int(verbose)
+        if unrestricted is None:
+            self.unrestricted = self.mult > 1
+        else:
+            self.unrestricted = unrestricted
+        self.grid_level = int(grid_level)
+        self.pruning = pruning.lower()
+
+        self.chkfile = None
+        self.out_fn = "pyscf.out"
+
+        self.use_gpu = use_gpu
+
+        lib.num_threads(self.pal)
+
+    @staticmethod
+    def geom_from_fn(fn, **kwargs):
+        geom = geom_loader(fn)
+        geom.set_calculator(PySCF(**kwargs))
+        return geom
+
+    def set_scf_params(self, mf):
+        mf.conv_tol = 1e-8
+        mf.max_cycle = 150
+
+    def build_grid(self, mf):
+        mf.grids.level = self.grid_level
+        mf.grids.prune = self.pruning_method[self.pruning]
+        mf.grids.build()
+
+    def prepare_mf(self, mf):
+        # Method can be overriden in a subclass to modify the mf object.
+        if self.use_gpu:
+            return mf.to_gpu()
+        else:
+            return mf
+
+    def get_driver(self, step, mol=None, mf=None):
+        def _get_driver():
+            return self.drivers[(step, self.unrestricted)]
+
+        if mol and (step == "dft"):
+            driver = _get_driver()
+            mf = driver(mol)
+            mf.xc = self.xc
+            self.set_scf_params(mf)
+            self.build_grid(mf)
+            mf = self.prepare_mf(mf)
+        elif mol and (step == "scf"):
+            driver = _get_driver()
+            mf = driver(mol)
+            self.set_scf_params(mf)
+            mf = self.prepare_mf(mf)
+        elif mf and (step == "mp2"):
+            mp2_mf = _get_driver()
+            mf = mp2_mf(mf)
+        elif mf and (step == "tddft"):
+            mf = pyscf.tddft.TDDFT(mf)
+            mf.nstates = self.nstates
+        elif mf and (step == "tda"):
+            mf = pyscf.tddft.TDA(mf)
+            mf.nstates = self.nstates
+        else:
+            raise Exception("Unknown method '{step}'!")
+        return mf
+
+    def prepare_mol(self, atoms, coords, build=True):
+        mol = gto.Mole()
+        mol.atom = [(atom, c) for atom, c in zip(atoms, coords.reshape(-1, 3))]
+        mol.basis = self.basis
+        mol.unit = "Bohr"
+        mol.charge = self.charge
+        mol.spin = self.mult - 1
+        mol.symmetry = False
+        mol.verbose = self.verbose
+        # Personally, I patched mole.py so it doesn't print
+        # messages regarding the output-file for verbose > QUIET.
+        # Just uncomment the lines after
+        #   if self.verbose > logger.QUIET:
+        #       ...
+        # in 'mole.Mole.build'. Around line 2046 for pyscf 1.6.5.
+        # Search for "output file" in gto/mole.py
+        # Search for "Large deviations found" in scf/{uhf,dhf,ghf}.py
+        mol.output = self.make_fn(self.out_fn)
+        mol.max_memory = self.mem * self.pal
+        if build:
+            mol.build(parse_arg=False)
+        return mol
+
+    def prepare_input(self, atoms, coords, build=True):
+        mol = self.prepare_mol(atoms, coords, build=build)
+        assert mol._built, "Please call mol.build(parse_arg=False)!"
+        return mol
+
+    def store_and_track(self, results, func, atoms, coords, **prepare_kwargs):
+        if self.track:
+            self.store_overlap_data(atoms, coords)
+            if self.track_root():
+                # Redo the calculation with the updated root
+                results = func(atoms, coords, **prepare_kwargs)
+        results["all_energies"] = self.parse_all_energies()
+        return results
+
+    def get_energy(self, atoms, coords, **prepare_kwargs):
+        point_charges = prepare_kwargs.get("point_charges", None)
+
+        mol = self.prepare_input(atoms, coords)
+        mf = self.run(mol, point_charges=point_charges)
+        results = {
+            "energy": mf.e_tot,
+        }
+        results = self.store_and_track(
+            results, self.get_energy, atoms, coords, **prepare_kwargs
+        )
+        return results
+
+    def get_forces(self, atoms, coords, **prepare_kwargs):
+        point_charges = prepare_kwargs.get("point_charges", None)
+
+        mol = self.prepare_input(atoms, coords)
+        mf = self.run(mol, point_charges=point_charges)
+        grad_driver = mf.Gradients()
+        if self.root:
+            grad_driver.state = self.root
+        gradient = grad_driver.kernel()
+        self.log("Completed gradient step")
+
+        try:
+            e_tot = mf._scf.e_tot
+        except AttributeError:
+            e_tot = mf.e_tot
+
+        results = {
+            "energy": e_tot,
+            "forces": -gradient.flatten(),
+        }
+        results = self.store_and_track(
+            results, self.get_forces, atoms, coords, **prepare_kwargs
+        )
+        return results
+
+    def get_hessian(self, atoms, coords, **prepare_kwargs):
+        point_charges = prepare_kwargs.get("point_charges", None)
+
+        mol = self.prepare_input(atoms, coords)
+        mf = self.run(mol, point_charges=point_charges)
+        H = mf.Hessian().kernel()
+
+        # The returned hessian is 4d ... ok. This probably serves a purpose
+        # that I don't understand. We transform H to a nice, simple 2d array.
+        H = np.hstack(np.concatenate(H, axis=1))
+        results = {
+            "energy": mf.e_tot,
+            "hessian": H,
+        }
+        # results = self.store_and_track(
+        # results, self.get_hessian, atoms, coords, **prepare_kwargs
+        # )
+        return results
+
+    def run_calculation(self, atoms, coords, **prepare_kwargs):
+        return self.get_energy(atoms, coords, **prepare_kwargs)
+
+    def run(self, mol, point_charges=None):
+        steps = self.multisteps[self.method]
+        self.log(f"Running steps '{steps}' for method {self.method}")
+        for i, step in enumerate(steps):
+            if i == 0:
+                mf = self.get_driver(step, mol=mol)
+                assert step in ("scf", "dft")
+                if self.chkfile:
+                    # Copy old chkfile to new chkfile
+                    new_chkfile = self.make_fn("chkfile", return_str=True)
+                    shutil.copy(self.chkfile, new_chkfile)
+                    self.chkfile = new_chkfile
+                    mf.chkfile = self.chkfile
+                    mf.init_guess = "chkfile"
+                    self.log(
+                        f"Using '{self.chkfile}' as initial guess for {step} calculation."
+                    )
+                if self.auxbasis:
+                    mf = mf.density_fit(auxbasis=self.auxbasis)
+                    self.log(f"Using density fitting with auxbasis {self.auxbasis}.")
+
+                if point_charges is not None:
+                    mf = qmmm.mm_charge(mf, point_charges[:, :3], point_charges[:, 3])
+                    self.log(
+                        f"Added {len(point_charges)} point charges with "
+                        f"sum(q)={sum(point_charges[:,3]):.4f}."
+                    )
+            else:
+                mf = self.get_driver(step, mf=prev_mf)  # noqa: F821
+
+            if self.keep_chk and (self.chkfile is None) and (step in ("dft", "scf")):
+                self.chkfile = self.make_fn("chkfile", return_str=True)
+                try:
+                    os.remove(self.chkfile)
+                except FileNotFoundError:
+                    self.log(f"Tried to remove '{self.chkfile}'. It doesn't exist.")
+                self.log(f"Created chkfile '{self.chkfile}'")
+                mf.chkfile = self.chkfile
+            mf.kernel()
+            self.log(f"Completed {step} step")
+            prev_mf = mf
+
+        # Keep mf and dump mol
+        # save_mol(mol, self.make_fn("mol.chk"))
+        self.mf = mf.reset()  # release integrals and other temporary intermediates.
+        self.calc_counter += 1
+
+        return mf
+
+    def parse_all_energies(self, exc_mf=None):
+        if exc_mf is None:
+            exc_mf = self.mf
+
+        try:
+            gs_energy = exc_mf._scf.e_tot
+            exc_energies = exc_mf.e_tot
+            all_energies = np.zeros(exc_energies.size + 1)
+            all_energies[0] = gs_energy
+            all_energies[1:] = exc_energies
+        except AttributeError:
+            gs_energy = exc_mf.e_tot
+            all_energies = np.array((gs_energy,))
+
+        return all_energies
+
+    def prepare_overlap_data(self, path):
+        gs_mf = self.mf._scf
+        exc_mf = self.mf
+
+        C = gs_mf.mo_coeff
+
+        first_Y = exc_mf.xy[0][1]
+        # In TDA calculations Y is just the integer 0.
+        if isinstance(first_Y, int) and (first_Y == 0):
+            X = np.array([state[0] for state in exc_mf.xy])
+            Y = np.zeros_like(X)
+        # In TD-DFT calculations the Y vectors is also present
+        else:
+            # Shape = (nstates, 2 (X,Y), occ, virt)
+            ci_coeffs = np.array(exc_mf.xy)
+            X = ci_coeffs[:, 0]
+            Y = ci_coeffs[:, 1]
+
+        all_energies = self.parse_all_energies(exc_mf)
+        return C, X, Y, all_energies
+
+    def parse_charges(self):
+        results = self.mf.analyze(with_meta_lowdin=False)
+        # Mulliken charges
+        charges = results[0][1]
+        return charges
+
+    def get_chkfiles(self):
+        return {
+            "chkfile": self.chkfile,
+        }
+
+    def set_chkfiles(self, chkfiles):
+        try:
+            chkfile = chkfiles["chkfile"]
+            self.chkfile = chkfile
+            self.log(f"Set chkfile '{chkfile}' as chkfile.")
+        except KeyError:
+            self.log("Found no chkfile information in chkfiles!")
+
+    def __str__(self):
+        return f"PySCF({self.name})"

pysisyphus/calculators/PyXTB.py

@@ -0,0 +1,73 @@
+import numpy as np
+try:
+    from xtb.interface import Environment, Param, Calculator as XTBCalculator
+    from xtb.libxtb import VERBOSITY_MINIMAL, VERBOSITY_FULL, VERBOSITY_MUTED
+except ModuleNotFoundError:
+    can_pyxtb = False
+
+from pysisyphus.calculators.Calculator import Calculator
+from pysisyphus.elem_data import ATOMIC_NUMBERS
+
+
+class PyXTB(Calculator):
+
+    def __init__(self, *args, gfn=2, acc=None, verbosity=0, keep_calculator=False, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.env = Environment()
+        self.gfn = str(gfn)
+        self.acc = acc
+        avail_verbosities = {
+            0: VERBOSITY_MUTED,
+            1: VERBOSITY_MINIMAL,
+            2: VERBOSITY_FULL,
+        }
+        self.verbosity = avail_verbosities[verbosity]
+        avail_params = {
+            "0": Param.GFN0xTB,
+            "1": Param.GFN1xTB,
+            "2": Param.GFN2xTB,
+            "ff": Param.GFNFF,
+        }
+        self.keep_calculator = keep_calculator
+        self._calculator = None
+
+        self.param = avail_params[self.gfn]
+        self.uhf = self.mult - 1
+
+    def get_calculator(self, atoms, coords):
+        # Reuse old calculator and updatec coordinates
+        if self._calculator:
+            self._calculator.update(coords.reshape(-1, 3))
+            return self._calculator
+
+        # Create new calculator
+        numbers = np.array([ATOMIC_NUMBERS[atom.lower()] for atom in atoms])
+        calc = XTBCalculator(
+            self.param, numbers, coords.reshape(-1, 3), charge=self.charge, uhf=self.uhf
+        )
+        calc.set_verbosity(self.verbosity)
+        if self.acc:
+            calc.set_accuracy(self.acc)
+
+        # Keep calculator, if requested
+        if self.keep_calculator and (self._calculator is None):
+            self._calculator = calc
+        return calc
+
+    def get_energy(self, atoms, coords, **prepare_kwargs):
+        calc = self.get_calculator(atoms, coords)
+        xtb_result = calc.singlepoint(copy=True)
+        results = {
+            "energy": xtb_result.get_energy(),
+        }
+        return results
+
+    def get_forces(self, atoms, coords, **prepare_kwargs):
+        calc = self.get_calculator(atoms, coords)
+        xtb_result = calc.singlepoint(copy=True)
+        results = {
+            "energy": xtb_result.get_energy(),
+            "forces": -xtb_result.get_gradient().flatten(),
+        }
+        return results