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

pysisyphus/io/pdb.py

@@ -0,0 +1,269 @@
+from collections import OrderedDict
+import re
+import struct
+import textwrap
+
+from jinja2 import Template
+import numpy as np
+
+from pysisyphus.elem_data import KNOWN_ATOMS
+from pysisyphus.constants import ANG2BOHR
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import chunks, file_or_str
+from pysisyphus.intcoords.setup_fast import find_bonds
+
+
+AMINOS = (
+    "ALA CYS ASP GLU PHE GLY HIS ILE LYS LEU"
+    "MET ASN PRO GLN ARG SER THR VAL TRP TYR".split()
+)
+
+
+def get_parser(widths):
+    fmt = " ".join("{}{}".format(width, "s") for width in widths)
+    fieldstruct = struct.Struct(fmt)
+    parse = lambda line: tuple(s.decode() for s in fieldstruct.unpack(line.encode()))
+    return parse
+
+
+STRIP_RE = re.compile(r"[\d\s]*")  # To remove numbers and whitespace
+NAME_MAP = {
+    "hb": "H",
+    "he": "H",
+    "hh": "H",
+    "hd": "H",
+    "hg": "H",
+    "so": "Na",
+}
+FULL_NAME = {
+    " sod",
+    " cla",
+    " cal",
+}
+
+
+def parse_atom_name(name):
+    org_name = name
+    assert len(name) == 4
+    name = name.lower()
+    # Cases like " SOD" require special handling. Sticking to the PDB specification (!)
+    # and using only the first two characters would result in S (sulphur).
+    use_full_name = name in FULL_NAME
+    if not use_full_name:
+        name = name[:2]
+    stripped = STRIP_RE.sub("", name).lower()
+    if use_full_name:
+        stripped = stripped[:2]
+    try:
+        mapped = NAME_MAP[stripped]
+    except KeyError:
+        assert stripped in KNOWN_ATOMS, f"Could not parse atom name '{org_name}'"
+        mapped = stripped
+    return mapped.capitalize()
+
+
+@file_or_str(".pdb")
+def parse_pdb(text):
+    atm_lines = list()
+    for line in text.split("\n"):
+        if line.startswith("HETATM") or line.startswith("ATOM"):
+            atm_lines.append(line.strip())
+            continue
+        elif line.startswith("MASTER"):
+            master_line = line.strip()
+
+    """
+        0  Record name
+        1  serial
+        2  name
+        3  altLoc
+        4  ResName
+        5  chainID
+        6  resSeq
+        7  iCode
+        8  x
+        9  y
+        10 z
+        11 occupancy
+        12 tempFactor
+        13 element
+        (14 charge), not parsed
+
+        See https://stackoverflow.com/questions/4914008
+    """
+    atm_widths = (6, 6, 4, 1, 4, 1, 4, 1, 11, 8, 8, 6, 6, 12)
+    atm_parse = get_parser(atm_widths)
+
+    """
+        0  Record name
+        1  Num. of REMARK
+        2  "0"
+        3  Num. of HET
+        4  Num. of HELIX
+        5  Num. of SHEET
+        6  deprecated
+        7  Num. of SITE
+        8  numXform
+        9  Num. of atomic coordinates (HETATM + ATOMS)
+        10 Num. of CONECT
+        11 Num. of SEQRES
+    """
+    master_widths = (6, 9, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5)
+    master_parse = get_parser(master_widths)
+    try:
+        master_fields = master_parse(master_line)
+    except UnboundLocalError:
+        master_fields = None
+
+    atoms = list()
+    coords = list()
+    fragments = {}
+    atom_map = dict()
+
+    for i, line in enumerate(atm_lines):
+        # Charge is not considered right now ...
+        fields = atm_parse(f"{line[:78]: <78}")
+        res_name = fields[4].strip()
+        chain = fields[5].strip()
+        res_seq = int(fields[6])
+        frag = f"{chain}_{res_name}{res_seq}"
+
+        xyz = np.array(fields[8:11], dtype=float)
+        atom = fields[13].strip()
+        if not atom.lower() in KNOWN_ATOMS:
+            name = fields[2]
+            atom = parse_atom_name(name)
+        atoms.append(atom)
+        coords.append(xyz)
+        id_ = int(fields[1])
+        atom_map[id_] = i
+
+        frag = fragments.setdefault(frag, list())
+        frag.append(i)
+
+    # Verification using MASTER record, if present.
+    if master_fields is not None:
+        num_coord = int(master_fields[9])
+        try:
+            assert len(atoms) == num_coord
+        except AssertionError as err:
+            if len(atoms) < 99_999:  # See, e.g., 1HTQ
+                raise err
+
+    coords = np.array(coords, dtype=float) * ANG2BOHR
+    return atoms, coords.flatten(), fragments, atom_map
+
+
+def geom_from_pdb(fn, **kwargs):
+    atoms, coords, fragments, _ = parse_pdb(fn)
+    kwargs["fragments"] = fragments
+    geom = Geometry(atoms, coords.flatten(), **kwargs)
+    return geom
+
+
+def get_conect_lines(atoms, coords):
+    bonds = find_bonds(atoms, coords)
+    # PDB indexing is 1-based
+    bonds += 1
+    # Bring in a suitable order for CONECT entries
+    bonds.sort(axis=1)
+    bonds = sorted(bonds, key=lambda row: row[0])
+    conect = OrderedDict()
+    for from_, to_ in bonds:
+        conect.setdefault(from_, list()).append(to_)
+
+    conect_lines = list()
+    for from_, to_ in conect.items():
+        to_iter = chunks(sorted(to_), 4)
+        for to_chunk in to_iter:
+            conect_lines.append([from_] + to_chunk)
+    return conect_lines
+
+
+def atoms_coords_to_pdb_str(atoms, coords, fragments=None, resname="", conect=True):
+    coords3d = coords.reshape(-1, 3)
+    coords3d_ang = coords3d / ANG2BOHR
+
+    coord_fmt = "{: >8.3f}" * 3
+    #            serial  name  altLoc      chainID     iCode
+    #                                resName     resSeq
+    hetatm_fmt = "HETATM{: >5d} {: >4}{: >1}{: >3}{: >1} {: >4d}{: >1}   "
+    #      xyz
+    hetatm_fmt += coord_fmt
+    #      occupancy tempFactor            atom
+    hetatm_fmt += "{: >6.2f}{: >6.2f}" + 10 * " " + "{: >2s}"
+    ter_fmt = "TER   {: >5d}      {: >3s} {:1s}{: >4d}"
+
+    if fragments is None:
+        fragments = [
+            range(len(atoms)),
+        ]
+
+    # Fixed for now
+    altLoc = ""
+    resName = resname
+    chainID = ""
+    iCode = ""
+    occupancy = 1.0
+    tempFactor = 0.0
+
+    lines = list()
+    serial = 1
+    for resSeq, fragment in enumerate(fragments, 1):
+        for id_ in fragment:
+            name = atoms[id_]
+            xyz = coords3d_ang[id_]
+            line = hetatm_fmt.format(
+                serial,
+                name,
+                altLoc,
+                resName,
+                chainID,
+                resSeq,
+                iCode,
+                *xyz,
+                occupancy,
+                tempFactor,
+                name,
+            )
+            lines.append(line)
+            serial += 1
+        lines.append(ter_fmt.format(serial, resName, chainID, resSeq))
+
+    def fmt_conect_line(conect_line):
+        return "CONECT" + ("{: >5d}" * len(conect_line)).format(*conect_line)
+
+    if conect:
+        conect_lines = get_conect_lines(atoms, coords)
+    else:
+        conect_lines = []
+    conect_str = "\n".join([fmt_conect_line(cl) for cl in conect_lines])
+
+    pdb_tpl = Template(
+        textwrap.dedent(
+            """\
+    REMARK 1 Created by pysisyphus
+    {%+ for line in lines -%}
+    {{ line }}
+    {%+ endfor -%}
+    {{- conect_str }}
+    END"""
+        )
+    )
+
+    pdb_str = pdb_tpl.render(lines=lines, conect_str=conect_str)
+    return pdb_str
+
+
+def geom_to_pdb_str(geom, detect_fragments=False, **kwargs):
+    fragments = None
+    if detect_fragments:
+        try:
+            fragments = geom.internal.fragments
+        except AttributeError:
+            geom_ = geom.copy(coord_type="redund")
+            fragments = geom_.internal.fragments
+    pdb_str = atoms_coords_to_pdb_str(
+        geom.atoms, geom.coords3d, fragments=fragments, **kwargs
+    )
+    return pdb_str

