stjames 0.0.40__py3-none-any.whl → 0.0.41__py3-none-any.whl

stjames/__init__.py

@@ -1,6 +1,8 @@
 # ruff: noqa: I001
 
 from .calculation import *
+from .atom import *
+from .periodic_cell import *
 from .molecule import *
 from .workflows import *
 
@@ -21,3 +23,4 @@ from .mode import *
 from .status import *
 from .constraint import *
 from .message import *
+from .types import *

stjames/atom.py

@@ -0,0 +1,66 @@
+from typing import Self, Sequence
+
+from pydantic import NonNegativeInt
+
+from .base import Base
+from .data import ELEMENT_SYMBOL, SYMBOL_ELEMENT
+from .types import Vector3D
+
+
+class Atom(Base):
+    atomic_number: NonNegativeInt
+    position: Vector3D  # in Å
+
+    def __repr__(self) -> str:
+        """
+        >>> Atom(atomic_number=2, position=[0, 1, 2])
+        Atom(2, [0.00000, 1.00000, 2.00000])
+        """
+        x, y, z = self.position
+        return f"Atom({self.atomic_number}, [{x:.5f}, {y:.5f}, {z:.5f}])"
+
+    def __str__(self) -> str:
+        """
+        >>> str(Atom(atomic_number=2, position=[0, 1, 2]))
+        'He    0.0000000000    1.0000000000    2.0000000000'
+        """
+        x, y, z = self.position
+        return f"{self.atomic_symbol:2} {x:15.10f} {y:15.10f} {z:15.10f}"
+
+    @property
+    def atomic_symbol(self) -> str:
+        """
+        >>> Atom(atomic_number=2, position=[0, 1, 2]).atomic_symbol
+        'He'
+        """
+        return ELEMENT_SYMBOL[self.atomic_number]
+
+    def edited(self, atomic_number: int | None = None, position: Sequence[float] | None = None) -> Self:
+        """
+        Create a new Atom with the specified changes.
+
+        >>> a = Atom(atomic_number=2, position=[0, 1, 2])
+        >>> a2 = a.edited(3)
+        >>> a is a2
+        False
+        >>> a2
+        Atom(3, [0.00000, 1.00000, 2.00000])
+        """
+        if atomic_number is None:
+            atomic_number = self.atomic_number
+        if position is None:
+            position = list(self.position)
+
+        return self.__class__(atomic_number=atomic_number, position=position)
+
+    @classmethod
+    def from_xyz(cls: type[Self], xyz_line: str) -> Self:
+        """
+        >>> Atom.from_xyz("H 0 0 0")
+        Atom(1, [0.00000, 0.00000, 0.00000])
+        """
+        name, *xyz = xyz_line.split()
+        symbol = int(name) if name.isdigit() else SYMBOL_ELEMENT[name]
+        if not len(xyz) == 3:
+            raise ValueError("XYZ file should have 3 coordinates per atom")
+        return cls(atomic_number=symbol, position=xyz)

stjames/base.py

@@ -1,42 +1,42 @@
 from enum import Enum
-from typing import Annotated, Hashable, TypeVar
+from typing import Annotated, Any, Hashable, TypeVar
 
 import numpy as np
 import pydantic
 
-T = TypeVar("T")
+_T = TypeVar("_T")
 
 
 class Base(pydantic.BaseModel):
     @pydantic.field_validator("*", mode="before")
     @classmethod
-    def coerce_numpy(cls, val: T) -> T | list:
+    def coerce_numpy(cls, val: _T) -> _T | list[Any]:
         if isinstance(val, np.ndarray):
-            return val.tolist()
-        else:
-            return val
+            return val.tolist()  # type: ignore [no-any-return, unused-ignore]
+
+        return val
 
 
 class LowercaseStrEnum(str, Enum):
     """Enum where hyphens, underscores, and case are ignored."""
 
     @classmethod
-    def _missing_(cls, value: str) -> str | None:  # type: ignore
-        # Type note: technically breaking Liskov, value: object in Enum
+    def _missing_(cls, value: object) -> str | None:
         for member in cls:
-            if member.lower().replace("-", "").replace("_", "") == value.lower().replace("-", "").replace("_", ""):
-                return member
+            if isinstance(value, str):
+                if member.lower().replace("-", "").replace("_", "") == value.lower().replace("-", "").replace("_", ""):
+                    return member
         return None
 
 
 # cf. https://github.com/pydantic/pydantic-core/pull/820#issuecomment-1670475909
