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

pysisyphus/io/crd.py

@@ -0,0 +1,101 @@
+import re
+
+import jinja2
+import numpy as np
+
+from pysisyphus.constants import ANG2BOHR
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import file_or_str
+
+
+@file_or_str(".crd")
+def parse_crd(crd):
+    lines = [l.strip() for l in crd.strip().split("\n") if not l.startswith("*")]
+    expect_line = lines.pop(0)
+    expect_num = int(expect_line.split()[0])
+    coords3d = np.zeros((expect_num, 3), dtype=float)
+    atoms = list()
+    re_ = re.compile(r"\d+")
+    for i, line in enumerate(lines):
+        _, _, res, atom, x, y, z, *_ = line.split()
+        coords3d[i] = (x, y, z)
+        atom = re_.sub("", atom)  # Delte number
+        atoms.append(atom)
+    assert len(atoms) == expect_num
+    coords3d *= ANG2BOHR
+    return atoms, coords3d.flatten()
+
+
+def geom_from_crd(fn, **kwargs):
+    atoms, coords = parse_crd(fn)
+    return Geometry(atoms, coords, **kwargs)
+
+
+# I10   I10 2X A8 2X A8       3F20.10     2X A8 2X A8 F20.10
+CRD_TPL = """* Created by pysisyphus
+*
+{{ i10(atomnum) }}  EXT
+{% for atom, x, y, z in atoms_xyz -%}
+ {{ i10(loop.index) }}{{ i10(resno) }}  {{ a8(res) }}  {{ a8(atom) }}{{ f20(x) }}{{ f20(y) }}{{ f20(z) }}  {{ a8(segid) }}  {{ a8(resid) }}{{ f20(0.0) }}
+{% endfor %}
+"""
+
+
+def i10(int_):
+    return f"{int_:10d}"
+
+
+def f20(float_):
+    return f"{float_:>20.10f}"
+
+
+def a8(str_):
+    return f"{str(str_):<8s}"
+
+
+def atoms_coords_to_crd_str(
+    atoms,
+    coords,
+    resno=1,
+    res="UNL1",
+    segid=None,
+    resid=1,
+    ref_atoms=None,
+    del_atoms=None,
+):
+    if segid is None:
+        segid = res
+    if del_atoms is None:
+        del_atoms = []
+
+    env = jinja2.Environment(loader=jinja2.BaseLoader)
+    env.globals.update(i10=i10, f20=f20, a8=a8)
+    tpl = env.from_string(CRD_TPL)
+
+    if ref_atoms is not None:
+        atoms_for_names = ref_atoms
+    else:
+        atoms_for_names = atoms
+    counter = dict()
+    atom_names = list()
+    for atom in atoms_for_names:
+        counter.setdefault(atom, 1)
+        atom_names.append(f"{atom}{counter[atom]}")
+        counter[atom] += 1
+    atom_names = [an for i, an in enumerate(atom_names) if i not in del_atoms]
+
+    coords3d = coords.reshape(-1, 3) / ANG2BOHR
+    atoms_xyz = list(zip(atom_names, *coords3d.T))
+    rendered = tpl.render(
+        atomnum=len(atoms_xyz),
+        atoms_xyz=atoms_xyz,
+        resno=resno,
+        res=res,
+        resid=resid,
+        segid=segid,
+    )
+    return rendered
+
+
+def geom_to_crd_str(geom, **kwargs):
+    return atoms_coords_to_crd_str(geom.atoms, geom.cart_coords, **kwargs)

pysisyphus/io/cube.py