pysisyphus/io/psf.py

@@ -0,0 +1,79 @@
+import pyparsing as pp
+
+from pysisyphus.helpers_pure import file_or_str
+
+
+@file_or_str(".psf")
+def parse_psf(text):
+    int_ = pp.common.integer
+    real = pp.common.sci_real
+    psf_style_comment = pp.Regex(r"\*.*")
+
+    psf = pp.CaselessLiteral("PSF")
+    token = pp.ZeroOrMore(
+        pp.CaselessLiteral("EXT")
+        | pp.CaselessLiteral("CMAP")
+        | pp.CaselessLiteral("XPLOR")
+        | pp.CaselessLiteral("AUTOG")
+    ).set_results_name("token")
+    ntitle = int_ + pp.CaselessLiteral("!NTITLE")
+    natom = int_.set_results_name("natom") + pp.CaselessLiteral("!NATOM")
+
+    # Atom records
+    atom_name = pp.Word(pp.alphanums)
+    zero = pp.Literal("0")
+    atom_record = pp.Group(
+        int_.set_results_name("id")
+        + pp.Word(pp.alphas).set_results_name("segment")
+        + int_.set_results_name("resid")
+        + pp.Word(pp.alphanums).set_results_name("resname")  # VAL, ALA, TIP3, etc.
+        + atom_name.set_results_name("atom_name")
+        + atom_name.set_results_name("atom_type")
+        + real.set_results_name("charge")
+        + real.set_results_name("mass")
+        + zero
+    )
+
+    new_section = int_ + pp.Literal("!")
+
+    def get_section(lhs, rhs=None):
+        parser = int_.set_results_name("nterm") + pp.CaselessLiteral(lhs)
+
+        if rhs is not None:
+            parser += pp.Literal(":") + pp.CaselessLiteral(rhs)
+        # Negative lookahead prevents matching 'nterm' of the next section
+        parser += pp.Group(pp.ZeroOrMore(~new_section + int_)).set_results_name("inds")
+        parser = pp.Group(parser)
+        return parser
+
+    nbond = get_section("!NBOND", "bonds").set_results_name("nbond")
+    ntheta = get_section("!NTHETA", "angles").set_results_name("ntheta")
+    nphi = get_section("!NPHI", "dihedrals").set_results_name("nphi")
+    nimphi = get_section("!NIMPHI", "impropers").set_results_name("nimphi")
+    ndon = get_section("!NDON", "donors").set_results_name("ndon")
+    nacc = get_section("!NACC", "acceptors").set_results_name("nacc")
+    nnb = get_section("!NNB").set_results_name("nnb")
+    ncrterm = get_section("!NCRTERM", "cross-terms").set_results_name("ncrterm")
+
+    parser = (
+        psf
+        + token
+        + ntitle
+        + natom
+        + pp.OneOrMore(atom_record).set_results_name("atoms")
+        + nbond  # bonds
+        + ntheta  # bends
+        + nphi  # dihedrals
+        + nimphi  # impropers
+        + ndon  # donors
+        + nacc  # acceptors
+        + nnb  # Non bonding something?
+        # 0 0 !NGRP NST2
+        # 0 0 !NUMLPH NUMLPH
+        # + ncrterm
+    )
+    parser.ignore(psf_style_comment)
+
+    result = parser.parseString(text)
+    as_dict = result.asDict()
+    return as_dict