-H = TypeVar("H", bound=Hashable)
+_H = TypeVar("_H", bound=Hashable)
 
 
-def _validate_unique_list(v: list[H]) -> list[H]:
+def _validate_unique_list(v: list[_H]) -> list[_H]:
     if len(v) != len(set(v)):
         raise ValueError("this list must be unique, and isn't!")
     return v
 
 
-UniqueList = Annotated[list[H], pydantic.AfterValidator(_validate_unique_list)]
+UniqueList = Annotated[list[_H], pydantic.AfterValidator(_validate_unique_list)]

stjames/basis_set.py

@@ -1,10 +1,6 @@
-import pydantic
-from pydantic import PositiveFloat, PositiveInt
+from typing import Optional, Self
 
-try:
-    from typing import Optional, Self
-except ImportError:
-    from typing_extensions import Optional, Self
+from pydantic import PositiveFloat, PositiveInt, model_validator
 
 from .base import Base
 
@@ -14,7 +10,7 @@ class BasisSetOverride(Base):
     atomic_numbers: Optional[list[PositiveInt]] = None
     atoms: Optional[list[PositiveInt]] = None  # 1-indexed
 
-    @pydantic.model_validator(mode="after")
+    @model_validator(mode="after")
     def check_override(self) -> Self:
         # ^ is xor
         assert (self.atomic_numbers is not None) ^ (self.atoms is not None), "Exactly one of ``atomic_numbers`` or ``atoms`` must be specified!"

stjames/calculation.py

@@ -5,6 +5,7 @@ from .message import Message
 from .molecule import Molecule
 from .settings import Settings
 from .status import Status
+from .types import UUID
 
 
 class StJamesVersion(LowercaseStrEnum):
@@ -29,6 +30,7 @@ class Calculation(Base):
     messages: list[Message] = []
 
     engine: Optional[str] = "peregrine"
+    uuids: list[UUID | None] | None = None
 
     # not to be changed by end users, diff. versions will have diff. defaults
     json_format: str = StJamesVersion.V0

stjames/constraint.py

@@ -1,3 +1,5 @@
+from pydantic import PositiveFloat, PositiveInt
+
 from .base import Base, LowercaseStrEnum
 
 
@@ -10,7 +12,25 @@ class ConstraintType(LowercaseStrEnum):
 
 
 class Constraint(Base):
-    """Represents a single constraint."""
+    """Represents a single (absolute) constraint."""
 
     constraint_type: ConstraintType
-    atoms: list[int]  # 1-indexed
+    atoms: list[PositiveInt]  # 1-indexed
+
+
+class PairwiseHarmonicConstraint(Base):
+    """
+    Represents a harmonic constraint, with a characteristic spring constant.
+    """
+
+    atoms: tuple[PositiveInt, PositiveInt]  # 1-indexed
+    spring_constant: PositiveFloat  # kcal/mol / Å**2
+
+
+class SphericalHarmonicConstraint(Base):
+    """
+    Represents a spherical harmonic constraint to keep a system near the origin.
+    """
+
+    confining_radius: PositiveFloat
+    confining_force_constant: PositiveFloat = 10  # kcal/mol / Å**2

stjames/data/__init__.py

@@ -0,0 +1 @@
+from .elements import *

stjames/data/elements.py

@@ -0,0 +1,27 @@
+"""Read elemental data from files."""
+
+import json
+from collections import namedtuple
+from importlib import resources
+
+data_dir = resources.files("stjames").joinpath("data")
+
+with data_dir.joinpath("symbol_element.json").open() as f:
+    SYMBOL_ELEMENT: dict[str, int] = json.loads(f.read())
+
+ELEMENT_SYMBOL = {v: k for k, v in SYMBOL_ELEMENT.items()}
+
+Isotope = namedtuple("Isotope", ["relative_atomic_mass", "isotopic_composition", "standard_atomic_weight"])
+with data_dir.joinpath("nist_isotopes.json").open() as f:
+    d = json.loads(f.read())
+
+    ISOTOPES: dict[int, dict[int, Isotope]] = {
+        int(k): {
+            int(kk): Isotope(*vv)
+            for kk, vv in v.items()  # stay open
+        }
+        for k, v in d.items()
+    }
+
+with data_dir.joinpath("bragg_radii.json").open() as f:
+    BRAGG_RADII: dict[int, float] = json.loads(f.read())

stjames/data/read_nist_isotopes.py