@@ -0,0 +1,220 @@
+from dataclasses import dataclass
+from math import ceil
+import tempfile
+from typing import Tuple
+
+import jinja2
+import numpy as np
+from numpy.typing import NDArray
+import pyparsing as pp
+
+from pysisyphus.elem_data import ATOMIC_NUMBERS, INV_ATOMIC_NUMBERS
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import file_or_str
+from pysisyphus.wrapper.jmol import view_cdd_cube
+
+
+def get_grid(coords3d, num=10, offset=3.0):
+    minx, miny, minz = coords3d.min(axis=0) - offset
+    maxx, maxy, maxz = coords3d.max(axis=0) + offset
+    X, Y, Z = np.mgrid[
+        minx : maxx : num * 1j,
+        miny : maxy : num * 1j,
+        minz : maxz : num * 1j,
+    ]
+    xyz = np.stack((X.flatten(), Y.flatten(), Z.flatten()), axis=1)
+    spacing = np.array((maxx - minx, maxy - miny, maxz - minz)) / (num - 1)
+    return xyz, spacing, (num, num, num)
+
+
+def get_grid_with_spacing(coords3d, spacing=0.30, margin=3.0):
+    minx, miny, minz = coords3d.min(axis=0) - margin
+    maxx, maxy, maxz = coords3d.max(axis=0) + margin
+    dx = maxx - minx
+    dy = maxy - miny
+    dz = maxz - minz
+
+    nx = ceil(dx / spacing)
+    ny = ceil(dy / spacing)
+    nz = ceil(dz / spacing)
+    X, Y, Z = np.mgrid[
+        minx : maxx : nx * 1j,
+        miny : maxy : ny * 1j,
+        minz : maxz : nz * 1j,
+    ]
+    xyz = np.stack((X.flatten(), Y.flatten(), Z.flatten()), axis=1)
+    act_spacing = np.array(
+        ((maxx - minx) / (nx - 1), (maxy - miny) / (ny - 1), (maxz - minz) / (nz - 1))
+    )
+    return xyz, act_spacing, (nx, ny, nz)
+
+
+@dataclass
+class Cube:
+    atoms: Tuple
+    coords3d: NDArray
+    origin: NDArray
+    npoints: NDArray
+    axes: NDArray
+    vol_data: NDArray
+    comment1: str = "Generated by pysisyphus"
+    comment2: str = ""
+
+    @staticmethod
+    def from_file(fn):
+        return parse_cube(fn)
+
+    @staticmethod
+    def from_wf_and_grid(wf, vol_data, origin, axes):
+        assert vol_data.ndim == 3, "Volume data 'vol_data' must be 3d!"
+        return Cube(
+            atoms=wf.atoms,
+            coords3d=wf.coords3d,
+            origin=origin,
+            npoints=vol_data.shape,
+            axes=axes,
+            vol_data=vol_data,
+        )
+
+    def to_str(self):
+        return cube_to_str(
+            self.atoms,
+            self.coords3d,
+            self.vol_data.reshape(*self.npoints),
+            self.origin,
+            self.axes,
+        )
+
+    def write(self, fn):
+        with open(fn, "w") as handle:
+            handle.write(self.to_str())
+
+    def view_cdd(self, **kwargs):
+        with tempfile.NamedTemporaryFile("w", suffix=".cube") as tmp_cube:
+            self.write(tmp_cube.name)
+            view_cdd_cube(tmp_cube.name, **kwargs)
+
+
+@file_or_str(".cube", ".cub")
+def parse_cube(text):
+    int_ = pp.common.integer
+    sci_real = pp.common.sci_real
+
+    def get_line_word(*args):
+        return pp.Word(*args).setWhitespaceChars("\n")
+
+    comment = get_line_word(pp.printables + " \t")
+    cart_vec = pp.Group(sci_real + sci_real + sci_real)
+    axis = pp.Group(int_ + cart_vec)
+    atom_line = pp.Group(
+        int_.set_results_name("atomic_num")
+        + sci_real.set_results_name("charge")
+        + cart_vec.set_results_name("coords")
+        + pp.LineEnd()
+    )
+
+    parser = (
+        comment.set_results_name("comment1")
+        + comment.set_results_name("comment2")
+        + int_.set_results_name("atom_num")
+        + cart_vec.set_results_name("origin")
+        + axis.set_results_name("axis1")
+        + axis.set_results_name("axis2")
+        + axis.set_results_name("axis3")
+        + pp.ZeroOrMore(atom_line).set_results_name("atom_lines")
+        + pp.ZeroOrMore(sci_real).set_results_name("vol_data")
+    )
+    res = parser.parseString(text).as_dict()
+
+    comment1 = " ".join(res["comment1"])
+    comment2 = " ".join(res["comment2"])
+    atom_num = res["atom_num"]
+    atom_lines = res["atom_lines"]
+    assert atom_num == len(atom_lines)
+
+    atom_nums, coords3d = zip(*[(al["atomic_num"], al["coords"]) for al in atom_lines])
+    atoms = [INV_ATOMIC_NUMBERS[an].capitalize() for an in atom_nums]
+    coords3d = np.array(coords3d)
+    origin = np.array(res["origin"])
+
+    npoints, axes = zip(*[res[k] for k in ("axis1", "axis2", "axis3")])
+    npoints = np.array(npoints, int)
+    assert all(npoints > 0), "Only Bohr are supported!"
+    axes = np.array(axes, float)
+    vol_data = np.array(res["vol_data"]).reshape(*npoints)
+
+    cube = Cube(
+        atoms=atoms,
+        coords3d=coords3d,
+        origin=origin,
+        npoints=npoints,
+        axes=axes,
+        vol_data=vol_data,
+        comment1=comment1,
+        comment2=comment2,
+    )
+
+    return cube
+
+
+@file_or_str(".cube", ".cub")
+def geom_from_cube(text, **kwargs):
+    cube = parse_cube(text)
+    geom = Geometry(cube.atoms, cube.coords3d, **kwargs)
+    return geom
+
+
+CUBE_TPL = """{{ comment1 }}
+{{ comment2 }}
+{{ ifmt(atom_nums|length) }} {{ fmt(origin_x) }} {{ fmt(origin_y) }} {{ fmt(origin_z) }}
+{% for n_points, (x, y, z) in zip(axes_npoints, axes) %}
+{{ ifmt(n_points) }} {{ fmt(x) }} {{ fmt(y) }} {{ fmt(z) }}
+{% endfor %}
+{% for atomic_num, (x, y, z) in zip(atom_nums, coords3d) %}
+{{ ifmt(atomic_num) }} 0.000000 {{ fmt(x) }} {{ fmt(y) }} {{ fmt(z) }}
+{% endfor %}
+{{ grid_str }}
+"""
+
+
+def cube_to_str(atoms, coords3d, vol_data, origin, axes):
+    assert vol_data.ndim == 3
+    env = jinja2.Environment(trim_blocks=True, lstrip_blocks=True)
+    env.globals.update(zip=zip)
+    tpl = env.from_string(CUBE_TPL)
+
+    org_x, org_y, org_z = origin
+    axes_npoints = vol_data.shape
+    atom_nums = [ATOMIC_NUMBERS[atom.lower()] for atom in atoms]
+    nx, ny, nz = axes_npoints
+    grid_str = ""
+    for x in range(nx):
+        for y in range(ny):
+            for z in range(nz):
+                d = "{: >12.6e} ".format(vol_data[x, y, z])
+                grid_str += d
+                if z % 6 == 5:
+                    grid_str += "\n"
+            grid_str += "\n"
+
+    def fmt(num):
+        return f"{num: >12.6f}"
+
+    def ifmt(num):
+        return f"{num: >5d}"
+
+    rendered = tpl.render(
+        fmt=fmt,
+        ifmt=ifmt,
+        origin_x=org_x,
+        origin_y=org_y,
+        origin_z=org_z,
+        comment1="Generated by pysisyphus",
+        comment2=f"Total of {nx*ny*nz} points",
+        axes=axes,
+        axes_npoints=axes_npoints,
+        atom_nums=atom_nums,
+        coords3d=coords3d,
+        grid_str=grid_str,
+    )
+    return rendered