pysisyphus/io/pubchem.py

@@ -0,0 +1,31 @@
+import json
+import urllib.request
+from urllib.error import HTTPError
+
+from pysisyphus.io.sdf import geom_from_sdf
+
+
+def cid_from_name(name):
+    url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/name/{name}/property/IUPACName/JSON"
+    with urllib.request.urlopen(url) as handle:
+        text = handle.read()
+    json_ = json.loads(text)
+    cid = json_["PropertyTable"]["Properties"][0]["CID"]
+    return cid
+
+
+def sdf_from_cid(cid):
+    url = f"https://pubchem.ncbi.nlm.nih.gov/rest/pug/compound/cid/{cid}/SDF?record_type=3d"
+    with urllib.request.urlopen(url) as handle:
+        sdf = handle.read().decode("utf-8")
+    return sdf
+
+
+def geom_from_pubchem_name(name, **kwargs):
+    try:
+        cid = cid_from_name(name)
+        sdf = sdf_from_cid(cid)
+        geom = geom_from_sdf(sdf, **kwargs)
+    except HTTPError:
+        geom = None
+    return geom

pysisyphus/io/qcschema.py

@@ -0,0 +1,34 @@
+import functools
+import json
+from pathlib import Path
+from typing import Dict
+
+import numpy as np
+
+from pysisyphus.Geometry import Geometry
+
+
+@functools.singledispatch
+def geom_from_qcschema(qcschema: Dict, **geom_kwargs):
+    mol = qcschema["molecule"]
+    coords = np.array((mol["geometry"]))
+    atoms = mol["symbols"]
+    return Geometry(atoms, coords, **geom_kwargs)
+
+
+# These definitions may seem a bit wild, but I was not able to make
+# singledispatch work with 'file_or_str'.
+
+
+@geom_from_qcschema.register
+def _(text: str, **geom_kwargs):
+    if (p := Path(text).exists()):
+        with open(p, "r") as handle:
+            text = handle.read()
+    qcschema = json.loads(text)
+    return geom_from_qcschema(qcschema, **geom_kwargs)
+
+
+@geom_from_qcschema.register
+def _(path: Path, **geom_kwargs):
+    return geom_from_qcschema(str(path), **geom_kwargs)