@@ -0,0 +1,116 @@
+"""
+Read the NIST isotopes data file and write it to a JSON file.
+
+NIST Isotopes data from:
+https://physics.nist.gov/cgi-bin/Compositions/stand_alone.pl?ele=&all=all&ascii=ascii2
+"""
+
+import json
+from collections import defaultdict
+from importlib import resources
+from typing import Callable, TypeVar
+
+data_dir = resources.files("stjames").joinpath("data")
+
+_T = TypeVar("_T")
+
+
+def process_line(line: str, fmt: Callable[[str], _T] = str) -> _T:  # type: ignore[assignment]
+    """
+    Process a line from the NIST data file.
+
+    :param line: line to process
+    :param fmt: function to format the value
+    >>> process_line("Atomic Number = 1", int)
+    1
+    """
+    return fmt(line.split("=")[-1].strip())
+
+
+def fmt_float(val: str) -> float:
+    """
+    Format a float from the NIST data file.
+
+    >>> fmt_float(" 1.00784(7)")
+    1.00784
+    """
+    return float(val.strip().split("(")[0])
+
+
+def fmt_maybe_list(val: str) -> float:
+    """
+    Format a float or list of floats from the NIST data file.
+
+    Only the first value is returned.
+
+    >>> fmt_maybe_list("1.00784(7)")
+    1.00784
+    >>> fmt_maybe_list(" [1.00784,1.00811]")
+    1.00784
+    >>> fmt_maybe_list(" [98]")
+    98.0
+    """
+    val = val.strip()
+    if val.startswith("["):
+        val = val[1:-1].split(",")[0]
+    return fmt_float(val)
+
+
+def process_chunk(chunk: str) -> tuple[int, int, tuple[float, float, float]]:
+    r"""
+    Atomic Number, Mass Number, (Relative Atomic Mass, Isotopic Composition, Standard Atomic Weight)
+
+    >>> process_chunk('''\
+    ... Atomic Number = 1
+    ... Atomic Symbol = H
+    ... Mass Number = 1
+    ... Relative Atomic Mass = 1.00784(7)
+    ... Isotopic Composition = 0.999885(70)
+    ... Standard Atomic Weight = [1.00784,1.00811]
+    ... Notes = m
+    ... ''')
+    (1, 1, (1.00784, 0.999885, 1.00784))
+    """
+    lines = chunk.splitlines()
+
+    atomic_number = process_line(lines[0], int)
+    _atomic_symbol = process_line(lines[1], str)
+    mass_number = process_line(lines[2], int)
+    relative_atomic_mass = process_line(lines[3], fmt_float)
+    try:
+        isotopic_composition = process_line(lines[4], fmt_float)
+    except ValueError:
+        isotopic_composition = 0
+    try:
+        standard_atomic_weight = process_line(lines[5], fmt_maybe_list)
+    except ValueError:
+        standard_atomic_weight = relative_atomic_mass
+
+    return atomic_number, mass_number, (relative_atomic_mass, isotopic_composition, standard_atomic_weight)
+
+
+def read_nist_isotopes() -> dict[int, dict[int, tuple[float, float, float]]]:
+    """
+    Read the NIST data file and write it to a JSON file.
+
+    {Atomic Number: {Mass Number, (Relative Atomic Mass, Isotopic Composition, Standard Atomic Weight)}}
+    """
+    with data_dir.joinpath("nist_isotopes.txt").open() as f:
+        next(f), next(f)  # Skip the first two lines
+        nist_isotopes = f.read()
+
+    isotopes: dict[int, dict[int, tuple[float, float, float]]] = defaultdict(dict)
+    for chunk in nist_isotopes.split("\n\n"):
+        atomic_number, mass_number, values = process_chunk(chunk)
+        isotopes[atomic_number][mass_number] = values
+
+    with open("nist_isotopes.json", "w") as f:
+        json.dump(isotopes, f)
+
+    return isotopes
+
+
+if __name__ == "__main__":
+    from pprint import pprint
+
+    pprint(read_nist_isotopes())

stjames/method.py

@@ -29,10 +29,43 @@ class Method(LowercaseStrEnum):
 
     AIMNET2_WB97MD3 = "aimnet2_wb97md3"
 
+    GFN_FF = "gfn_ff"
     GFN0_XTB = "gfn0_xtb"
     GFN1_XTB = "gfn1_xtb"
     GFN2_XTB = "gfn2_xtb"
-    GFN_FF = "gfn_ff"
 
     # this was going to be removed, but Jonathon wrote such a nice basis set test... it's off the front end.
     BP86 = "bp86"
+
+
+MLFF = [
+    Method.AIMNET2_WB97MD3,
+]
+
+XTB_METHODS = [
+    Method.GFN_FF,
+    Method.GFN0_XTB,
+    Method.GFN1_XTB,
+    Method.GFN2_XTB,
+]
+
+COMPOSITE_METHODS = [
+    Method.HF3C,
+    Method.B973C,
+    Method.R2SCAN3C,
+    Method.WB97X3C,
+]
+
+PREPACKAGED_METHODS = [
+    *MLFF,
+    *XTB_METHODS,
+    *COMPOSITE_METHODS,
+]
+
+METHODS_WITH_CORRECTION = [
+    Method.WB97XD3,
+    Method.WB97XV,
+    Method.WB97MV,
+    Method.WB97MD3BJ,
+    Method.DSDBLYPD3BJ,
+]