pysisyphus/io/fchk.py

@@ -0,0 +1,184 @@
+import re
+
+import numpy as np
+
+from pysisyphus.elem_data import INV_ATOMIC_NUMBERS
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import file_or_str
+from pysisyphus.wavefunction.shells import Shell, FCHKShells
+from pysisyphus.wavefunction import Wavefunction
+from pysisyphus.wavefunction.helpers import BFType
+
+
+@file_or_str(".fchk")
+def parse_fchk(text):
+    header_re = r"\s+".join(
+        (
+            r"(?P<keyword>[\w\-\s/\(\)=\*]+)",
+            r"(?P<kind>[HIRCL])",
+            r"(?P<length>N=)?",
+            r"(?P<num>[\d\-\+\.E]+)",
+        )
+    )
+    header_re = re.compile(header_re)
+
+    cur_keyword = None
+    lines = text.strip().split("\n")
+    # title1, titl2 are not used
+    _, _, *lines = lines
+
+    conv_funcs = {
+        "R": lambda line: [float(_) for _ in line.split()],  # real
+        "I": lambda line: [int(_) for _ in line.split()],  # integer
+        "C": lambda line: [line],  # character
+        "H": lambda line: [bool(_) for _ in line.split()],  # logical
+        "L": lambda line: [bool(_) for _ in line.split()],  # logical
+    }
+
+    data = {}
+    for line in lines:
+        line = line.strip()
+        if mobj := header_re.match(line):
+            cur_keyword = mobj.group("keyword").strip()
+            conv_func = conv_funcs[mobj.group("kind")]
+            if mobj.group("length"):
+                data[cur_keyword] = list()
+            else:
+                data[cur_keyword] = conv_func(mobj.group("num"))[0]
+            continue
+        data[cur_keyword].extend(conv_func(line))
+    return data
+
+
+@file_or_str(".fchk")
+def shells_from_fchk(text):
+    data = parse_fchk(text)
+
+    def to_arr(key, dtype):
+        return np.array(data[key], dtype=dtype)
+
+    """
+    These coordinates may not coincide with the basis functions?! Nonetheless,
+    basis functions are usually atom-centered. But just to be sure we use the
+    entries from 'Coordinates of each shell'.
+    # coords3d = to_arr("Current cartesian coordinates", float).reshape(-1, 3)
+    """
+    atomic_nums = to_arr("Atomic numbers", int)
+
+    centers = to_arr("Shell to atom map", int) - 1
+    ang_moms = to_arr("Shell types", int)
+    prims_per_shell = to_arr("Number of primitives per shell", int)
+
+    prim_exponents = to_arr("Primitive exponents", float)
+    prim_coeffs = to_arr("Contraction coefficients", float)
+    try:
+        p_prim_coeffs = to_arr("P(S=P) Contraction coefficients", float)
+    except KeyError:
+        p_prim_coeffs = np.zeros_like(prim_coeffs)
+    shell_coords3d = to_arr("Coordinates of each shell", float).reshape(-1, 3)
+
+    _shells = list()
+    prim_ind = 0
+    for i, (center_ind, L, prim_num) in enumerate(
+        zip(centers, ang_moms, prims_per_shell)
+    ):
+        # 0: s, 1: p, -1: sp, 2: 6d, -2: 5d, 3: 10f, -3: 7f, etc...
+        assert L <= 1, "Only spherical basis functions are supported."
+        # Convert SP-shell with L = -1 to a s-shell with L = 0.
+        # The p-shell is also created later.
+        L, is_sp = (0, True) if (L == -1) else (L, False)
+
+        coeffs = list()
+        p_coeffs = list()
+        exps = list()
+        for _ in range(prim_num):
+            coeffs.append(prim_coeffs[prim_ind])
+            p_coeffs.append(p_prim_coeffs[prim_ind])
+            exps.append(prim_exponents[prim_ind])
+            prim_ind += 1
+        atomic_num = atomic_nums[center_ind]
+        center = shell_coords3d[i]
+
+        def append_shell(coeffs, L=L):
+            shell = Shell(
+                L=abs(L),
+                atomic_num=atomic_num,
+                center=center,
+                coeffs=coeffs,
+                exps=exps,
+                center_ind=center_ind,
+            )
+            _shells.append(shell)
+
+        # Create shell with actual L.
+        # If we deal with a SP shell we also create the P shell below.
+        append_shell(coeffs)
+        if is_sp:
+            # Explicitly create p-shell
+            append_shell(p_coeffs, L=1)
+    shells = FCHKShells(_shells)
+    return shells
+
+
+def atoms_from_data(data):
+    atomic_numbers = data["Atomic numbers"]
+    atoms = tuple([INV_ATOMIC_NUMBERS[Z] for Z in atomic_numbers])
+    return atoms
+
+
+@file_or_str(".fchk")
+def wavefunction_from_fchk(text):
+    data = parse_fchk(text)
+
+    charge = data["Charge"]
+    mult = data["Multiplicity"]
+    atoms = atoms_from_data(data)
+    coords = data["Current cartesian coordinates"]
+    unrestricted = "Beta Orbital Energies" in data
+
+    def mos_for_spin(spin):
+        mo_ens = data[f"{spin} Orbital Energies"]
+        mo_num = len(mo_ens)
+        mo_coeffs = np.array(data[f"{spin} MO coefficients"])
+        mo_coeffs = mo_coeffs.reshape(mo_num, mo_num)
+        return mo_coeffs
+
+    C_a = mos_for_spin("Alpha").T
+    occ_a = data["Number of alpha electrons"]
+    if unrestricted:
+        C_b = mos_for_spin("Beta")
+        occ_b = data["Number of beta electrons"]
+    else:
+        C_b = C_a.copy()
+        occ_b = occ_a
+    C = np.stack((C_a, C_b))
+    shells = shells_from_fchk(text)
+    return Wavefunction(
+        atoms=atoms,
+        coords=coords,
+        charge=charge,
+        mult=mult,
+        unrestricted=unrestricted,
+        occ=(occ_a, occ_b),
+        C=C,
+        bf_type=BFType.PURE_SPHERICAL,
+        shells=shells,
+    )
+
+
+@file_or_str(".fchk")
+def geom_from_fchk(text, **geom_kwargs):
+    data = parse_fchk(text)
+    atoms = atoms_from_data(data)
+
+    try:
+        coords = data["Opt point       1 Geometries"]
+    except KeyError:
+        coords = data["Current cartesian coordinates"]
+    coords = np.array(coords)
+    coords = coords.reshape(-1, 3 * len(atoms))
+
+    geoms = [Geometry(atoms, coords_, **geom_kwargs) for coords_ in coords]
+    if len(geoms) == 1:
+        geoms = geoms[0]
+    return geoms