pysisyphus/io/sdf.py

@@ -0,0 +1,29 @@
+import numpy as np
+
+from pysisyphus.constants import BOHR2ANG
+from pysisyphus.Geometry import Geometry
+
+
+def parse_coord_line(line):
+    x, y, z, atom, *_ = line.strip().split()
+    return (x, y, z), atom
+
+
+def parse_sdf(text):
+    lines = text.split("\n")
+    # title, program, comment = lines[:3]
+    count = lines[3]
+    atoms, *_ = count.split()
+    atoms = int(atoms)
+
+    coord_lines = lines[4:4+int(atoms)]
+    coords, atoms = zip(*[parse_coord_line(cl) for cl in coord_lines])
+    coords = np.array(coords, dtype=float)
+    coords /= BOHR2ANG
+    return atoms, coords
+
+
+def geom_from_sdf(text, **kwargs):
+    atoms, coords = parse_sdf(text)
+    geom = Geometry(atoms, coords, **kwargs)
+    return geom

pysisyphus/io/xyz.py

@@ -0,0 +1,61 @@
+import numpy as np
+
+from pysisyphus.constants import ANG2BOHR
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import file_or_str
+from pysisyphus.xyzloader import split_xyz_str
+
+
+@file_or_str(".xyz", "_trj.xyz")
+def parse_xyz(xyz_str, with_comment=False):
+    lines = iter(xyz_str.strip().split("\n"))
+    atoms_coords = list()
+    comments = list()
+    for line in lines:
+        line = line.strip()
+        # Skip emtpy lines
+        if not line:
+            continue
+        # A new xyz block starts with an integer
+        try:
+            expect = int(line)
+        # We only skip empty lines, otherwise we leave the loop.
+        except ValueError:
+            break
+        atoms = list()
+        coords3d = np.zeros((expect, 3), dtype=float)
+        # Advance iterator over xyz definition
+        comment = next(lines)
+        comments.append(comment if with_comment else None)
+        for i in range(expect):
+            atom, *xyz_ = next(lines).split()
+            atoms.append(atom)
+            coords3d[i] = xyz_
+        assert i == expect - 1
+        assert len(atoms) == expect
+        coords3d *= ANG2BOHR
+        atoms_coords.append((atoms, coords3d.flatten()))
+    return atoms_coords, comments
+
+
+
+def geoms_from_xyz(fn, **kwargs):
+    atoms_coords, comments = parse_xyz(fn, with_comment=True)
+    geoms = [
+        Geometry(atoms, coords.flatten(), comment=comment, **kwargs)
+        for (atoms, coords), comment in zip(atoms_coords, comments)
+    ]
+    return geoms
+
+
+def geom_from_xyz(fn, **kwargs):
+    return geoms_from_xyz(fn, **kwargs)[0]
+
+
+def geoms_from_inline_xyz(inline_xyz, **kwargs):
+    atoms_coords = split_xyz_str(inline_xyz)
+    # We excpect the coordinates to be given in Angstrom
+    geoms = [
+        Geometry(atoms, coords * ANG2BOHR, **kwargs) for atoms, coords in atoms_coords
+    ]
+    return geoms

pysisyphus/io/zmat.py