stjames/molecule.py

@@ -1,24 +1,26 @@
+from pathlib import Path
+from typing import Iterable, Optional, Self
+
 import pydantic
 from pydantic import NonNegativeInt, PositiveInt
 
-try:
-    from typing import Optional, Self
-except ImportError:
-    from typing_extensions import Optional, Self
-
+from .atom import Atom
 from .base import Base
+from .periodic_cell import PeriodicCell
+from .types import Matrix3x3, Vector3D, Vector3DPerAtom
+
+
+class MoleculeReadError(RuntimeError):
+    pass
 
 
 class VibrationalMode(Base):
     frequency: float  # in cm-1
-    reduced_mass: float
-    force_constant: float
-    displacements: list[list[float]]
-
+    reduced_mass: float  # amu
 
-class Atom(Base):
-    atomic_number: NonNegativeInt
-    position: list[float]  # in Å
+    # todo - check units here?
+    force_constant: float
+    displacements: Vector3DPerAtom
 
 
 class Molecule(Base):
@@ -26,6 +28,9 @@ class Molecule(Base):
     multiplicity: PositiveInt
     atoms: list[Atom]
 
+    # for periodic boundary conditions
+    cell: Optional[PeriodicCell] = None
+
     energy: Optional[float] = None  # in Hartree
     scf_iterations: Optional[NonNegativeInt] = None
     scf_completed: Optional[bool] = None
@@ -33,11 +38,14 @@ class Molecule(Base):
 
     homo_lumo_gap: Optional[float] = None  # in eV
 
-    gradient: Optional[list[list[float]]] = None  # Hartree/Bohr
+    gradient: Optional[Vector3DPerAtom] = None  # Hartree/Å
+    stress: Optional[Matrix3x3] = None  # Hartree/Å
+
+    velocities: Optional[Vector3DPerAtom] = None  # Å/fs
 
     mulliken_charges: Optional[list[float]] = None
     mulliken_spin_densities: Optional[list[float]] = None
-    dipole: Optional[list[float]] = None  # in Debye
+    dipole: Optional[Vector3D] = None  # in Debye
 
     vibrational_modes: Optional[list[VibrationalMode]] = None
 
@@ -49,8 +57,18 @@ class Molecule(Base):
     def __len__(self) -> int:
         return len(self.atoms)
 
+    def distance(self, atom1: PositiveInt, atom2: PositiveInt) -> float:
+        r"""
+        Get the distance between atoms.
+
+        >>> mol = Molecule.from_xyz("H 0 1 0\nH 0 0 1")
+        >>> mol.distance(1, 2)
+        1.4142135623730951
+        """
+        return sum((q2 - q1) ** 2 for q1, q2 in zip(self.atoms[atom1 - 1].position, self.atoms[atom2 - 1].position)) ** 0.5  # type: ignore [no-any-return,unused-ignore]
+
     @property
-    def coordinates(self) -> list[list[float]]:
+    def coordinates(self) -> Vector3DPerAtom:
         return [a.position for a in self.atoms]
 
     @property
@@ -92,3 +110,83 @@ class Molecule(Base):
             )
 
         return self