pysisyphus/io/hdf5.py

@@ -0,0 +1,49 @@
+import h5py
+
+
+def init_h5_group(f, group_name, data_model):
+    """Create group with given name and data model."""
+    group = f.create_group(group_name)
+    # Create (resizable) datasets by using None in maxshape
+    for key, shape in data_model.items():
+        maxshape = (None,) if (len(shape) == 1) else (None, *shape[1:])
+        group.create_dataset(key, shape, maxshape=maxshape)
+
+
+def get_h5_group(fn, group_name, data_model=None, reset=False):
+    """Return (and create if neccesary) group with given name and
+    data model."""
+
+    f = h5py.File(fn, mode="a")
+    # Shortcut
+    if data_model is None:
+        return f[group_name]
+
+    if group_name not in f:
+        init_h5_group(f, group_name, data_model)
+    group = f[group_name]
+
+    # Check compatibility of data_model and group. If they aren't compatible
+    # recreate the group with the proper shapes.
+    try:
+        compatible = [group[key].shape == shape for key, shape in data_model.items()]
+    except KeyError:
+        compatible = (False,)
+    compatible = all(compatible)
+    if (not compatible) or reset:
+        del f[group_name]
+        init_h5_group(f, group_name, data_model)
+    group = f[group_name]
+
+    return group
+
+
+def resize_h5_group(group, max_cycles):
+    """Increase size of first dimension of datasets in the given group."""
+    for key, dataset in group.items():
+        # No need to resize scalar datasets
+        if dataset.shape == ():
+            continue
+        new_shape = list(dataset.shape).copy()
+        new_shape[0] = max_cycles
+        dataset.resize(new_shape)