@@ -0,0 +1,175 @@
+from collections import namedtuple
+
+import numpy as np
+
+from pysisyphus.constants import ANG2BOHR as ANG2BOHR
+from pysisyphus.Geometry import Geometry
+
+
+ZLine = namedtuple(
+    "ZLine", "atom rind r aind a dind d", defaults=(None, None, None, None, None, None)
+)
+
+
+def geom_from_zmat(
+    zmat,
+    atoms=None,
+    coords3d=None,
+    geom=None,
+    start_at=None,
+    drop_dummy=True,
+    **geom_kwargs
+):
+    """Adapted from https://github.com/robashaw/geomConvert by Robert Shaw."""
+
+    if isinstance(zmat, str):
+        zmat = zmat_from_str(zmat)
+
+    zmat_atoms = [zline.atom for zline in zmat]
+    # Extend supplied geometry by zmat
+    if geom is not None:
+        atoms = geom.atoms
+        coords3d = geom.coords3d
+        start_at = len(geom.atoms)
+
+    if atoms is not None:
+        atoms = list(atoms) + zmat_atoms
+    else:
+        atoms = zmat_atoms
+
+    # Grow coordindate array and assign old coordinates
+    if coords3d is not None:
+        _coords3d = np.zeros((len(coords3d) + len(zmat), 3))
+        _coords3d[: len(coords3d)] = coords3d
+        coords3d = _coords3d
+    else:
+        coords3d = np.zeros((len(zmat), 3), dtype=float)
+
+    if atoms or coords3d:
+        assert len(coords3d) == len(atoms)
+
+    if start_at is None:
+        start_at = 0
+
+    for i, zline in enumerate(zmat, start_at):
+        assert all(
+            [
+                (ind is None) or (ind >= 0)
+                for ind in (zline.rind, zline.aind, zline.dind)
+            ]
+        ), "Found invalid atom index. Atom indices start with 1, not 0!"
+
+        r = zline.r
+        # First atom is placed at the origin
+        if i == 0:
+            continue
+        # Bond along x-axis
+        elif i == 1:
+            coords3d[i, 0] = r
+        # Angle in xy-plane from polar coordinates
+        elif i == 2:
+            r"""
+            M       P <- add
+             \     /
+              u   v
+               \ /
+                O
+            """
+            theta = np.deg2rad(zline.a)
+            # Center
+            O = coords3d[zline.rind]
+            # Bond, pointing away from O to M
+            u = coords3d[zline.aind] - O
+            # Direction of u along x axis (left/right)
+            sign = np.sign(u[0])
+            # Polar coordinates
+            x = r * np.cos(theta)
+            y = r * np.sin(theta)
+            # Translate from center with correct orientation
+            coords3d[i] = O + (sign * x, sign * y, 0.0)
+        # Dihedral in xyz-space from spherical coordinates
+        else:
+            theta, phi = np.deg2rad((zline.a, zline.d))
+
+            sin_theta = np.sin(theta)
+            cos_theta = np.cos(theta)
+            sin_phi = np.sin(phi)
+            cos_phi = np.cos(phi)
+
+            x = r * cos_theta
+            y = r * sin_theta * cos_phi
+            z = r * sin_theta * sin_phi
+
+            r"""
+            M <- add      N
+             \           /
+              u         v
+               \       /
+                O--w--P
+            """
+
+            O = coords3d[zline.rind]
+            P = coords3d[zline.aind]
+            N = coords3d[zline.dind]
+
+            # Local axis system
+            v_ = P - N
+            v = v_ / np.linalg.norm(v_)
+            w_ = O - P
+            w = w_ / np.linalg.norm(w_)
+            a = np.cross(v, w)
+            a /= np.linalg.norm(a)
+            b = np.cross(a, w)
+            b /= np.linalg.norm(b)
+            coords3d[i] = O - w * x + b * y + a * z
+
+    if drop_dummy:
+        atoms_ = list()
+        coords3d_ = list()
+        for atom, xyz in zip(atoms, coords3d):
+            if atom.lower() == "x":
+                continue
+            atoms_.append(atom)
+            coords3d_.append(xyz)
+        atoms = atoms_
+        coords3d = np.array(coords3d_)
+
+    geom = Geometry(atoms, coords3d, **geom_kwargs)
+    return geom
+
+
+def geom_from_zmat_str(text, coord_type="cart", coord_kwargs=None):
+    zmat = zmat_from_str(text)
+    return geom_from_zmat(zmat, coord_type=coord_type, coord_kwargs=coord_kwargs)
+
+
+def zmat_from_str(text):
+    def dec(str_):
+        return int(str_) - 1
+
+    def to_bohr(str_):
+        return float(str_) * ANG2BOHR
+
+    def convert(items):
+        funcs = (str, dec, to_bohr, dec, float, dec, float)
+        return [f(item) for f, item in zip(funcs, items)]
+
+    zmat = list()
+    for line in text.strip().split("\n"):
+        line = line.strip()
+        if line.startswith("#"):
+            continue
+        zmat.append(ZLine(*convert(line.split())))
+    return zmat
+
+
+def zmat_from_fn(fn):
+    with open(fn) as handle:
+        text = handle.read()
+    return zmat_from_str(text)
+
+
+def geom_from_zmat_fn(fn, **geom_kwargs):
+    zmat = zmat_from_fn(fn)
+    geom = geom_from_zmat(zmat, **geom_kwargs)
+    return geom