+
+    @classmethod
+    def from_file(cls: type[Self], filename: Path | str, format: str | None = None, charge: int = 0, multiplicity: PositiveInt = 1) -> Self:
+        r"""
+        Read a molecule from a file.
+
+        >>> import tempfile
+        >>> with tempfile.NamedTemporaryFile("w+", suffix=".xyz") as f:
+        ...    _ = f.write("2\nComment\nH 0 0 0\nF 0 0 1")
+        ...    _ = f.seek(0)
+        ...    mol = Molecule.from_file(f.name)
+        >>> print(mol.to_xyz())
+        2
+        <BLANKLINE>
+        H     0.0000000000    0.0000000000    0.0000000000
+        F     0.0000000000    0.0000000000    1.0000000000
+        """
+        filename = Path(filename)
+        if not format:
+            format = filename.suffix[1:]
+
+        with open(filename) as f:
+            match format:
+                case "xyz":
+                    return cls.from_xyz_lines(f.readlines(), charge=charge, multiplicity=multiplicity)
+                case _:
+                    raise ValueError(f"Unsupported {format=}")
+
+    @classmethod
+    def from_xyz(cls: type[Self], xyz: str, charge: int = 0, multiplicity: PositiveInt = 1) -> Self:
+        r"""
+        Generate a Molecule from an XYZ string.
+
+        Note: only supports single molecule inputs.
+
+        >>> len(Molecule.from_xyz("2\nComment\nH 0 0 0\nH 0 0 1"))
+        2
+        """
+        return cls.from_xyz_lines(xyz.strip().splitlines(), charge=charge, multiplicity=multiplicity)
+
+    @classmethod
+    def from_xyz_lines(cls: type[Self], lines: Iterable[str], charge: int = 0, multiplicity: PositiveInt = 1) -> Self:
+        lines = list(lines)
+        if len(lines[0].split()) == 1:
+            natoms = lines[0].strip()
+            if not natoms.isdigit() or (int(lines[0]) != len(lines) - 2):
+                raise MoleculeReadError(f"First line of XYZ file should be the number of atoms, got: {lines[0]} != {len(lines) - 2}")
+            lines = lines[2:]
+
+        try:
+            return cls(atoms=[Atom.from_xyz(line) for line in lines], charge=charge, multiplicity=multiplicity)
+        except Exception as e:
+            raise MoleculeReadError("Error reading molecule from xyz") from e
+
+    def to_xyz(self, comment: str = "", out_file: Path | str | None = None) -> str:
+        r"""
+        Generate an XYZ string.
+
+        >>> mol = Molecule.from_xyz("2\nComment\nH 0 1 2\nF 1 2 3")
+        >>> print(mol.to_xyz(comment="HF"))
+        2
+        HF
+        H     0.0000000000    1.0000000000    2.0000000000
+        F     1.0000000000    2.0000000000    3.0000000000
+        >>> import tempfile
+        >>> with tempfile.TemporaryDirectory() as directory:
+        ...     file = Path(directory) / "mol.xyz"
+        ...     out = mol.to_xyz(comment="HF", out_file=file)
+        ...     with file.open() as f:
+        ...         Molecule.from_xyz(f.read()).to_xyz("HF") == out
+        True
+        """
+        geom = "\n".join(map(str, self.atoms))
+        out = f"{len(self)}\n{comment}\n{geom}"
+
+        if out_file:
+            with Path(out_file).open("w") as f:
+                f.write(out)
+
+        return out

stjames/opt_settings.py

@@ -1,4 +1,6 @@
-from pydantic import Field, PositiveFloat, PositiveInt
+from typing import Sequence
+
+from pydantic import PositiveFloat, PositiveInt
 
 from .base import Base
 from .constraint import Constraint
@@ -8,9 +10,12 @@ class OptimizationSettings(Base):
     max_steps: PositiveInt = 250
     transition_state: bool = False
 
-    # when are we converged?
-    max_gradient_threshold: PositiveFloat = 4.5e-4
-    rms_gradient_threshold: PositiveFloat = 3.0e-4
+    # when are we converged? (Hartree and Hartree/Å)
+    max_gradient_threshold: PositiveFloat = 7e-4
+    rms_gradient_threshold: PositiveFloat = 6e-4
     energy_threshold: PositiveFloat = 1e-6
 
-    constraints: list[Constraint] = Field(default_factory=list)
+    # for periodic systems only
+    optimize_cell: bool = False
+
+    constraints: Sequence[Constraint] = tuple()

stjames/periodic_cell.py

@@ -0,0 +1,34 @@
+from typing import TypeAlias
+
+import numpy as np
+import pydantic
+
+from .base import Base
+from .types import Matrix3x3
+
+Bool3: TypeAlias = tuple[bool, bool, bool]
+
+
+class PeriodicCell(Base):
+    lattice_vectors: Matrix3x3
+    is_periodic: Bool3 = (True, True, True)
+
+    @pydantic.field_validator("lattice_vectors")
+    @classmethod
+    def check_tensor_3D(cls, v: Matrix3x3) -> Matrix3x3:
+        if len(v) != 3 or any(len(row) != 3 for row in v):
+            raise ValueError("Cell tensor must be a 3x3 list of floats")
+
+        return v
+
+    @pydantic.field_validator("is_periodic")
+    @classmethod
+    def check_pbc(cls, v: Bool3) -> Bool3:
+        if not any(v):
+            raise ValueError("For periodic boundary conditions, at least one dimension must be periodic!")
+        return v
+
+    @pydantic.computed_field  # type: ignore[misc, prop-decorator, unused-ignore]
+    @property
+    def volume(self) -> float:
+        return float(np.abs(np.linalg.det(np.array(self.lattice_vectors))))