pysisyphus/io/hessian.py

@@ -0,0 +1,72 @@
+import numpy as np 
+import h5py
+
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import eigval_to_wavenumber
+
+
+def save_hessian(h5_fn, geom, cart_hessian=None, energy=None, mult=None):
+    if cart_hessian is None:
+        cart_hessian = geom.cart_hessian
+
+    if energy is None:
+        energy = geom.energy
+
+    if mult is None:
+        mult = geom.calculator.mult
+
+    if len(geom.atoms) > 1:
+        proj_hessian = geom.eckart_projection(geom.mass_weigh_hessian(cart_hessian))
+    else:
+        proj_hessian = cart_hessian
+    eigvals, _ = np.linalg.eigh(proj_hessian)
+    vibfreqs = eigval_to_wavenumber(eigvals)
+
+    masses = geom.masses
+    atoms = geom.atoms
+    coords3d = geom.coords3d
+
+    # with h5py.File(h5_fn, "w") as handle:
+    #     handle.create_dataset("hessian", data=cart_hessian)
+    #     handle.create_dataset("vibfreqs", data=vibfreqs)
+    #     handle.create_dataset("masses", data=masses)
+    #     handle.create_dataset("coords3d", data=coords3d)
+
+    #     try:
+    #         handle.attrs["atoms"] = [atom.lower() for atom in atoms]
+    #     except OSError:
+    #         handle.create_dataset("atoms", data=[atom.lower() for atom in atoms])
+    #     handle.attrs["energy"] = energy
+    #     handle.attrs["mult"] = mult
+
+
+def save_third_deriv(h5_fn, geom, third_deriv_result, H_mw, H_proj):
+    # with h5py.File(h5_fn, "w") as handle:
+    #     for key, value in third_deriv_result._asdict().items():
+    #         handle.create_dataset(key, data=value)
+
+    #     handle.create_dataset("coords3d", data=geom.coords3d)
+    #     handle.create_dataset("masses", data=geom.masses)
+    #     handle.create_dataset("H_mw", data=H_mw)
+    #     handle.create_dataset("H_proj", data=H_proj)
+    #     try:
+    #         handle.attrs["atoms"] = [atom.lower() for atom in atoms]
+    #     except OSError:
+    #         handle.create_dataset("atoms", data=[atom.lower() for atom in atoms])
+    pass
+
+
+def geom_from_hessian(h5_fn, **geom_kwargs):
+    with h5py.File(h5_fn, "r") as handle:
+        try:
+            atoms = [atom.capitalize() for atom in handle.attrs["atoms"]]
+        except (KeyError, TypeError):
+            atoms = [atom.capitalize() for atom in handle["atoms"][:]]
+        coords3d = handle["coords3d"][:]
+        energy = handle.attrs["energy"]
+        cart_hessian = handle["hessian"][:]
+
+    geom = Geometry(atoms=atoms, coords=coords3d, **geom_kwargs)
+    geom.cart_hessian = cart_hessian
+    geom.energy = energy
+    return geom