diff --git a/.bumpversion.cfg b/.bumpversion.cfg index db41815..4068a9d 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,6 +1,8 @@ [bumpversion] -current_version = 0.11.0 +current_version = 0.12.0 commit = True tag = True [bumpversion:file:mudslide/version.py] + +[bumpversion:file:docs/source/conf.py] diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml index 3a5efad..cc4fb0b 100644 --- a/.github/workflows/python-package.yml +++ b/.github/workflows/python-package.yml @@ -10,8 +10,26 @@ on: branches: [ master ] jobs: - build: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e ".[dev]" + - name: Lint with pylint + run: | + pylint mudslide/ --fail-under=9.5 + - name: Type check with mypy + run: | + mypy mudslide/ + test: runs-on: ubuntu-latest strategy: fail-fast: false @@ -19,17 +37,17 @@ jobs: python-version: ["3.9", "3.10", "3.11", "3.12", "3.13", "3.14"] steps: - - uses: actions/checkout@v3 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v3 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install pytest + python -m pip install pytest pytest-cov if [ -f requirements.txt ]; then pip install -r requirements.txt; fi python -m pip install -e . - name: Test with pytest run: | - pytest + pytest --cov=mudslide --cov-report=term-missing diff --git a/.gitignore b/.gitignore index 95dfdb8..a09fdbb 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,8 @@ _build _static _templates +dist/ test/checks .claude .vscode +mudslide.report diff --git a/.pylintrc b/.pylintrc index 8e8f95d..a96a70b 100644 --- a/.pylintrc +++ b/.pylintrc @@ -99,9 +99,7 @@ recursive=no # source root. source-roots= -# When enabled, pylint would attempt to guess common misconfiguration and emit -# user-friendly hints instead of false-positive error messages. -suggestion-mode=yes +# suggestion-mode was removed in pylint 4.x # Allow loading of arbitrary C extensions. Extensions are imported into the # active Python interpreter and may run arbitrary code. @@ -200,11 +198,34 @@ good-names=i, v, V, p, - P + P, + A, + B, + C, + D, + F, + G, + H, + L, + M, + N, + Q, + R, + T, + U, + W # Good variable names regexes, separated by a comma. If names match any regex, # they will always be accepted -good-names-rgxs= +# Scientific/mathematical naming conventions: +# ^[A-Z][a-z0-9]+$ - PascalWord: Hbar, Trace, Reff, Peff, Sij, V12 +# ^[A-Z]+[0-9]+$ - UPPER+digits: W00, H0, TV00, E1 +# ^[A-Z]{2}$ - two-letter uppercase: HI, AA +# ^[A-Z]+_\w+$ - UPPER_word: NAC_matrix, M_inv, V_nuc, T_window +# ^[a-z_]+[A-Z]\w*$ - mixed_Case: dV, delR, ddP, kbT2, compute_delF, avg_KE +# ^[A-Z]\w*_$ - trailing underscore: ElectronicModel_, DiabaticModel_ +# ^[A-Z][a-z]+[A-Z]\w*$ - PascalCamel: LmR, LpR +good-names-rgxs=^[A-Z][a-z0-9]+$,^[A-Z]+[0-9]+$,^[A-Z]{2}$,^[A-Z]+_\w+$,^[a-z_]+[A-Z]\w*$,^[A-Z]\w*_$,^[A-Z][a-z]+[A-Z]\w*$ # Include a hint for the correct naming format with invalid-name. include-naming-hint=no @@ -440,7 +461,9 @@ disable=raw-checker-failed, deprecated-pragma, use-symbolic-message-instead, use-implicit-booleaness-not-comparison-to-string, - use-implicit-booleaness-not-comparison-to-zero + use-implicit-booleaness-not-comparison-to-zero, + too-many-positional-arguments, + too-few-public-methods # Enable the message, report, category or checker with the given id(s). You can # either give multiple identifier separated by comma (,) or put this option diff --git a/README.md b/README.md index 5b8af01..b246e4b 100644 --- a/README.md +++ b/README.md @@ -1,28 +1,38 @@ -# Fewest Switches Surface Hopping [![Build Status](https://github.com/smparker/mudslide/actions/workflows/python-package.yml/badge.svg)](https://github.com/smparker/mudslide/actions/workflows/python-package.yml) [![Documentation Status](https://readthedocs.org/projects/mudslide/badge/?version=latest)](https://mudslide.readthedocs.io/en/latest/?badge=latest) -Python implementation of Tully's Fewest Switches Surface Hopping (FSSH) for model problems including -a propagator and an implementation of Tully's model problems described in Tully, J.C. _J. Chem. Phys._ (1990) **93** 1061. -The current implementation works for diabatic as well as ab initio models, with two or more electronic states, and with one or more -dimensional potentials. +# Mudslide [![Build Status](https://github.com/smparker/mudslide/actions/workflows/python-package.yml/badge.svg)](https://github.com/smparker/mudslide/actions/workflows/python-package.yml) [![Documentation Status](https://readthedocs.org/projects/mudslide/badge/?version=latest)](https://mudslide.readthedocs.io/en/latest/?badge=latest) +A Python library for nonadiabatic molecular dynamics, implementing Fewest Switches Surface Hopping (FSSH) +and related methods. Includes Tully's model problems (Tully, J.C. _J. Chem. Phys._ (1990) **93** 1061) +as well as interfaces to ab initio electronic structure codes. Supports diabatic and adiabatic models +with two or more electronic states and one or more dimensional potentials. ## Contents * `mudslide` package that contains - - implementation of all surface hopping methods + - nonadiabatic and adiabatic dynamics methods - `SurfaceHoppingMD` - Standard FSSH implementation - `Ehrenfest` - Ehrenfest dynamics - - `AugmentedFSSH` - Augmented FSSH implementation + - `AugmentedFSSH` - Augmented FSSH (A-FSSH) implementation - `EvenSamplingTrajectory` - FSSH with even sampling of phase space - - collection of 1D models + - `AdiabaticMD` - Adiabatic (ground state) molecular dynamics + - collection of 1D model potentials - `TullySimpleAvoidedCrossing` - `TullyDualAvoidedCrossing` - `TullyExtendedCouplingReflection` - `SuperExchange` - `SubotnikModelX` - `SubotnikModelS` + - `SubotnikModelW` + - `SubotnikModelZ` - `ShinMetiu` - - some 2D models + - `LinearVibronic` + - 2D models - `Subotnik2D` + - ab initio interfaces + - `TMModel` - Turbomole interface for TDDFT-based NAMD + - `OpenMM` - OpenMM interface for classical MD + - `QMMM` - QM/MM combining Turbomole and OpenMM + - `HarmonicModel` - Harmonic approximation from Hessian data * `mudslide` script that runs simple model trajectories -* `mudslide-surface` script that prints 1D surface and couplings +* `mud` script providing a unified CLI with subcommands +* `mudslide-surface` script that prints 1D surfaces and couplings ## Requirements * numpy @@ -85,10 +95,9 @@ will run 4 scattering simulations with a particle starting in the ground state ( * `total_time` - total simulation length (default: 2 * abs(position/velocity)) * `samples` - number of trajectories to run (default: 2000) * `seed` - random seed for trajectories (defaults however numpy does) -* `propagator` - method used to propagate electronic wavefunction - * "exponential" (default) - apply exponentiated Hamiltonian via diagonalization - * "ode" - scipy's ODE integrator -* `nprocs` - number of processes over which to parallelize trajectories (default: 1) +* `electronic_integration` - method used to propagate electronic wavefunction + * "exp" (default) - apply exponentiated Hamiltonian via diagonalization + * "linear-rk4" - interpolated RK4 integration * `outcome_type` - how to count statistics at the end of a trajectory * "state" (default) - use the state attribute of the simulation only * "populations" - use the diagonals of the density matrix @@ -105,15 +114,17 @@ and should implement ### compute() function The `compute()` function needs to have the following signature: - def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any None, reference: Any = None) -> None + def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, reference: Any = None) -> None The `X` input to the `compute()` function is an array of the positions. All other inputs are ignored for now but will eventually be used to allow the trajectory to enumerate precisely which quantities are desired at each call. At the end of the `compute()` function, the object must store -* `self.hamiltonian` - An `nstates x nstates` array of the Hamiltonian (`nstates` is the number of electronic states) -* `self.force` - An `nstates x ndim` array of the force on each PES (`ndim` is the number of classical degrees of freedom) -* `self.derivative_coupling` - An `nstates x nstates x ndim` array where `self.derivative_coupling[i,j,:]` contains . +* `self._hamiltonian` - An `nstates x nstates` array of the Hamiltonian (`nstates` is the number of electronic states) +* `self._force` - An `nstates x ndim` array of the force on each PES (`ndim` is the number of classical degrees of freedom) +* `self._derivative_coupling` - An `nstates x nstates x ndim` array where `derivative_coupling[i,j,:]` contains +* `self._forces_available` - A boolean array of length `nstates` indicating which forces were computed +* `self._derivative_couplings_available` - An `nstates x nstates` boolean array indicating which couplings were computed See the file `mudslide/turbomole_model.py` for an example of a standalone ab initio model. @@ -143,12 +154,6 @@ For batch runs, one must tell `BatchedTraj` how to decide on new initial conditi and how to decide when a trajectory has finished. The basic requirements for each of those is simple. -The structure of these classes is somewhat strange because of the limitations of -multiprocessing in python. To make use of multiprocessing, every object -must be able to be `pickle`d, meaning that multiprocessing inherits all the -same limitations. As a result, when using multiprocessing, the trajectory generator class must -be fully defined in the default namespace. - ### Generating initial conditions This should be a generator function that accepts a number of samples and returns a dictionary with starting conditions filled in, e.g., diff --git a/docs/source/conf.py b/docs/source/conf.py index 98a5e8b..983359d 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -21,8 +21,9 @@ author = 'Shane M. Parker' # The full version, including alpha/beta/rc tags -version = '0.11.0' -release = '0.11.0' +from mudslide.version import __version__ +version = __version__ +release = __version__ master_doc = 'index' diff --git a/docs/source/developing.rst b/docs/source/developing.rst new file mode 100644 index 0000000..03dcb55 --- /dev/null +++ b/docs/source/developing.rst @@ -0,0 +1,110 @@ +Developer Guide +==================================== + +This page describes how to set up a development environment +and run the checks that are required to pass CI. + +Setting Up a Development Environment +------------------------------------- +Clone the repository and install in editable mode with the +development dependencies: + +.. code-block:: shell + + $ git clone https://github.com/smparker/mudslide.git + $ cd mudslide + $ pip install -e ".[dev]" + +This installs the package along with ``pytest``, ``mypy``, ``pylint``, +and ``yapf``. + +Running the CI Checks Locally +------------------------------ +The GitHub Actions CI runs three checks on every push and pull request: +tests, linting, and type checking. You should run all three locally +before pushing to make sure CI will pass. + +Tests (pytest) +^^^^^^^^^^^^^^ +Run the full test suite with: + +.. code-block:: shell + + $ pytest + +All tests must pass. To run a single test file: + +.. code-block:: shell + + $ pytest test/test_math.py + +Linting (pylint) +^^^^^^^^^^^^^^^^^ +Run pylint with: + +.. code-block:: shell + + $ pylint mudslide/ + +CI requires a minimum score of **9.5/10**. The project's ``.pylintrc`` +file configures allowed variable names, disabled checks, and other +settings. If pylint reports a score below 9.5, you can see individual +messages to understand what needs to be fixed. + +You can also check a specific file: + +.. code-block:: shell + + $ pylint mudslide/batch.py + +Type Checking (mypy) +^^^^^^^^^^^^^^^^^^^^^ +Run mypy with: + +.. code-block:: shell + + $ mypy mudslide/ + +This must exit with **zero errors**. The mypy configuration in +``pyproject.toml`` enforces strict settings including +``disallow_untyped_defs`` and ``disallow_incomplete_defs``, so all +functions must have type annotations. + +If you are adding new code, make sure to include type annotations on +all function signatures. Common patterns in the codebase: + +.. code-block:: python + + import numpy as np + from numpy.typing import ArrayLike, NDArray + + def compute_energy(x: ArrayLike) -> NDArray: + ... + +Code Formatting (yapf) +^^^^^^^^^^^^^^^^^^^^^^^ +The project uses ``yapf`` with a custom style. To format +a file in place: + +.. code-block:: shell + + $ yapf -i mudslide/batch.py + +To format all source files: + +.. code-block:: shell + + $ yapf -i mudslide/*.py + +Formatting is not currently enforced in CI, but consistent formatting +is expected for contributions. + +Running All Checks +^^^^^^^^^^^^^^^^^^^ +To run all three CI checks in sequence: + +.. code-block:: shell + + $ pytest && pylint mudslide/ --fail-under=9.5 && mypy mudslide/ + +If all three commands succeed, your changes should pass CI. diff --git a/docs/source/index.rst b/docs/source/index.rst index e991e45..5ba61f2 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -15,6 +15,7 @@ Mudslide! components usage library + developing references Indices and tables diff --git a/mudslide/__init__.py b/mudslide/__init__.py index 610b1d1..c530840 100644 --- a/mudslide/__init__.py +++ b/mudslide/__init__.py @@ -7,6 +7,11 @@ from .version import __version__ from .header import print_header +from .exceptions import ( + MudslideError, ConfigurationError, ExternalCodeError, + ConvergenceError, ComputeError, MissingDataError, + MissingForceError, MissingCouplingError +) from . import units from . import models @@ -18,6 +23,7 @@ from . import surface from . import turbo_make_harmonic +from .trajectory_md import * from .surface_hopping_md import * from .adiabatic_propagator import * from .adiabatic_md import * @@ -26,4 +32,3 @@ from .afssh import * from .batch import * from .tracer import * - diff --git a/mudslide/__main__.py b/mudslide/__main__.py index 5741a92..f5230cf 100644 --- a/mudslide/__main__.py +++ b/mudslide/__main__.py @@ -1,12 +1,14 @@ # -*- coding: utf-8 -*- """Code for the mudslide runtime""" -from __future__ import print_function, division - -import numpy as np +from __future__ import annotations +import argparse as ap import pickle import sys +from typing import Any + +import numpy as np from .surface_hopping_md import SurfaceHoppingMD from .even_sampling import EvenSamplingTrajectory @@ -15,29 +17,42 @@ from .batch import TrajGenConst, TrajGenNormal, BatchedTraj from .tracer import TraceManager from .models import scattering_models as models +from .exceptions import ConfigurationError from .version import __version__, get_version_info -import argparse as ap - -from typing import Any - # Add a method into this dictionary to register it with argparse methods = { - "fssh": lambda *args, **kwargs: SurfaceHoppingMD(*args, hopping_method="instantaneous", **kwargs), - "cumulative-sh": lambda *args, **kwargs: SurfaceHoppingMD(*args, hopping_method="cumulative", **kwargs), - "ehrenfest": Ehrenfest, - "afssh": AugmentedFSSH, - "even-sampling": EvenSamplingTrajectory + "fssh": + lambda *args, **kwargs: SurfaceHoppingMD( + *args, hopping_method="instantaneous", **kwargs), + "cumulative-sh": + lambda *args, **kwargs: SurfaceHoppingMD( + *args, hopping_method="cumulative", **kwargs), + "ehrenfest": + Ehrenfest, + "afssh": + AugmentedFSSH, + "even-sampling": + EvenSamplingTrajectory } -def main(argv=None, file=sys.stdout) -> None: +def main(argv: list[str] | None = None, file: Any = sys.stdout) -> None: + """CLI entry point for running scattering model surface hopping simulations.""" parser = ap.ArgumentParser(description="Mudslide test driver", epilog=get_version_info(), formatter_class=ap.RawDescriptionHelpFormatter) - parser.add_argument('-v', '--version', action='version', version=get_version_info()) + parser.add_argument('-v', + '--version', + action='version', + version=get_version_info()) - parser.add_argument('-a', '--method', default="fssh", choices=methods.keys(), type=str.lower, help="Variant of SH") + parser.add_argument('-a', + '--method', + default="fssh", + choices=methods.keys(), + type=str.lower, + help="Variant of SH") parser.add_argument('-m', '--model', default='simple', @@ -55,62 +70,119 @@ def main(argv=None, file=sys.stdout) -> None: nargs=2, type=float, help="range of momenta to consider (%(default)s)") - parser.add_argument('-n', '--nk', default=20, type=int, help="number of momenta to compute (%(default)d)") + parser.add_argument('-n', + '--nk', + default=20, + type=int, + help="number of momenta to compute (%(default)d)") parser.add_argument('-l', '--kspacing', default="linear", type=str, choices=('linear', 'log'), help="linear or log spacing for momenta (%(default)s)") - parser.add_argument('-K', - '--ksampling', - default="none", - type=str, - choices=('none', 'normal'), - help="how to sample momenta for a set of simulations (%(default)s)") - parser.add_argument('-f', - '--normal', - default=20, + parser.add_argument( + '-K', + '--ksampling', + default="none", + type=str, + choices=('none', 'normal'), + help="how to sample momenta for a set of simulations (%(default)s)") + parser.add_argument( + '-f', + '--normal', + default=20, + type=float, + help= + "standard deviation as a proportion of inverse momentum for normal samping (%(default)s)" + ) + parser.add_argument('-s', + '--samples', + default=200, + type=int, + help="number of samples (%(default)d)") + parser.add_argument( + '--sample-stack', + default=[10], + nargs='*', + type=int, + help= + "number of samples at each sampling depth for even sampling algorithm (%(default)s)" + ) + parser.add_argument('-M', + '--mass', + default=2000.0, type=float, - help="standard deviation as a proportion of inverse momentum for normal samping (%(default)s)") - parser.add_argument('-s', '--samples', default=200, type=int, help="number of samples (%(default)d)") - parser.add_argument('--sample-stack', - default=[10], - nargs='*', + help="particle mass (%(default)s)") + parser.add_argument('-t', + '--dt', + default=20.0, + type=float, + help="time step in a.u.(%(default)s)") + parser.add_argument('-y', + '--scale_dt', + dest="scale_dt", + action="store_true", + help="use dt=[dt]/k (%(default)s)") + parser.add_argument('-T', + '--nt', + default=50000, type=int, - help="number of samples at each sampling depth for even sampling algorithm (%(default)s)") - parser.add_argument('-j', '--nprocs', default=1, type=int, help="number of processors (%(default)d)") - parser.add_argument('-M', '--mass', default=2000.0, type=float, help="particle mass (%(default)s)") - parser.add_argument('-t', '--dt', default=20.0, type=float, help="time step in a.u.(%(default)s)") - parser.add_argument('-y', '--scale_dt', dest="scale_dt", action="store_true", help="use dt=[dt]/k (%(default)s)") - parser.add_argument('-T', '--nt', default=50000, type=int, help="max number of steps (%(default)s)") - parser.add_argument('-e', '--every', default=1, type=int, help="store a snapshot every nth step (%(default)s)") - parser.add_argument('-x', '--position', default=-10.0, type=float, help="starting position (%(default)s)") - parser.add_argument('-b', '--bounds', default=5.0, type=float, help="bounding box to end simulation (%(default)s)") - parser.add_argument('-p', - '--probability', - choices=["tully", "poisson"], - default="tully", - type=str, - help="how to determine hopping probabilities from gk->n * dt (%(default)s)") + help="max number of steps (%(default)s)") + parser.add_argument('-e', + '--every', + default=1, + type=int, + help="store a snapshot every nth step (%(default)s)") + parser.add_argument('-x', + '--position', + default=-10.0, + type=float, + help="starting position (%(default)s)") + parser.add_argument('-b', + '--bounds', + default=5.0, + type=float, + help="bounding box to end simulation (%(default)s)") + parser.add_argument( + '-p', + '--probability', + choices=["tully", "poisson"], + default="tully", + type=str, + help= + "how to determine hopping probabilities from gk->n * dt (%(default)s)") parser.add_argument('-o', '--output', default="averaged", type=str, - choices=('averaged', 'single', 'pickle', 'swarm', 'hack'), + choices=('averaged', 'single', 'pickle', 'swarm', + 'hack'), help="what to produce as output (%(default)s)") parser.add_argument('-O', '--outfile', default="sh.pickle", type=str, help="name of pickled file to produce (%(default)s)") - parser.add_argument('-z', '--seed', default=None, type=int, help="random seed (None)") - parser.add_argument("--log", choices=["memory", "yaml"], default="memory", help="how to store trajectory data") - parser.add_argument('--logdir', default="", type=str, help="directory to put log results (%(default)s)") - parser.add_argument('--published', - dest="published", - action="store_true", - help="override ranges to use those found in relevant papers (%(default)s)") + parser.add_argument('-z', + '--seed', + default=None, + type=int, + help="random seed (None)") + parser.add_argument("--log", + choices=["memory", "yaml"], + default="memory", + help="how to store trajectory data") + parser.add_argument('--logdir', + default="", + type=str, + help="directory to put log results (%(default)s)") + parser.add_argument( + '--published', + dest="published", + action="store_true", + help= + "override ranges to use those found in relevant papers (%(default)s)") args = parser.parse_args(argv) @@ -119,26 +191,29 @@ def main(argv=None, file=sys.stdout) -> None: nk = args.nk min_k, max_k = args.krange - if (args.published): # hack spacing to resemble Tully's - if (args.model == "simple"): + if args.published: # hack spacing to resemble Tully's + if args.model == "simple": min_k, max_k = 1.0, 35.0 - elif (args.model == "dual"): - min_k, max_k = np.log10(np.sqrt(2.0 * args.mass * np.exp(-4.0))), np.log10( - np.sqrt(2.0 * args.mass * np.exp(1.0))) - elif (args.model == "extended"): + elif args.model == "dual": + min_k, max_k = np.log10(np.sqrt( + 2.0 * args.mass * np.exp(-4.0))), np.log10( + np.sqrt(2.0 * args.mass * np.exp(1.0))) + elif args.model == "extended": min_k, max_k = 1.0, 35.0 - elif (args.model == "super"): + elif args.model == "super": min_k, max_k = 0.5, 20.0 else: - print("Warning! Published option chosen but no available bounds! Using inputs.", file=sys.stderr) + print( + "Warning! Published option chosen but no available bounds! Using inputs.", + file=sys.stderr) - kpoints = [] + kpoints: np.ndarray if args.kspacing == "linear": kpoints = np.linspace(min_k, max_k, nk) elif args.kspacing == "log": kpoints = np.logspace(min_k, max_k, nk) else: - raise Exception("Unrecognized type of spacing") + raise ConfigurationError("Unrecognized type of spacing") trajectory_type = methods[args.method] @@ -149,11 +224,11 @@ def main(argv=None, file=sys.stdout) -> None: all_results = [] - if (args.output == "averaged" or args.output == "pickle"): + if args.output in ("averaged", "pickle"): print("# momentum ", end='', file=file) for ist in range(model.nstates): for d in ["reflected", "transmitted"]: - print("%d_%s" % (ist, d), end=' ', file=file) + print(f"{ist}_{d}", end=' ', file=file) print(file=file) for k in kpoints: @@ -162,7 +237,11 @@ def main(argv=None, file=sys.stdout) -> None: if args.ksampling == "none": traj_gen = TrajGenConst(args.position, v, 0, seed=args.seed) elif args.ksampling == "normal": - traj_gen = TrajGenNormal(args.position, v, 0, sigma=args.normal / v, seed=args.seed) + traj_gen = TrajGenNormal(args.position, + v, + 0, + sigma=args.normal / v, + seed=args.seed) dt = (args.dt / k) if args.scale_dt else args.dt @@ -170,11 +249,12 @@ def main(argv=None, file=sys.stdout) -> None: traj_gen, trajectory_type=trajectory_type, samples=args.samples, - nprocs=args.nprocs, dt=dt, max_steps=args.nt, - bounds=[ -abs(args.bounds), abs(args.bounds) ], - tracemanager=TraceManager(trace_type, trace_kwargs=trace_options), + bounds=[-abs(args.bounds), + abs(args.bounds)], + tracemanager=TraceManager( + trace_type, trace_kwargs=trace_options), trace_every=args.every, spawn_stack=args.sample_stack, electronic_integration=args.electronic, @@ -183,11 +263,11 @@ def main(argv=None, file=sys.stdout) -> None: results = fssh.compute() outcomes = results.outcomes - if (args.output == "single"): + if args.output == "single": results.traces[0].print(file=file) - elif (args.output == "swarm"): + elif args.output == "swarm": maxsteps = max([len(t) for t in results.traces]) - outfiles = ["state_%d.trace" % i for i in range(model.nstates)] + outfiles = [f"state_{i}.trace" for i in range(model.nstates)] fils = [open(o, "w") for o in outfiles] for i in range(maxsteps): nswarm = [0 for x in fils] @@ -202,21 +282,23 @@ def main(argv=None, file=sys.stdout) -> None: for ist in range(model.nstates): if nswarm[ist] == 0: - print("%12.6f" % -9999999, file=fils[ist]) + print(f"{-9999999:12.6f}", file=fils[ist]) print(file=fils[ist]) print(file=fils[ist]) for f in fils: f.close() - elif (args.output == "averaged" or args.output == "pickle"): - print("%12.6f %s" % (k, " ".join(["%12.6f" % x for x in np.nditer(outcomes)])), file=file) - if (args.output == "pickle"): # save results for later processing + elif args.output in ("averaged", "pickle"): + print(f"{k:12.6f} {' '.join(f'{float(x):12.6f}' for x in outcomes.flat)}", + file=file) + if args.output == "pickle": # save results for later processing all_results.append((k, results)) - elif (args.output == "hack"): + elif args.output == "hack": print("Hack something here, if you like.", file=file) else: - print("Not printing results. This is probably not what you wanted!", file=file) + print("Not printing results. This is probably not what you wanted!", + file=file) - if (len(all_results) > 0): + if len(all_results) > 0: pickle.dump(all_results, open(args.outfile, "wb")) diff --git a/mudslide/adiabatic_md.py b/mudslide/adiabatic_md.py index e5e69ae..997c18d 100644 --- a/mudslide/adiabatic_md.py +++ b/mudslide/adiabatic_md.py @@ -5,137 +5,46 @@ similar to ground state molecular dynamics simulations. """ -from typing import Dict, Any -import copy as cp +from __future__ import annotations + +from typing import Dict, Any, TYPE_CHECKING import numpy as np -from numpy.typing import ArrayLike -from .util import check_options -from .constants import boltzmann -from .tracer import Trace +from .trajectory_md import TrajectoryMD +from .propagator import Propagator_ from .adiabatic_propagator import AdiabaticPropagator -class AdiabaticMD: +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + + +class AdiabaticMD(TrajectoryMD): """Class to propagate a single adiabatic trajectory, like ground state MD. This class handles the propagation of molecular dynamics trajectories in the adiabatic regime, similar to ground state molecular dynamics. """ - recognized_options = [ - "dt", "t0", "trace_every", "remove_com_every", "remove_angular_momentum_every", - "max_steps", "max_time", "bounds", "propagator", "seed_sequence", "electronics", - "outcome_type", "weight", "last_velocity", "previous_steps", "restarting" - ] - - def __init__(self, - model: Any, - x0: ArrayLike, - v0: ArrayLike, - tracer: Any = None, - queue: Any = None, - strict_option_check: bool = True, - **options: Any): - """Initialize the AdiabaticMD class. + def make_propagator(self, model: ElectronicModel_, + options: Dict[str, Any]) -> Propagator_: + """Create the adiabatic propagator. Parameters ---------- model : Any Model object defining problem. - x0 : ArrayLike - Initial position. - v0 : ArrayLike - Initial velocity. - tracer : Any, optional - Spawn from TraceManager to collect results. - queue : Any, optional - Trajectory queue. - strict_option_check : bool, optional - Whether to strictly check options. - **options : Any - Additional options for the simulation. Recognized options are: + options : Dict[str, Any] + Options dictionary. - dt : float, optional - Time step for nuclear propagation (in atomic units). Required. - t0 : float, optional - Initial time. Default is 0.0. - trace_every : int, optional - Interval (in steps) at which to record trajectory data. Default is 1. - remove_com_every : int, optional - Interval for removing center-of-mass motion. Default is 0 (never). - remove_angular_momentum_every : int, optional - Interval for removing angular momentum. Default is 0 (never). - max_steps : int, optional - Maximum number of steps. Default is 100000. - max_time : float, optional - Maximum simulation time. Default is 1e25. - bounds : tuple or list, optional - Tuple or list of (lower, upper) bounds for the simulation box. Default is None. - propagator : str or dict, optional - The propagator to use for nuclear motion. Can be a string (e.g., 'VV') or a - dictionary with more options. Default is 'VV'. - seed_sequence : int or numpy.random.SeedSequence, optional - Seed or SeedSequence for random number generation. Default is None. - electronics : object, optional - Initial electronic state object. Default is None. - outcome_type : str, optional - Type of outcome to record (e.g., 'state'). Default is 'state'. - weight : float, optional - Statistical weight of the trajectory. Default is 1.0. - restarting : bool, optional - Whether this is a restarted trajectory. Default is False. + Returns + ------- + Propagator_ + Adiabatic propagator instance. """ - check_options(options, self.recognized_options, strict=strict_option_check) - - self.model = model - self.tracer = Trace(tracer) - self.queue: Any = queue - self.mass = model.mass - self.position = np.array(x0, dtype=np.float64).reshape(model.ndof) - self.velocity = np.array(v0, dtype=np.float64).reshape(model.ndof) - self.last_position = np.zeros_like(self.position, dtype=np.float64) - self.last_velocity = np.zeros_like(self.velocity, dtype=np.float64) - if "last_velocity" in options: - self.last_velocity[:] = options["last_velocity"] - - # function duration_initialize should get us ready to for future continue_simulating calls - # that decide whether the simulation has finished - if "duration" in options: - self.duration = options["duration"] - else: - self.duration_initialize(options) - - # fixed initial parameters - self.time = np.longdouble(options.get("t0", 0.0)) - self.nsteps = int(options.get("previous_steps", 0)) - self.max_steps = int(options.get("max_steps", 100000)) - self.max_time = float(options.get("max_time", 1e25)) - self.trace_every = int(options.get("trace_every", 1)) - - self.propagator = AdiabaticPropagator(self.model, options.get("propagator", "VV")) - - self.remove_com_every = int(options.get("remove_com_every", 0)) - self.remove_angular_momentum_every = int(options.get("remove_angular_momentum_every", 0)) - - # read out of options - self.dt = float(options["dt"]) - self.outcome_type = options.get("outcome_type", "state") - - ss = options.get("seed_sequence", None) - self.seed_sequence = ss if isinstance(ss, np.random.SeedSequence) \ - else np.random.SeedSequence(ss) - self.random_state = np.random.default_rng(self.seed_sequence) - - self.electronics = options.get("electronics", None) - self.last_electronics = None - - self.weight = np.float64(options.get("weight", 1.0)) - - self.restarting = options.get("restarting", False) - self.force_quit = False + return AdiabaticPropagator(model, options.get("propagator", "VV")) # type: ignore[return-value] @classmethod - def restart(cls, model, log, **options) -> 'AdiabaticMD': + def restart(cls, model: ElectronicModel_, log: Any, **options: Any) -> 'AdiabaticMD': """Restart trajectory from log. Parameters @@ -179,258 +88,40 @@ def restart(cls, model, log, **options) -> 'AdiabaticMD': restarting=True, **options) - def update_weight(self, weight: np.float64) -> None: - """Update weight held by trajectory and by trace. - - Parameters - ---------- - weight : np.float64 - New weight value to set. - """ - self.weight = weight - self.tracer.weight = weight - - if self.weight == 0.0: - self.force_quit = True - - def __deepcopy__(self, memo: Any) -> 'AdiabaticMD': - """Override deepcopy. - - Parameters - ---------- - memo : Any - Memo dictionary for deepcopy. - - Returns - ------- - AdiabaticMD - Deep copy of the current object. - """ - cls = self.__class__ - result = cls.__new__(cls) - memo[id(self)] = result - shallow_only = ["queue"] - for k, v in self.__dict__.items(): - setattr(result, k, cp.deepcopy(v, memo) if k not in shallow_only else cp.copy(v)) - return result - - def clone(self) -> 'AdiabaticMD': - """Clone existing trajectory for spawning. - - Returns - ------- - AdiabaticMD - Copy of current object. - """ - return cp.deepcopy(self) - - def random(self) -> np.float64: - """Get random number for hopping decisions. - - Returns - ------- - np.float64 - Uniform random number between 0 and 1. - """ - return self.random_state.uniform() - - def currently_interacting(self) -> bool: - """Determine whether trajectory is currently inside an interaction region. - - Returns - ------- - bool - True if trajectory is inside interaction region, False otherwise. - """ - if self.duration["box_bounds"] is None: - return False - return np.all(self.duration["box_bounds"][0] < self.position) and np.all( - self.position < self.duration["box_bounds"][1]) - - def duration_initialize(self, options: Dict[str, Any]) -> None: - """Initialize variables related to continue_simulating. - - Parameters - ---------- - options : Dict[str, Any] - Dictionary with options. - """ - - duration = {} # type: Dict[str, Any] - duration['found_box'] = False - - bounds = options.get('bounds', None) - if bounds: - b0 = np.array(bounds[0], dtype=np.float64) - b1 = np.array(bounds[1], dtype=np.float64) - duration["box_bounds"] = (b0, b1) - else: - duration["box_bounds"] = None - - self.duration = duration - - def continue_simulating(self) -> bool: - """Decide whether trajectory should keep running. - - Returns - ------- - bool - True if trajectory should keep running, False if it should finish. - """ - if self.force_quit: # pylint: disable=no-else-return - return False - elif self.max_steps >= 0 and self.nsteps >= self.max_steps: - return False - elif self.time >= self.max_time or np.isclose( - self.time, self.max_time, atol=1e-8, rtol=0.0): - return False - elif self.duration["found_box"]: - return self.currently_interacting() - else: - if self.currently_interacting(): - self.duration["found_box"] = True - return True - - def trace(self, force: bool = False) -> None: - """Add results from current time point to tracing function. - - Only adds snapshot if nsteps%trace_every == 0, unless force=True. - - Parameters - ---------- - force : bool, optional - Force snapshot regardless of trace_every interval. - """ - if force or (self.nsteps % self.trace_every) == 0: - self.tracer.collect(self.snapshot()) - - def snapshot(self) -> Dict[str, Any]: - """Collect data from run for logging. - - Returns - ------- - Dict[str, Any] - Dictionary with all data from current time step. - """ - out = { - "time": self.time, - "position": self.position.tolist(), - "velocity": self.velocity.tolist(), - "potential": self.potential_energy().item(), - "kinetic": self.kinetic_energy().item(), - "energy": self.total_energy().item(), - "temperature": 2 * self.kinetic_energy() / ( boltzmann * self.model.ndof), - "electronics": self.electronics.as_dict() - } - return out - - def kinetic_energy(self) -> np.float64: - """Calculate kinetic energy. - - Returns - ------- - np.float64 - Kinetic energy. - """ - return 0.5 * np.einsum('m,m,m', self.mass, self.velocity, self.velocity) - - def potential_energy(self, electronics: 'ElectronicModel_' = None) -> np.floating: + def potential_energy(self, + electronics: ElectronicModel_ | None = None + ) -> float: """Calculate potential energy. Parameters ---------- - electronics : ElectronicModel, optional + electronics : ElectronicModel_, optional Electronic states from current step. Returns ------- - np.floating + float Potential energy. """ if electronics is None: electronics = self.electronics + assert electronics is not None return electronics.energies[0] - def total_energy(self, electronics: 'ElectronicModel_' = None) -> np.floating: - """Calculate total energy (kinetic + potential). - - Parameters - ---------- - electronics : ElectronicModel, optional - Electronic states from current step. - - Returns - ------- - np.floating - Total energy. - """ - potential = self.potential_energy(electronics) - kinetic = self.kinetic_energy() - return potential + kinetic - - def force(self, electronics: 'ElectronicModel_' = None) -> ArrayLike: - """Compute force on active state. + def force(self, electronics: ElectronicModel_ | None = None) -> np.ndarray: + """Compute force on ground state. Parameters ---------- - electronics : 'ElectronicModel', optional + electronics : ElectronicModel_, optional ElectronicStates from current step. Returns ------- - ArrayLike - Force on active electronic state. + np.ndarray + Force on ground electronic state. """ if electronics is None: electronics = self.electronics + assert electronics is not None return electronics.force(0) - - def mode_kinetic_energy(self, direction: ArrayLike) -> np.float64: - """Calculate kinetic energy along given momentum mode. - - Parameters - ---------- - direction : ArrayLike - Array defining direction. - - Returns - ------- - np.float64 - Kinetic energy along specified direction. - """ - u = direction / np.linalg.norm(direction) - momentum = self.velocity * self.mass - component = np.dot(u, momentum) * u - return 0.5 * np.einsum('m,m,m', 1.0 / self.mass, component, component) - - def simulate(self) -> 'Trace': - """Run the simulation. - - Returns - ------- - Trace - Trace of trajectory. - """ - - if not self.continue_simulating(): - return self.tracer - - if self.electronics is None: - self.electronics = self.model.update(self.position) - - if not self.restarting: - self.trace() - - # propagation - while True: - self.propagator(self, 1) # pylint: disable=not-callable - - # ending condition - if not self.continue_simulating(): - break - - self.trace() - - self.trace(force=True) - - return self.tracer diff --git a/mudslide/adiabatic_propagator.py b/mudslide/adiabatic_propagator.py index 3bce3a4..cc1670f 100644 --- a/mudslide/adiabatic_propagator.py +++ b/mudslide/adiabatic_propagator.py @@ -6,13 +6,21 @@ """ # pylint: disable=too-few-public-methods, too-many-arguments, invalid-name -from typing import Any, Dict +from __future__ import annotations + +from typing import Any, Dict, TYPE_CHECKING import numpy as np from .constants import fs_to_au, boltzmann, amu_to_au +from .exceptions import ConfigurationError from .propagator import Propagator_ -from .util import remove_center_of_mass_motion, remove_angular_momentum +from .util import remove_center_of_mass_motion, remove_angular_momentum, check_options + +if TYPE_CHECKING: + from .adiabatic_md import AdiabaticMD + from .models.electronics import ElectronicModel_ + class VVPropagator(Propagator_): """Velocity Verlet propagator. @@ -26,6 +34,8 @@ class VVPropagator(Propagator_): Additional options for the propagator. """ + recognized_options: list[str] = ["ndof"] + def __init__(self, **options: Any) -> None: """Initialize the Velocity Verlet propagator. @@ -34,6 +44,7 @@ def __init__(self, **options: Any) -> None: **options : Any Additional options for the propagator. """ + check_options(options, self.recognized_options) super().__init__() def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: @@ -55,7 +66,8 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: # calculate electronics at new position traj.last_electronics = traj.electronics - traj.electronics = traj.model.update(traj.position, electronics=traj.electronics) + traj.electronics = traj.model.update(traj.position, + electronics=traj.electronics) # update velocity last_acceleration = acceleration @@ -65,10 +77,11 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: traj.velocity += 0.5 * (last_acceleration + acceleration) * dt # optionally remove COM motion and total angular momentum - if traj.remove_com_every > 0 and (traj.nsteps % traj.remove_com_every) == 0: + if traj.remove_com_every > 0 and (traj.nsteps % + traj.remove_com_every) == 0: d = traj.model.dimensionality v = traj.velocity.reshape(d) - m = traj.mass.reshape(d)[:,0] + m = traj.mass.reshape(d)[:, 0] vnew = remove_center_of_mass_motion(v, m) traj.velocity = vnew.reshape(traj.velocity.shape) @@ -76,7 +89,7 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: (traj.nsteps % traj.remove_angular_momentum_every) == 0: d = traj.model.dimensionality v = traj.velocity.reshape(d) - m = traj.mass.reshape(d)[:,0] + m = traj.mass.reshape(d)[:, 0] x = traj.position.reshape(d) vnew = remove_angular_momentum(v, m, x) traj.velocity = vnew.reshape(traj.velocity.shape) @@ -84,6 +97,7 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: traj.time += dt traj.nsteps += 1 + class NoseHooverChainPropagator(Propagator_): """Nose-Hoover Chain thermostat propagator. @@ -112,8 +126,13 @@ class NoseHooverChainPropagator(Propagator_): Molecular Physics, 87, 1117-1157 (1996) """ - def __init__(self, temperature: np.float64, timescale: np.float64 = 1e2 * fs_to_au, - ndof: int = 3, nchains: int = 3, nys: int = 3, nc: int = 1): + def __init__(self, + temperature: float, + timescale: float = 1e2 * fs_to_au, + ndof: int = 3, + nchains: int = 3, + nys: int = 3, + nc: int = 1): """Initialize the Nose-Hoover Chain thermostat. Parameters @@ -141,7 +160,7 @@ def __init__(self, temperature: np.float64, timescale: np.float64 = 1e2 * fs_to_ assert self.temperature > 0.0 assert self.timescale > 0.0 assert self.nchains >= 1 - assert self.nys in (3,5) + assert self.nys in (3, 5) assert self.nc >= 1 self.nh_position = np.zeros(nchains, dtype=np.float64) @@ -151,13 +170,13 @@ def __init__(self, temperature: np.float64, timescale: np.float64 = 1e2 * fs_to_ self.nh_mass[0] *= ndof if nys == 3: - tmp = 1 / (2 - 2**(1./3)) - self.w = np.array([tmp, 1 - 2*tmp, tmp]) / nc + tmp = 1 / (2 - 2**(1. / 3)) + self.w = np.array([tmp, 1 - 2 * tmp, tmp]) / nc elif nys == 5: - tmp = 1 / (4 - 4**(1./3)) - self.w = np.array([tmp, tmp, 1 - 4*tmp, tmp, tmp]) / nc + tmp = 1 / (4 - 4**(1. / 3)) + self.w = np.array([tmp, tmp, 1 - 4 * tmp, tmp, tmp]) / nc else: - raise ValueError("nys must be either 3 or 5") + raise ConfigurationError("nys must be either 3 or 5") self.G = np.zeros(nchains) @@ -167,7 +186,7 @@ def __init__(self, temperature: np.float64, timescale: np.float64 = 1e2 * fs_to_ print(f" Timescale: {timescale / fs_to_au:.2f} fs") print(f" Thermostat mass: {self.nh_mass / amu_to_au} amu") - def nhc_step(self, velocity, mass, dt: float): + def nhc_step(self, velocity: np.ndarray, mass: np.ndarray, dt: float) -> float: """Move forward one step in the extended system variables. Parameters @@ -200,9 +219,9 @@ def nhc_step(self, velocity, mass, dt: float): for _ in range(self.nc): for iys in range(self.nys): wdt = self.w[iys] * dt - V[M-1] += G[M-1] * wdt / 4.0 + V[M - 1] += G[M - 1] * wdt / 4.0 - for kk in range(0, M-1): + for kk in range(0, M - 1): AA = np.exp(-wdt * V[M - 1 - kk] / 8.0) V[M - 2 - kk] *= AA * AA V[M - 2 - kk] += 0.25 * wdt * G[M - 2 - kk] * AA @@ -222,7 +241,8 @@ def nhc_step(self, velocity, mass, dt: float): AA = np.exp(-0.125 * wdt * V[kk + 1]) V[kk] = V[kk] * AA * AA \ + 0.25 * wdt * G[kk] * AA - G[kk+1] = (self.nh_mass[kk] * V[kk]**2 - kt) / self.nh_mass[kk + 1] + G[kk + 1] = (self.nh_mass[kk] * V[kk]**2 - + kt) / self.nh_mass[kk + 1] V[M - 1] += 0.25 * G[M - 1] * wdt return scale @@ -253,21 +273,23 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: # calculate electronics at new position traj.last_electronics = traj.electronics - traj.electronics = traj.model.update(traj.position, electronics=traj.electronics) + traj.electronics = traj.model.update(traj.position, + electronics=traj.electronics) acceleration = traj.force(traj.electronics) / traj.mass - v2p = v1 + 0.5 * dt * acceleration + v2p = v1 + 0.5 * dt * acceleration vscale = self.nhc_step(v2p, traj.mass, dt) traj.last_velocity = traj.velocity traj.velocity = v2p * vscale # optionally remove COM motion and total angular momentum - if traj.remove_com_every > 0 and (traj.nsteps % traj.remove_com_every) == 0: + if traj.remove_com_every > 0 and (traj.nsteps % + traj.remove_com_every) == 0: d = traj.model.dimensionality v = traj.velocity.reshape(d) - m = traj.mass.reshape(d)[:,0] + m = traj.mass.reshape(d)[:, 0] vnew = remove_center_of_mass_motion(v, m) traj.velocity = vnew.reshape(traj.velocity.shape) @@ -275,7 +297,7 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: (traj.nsteps % traj.remove_angular_momentum_every) == 0: d = traj.model.dimensionality v = traj.velocity.reshape(d) - m = traj.mass.reshape(d)[:,0] + m = traj.mass.reshape(d)[:, 0] x = traj.position.reshape(d) vnew = remove_angular_momentum(v, m, x) traj.velocity = vnew.reshape(traj.velocity.shape) @@ -283,6 +305,7 @@ def __call__(self, traj: 'AdiabaticMD', nsteps: int) -> None: traj.time += dt traj.nsteps += 1 + class AdiabaticPropagator: """Factory class for creating propagator objects. @@ -306,7 +329,8 @@ class AdiabaticPropagator: ValueError If the propagator type is unknown or if the propagator options are invalid. """ - def __new__(cls, model: Any, prop_options: Any = "vv") -> Propagator_: + + def __new__(cls, model: ElectronicModel_, prop_options: Any = "vv") -> Propagator_: # type: ignore[misc] """Create a new propagator instance. Parameters @@ -327,7 +351,7 @@ def __new__(cls, model: Any, prop_options: Any = "vv") -> Propagator_: If the propagator type is unknown or if the propagator options are invalid. """ if isinstance(prop_options, str): - prop_options = { "type": prop_options.lower() } + prop_options = {"type": prop_options.lower()} prop_options["ndof"] = model.ndof @@ -338,6 +362,6 @@ def __new__(cls, model: Any, prop_options: Any = "vv") -> Propagator_: elif prop_type in ["nh", "nhc", "nose-hoover"]: return NoseHooverChainPropagator(**prop_options) else: - raise ValueError(f"Unknown propagator type: {prop_type}") + raise ConfigurationError(f"Unknown propagator type: {prop_type}") - raise ValueError(f"Unknown propagator options: {prop_options}") + raise ConfigurationError(f"Unknown propagator options: {prop_options}") diff --git a/mudslide/afssh.py b/mudslide/afssh.py index cd61d93..7679a8e 100644 --- a/mudslide/afssh.py +++ b/mudslide/afssh.py @@ -1,17 +1,28 @@ # -*- coding: utf-8 -*- """Propagating Augmented-FSSH (A-FSSH) trajectories.""" -from typing import Any, Union +from __future__ import annotations + +from typing import Any, Union, TYPE_CHECKING import numpy as np from numpy.typing import ArrayLike +from .exceptions import ComputeError, ConfigurationError from .util import is_string from .math import poisson_prob_scale from .propagation import rk4 from .surface_hopping_md import SurfaceHoppingMD from .propagator import Propagator_ +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + from .surface_hopping_propagator import SHPropagator + + +IMAGINARY_NORM_TOLERANCE: float = 1e-10 +ZERO_DIVISION_FLOOR: float = 1e-10 + class AFSSHVVPropagator(Propagator_): """Surface Hopping Velocity Verlet propagator.""" @@ -26,12 +37,12 @@ def __init__(self, **options: Any) -> None: """ super().__init__() - def __call__(self, traj: 'SurfaceHoppingMD', nsteps: int) -> None: + def __call__(self, traj: 'AugmentedFSSH', nsteps: int) -> None: # type: ignore[override] """Propagate trajectory using Surface Hopping Velocity Verlet algorithm. Parameters ---------- - traj : SurfaceHoppingMD + traj : AugmentedFSSH Trajectory object to propagate. nsteps : int Number of steps to propagate. @@ -40,40 +51,46 @@ def __call__(self, traj: 'SurfaceHoppingMD', nsteps: int) -> None: # first update nuclear coordinates for _ in range(nsteps): # Advance position using Velocity Verlet - acceleration = traj._force(traj.electronics) / traj.mass + acceleration = traj.force(traj.electronics) / traj.mass traj.last_position = traj.position traj.position += traj.velocity * dt + 0.5 * acceleration * dt * dt - traj.advance_delR(traj.last_electronics, traj.electronics) + traj.advance_delR(traj.last_electronics, traj.electronics) # type: ignore[arg-type] # calculate electronics at new position traj.last_electronics, traj.electronics = traj.electronics, traj.model.update( - traj.position, electronics=traj.electronics, - gradients=traj.needed_gradients(), couplings=traj.needed_couplings()) + traj.position, + electronics=traj.electronics, + gradients=traj.needed_gradients(), + couplings=traj.needed_couplings()) # Update velocity using Velocity Verlet - last_acceleration = traj._force(traj.last_electronics) / traj.mass - this_acceleration = traj._force(traj.electronics) / traj.mass + last_acceleration = traj.force(traj.last_electronics) / traj.mass + this_acceleration = traj.force(traj.electronics) / traj.mass traj.last_velocity = traj.velocity traj.velocity += 0.5 * (last_acceleration + this_acceleration) * dt + assert traj.last_electronics is not None + assert traj.electronics is not None traj.advance_delP(traj.last_electronics, traj.electronics) # now propagate the electronic wavefunction to the new time - traj.propagate_electronics(traj.last_electronics, traj.electronics, dt) + traj.propagate_electronics(traj.last_electronics, traj.electronics, + dt) traj.surface_hopping(traj.last_electronics, traj.electronics) traj.time += dt traj.nsteps += 1 -class AFSSHPropagator(Propagator_): + +class AFSSHPropagator(Propagator_): # pylint: disable=abstract-method """Surface Hopping propagator factory. This class serves as a factory for creating different types of propagators used in adiabatic FSSH molecular dynamics simulations. """ - def __new__(cls, model: Any, prop_options: Any = "vv") -> 'SHPropagator': + def __new__(cls, model: ElectronicModel_, prop_options: Any = "vv") -> Propagator_: # type: ignore[misc] """Create a new surface hopping propagator instance. Parameters @@ -98,13 +115,14 @@ def __new__(cls, model: Any, prop_options: Any = "vv") -> 'SHPropagator': if is_string(prop_options): prop_options = {"type": prop_options} elif not isinstance(prop_options, dict): - raise Exception("prop_options must be a string or a dictionary") + raise ConfigurationError("prop_options must be a string or a dictionary") proptype = prop_options.get("type", "vv") if proptype.lower() == "vv": return AFSSHVVPropagator(**prop_options) - else: - raise ValueError(f"Unrecognized surface hopping propagator type: {proptype}.") + raise ConfigurationError( + f"Unrecognized surface hopping propagator type: {proptype}.") + class AugmentedFSSH(SurfaceHoppingMD): """Augmented-FSSH (A-FSSH) dynamics, by Subotnik and coworkers. @@ -112,24 +130,39 @@ class AugmentedFSSH(SurfaceHoppingMD): Initial implementation based on original paper: Subotnik, Shenvi JCP 134, 024105 (2011); doi: 10.1063/1.3506779 """ + + recognized_options = SurfaceHoppingMD.recognized_options + [ + "augmented_integration" + ] + def __init__(self, *args: Any, **options: Any): - options['hopping_method'] = 'instantaneous' # force instantaneous hopping + options[ + 'hopping_method'] = 'instantaneous' # force instantaneous hopping SurfaceHoppingMD.__init__(self, *args, **options) - self.augmented_integration = options.get("augmented_integration", self.electronic_integration).lower() + self.augmented_integration = options.get( + "augmented_integration", self.electronic_integration).lower() - self.delR = np.zeros([self.model.ndof, self.model.nstates, self.model.nstates], - dtype=np.complex128) - self.delP = np.zeros([self.model.ndof, self.model.nstates, self.model.nstates], - dtype=np.complex128) + self.delR = np.zeros( + [self.model.ndof, self.model.nstates, self.model.nstates], + dtype=np.complex128) + self.delP = np.zeros( + [self.model.ndof, self.model.nstates, self.model.nstates], + dtype=np.complex128) - self.propagator = AFSSHPropagator(self.model, "vv") + self.propagator: Propagator_ = AFSSHPropagator(self.model, "vv") # type: ignore[assignment] - def needed_gradients(self): - """A-FSSH needs all forces for force_matrix computation.""" + def needed_gradients(self) -> list[int] | None: + """A-FSSH needs all forces for force_matrix computation. + + Returns + ------- + None + None means all state gradients are needed. + """ return None - def compute_delF(self, this_electronics): + def compute_delF(self, this_electronics: ElectronicModel_) -> np.ndarray: """Compute the difference in forces between states. Parameters @@ -143,12 +176,12 @@ def compute_delF(self, this_electronics): Matrix of force differences between states. """ delF = np.copy(this_electronics.force_matrix) - F0 = self._force(this_electronics) + F0 = self.force(this_electronics) for i in range(self.model.nstates): - delF[i,i,:] -= F0 + delF[i, i, :] -= F0 return delF - def advance_delR(self, last_electronics, this_electronics): + def advance_delR(self, last_electronics: ElectronicModel_, this_electronics: ElectronicModel_) -> None: """Propagate delR using Eq. (29) from Subotnik 2011 JCP. Parameters @@ -162,7 +195,7 @@ def advance_delR(self, last_electronics, this_electronics): H = self.hamiltonian_propagator(last_electronics, this_electronics) delV = np.zeros_like(self.delP) for x in range(self.delP.shape[0]): - delV[x,:,:] = self.delP[x,:,:] / self.mass[x] + delV[x, :, :] = self.delP[x, :, :] / self.mass[x] if self.augmented_integration == "exp": eps, co = np.linalg.eigh(H) @@ -173,7 +206,8 @@ def advance_delR(self, last_electronics, this_electronics): Rt = (RR + delV * dt) * expiht self.delR = np.einsum("pi,xij,qj->xpq", co, Rt, co.conj()) elif self.augmented_integration == "rk4": - def ydot(RR: ArrayLike, t: np.floating) -> ArrayLike: + + def ydot(RR: np.ndarray, t: float) -> np.ndarray: assert t >= 0.0 and t <= dt HR = np.einsum("pr,xrq->xpq", H, RR) RH = np.einsum("xpr,rq->xpq", RR, H) @@ -184,9 +218,11 @@ def ydot(RR: ArrayLike, t: np.floating) -> ArrayLike: Rt = rk4(self.delR, ydot, 0.0, dt, nsteps) self.delR = Rt else: - raise Exception("Unrecognized propagate delR method") + raise ConfigurationError( + f"Unrecognized augmented integration method: {self.augmented_integration}" + ) - def advance_delP(self, last_electronics, this_electronics): + def advance_delP(self, last_electronics: ElectronicModel_, this_electronics: ElectronicModel_) -> None: """Propagate delP using Eq. (31) from Subotnik JCP 2011. Parameters @@ -204,7 +240,7 @@ def advance_delP(self, last_electronics, this_electronics): eps, co = np.linalg.eigh(H) expiht = np.exp(-1j * dt * np.subtract.outer(eps, eps)) - eee = np.subtract.outer(2*eps, np.add.outer(eps,eps)) + eee = np.subtract.outer(2 * eps, np.add.outer(eps, eps)) poiss = -poisson_prob_scale(1j * eee * dt) * dt poiss_star = -poisson_prob_scale(-1j * eee * dt) * dt @@ -212,15 +248,18 @@ def advance_delP(self, last_electronics, this_electronics): PP = np.einsum("pi,xpq,qj->xij", co.conj(), self.delP, co) rho = np.einsum("pi,pq,qj->ij", co.conj(), self.rho, co) - FF = np.einsum("xik,kj,jik->xij", delF, rho, poiss) + np.einsum("ik,xkj,ijk->xij", rho, delF, poiss_star) + FF = np.einsum("xik,kj,jik->xij", delF, rho, poiss) + np.einsum( + "ik,xkj,ijk->xij", rho, delF, poiss_star) FF *= -0.5 Pt = (PP + FF) * expiht self.delP = np.einsum("pi,xij,qj->xpq", co, Pt, co.conj()) elif self.augmented_integration == "rk4": - dFrho_comm = np.einsum("prx,rq->xpq", delF, self.rho) + np.einsum("pr,rqx->xpq", self.rho, delF) + dFrho_comm = np.einsum("prx,rq->xpq", delF, self.rho) + np.einsum( + "pr,rqx->xpq", self.rho, delF) dFrho_comm *= 0.5 - def ydot(PP: ArrayLike, t: np.floating) -> ArrayLike: + + def ydot(PP: np.ndarray, t: float) -> np.ndarray: assert t >= 0.0 and t <= dt HP = np.einsum("pr,xrq->xpq", H, PP) PH = np.einsum("xpr,rq->xpq", PP, H) @@ -231,10 +270,16 @@ def ydot(PP: ArrayLike, t: np.floating) -> ArrayLike: Pt = rk4(self.delP, ydot, 0.0, dt, nsteps) self.delP = Pt else: - raise Exception("Unrecognized propagate delP method") + raise ConfigurationError( + f"Unrecognized augmented integration method: {self.augmented_integration}" + ) return - def direction_of_rescale(self, source: int, target: int, electronics: 'ElectronicModel_'=None) -> np.ndarray: + def direction_of_rescale( + self, + source: int, + target: int, + electronics: ElectronicModel_ | None = None) -> np.ndarray: """Return direction in which to rescale momentum. In Subotnik JCP 2011, they suggest to use the difference between the momenta on delP. @@ -254,10 +299,14 @@ def direction_of_rescale(self, source: int, target: int, electronics: 'Electroni Unit vector pointing in direction of rescale. """ out = self.delP[:, source, source] - self.delP[:, target, target] - assert np.linalg.norm(np.imag(out)) < 1e-8 + if np.linalg.norm(np.imag(out)) >= IMAGINARY_NORM_TOLERANCE: + raise ComputeError( + "Rescale direction has unexpectedly large imaginary component: " + f"{np.linalg.norm(np.imag(out)):.2e}") return np.real(out) - def gamma_collapse(self, electronics: 'ElectronicModel_'=None) -> np.ndarray: + def gamma_collapse(self, + electronics: ElectronicModel_ | None = None) -> np.ndarray: """Compute probability of collapse to each electronic state. Uses Eq. (55) in Subotnik JCP 2011. This formula has some major problems @@ -277,26 +326,29 @@ def gamma_collapse(self, electronics: 'ElectronicModel_'=None) -> np.ndarray: ndof = self.model.ndof out = np.zeros(nst, dtype=np.float64) - def shifted_diagonal(X, k: int) -> np.ndarray: + def shifted_diagonal(X: np.ndarray, k: int) -> np.ndarray: out = np.zeros([nst, ndof]) for i in range(nst): - out[i,:] = np.real(X[:,k,k] - X[:,i,i]) + out[i, :] = np.real(X[:, k, k] - X[:, i, i]) return out ddR = shifted_diagonal(self.delR, self.state) ddP = shifted_diagonal(self.delP, self.state) - ddP = np.where(np.abs(ddP) == 0.0, 1e-10, ddP) - ddF = shifted_diagonal(np.einsum("pqx->xpq", electronics.force_matrix), self.state) - ddR = ddR * np.sign(ddR/ddP) + ddP = np.where(np.abs(ddP) == 0.0, ZERO_DIVISION_FLOOR, ddP) + assert electronics is not None + ddF = shifted_diagonal(np.einsum("pqx->xpq", electronics.force_matrix), + self.state) + ddR = ddR * np.sign(ddR / ddP) for i in range(nst): - out[i] = np.dot(ddF[i,:], ddR[i,:]) + out[i] = np.dot(ddF[i, :], ddR[i, :]) - out[self.state] = 0.0 # zero out self collapse for safety + out[self.state] = 0.0 # zero out self collapse for safety return 0.5 * out * self.dt - def surface_hopping(self, last_electronics: 'ElectronicModel_', this_electronics: 'ElectronicModel_') -> None: + def surface_hopping(self, last_electronics: ElectronicModel_, + this_electronics: ElectronicModel_) -> None: """Specialized version of surface_hopping that handles collapsing. Parameters @@ -306,7 +358,8 @@ def surface_hopping(self, last_electronics: 'ElectronicModel_', this_electronics this_electronics : ElectronicModel ElectronicStates from current step. """ - SurfaceHoppingMD.surface_hopping(self, last_electronics, this_electronics) + SurfaceHoppingMD.surface_hopping(self, last_electronics, + this_electronics) gamma = self.gamma_collapse(this_electronics) @@ -318,24 +371,29 @@ def surface_hopping(self, last_electronics: 'ElectronicModel_', this_electronics eta[i] = e if e < gamma[i]: - assert self.model.nstates == 2 + if self.model.nstates != 2: + raise NotImplementedError( + "A-FSSH collapse is only implemented for 2-state systems" + ) # reset the density matrix - self.rho[:,:] = 0.0 + self.rho[:, :] = 0.0 self.rho[self.state, self.state] = 1.0 # reset delR and delP - self.delR[:,:,:] = 0.0 - self.delP[:,:,:] = 0.0 - - self.tracer.record_event({ - "time" : self.time, - "removed" : i, - "gamma" : gamma[i], - "eta" : eta - }, event_type="collapse") - - def hop_update(self, hop_from, hop_to): + self.delR[:, :, :] = 0.0 + self.delP[:, :, :] = 0.0 + + self.tracer.record_event( + { + "time": self.time, + "removed": i, + "gamma": gamma[i], + "eta": eta + }, + event_type="collapse") + + def hop_update(self, hop_from: int, hop_to: int) -> None: """Shift delR and delP after hops. Parameters @@ -349,5 +407,5 @@ def hop_update(self, hop_from, hop_to): dPb = self.delP[:, hop_to, hop_to] for i in range(self.model.nstates): - self.delR[:,i,i] -= dRb - self.delP[:,i,i] -= dPb + self.delR[:, i, i] -= dRb + self.delP[:, i, i] -= dPb diff --git a/mudslide/batch.py b/mudslide/batch.py index 06291fd..4f8bd76 100644 --- a/mudslide/batch.py +++ b/mudslide/batch.py @@ -1,18 +1,21 @@ # -*- coding: utf-8 -*- """Code for running batches of trajectories.""" -from typing import Any, Iterator, Tuple +from __future__ import annotations + +from typing import Any, Iterator, Tuple, TYPE_CHECKING import logging import queue -import sys import numpy as np from numpy.typing import ArrayLike from .constants import boltzmann -from .exceptions import StillInteracting from .tracer import TraceManager +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + logger = logging.getLogger("mudslide") @@ -21,9 +24,9 @@ class TrajGenConst: Parameters ---------- - position : ArrayLike + position : np.ndarray Initial position. - velocity : ArrayLike + velocity : np.ndarray Initial velocity. initial_state : Any Initial state specification, should be either an integer or "ground". @@ -36,7 +39,11 @@ class TrajGenConst: Generator yielding tuples of (position, velocity, initial_state, params). """ - def __init__(self, position: ArrayLike, velocity: ArrayLike, initial_state: Any, seed: Any = None): + def __init__(self, + position: np.ndarray, + velocity: np.ndarray, + initial_state: Any, + seed: Any = None): self.position = position self.velocity = velocity self.initial_state = initial_state @@ -57,7 +64,9 @@ def __call__(self, nsamples: int) -> Iterator: """ seedseqs = self.seed_sequence.spawn(nsamples) for i in range(nsamples): - yield (self.position, self.velocity, self.initial_state, {"seed_sequence": seedseqs[i]}) + yield (self.position, self.velocity, self.initial_state, { + "seed_sequence": seedseqs[i] + }) class TrajGenNormal: @@ -65,13 +74,13 @@ class TrajGenNormal: Parameters ---------- - position : ArrayLike + position : np.ndarray Center of normal distribution for position. - velocity : ArrayLike + velocity : np.ndarray Center of normal distribution for velocity. initial_state : Any Initial state designation. - sigma : ArrayLike + sigma : np.ndarray Standard deviation of distribution. seed : Any, optional Initial seed to give to trajectory, by default None. @@ -80,10 +89,10 @@ class TrajGenNormal: """ def __init__(self, - position: ArrayLike, - velocity: ArrayLike, + position: np.ndarray, + velocity: np.ndarray, initial_state: Any, - sigma: ArrayLike, + sigma: np.ndarray, seed: Any = None, seed_traj: Any = None): self.position = position @@ -94,12 +103,12 @@ def __init__(self, self.seed_sequence = np.random.SeedSequence(seed) self.random_state = np.random.default_rng(seed_traj) - def vskip(self, vtest: float) -> bool: + def vskip(self, vtest: np.ndarray) -> bool: """Determine whether to skip given velocity. Parameters ---------- - vtest : float + vtest : np.ndarray Velocity to test. Returns @@ -107,7 +116,7 @@ def vskip(self, vtest: float) -> bool: bool True if velocity should be skipped, False otherwise. """ - return np.any(vtest < 0.0) + return bool(np.any(vtest < 0.0)) def __call__(self, nsamples: int) -> Iterator: """Generate initial conditions. @@ -137,9 +146,9 @@ class TrajGenBoltzmann: Parameters ---------- - position : ArrayLike + position : np.ndarray Initial positions. - mass : ArrayLike + mass : np.ndarray Array of particle masses. temperature : float Initial temperature to determine velocities. @@ -154,8 +163,8 @@ class TrajGenBoltzmann: """ def __init__(self, - position: ArrayLike, - mass: ArrayLike, + position: np.ndarray, + mass: np.ndarray, temperature: float, initial_state: Any, scale: bool = True, @@ -227,13 +236,17 @@ class BatchedTraj: | key | default | ---------------------|----------------------------| | t0 | 0.0 | - | nprocs | 1 | | seed | None (date) | """ - batch_only_options = [ "samples", "nprocs" ] + batch_only_options = ["samples"] - def __init__(self, model: 'ElectronicModel_', traj_gen: Any, trajectory_type: Any, tracemanager: Any = None, **inp: Any): + def __init__(self, + model: ElectronicModel_, + traj_gen: Any, + trajectory_type: Any, + tracemanager: Any = None, + **inp: Any): self.model = model if tracemanager is None: self.tracemanager = TraceManager() @@ -244,8 +257,7 @@ def __init__(self, model: 'ElectronicModel_', traj_gen: Any, trajectory_type: An self.batch_options = {} # statistical parameters - self.batch_options["samples"] = inp.get("samples", 2000) - self.batch_options["nprocs"] = inp.get("nprocs", 1) + self.batch_options["samples"] = inp.get("samples", 2000) # other options get copied over self.traj_options = {} @@ -261,21 +273,9 @@ def compute(self) -> TraceManager: TraceManager Object containing the results. """ - # for now, define four possible outcomes of the simulation nsamples = self.batch_options["samples"] - nprocs = self.batch_options["nprocs"] - - if nprocs > 1: - logger.warning('nprocs {} specified, but parallelism is not currently handled'.format(nprocs)) traj_queue: Any = queue.Queue() - results_queue: Any = queue.Queue() - - #traj_queue = mp.JoinableQueue() - #results_queue = mp.Queue() - #procs = [ mp.Process(target=traj_runner, args=(traj_queue, results_queue, )) for p in range(nprocs) ] - #for p in procs: - # p.start() for x0, v0, initial, params in self.traj_gen(nsamples): traj_input = self.traj_options @@ -292,33 +292,7 @@ def compute(self) -> TraceManager: while not traj_queue.empty(): traj = traj_queue.get() results = traj.simulate() - results_queue.put(results) - - #traj_queue.join() - #for p in procs: - # p.terminate() - - while not results_queue.empty(): - r = results_queue.get() - self.tracemanager.merge_tracer(r) + self.tracemanager.merge_tracer(results) self.tracemanager.outcomes = self.tracemanager.outcome() return self.tracemanager - - -def traj_runner(traj_queue: Any, results_queue: Any) -> None: - """Runner for computing jobs from queue. - - Parameters - ---------- - traj_queue : Any - Queue containing trajectories with a `simulate()` function. - results_queue : Any - Queue to store results of each call to `simulate()`. - """ - while True: - traj = traj_queue.get() - if traj is not None: - results = traj.simulate() - results_queue.put(results) - traj_queue.task_done() diff --git a/mudslide/collect.py b/mudslide/collect.py index 370db86..2df8055 100644 --- a/mudslide/collect.py +++ b/mudslide/collect.py @@ -1,25 +1,46 @@ # -*- coding: utf-8 -*- """CLI for collecting data from a trajectory""" +from __future__ import annotations + +from typing import Any + from .tracer import YAMLTrace -legend = { "t": "time", "k": "kinetic", "p": "potential", "e": "energy", "a": "active" } -legend_format = { "t": "12.8f", "k": "12.8f", "p": "12.8f", "e": "12.8f", "a": "12d" } +legend = { + "t": "time", + "k": "kinetic", + "p": "potential", + "e": "energy", + "a": "active" +} +legend_format = { + "t": "12.8f", + "k": "12.8f", + "p": "12.8f", + "e": "12.8f", + "a": "12d" +} + -def add_collect_parser(subparsers) -> None: +def add_collect_parser(subparsers: Any) -> None: """ Add a parser for the collect command to the subparsers :param subparsers: subparsers to add the parser to """ - parser = subparsers.add_parser('collect', help="Collect data from a trajectory") + parser = subparsers.add_parser('collect', + help="Collect data from a trajectory") parser.add_argument('logname', help="name of the trajectory to collect") - parser.add_argument('-k', '--keys', default="tkpea", + parser.add_argument('-k', + '--keys', + default="tkpea", help="keys to collect (default: %(default)s)") parser.set_defaults(func=collect_wrapper) -def collect(logname: str, keys: str="tkpea") -> None: + +def collect(logname: str, keys: str = "tkpea") -> None: """ Collect data from a trajectory @@ -31,9 +52,12 @@ def collect(logname: str, keys: str="tkpea") -> None: with open(logname + ".dat", "w", encoding="utf-8") as f: print("#", " ".join([f"{legend[k]:>12s}" for k in keys]), file=f) for snap in log: - print(" " + " ".join([f"{snap[legend[k]]:{legend_format[k]}}" for k in keys]), file=f) + print(" " + " ".join( + [f"{snap[legend[k]]:{legend_format[k]}}" for k in keys]), + file=f) + -def collect_wrapper(args) -> None: +def collect_wrapper(args: Any) -> None: """ Wrapper for collect """ diff --git a/mudslide/config.py b/mudslide/config.py index 5b29faf..08bc22b 100644 --- a/mudslide/config.py +++ b/mudslide/config.py @@ -5,6 +5,8 @@ ``$XDG_CONFIG_HOME/mudslide/config.yaml`` when that variable is set). """ +from __future__ import annotations + import os from pathlib import Path from typing import Any, Optional diff --git a/mudslide/ehrenfest.py b/mudslide/ehrenfest.py index 9359739..c043e02 100644 --- a/mudslide/ehrenfest.py +++ b/mudslide/ehrenfest.py @@ -1,13 +1,19 @@ # -*- coding: utf-8 -*- """Propagating Ehrenfest trajectories""" -from typing import Any +from __future__ import annotations + +from typing import Any, TYPE_CHECKING import numpy as np from numpy.typing import ArrayLike +from .trajectory_md import TrajectoryMD from .surface_hopping_md import SurfaceHoppingMD +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + class Ehrenfest(SurfaceHoppingMD): """Ehrenfest dynamics. @@ -16,7 +22,14 @@ class Ehrenfest(SurfaceHoppingMD): are treated quantum mechanically and the nuclear degrees of freedom are treated classically. The force on the nuclei is computed as the expectation value of the force operator over the electronic density matrix. + + Surface hopping options (hopping_probability, hopping_method, forced_hop_threshold, + zeta_list) are not used in Ehrenfest dynamics. """ + recognized_options = TrajectoryMD.recognized_options + [ + "electronic_integration", "max_electronic_dt", + "starting_electronic_intervals", "state0" + ] def __init__(self, *args: Any, **kwargs: Any): """Initialize Ehrenfest dynamics. @@ -28,13 +41,21 @@ def __init__(self, *args: Any, **kwargs: Any): **kwargs : Any Keyword arguments passed to SurfaceHoppingMD """ + kwargs.setdefault("outcome_type", "populations") SurfaceHoppingMD.__init__(self, *args, **kwargs) - def needed_gradients(self): - """Ehrenfest needs all forces since it sums over all states.""" + def needed_gradients(self) -> list[int] | None: + """Ehrenfest needs all forces since it sums over all states. + + Returns + ------- + None + None means all state gradients are needed. + """ return None - def potential_energy(self, electronics: 'ElectronicModel_' = None) -> np.floating: + def potential_energy(self, + electronics: ElectronicModel_ | None = None) -> float: """Calculate Ehrenfest potential energy. The potential energy is computed as the trace of the product of the @@ -52,9 +73,10 @@ def potential_energy(self, electronics: 'ElectronicModel_' = None) -> np.floatin """ if electronics is None: electronics = self.electronics + assert electronics is not None return np.real(np.trace(np.dot(self.rho, electronics.hamiltonian))) - def _force(self, electronics: 'ElectronicModel_' = None) -> ArrayLike: + def force(self, electronics: ElectronicModel_ | None = None) -> np.ndarray: """Calculate Ehrenfest force. The force is computed as the trace of the product of the density matrix @@ -72,14 +94,15 @@ def _force(self, electronics: 'ElectronicModel_' = None) -> ArrayLike: """ if electronics is None: electronics = self.electronics + assert electronics is not None out = np.zeros([electronics.ndof]) for i in range(electronics.nstates): - out += np.real(self.rho[i,i]) * electronics.force(i) + out += np.real(self.rho[i, i]) * electronics.force(i) return out - def surface_hopping(self, last_electronics: 'ElectronicModel_', - this_electronics: 'ElectronicModel_'): + def surface_hopping(self, last_electronics: ElectronicModel_, + this_electronics: ElectronicModel_) -> None: """Handle surface hopping. In Ehrenfest dynamics, surface hopping is not performed as the electronic diff --git a/mudslide/even_sampling.py b/mudslide/even_sampling.py index 557e2c4..c7631e5 100644 --- a/mudslide/even_sampling.py +++ b/mudslide/even_sampling.py @@ -4,17 +4,24 @@ This module implements trajectory surface hopping with even sampling of phase space. """ +from __future__ import annotations + from itertools import count -from typing import Optional, List, Any, Dict, Union +from typing import Optional, List, Any, Dict, Union, Iterator, TYPE_CHECKING import copy as cp import numpy as np from numpy.typing import ArrayLike +from .exceptions import ComputeError, ConfigurationError from .integration import quadrature from .surface_hopping_md import SurfaceHoppingMD +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + # pylint: disable=no-member + class SpawnStack: """Data structure to inform how new traces are spawned and weighted. @@ -36,7 +43,7 @@ def __init__(self, sample_stack: List, weight: float = 1.0): if sample_stack: weights = np.array([s["dw"] for s in sample_stack]) mw = np.ones(len(sample_stack)) - mw[1:] -= np.cumsum(weights[: len(weights) - 1]) + mw[1:] -= np.cumsum(weights[:len(weights) - 1]) self.marginal_weights = mw self.last_dw = weights[0] else: @@ -56,9 +63,10 @@ def zeta(self) -> float: """ return self.zeta_ - def next_zeta(self, - current_value: float, - random_state: Optional[np.random.RandomState] = None) -> float: + def next_zeta( + self, + current_value: float, + random_state: Optional[np.random.RandomState] = None) -> float: """Calculate the next zeta value. Parameters @@ -74,7 +82,8 @@ def next_zeta(self, Next zeta value for determining hops. """ if not self.sample_stack: - self.zeta_ = random_state.uniform() if random_state else np.random.uniform() + self.zeta_ = random_state.uniform( + ) if random_state else np.random.uniform() return self.zeta_ izeta = self.izeta @@ -84,13 +93,12 @@ def next_zeta(self, if izeta != self.izeta: # it means there was a hop, so update last_dw weights = np.array([s["dw"] for s in self.sample_stack]) - self.last_dw = np.sum(weights[self.izeta : izeta]) + self.last_dw = np.sum(weights[self.izeta:izeta]) # should actually probably be a merge operation self.last_stack = self.sample_stack[self.izeta] self.marginal_weight = (self.marginal_weights[izeta] - if izeta != len(self.sample_stack) - else 0.0) + if izeta != len(self.sample_stack) else 0.0) self.izeta = izeta if self.izeta < len(self.sample_stack): @@ -132,7 +140,8 @@ def spawn(self, reweight: float = 1.0) -> "SpawnStack": samp = self.last_stack dw = self.last_dw if dw == 0: - raise ValueError("What happened? A hop with no differential weight?") + raise ComputeError( + "What happened? A hop with no differential weight?") weight = self.base_weight * dw * reweight next_stack = samp["children"] else: @@ -164,15 +173,13 @@ def spawn_size(self) -> int: return samp["spawn_size"] return 1 - def append_layer( - self, - zetas: list, - dws: list, - stack=None, - node=None, - nodes=None, - adj_matrix=None - ): + def append_layer(self, + zetas: list, + dws: list, + stack: list | None = None, + node: int | None = None, + nodes: Iterator[int] | None = None, + adj_matrix: dict | None = None) -> None: """Append a layer to all leaves in the sample stack tree. A depth-first traversal of a sample_stack tree to append a layer to all leaves. @@ -203,7 +210,7 @@ def append_layer( stack = self.sample_stack if len(zetas) != len(dws): - raise ValueError("dimension of dws should be same as zetas") + raise ConfigurationError("dimension of dws should be same as zetas") l = len(stack) @@ -216,12 +223,18 @@ def append_layer( for i in range(l): adj_matrix[1].append(next(nodes)) else: + assert nodes is not None adj_matrix[node] = [] for i in range(l): adj_matrix[node].append(next(nodes)) if l == 0: for z, dw in zip(zetas, dws): - stack.append({"zeta": z, "dw": dw, "children": [], "spawn_size": 1}) + stack.append({ + "zeta": z, + "dw": dw, + "children": [], + "spawn_size": 1 + }) else: for i in range(l): self.append_layer( @@ -233,7 +246,7 @@ def append_layer( adj_matrix=adj_matrix, ) - def unpack(self, zeta_list, dw_list, stack=None, depth=0): + def unpack(self, zeta_list: list, dw_list: list, stack: list | dict | None = None, depth: int = 0) -> None: """Recursively unpack a sample stack. Fills in the zeta_list and dw_list with flattened lists of zeta values @@ -261,10 +274,12 @@ def unpack(self, zeta_list, dw_list, stack=None, depth=0): zeta_list.append((depth, stack["zeta"])) dw_list.append((depth, stack["dw"])) if stack["children"] != []: - self.unpack(zeta_list, dw_list, stack["children"], depth=depth + 1) - + self.unpack(zeta_list, + dw_list, + stack["children"], + depth=depth + 1) - def unravel(self): + def unravel(self) -> list: """Unravel the sample stack into points and weights. Calls unpack to recursively unpack a sample_stack and then unravels the list @@ -276,8 +291,8 @@ def unravel(self): List of tuples containing (points, weights) where points is a tuple of coordinates and weights is the product of weights at those coordinates. """ - zetas = [] - dws = [] + zetas: list[Any] = [] + dws: list[Any] = [] self.unpack(zeta_list=zetas, dw_list=dws) dim_list = [] @@ -285,8 +300,8 @@ def unravel(self): dim_list.append(tpl[0]) dim = max(dim_list) + 1 main_list = [] - coords = [] - weights = [] + coords: list[Any] = [] + weights: list[Any] = [] last_depth = 0 for i, tpl in enumerate(zetas): # When we recurse back up in depth, remove num_to_pop items from coords/weights. @@ -315,7 +330,9 @@ def unravel(self): main_list.append((tuple(coords), tuple(weights))) # Return product of weights - points_weights = [(points, np.prod(weights)) for (points, weights) in main_list] + points_weights = [ + (points, np.prod(weights)) for (points, weights) in main_list + ] return points_weights @@ -326,7 +343,7 @@ def from_quadrature( weight: float = 1.0, method: str = "gl", mcsamples: int = 1, - random_state: Optional[np.random.RandomState] = None, + random_state: Optional[np.random.RandomState | np.random.Generator] = None, ) -> "SpawnStack": """Create a SpawnStack from quadrature points. @@ -356,12 +373,15 @@ def from_quadrature( for ns in reversed(nsamples): leaves = cp.copy(forest) - samples, weights = quadrature(ns, 0.0, 1.0, method=method) # type: ignore + samples, weights = quadrature(ns, 0.0, 1.0, + method=method) # type: ignore spawnsize = spawn_size.pop(0) - forest = [ - {"zeta": s, "dw": dw, "children": cp.deepcopy(leaves), "spawn_size": spawnsize} - for s, dw in zip(samples, weights) - ] # type: ignore + forest = [{ + "zeta": s, + "dw": dw, + "children": cp.deepcopy(leaves), + "spawn_size": spawnsize + } for s, dw in zip(samples, weights)] # type: ignore[arg-type] return cls(forest, weight) @@ -387,9 +407,9 @@ class EvenSamplingTrajectory(SurfaceHoppingMD): """ recognized_options = (SurfaceHoppingMD.recognized_options + - ["spawn_stack", "quadrature", "mcsamples"]) + ["spawn_stack", "quadrature", "mcsamples"]) - def __init__(self, *args, **options): + def __init__(self, *args: Any, **options: Any) -> None: """Initialize the EvenSamplingTrajectory. Parameters @@ -415,14 +435,17 @@ def __init__(self, *args, **options): quadrature = options.get("quadrature", "gl") mcsamples = options.get("mcsamples", 1) self.spawn_stack = SpawnStack.from_quadrature( - ss, method=quadrature, mcsamples=mcsamples, random_state=self.random_state - ) + ss, + method=quadrature, + mcsamples=mcsamples, + random_state=self.random_state) else: self.spawn_stack = SpawnStack(ss) self.zeta = self.spawn_stack.next_zeta(0.0, self.random_state) - def clone(self, spawn_stack: Optional[Any] = None) -> "EvenSamplingTrajectory": + def clone(self, + spawn_stack: Optional[Any] = None) -> "EvenSamplingTrajectory": """Create a clone of the current trajectory. Parameters @@ -462,12 +485,12 @@ def clone(self, spawn_stack: Optional[Any] = None) -> "EvenSamplingTrajectory": ) return out - def hopper(self, gkndt: ArrayLike) -> List[Dict[str, Union[int, float]]]: + def hopper(self, gkndt: np.ndarray) -> List[Dict[str, Union[int, float]]]: """Determine whether and where to hop based on probabilities. Parameters ---------- - probs : ArrayLike + probs : np.ndarray Array of individual hopping probabilities. Returns @@ -482,36 +505,38 @@ def hopper(self, gkndt: ArrayLike) -> List[Dict[str, Union[int, float]]]: accumulated = 1 - (1 - accumulated) * np.exp(-gkdt) if accumulated > self.zeta: # then hop zeta = self.zeta - next_zeta = self.spawn_stack.next_zeta(accumulated, self.random_state) + next_zeta = self.spawn_stack.next_zeta(accumulated, + self.random_state) # where to hop hop_choice = gkndt / gkdt if self.spawn_stack.do_spawn(): nspawn = self.spawn_stack.spawn_size() spawn_weight = 1.0 / nspawn - targets = [ - { - "target": i, - "weight": hop_choice[i], - "zeta": zeta, - "prob": accumulated, - "stack": self.spawn_stack.spawn(spawn_weight * hop_choice[i]), - } - for i in range(self.model.nstates) - if i != self.state - for j in range(nspawn) - ] + targets = [{ + "target": + i, + "weight": + hop_choice[i], + "zeta": + zeta, + "prob": + accumulated, + "stack": + self.spawn_stack.spawn(spawn_weight * hop_choice[i]), + } for i in range(self.model.nstates) if i != self.state + for j in range(nspawn)] else: - target = self.random_state.choice(list(range(self.model.nstates)), p=hop_choice) - targets = [ - { - "target": target, - "weight": 1.0, - "zeta": zeta, - "prob": accumulated, - "stack": self.spawn_stack.spawn(), - } - ] + target = self.random_state.choice(list(range( + self.model.nstates)), + p=hop_choice) + targets = [{ + "target": target, + "weight": 1.0, + "zeta": zeta, + "prob": accumulated, + "stack": self.spawn_stack.spawn(), + }] # reset probabilities and random self.zeta = next_zeta @@ -524,7 +549,7 @@ def hopper(self, gkndt: ArrayLike) -> List[Dict[str, Union[int, float]]]: def hop_to_it(self, hop_targets: List[Dict[str, Any]], - electronics: 'ElectronicModel_' = None) -> None: + electronics: ElectronicModel_ | None = None) -> None: """Handle hopping by spawning new trajectories. This method spawns new trajectories instead of enacting hops directly. @@ -546,10 +571,13 @@ def hop_to_it(self, stack = hop["stack"] spawn = self.clone(stack) old_state = spawn.state - SurfaceHoppingMD.hop_to_it(spawn, [hop], electronics=spawn.electronics) + SurfaceHoppingMD.hop_to_it(spawn, [hop], + electronics=spawn.electronics) if spawn.state != old_state: - spawn.electronics.compute_additional(gradients=[spawn.state]) - spawn.time+= spawn.dt + assert spawn.electronics is not None + spawn.electronics.compute_additional( + gradients=[spawn.state]) + spawn.time += spawn.dt spawn.nsteps += 1 # trigger hop @@ -559,8 +587,11 @@ def hop_to_it(self, self.queue.put(spawn) self.update_weight(self.spawn_stack.weight()) else: - self.prob_cum = 0.0 + self.prob_cum = np.longdouble(0.0) old_state = self.state - SurfaceHoppingMD.hop_to_it(self, hop_targets, electronics=self.electronics) + SurfaceHoppingMD.hop_to_it(self, + hop_targets, + electronics=self.electronics) if self.state != old_state: + assert self.electronics is not None self.electronics.compute_additional(gradients=[self.state]) diff --git a/mudslide/exceptions.py b/mudslide/exceptions.py index f50319a..5b3d02c 100644 --- a/mudslide/exceptions.py +++ b/mudslide/exceptions.py @@ -1,10 +1,38 @@ # -*- coding: utf-8 -*- -"""Mudslide Exceptions""" +"""Mudslide exception hierarchy. +All mudslide-specific exceptions inherit from MudslideError, making it +possible to catch any library error with ``except MudslideError``. +""" -class StillInteracting(Exception): - """Exception class indicating that a simulation was terminated while still inside the - interaction region""" - def __init__(self) -> None: - Exception.__init__(self, "A simulation ended while still inside the interaction region.") +class MudslideError(Exception): + """Base class for all mudslide errors.""" + + +class ConfigurationError(MudslideError): + """Bad user input, invalid options, dimension mismatches, or unsupported features.""" + + +class ExternalCodeError(MudslideError): + """An external tool (e.g. Turbomole) crashed or is unavailable.""" + + +class ConvergenceError(ExternalCodeError): + """SCF or iterative solver convergence failure (recoverable subset of ExternalCodeError).""" + + +class ComputeError(MudslideError): + """Runtime numerical or algorithmic error during simulation.""" + + +class MissingDataError(MudslideError): + """Requested data (forces, couplings) was not computed.""" + + +class MissingForceError(MissingDataError): + """Requested force/gradient data was not computed.""" + + +class MissingCouplingError(MissingDataError): + """Requested derivative coupling data was not computed.""" diff --git a/mudslide/header.py b/mudslide/header.py index a3753a1..add31be 100644 --- a/mudslide/header.py +++ b/mudslide/header.py @@ -42,8 +42,8 @@ def print_header() -> None: border = "=" * width print(BANNER) - print(f" Nonadiabatic Molecular Dynamics") - print(f" with Trajectory Surface Hopping") + print(" Nonadiabatic Molecular Dynamics") + print(" with Trajectory Surface Hopping") print() print(border) for line in info_lines: diff --git a/mudslide/integration.py b/mudslide/integration.py index 2685ded..49e8efa 100644 --- a/mudslide/integration.py +++ b/mudslide/integration.py @@ -1,12 +1,19 @@ # -*- coding: utf-8 -*- """Quadrature implementations""" +from __future__ import annotations + from typing import Tuple import numpy as np from numpy.typing import ArrayLike -def clenshaw_curtis(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: +from .exceptions import ConfigurationError + + +def clenshaw_curtis(n: int, + a: float = -1.0, + b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: """ Computes the points and weights for a Clenshaw-Curtis integration from a to b. In other words, for the approximation to the integral @@ -44,7 +51,7 @@ def clenshaw_curtis(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, # sanity check imag_norm = np.linalg.norm(np.imag(wcc)) - assert imag_norm < 1e-14 + assert imag_norm < 10 * np.finfo(float).eps out = np.zeros(npoints) out[:nsegments] = np.real(wcc) @@ -55,7 +62,9 @@ def clenshaw_curtis(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, return xx, out -def midpoint(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: +def midpoint(n: int, + a: float = -1.0, + b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: """ Returns the points and weights for a midpoint integration from a to b. In other words, for the approximation to the integral @@ -70,7 +79,9 @@ def midpoint(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayL return points, weights -def trapezoid(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: +def trapezoid(n: int, + a: float = -1.0, + b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: """ Returns the points and weights for a trapezoid integration from a to b. In other words, for the approximation to the integral @@ -89,7 +100,9 @@ def trapezoid(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, Array return points, weights -def simpson(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: +def simpson(n: int, + a: float = -1.0, + b: float = 1.0) -> Tuple[ArrayLike, ArrayLike]: """ Returns the points and weights for a simpson rule integration from a to b. In other words, for the approximation to the integral @@ -99,8 +112,8 @@ def simpson(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayLi assert b > a and n > 1 if n % 2 != 1: - raise ValueError( - "Simpson's rule must be defined with an odd number of points") + raise ConfigurationError( + "Simpson's rule must be defined with an odd number of points") ninterval = n - 1 @@ -116,7 +129,9 @@ def simpson(n: int, a: float = -1.0, b: float = 1.0) -> Tuple[ArrayLike, ArrayLi return points, weights -def quadrature(n: int, a: float = -1.0, b: float = 1.0, +def quadrature(n: int, + a: float = -1.0, + b: float = 1.0, method: str = "gl") -> Tuple[ArrayLike, ArrayLike]: """ Returns a quadrature rule for the specified method and bounds @@ -135,5 +150,4 @@ def quadrature(n: int, a: float = -1.0, b: float = 1.0, return trapezoid(n, a, b) if method in ["simpson"]: return simpson(n, a, b) - else: - raise ValueError("Unrecognized quadrature choice") + raise ConfigurationError("Unrecognized quadrature choice") diff --git a/mudslide/io.py b/mudslide/io.py index b05a1a1..6c49cfc 100644 --- a/mudslide/io.py +++ b/mudslide/io.py @@ -1,9 +1,20 @@ # -*- coding: utf-8 -*- """Util functions""" +from __future__ import annotations + +from typing import Any, List, TextIO, TYPE_CHECKING + from .constants import bohr_to_angstrom -def write_xyz(coords, atom_types, file, comment=""): +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + + +def write_xyz(coords: Any, + atom_types: List[str], + file: TextIO, + comment: str = "") -> None: """Write coordinates to open file handle in XYZ format""" file.write(f"{len(coords)}\n") file.write(f"{comment}\n") @@ -11,9 +22,12 @@ def write_xyz(coords, atom_types, file, comment=""): acoords = coords * bohr_to_angstrom for atom, coord in zip(atom_types, acoords): atom = atom.capitalize() - file.write(f"{atom:3s} {coord[0]:20.12f} {coord[1]:20.12f} {coord[2]:20.12f}\n") + file.write( + f"{atom:3s} {coord[0]:20.12f} {coord[1]:20.12f} {coord[2]:20.12f}\n" + ) + -def write_trajectory_xyz(model, trace, filename, every=1): +def write_trajectory_xyz(model: ElectronicModel_, trace: Any, filename: str, every: int = 1) -> None: """Write trajectory to XYZ file""" natom, nd = model.dimensionality with open(filename, "w", encoding='utf-8') as file: @@ -22,5 +36,7 @@ def write_trajectory_xyz(model, trace, filename, every=1): continue desc = f"E={frame['energy']:g}; t={frame['time']:g}" coords = frame["position"].reshape(natom, nd) - atom_types = model.atom_types if model.atom_types is not None else ["X"] * natom + atom_types = model.atom_types if model.atom_types is not None else [ + "X" + ] * natom write_xyz(coords, atom_types, file, comment=desc) diff --git a/mudslide/math.py b/mudslide/math.py index ec963e7..44befa7 100644 --- a/mudslide/math.py +++ b/mudslide/math.py @@ -1,22 +1,25 @@ # -*- coding: utf-8 -*- """Math helper functions for molecular dynamics simulations.""" +from __future__ import annotations + from collections import deque import warnings import numpy as np from numpy.typing import ArrayLike +from .exceptions import ConfigurationError from .util import remove_center_of_mass_motion, remove_angular_momentum from .constants import boltzmann -def poisson_prob_scale(x: ArrayLike): +def poisson_prob_scale(x: np.ndarray) -> np.ndarray: """Compute (1 - exp(-x))/x for scaling Poisson probabilities. Parameters ---------- - x : ArrayLike + x : np.ndarray Input array of values. Returns @@ -27,27 +30,31 @@ def poisson_prob_scale(x: ArrayLike): """ with warnings.catch_warnings(): warnings.filterwarnings('ignore') - out = np.where(np.absolute(x) < 1e-3, - 1 - x/2 + x**2/6 - x**3/24, - -np.expm1(-x)/x) + out = np.where( + np.absolute(x) < 1e-3, 1 - x / 2 + x**2 / 6 - x**3 / 24, + -np.expm1(-x) / x) return out -def boltzmann_velocities(mass, temperature, remove_translation=True, - coords=None, remove_rotation=None, - scale=True, seed=None): +def boltzmann_velocities(mass: np.ndarray, + temperature: float, + remove_translation: bool = True, + coords: np.ndarray | None = None, + remove_rotation: bool | None = None, + scale: bool = True, + seed: int | None = None) -> np.ndarray: """Generate random velocities according to the Boltzmann distribution. Parameters ---------- - mass : ArrayLike + mass : np.ndarray Array of particle masses. temperature : float Target temperature for the velocity distribution. remove_translation : bool, optional Whether to remove center of mass translation from velocities. Default is True. - coords : ArrayLike or None, optional + coords : np.ndarray or None, optional Array of particle coordinates. Required if removing rotation. Default is None. remove_rotation : bool or None, optional @@ -74,12 +81,12 @@ def boltzmann_velocities(mass, temperature, remove_translation=True, if remove_rotation is None: remove_rotation = coords is not None elif remove_rotation and coords is None: - raise ValueError("Coordinates must be provided to remove rotation.") + raise ConfigurationError("Coordinates must be provided to remove rotation.") if remove_translation: v = p / mass v3 = v.reshape((-1, 3)) - M = mass.reshape((-1, 3))[:,0] # pylint: disable=invalid-name + M = mass.reshape((-1, 3))[:, 0] # pylint: disable=invalid-name v3_com = remove_center_of_mass_motion(v3, M) v = v3_com.reshape(mass.shape) @@ -87,7 +94,8 @@ def boltzmann_velocities(mass, temperature, remove_translation=True, if remove_rotation: v3 = v.reshape((-1, 3)) - M = mass.reshape((-1, 3))[:,0] + M = mass.reshape((-1, 3))[:, 0] + assert coords is not None coords3 = coords.reshape((-1, 3)) v3_am = remove_angular_momentum(v3, M, coords3) @@ -95,8 +103,8 @@ def boltzmann_velocities(mass, temperature, remove_translation=True, p = v * mass if scale: - avg_KE = 0.5 * np.sum(p**2 / mass) / mass.size # pylint: disable=invalid-name - kbT2 = 0.5 * kt # pylint: disable=invalid-name + avg_KE = 0.5 * np.sum(p**2 / mass) / mass.size # pylint: disable=invalid-name + kbT2 = 0.5 * kt # pylint: disable=invalid-name scal = np.sqrt(kbT2 / avg_KE) p *= scal @@ -128,7 +136,7 @@ class RollingAverage: The current sum of all values in the window. """ - def __init__(self, window_size=50): + def __init__(self, window_size: int = 50) -> None: """Initialize the RollingAverage calculator. Parameters @@ -138,10 +146,10 @@ def __init__(self, window_size=50): Default is 50. """ self.window_size = window_size - self.values = deque(maxlen=window_size) + self.values: deque[float] = deque(maxlen=window_size) self.sum = 0.0 - def insert(self, value): + def insert(self, value: float) -> None: """Add a new value to the rolling average. If the window is full, the oldest value is automatically removed. @@ -159,7 +167,7 @@ def insert(self, value): self.values.append(value) self.sum += value - def get_average(self): + def get_average(self) -> float: """Calculate and return the current rolling average. Returns @@ -171,7 +179,7 @@ def get_average(self): return 0.0 return self.sum / len(self.values) - def __len__(self): + def __len__(self) -> int: """Return the current number of values in the window. Returns diff --git a/mudslide/models/electronics.py b/mudslide/models/electronics.py index 907ee9c..e4458db 100644 --- a/mudslide/models/electronics.py +++ b/mudslide/models/electronics.py @@ -1,13 +1,17 @@ # -*- coding: utf-8 -*- """Handle storage and computation of electronic degrees of freedom""" +from __future__ import annotations + import copy as cp from typing import Tuple, Any, List import numpy as np from numpy.typing import ArrayLike -import math +from ..exceptions import (ComputeError, ConfigurationError, MissingCouplingError, + MissingForceError) + class ElectronicModel_: """Base class for handling electronic structure part of dynamics. @@ -31,20 +35,26 @@ class ElectronicModel_: The number of particles in the system. atom_types : List[str] The types of atoms in the system. - _position : "ArrayLike" + _position : np.ndarray The position of the system. - _hamiltonian : "ArrayLike" + _hamiltonian : np.ndarray The electronic Hamiltonian. - _force : "ArrayLike" + _force : np.ndarray The force on the system. - _forces_available : "ArrayLike" + _forces_available : np.ndarray A boolean array indicating which forces are available. - _derivative_coupling : "ArrayLike" + _derivative_coupling : np.ndarray The derivative coupling matrix. """ - def __init__(self, representation: str = "adiabatic", reference: Any = None, - nstates: int = 0, ndims: int = 1, nparticles: int = 1, ndof: int = None, - atom_types: List[str] = None): + + def __init__(self, + representation: str = "adiabatic", + reference: Any = None, + nstates: int = 0, + ndims: int = 1, + nparticles: int = 1, + ndof: int | None = None, + atom_types: List[str] | None = None): """Initialize the electronic model. Parameters @@ -75,16 +85,20 @@ def __init__(self, representation: str = "adiabatic", reference: Any = None, self._nstates = nstates self._representation = representation - self._position: "ArrayLike" + self._position: np.ndarray self._reference = reference - self._hamiltonian: "ArrayLike" - self._force: "ArrayLike" - self._forces_available: "ArrayLike" = np.zeros(self.nstates, dtype=bool) - self._derivative_coupling: "ArrayLike" - self._derivative_couplings_available: "ArrayLike" = np.zeros((self.nstates, self.nstates), dtype=bool) + self._hamiltonian: np.ndarray + self.energies: np.ndarray + self._force: np.ndarray + self._forces_available: np.ndarray = np.zeros(self.nstates, dtype=bool) + self._derivative_coupling: np.ndarray + self._derivative_couplings_available: np.ndarray = np.zeros( + (self.nstates, self.nstates), dtype=bool) + self._force_matrix: np.ndarray - self.atom_types: List[str] = atom_types + self.mass: np.ndarray + self.atom_types: List[str] | None = atom_types @property def ndof(self) -> int: @@ -142,51 +156,57 @@ def nstates(self) -> int: return self._nstates @property - def hamiltonian(self) -> "ArrayLike": + def hamiltonian(self) -> np.ndarray: """Get the electronic Hamiltonian. Returns ------- - ArrayLike + np.ndarray The electronic Hamiltonian matrix """ return self._hamiltonian - def force(self, state: int=0) -> "ArrayLike": + def force(self, state: int = 0) -> np.ndarray: """Return the force on a given state""" if not self._forces_available[state]: - raise ValueError("Force on state %d not available" % state) - return self._force[state,:] + raise MissingForceError(f"Force on state {state} not available") + return self._force[state, :] - def derivative_coupling(self, state1: int, state2: int) -> "ArrayLike": + def derivative_coupling(self, state1: int, state2: int) -> np.ndarray: """Return the derivative coupling between two states""" if not self._derivative_couplings_available[state1, state2]: - raise ValueError("Derivative coupling between states %d and %d not available" % (state1, state2)) + raise MissingCouplingError( + f"Derivative coupling between states {state1} and {state2} not available" + ) return self._derivative_coupling[state1, state2, :] @property - def derivative_coupling_tensor(self) -> "ArrayLike": + def derivative_coupling_tensor(self) -> np.ndarray: """Return the derivative coupling tensor""" if not np.all(self._derivative_couplings_available): false_indices = np.argwhere(~self._derivative_couplings_available) - print(f"Derivative couplings not available for state pairs: {false_indices.tolist()}") - print(f"Full availability matrix:\n{self._derivative_couplings_available}") - raise ValueError("All derivative couplings not available") + print( + f"Derivative couplings not available for state pairs: {false_indices.tolist()}" + ) + print( + f"Full availability matrix:\n{self._derivative_couplings_available}" + ) + raise MissingCouplingError("All derivative couplings not available") return self._derivative_coupling - def NAC_matrix(self, velocity: "ArrayLike") -> "ArrayLike": + def NAC_matrix(self, velocity: np.ndarray) -> np.ndarray: """Return the non-adiabatic coupling matrix for a given velocity vector """ if not np.all(self._derivative_couplings_available): - raise ValueError("NAC_matrix needs all derivative couplings") + raise MissingCouplingError("NAC_matrix needs all derivative couplings") return np.einsum("ijk,k->ij", self._derivative_coupling, velocity) @property - def force_matrix(self) -> "ArrayLike": + def force_matrix(self) -> np.ndarray: """Return the force matrix""" if not np.all(self._forces_available): - raise ValueError("Force matrix needs all forces") + raise MissingForceError("Force matrix needs all forces") return self._force_matrix def _needed_gradients(self, gradients: Any) -> List[int]: @@ -222,12 +242,18 @@ def _needed_couplings(self, couplings: Any) -> List[Tuple[int, int]]: Coupling pairs that still need to be computed. """ if couplings is None: - candidates = [(i, j) for i in range(self.nstates) for j in range(self.nstates)] + candidates = [ + (i, j) for i in range(self.nstates) for j in range(self.nstates) + ] else: candidates = couplings - return [(i, j) for (i, j) in candidates if not self._derivative_couplings_available[i, j]] + return [(i, j) + for (i, j) in candidates + if not self._derivative_couplings_available[i, j]] - def compute_additional(self, couplings: Any = None, gradients: Any = None) -> None: + def compute_additional(self, + couplings: Any = None, + gradients: Any = None) -> None: """Compute additional gradients/couplings at the current geometry. Checks what's already available and only computes what's missing. @@ -249,13 +275,19 @@ def compute_additional(self, couplings: Any = None, gradients: Any = None) -> No f"compute_additional not implemented; missing gradients={needed_g}, couplings={needed_c}" ) - def compute(self, X: "ArrayLike", couplings: Any = None, gradients: Any = None, reference: Any = None) -> None: + def compute(self, + X: np.ndarray, + couplings: Any = None, + gradients: Any = None, + reference: Any = None) -> None: """ - Central function for model objects. After the compute function exists, the following + Central function for model objects. After the compute function exits, the following data must be provided: - - self._hamiltonian -> n x n array containing electronic hamiltonian - - self.force -> n x ndof array containing the force on each diagonal - - self._derivative_coupling -> n x n x ndof array containing derivative couplings + - self._hamiltonian -> nstates x nstates array containing electronic hamiltonian + - self._force -> nstates x ndof array containing the force on each state + - self._derivative_coupling -> nstates x nstates x ndof array containing derivative couplings + - self._forces_available -> boolean array of length nstates + - self._derivative_couplings_available -> nstates x nstates boolean array Parameters ---------- @@ -273,7 +305,11 @@ def compute(self, X: "ArrayLike", couplings: Any = None, gradients: Any = None, """ raise NotImplementedError("ElectronicModel_ need a compute function") - def update(self, X: "ArrayLike", electronics: Any = None, couplings: Any = None, gradients: Any = None) -> 'ElectronicModel_': + def update(self, + X: np.ndarray, + electronics: Any = None, + couplings: Any = None, + gradients: Any = None) -> ElectronicModel_: """ Convenience function that copies the present object, updates the position, calls compute, and then returns the new object @@ -285,10 +321,13 @@ def update(self, X: "ArrayLike", electronics: Any = None, couplings: Any = None, else: reference = self._reference - out.compute(X, couplings=couplings, gradients=gradients, reference=reference) + out.compute(X, + couplings=couplings, + gradients=gradients, + reference=reference) return out - def clone(self): + def clone(self) -> ElectronicModel_: """Create a copy of the electronics object that owns its own resources, including disk @@ -302,7 +341,7 @@ def clone(self): """ return cp.deepcopy(self) - def as_dict(self): + def as_dict(self) -> dict[str, Any]: """Return a dictionary representation of the model""" out = { "nstates": self.nstates, @@ -313,7 +352,7 @@ def as_dict(self): "forces_available": self._forces_available.tolist() } - for key in [ "_derivative_coupling", "_force_matrix" ]: + for key in ["_derivative_coupling", "_force_matrix"]: if hasattr(self, key): out[key.lstrip('_')] = getattr(self, key).tolist() @@ -325,14 +364,22 @@ class DiabaticModel_(ElectronicModel_): To derive from DiabaticModel_, the following functions must be implemented: - - def V(self, X: ArrayLike) -> ArrayLike + - def V(self, X: np.ndarray) -> ArrayLike V(x) should return an ndarray of shape (nstates, nstates) - - def dV(self, X: ArrayLike) -> ArrayLike + - def dV(self, X: np.ndarray) -> ArrayLike dV(x) should return an ndarray of shape (nstates, nstates, ndof) """ - def __init__(self, representation: str = "adiabatic", reference: Any = None, - nstates:int = 0, ndof: int = 0): + #: Minimum energy gap threshold used when computing derivative couplings. + #: When the energy difference between two states is smaller than this value, + #: it is clamped to avoid division by near-zero. Override in subclasses if needed. + coupling_energy_threshold: float = 1.0e-14 + + def __init__(self, + representation: str = "adiabatic", + reference: Any = None, + nstates: int = 0, + ndof: int = 0): """Initialize the diabatic model. Parameters @@ -346,10 +393,17 @@ def __init__(self, representation: str = "adiabatic", reference: Any = None, ndof : int, optional Number of classical degrees of freedom, by default 0 """ - ElectronicModel_.__init__(self, representation=representation, reference=reference, - nstates=nstates, ndof=ndof) - - def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, reference: Any = None) -> None: + ElectronicModel_.__init__(self, + representation=representation, + reference=reference, + nstates=nstates, + ndof=ndof) + + def compute(self, + X: np.ndarray, + couplings: Any = None, + gradients: Any = None, + reference: Any = None) -> None: """Compute electronic properties at position X. Parameters @@ -365,15 +419,18 @@ def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, re """ self._position = X - self._reference, self._hamiltonian = self._compute_basis_states(self.V(X), reference=reference) + self._reference, self._hamiltonian = self._compute_basis_states( + self.V(X), reference=reference) dV = self.dV(X) - self._derivative_coupling = self._compute_derivative_coupling(self._reference, dV, np.diag(self._hamiltonian)) + self._derivative_coupling = self._compute_derivative_coupling( + self._reference, dV, np.diag(self._hamiltonian)) # Create new availability arrays to avoid sharing with shallow copies - self._derivative_couplings_available = np.zeros((self.nstates, self.nstates), dtype=bool) + self._derivative_couplings_available = np.zeros( + (self.nstates, self.nstates), dtype=bool) if couplings is None: - self._derivative_couplings_available[:,:] = True + self._derivative_couplings_available[:, :] = True else: for (i, j) in couplings: self._derivative_couplings_available[i, j] = True @@ -390,7 +447,9 @@ def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, re self._force_matrix = self._compute_force_matrix(dV, self._reference) - def compute_additional(self, couplings: Any = None, gradients: Any = None) -> None: + def compute_additional(self, + couplings: Any = None, + gradients: Any = None) -> None: """Compute additional gradients/couplings at the current geometry. Since diabatic models compute everything analytically, this just @@ -410,7 +469,10 @@ def compute_additional(self, couplings: Any = None, gradients: Any = None) -> No for (i, j) in needed_c: self._derivative_couplings_available[i, j] = True - def _compute_basis_states(self, V: ArrayLike, reference: Any = None) -> Tuple[ArrayLike, ArrayLike]: + def _compute_basis_states( + self, + V: np.ndarray, + reference: Any = None) -> Tuple[np.ndarray, np.ndarray]: """Computes coefficient matrix for basis states if a diabatic representation is chosen, no transformation takes place :param V: potential matrix @@ -421,18 +483,19 @@ def _compute_basis_states(self, V: ArrayLike, reference: Any = None) -> Tuple[Ar if reference is not None: try: for mo in range(self.nstates): - if (np.dot(coeff[:, mo], reference[:, mo]) < 0.0): + if np.dot(coeff[:, mo], reference[:, mo]) < 0.0: coeff[:, mo] *= -1.0 - except: - raise Exception("Failed to regularize new ElectronicStates from a reference object %s" % - (reference)) + except Exception as exc: + raise ComputeError( + f"Failed to regularize new ElectronicStates from a reference object {reference}" + ) from exc return (coeff, np.diag(energies)) elif self._representation == "diabatic": return (np.eye(self.nstates, dtype=np.float64), V) else: - raise Exception("Unrecognized run mode") + raise ConfigurationError("Unrecognized run mode") - def _compute_force(self, dV: ArrayLike, coeff: ArrayLike) -> ArrayLike: + def _compute_force(self, dV: np.ndarray, coeff: np.ndarray) -> np.ndarray: r""":math:`-\langle \phi_{\mbox{state}} | \nabla H | \phi_{\mbox{state}} \rangle`""" nst = self.nstates ndof = self.ndof @@ -444,23 +507,33 @@ def _compute_force(self, dV: ArrayLike, coeff: ArrayLike) -> ArrayLike: out[ist, :] += -np.einsum("i,ix->x", coeff[:, ist], half[:, ist, :]) return out - def _compute_force_matrix(self, dV: ArrayLike, coeff: ArrayLike) -> ArrayLike: + def _compute_force_matrix(self, dV: np.ndarray, + coeff: np.ndarray) -> np.ndarray: r"""returns :math:`F^\xi{ij} = \langle \phi_i | -\nabla_\xi H | \phi_j\rangle`""" out = -np.einsum("ip,xij,jq->pqx", coeff, dV, coeff) return out - def _compute_derivative_coupling(self, coeff: ArrayLike, dV: ArrayLike, energies: ArrayLike) -> ArrayLike: - r"""returns :math:`\phi_{i} | \nabla_\alpha \phi_{j} = d^\alpha_{ij}`""" + def _compute_derivative_coupling(self, coeff: np.ndarray, dV: np.ndarray, + energies: np.ndarray) -> np.ndarray: + r"""Compute derivative couplings :math:`d^\alpha_{ij} = \langle \phi_i | \nabla_\alpha \phi_j \rangle`. + + Uses the Hellmann-Feynman relation to compute derivative couplings from + the energy gap. When the energy gap between two states is smaller than + :attr:`coupling_energy_threshold`, the gap is clamped to avoid numerical + instability. + """ if self._representation == "diabatic": - return np.zeros([self.nstates, self.nstates, self.ndof], dtype=np.float64) + return np.zeros([self.nstates, self.nstates, self.ndof], + dtype=np.float64) out = np.einsum("ip,xij,jq->pqx", coeff, dV, coeff) + thresh = self.coupling_energy_threshold for j in range(self.nstates): for i in range(j): dE = energies[j] - energies[i] - if abs(dE) < 1.0e-10: - dE = np.copysign(1.0e-10, dE) + if abs(dE) < thresh: + dE = np.copysign(thresh, dE) out[i, j, :] /= dE out[j, i, :] /= -dE @@ -468,11 +541,15 @@ def _compute_derivative_coupling(self, coeff: ArrayLike, dV: ArrayLike, energies return out - def V(self, X: ArrayLike) -> ArrayLike: - raise NotImplementedError("Diabatic models must implement the function V") + def V(self, X: np.ndarray) -> np.ndarray: + """Return the diabatic potential matrix V(X).""" + raise NotImplementedError( + "Diabatic models must implement the function V") - def dV(self, X: ArrayLike) -> ArrayLike: - raise NotImplementedError("Diabatic models must implement the function dV") + def dV(self, X: np.ndarray) -> np.ndarray: + """Return the gradient of the diabatic potential matrix dV/dX.""" + raise NotImplementedError( + "Diabatic models must implement the function dV") class AdiabaticModel_(ElectronicModel_): @@ -482,8 +559,16 @@ class AdiabaticModel_(ElectronicModel_): many electronic states that are truncated to just a few. Sort of a truncated DiabaticModel_. """ - def __init__(self, representation: str = "adiabatic", reference: Any = None, - nstates:int = 0, ndof: int = 0): + #: Minimum energy gap threshold used when computing derivative couplings. + #: When the energy difference between two states is smaller than this value, + #: it is clamped to avoid division by near-zero. Override in subclasses if needed. + coupling_energy_threshold: float = 1.0e-14 + + def __init__(self, + representation: str = "adiabatic", + reference: Any = None, + nstates: int = 0, + ndof: int = 0): """Initialize the adiabatic model. Parameters @@ -503,11 +588,19 @@ def __init__(self, representation: str = "adiabatic", reference: Any = None, If representation is set to "diabatic" """ if representation == "diabatic": - raise Exception('Adiabatic models can only be run in adiabatic mode') - ElectronicModel_.__init__(self, representation=representation, reference=reference, - nstates=nstates, ndof=ndof) - - def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, reference: Any = None) -> None: + raise ConfigurationError( + 'Adiabatic models can only be run in adiabatic mode') + ElectronicModel_.__init__(self, + representation=representation, + reference=reference, + nstates=nstates, + ndof=ndof) + + def compute(self, + X: np.ndarray, + couplings: Any = None, + gradients: Any = None, + reference: Any = None) -> None: """Compute electronic properties at position X. Parameters @@ -523,15 +616,18 @@ def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, re """ self._position = X - self._reference, self._hamiltonian = self._compute_basis_states(self.V(X), reference=reference) + self._reference, self._hamiltonian = self._compute_basis_states( + self.V(X), reference=reference) dV = self.dV(X) - self._derivative_coupling = self._compute_derivative_coupling(self._reference, dV, np.diag(self._hamiltonian)) + self._derivative_coupling = self._compute_derivative_coupling( + self._reference, dV, np.diag(self._hamiltonian)) # Create new availability arrays to avoid sharing with shallow copies - self._derivative_couplings_available = np.zeros((self.nstates, self.nstates), dtype=bool) + self._derivative_couplings_available = np.zeros( + (self.nstates, self.nstates), dtype=bool) if couplings is None: - self._derivative_couplings_available[:,:] = True + self._derivative_couplings_available[:, :] = True else: for (i, j) in couplings: self._derivative_couplings_available[i, j] = True @@ -548,7 +644,9 @@ def compute(self, X: ArrayLike, couplings: Any = None, gradients: Any = None, re self._force_matrix = self._compute_force_matrix(dV, self._reference) - def compute_additional(self, couplings: Any = None, gradients: Any = None) -> None: + def compute_additional(self, + couplings: Any = None, + gradients: Any = None) -> None: """Compute additional gradients/couplings at the current geometry. Since adiabatic models compute everything analytically, this just @@ -568,7 +666,11 @@ def compute_additional(self, couplings: Any = None, gradients: Any = None) -> No for (i, j) in needed_c: self._derivative_couplings_available[i, j] = True - def update(self, X: ArrayLike, electronics: Any = None, couplings: Any = None, gradients: Any = None) -> 'AdiabaticModel_': + def update(self, + X: np.ndarray, + electronics: Any = None, + couplings: Any = None, + gradients: Any = None) -> AdiabaticModel_: """Update the model with new position and electronic information. Parameters @@ -591,10 +693,16 @@ def update(self, X: ArrayLike, electronics: Any = None, couplings: Any = None, g if electronics: self._reference = electronics._reference out._position = X - out.compute(X, couplings=couplings, gradients=gradients, reference=self._reference) + out.compute(X, + couplings=couplings, + gradients=gradients, + reference=self._reference) return out - def _compute_basis_states(self, V: ArrayLike, reference: Any = None) -> Tuple[ArrayLike, ArrayLike]: + def _compute_basis_states( + self, + V: np.ndarray, + reference: Any = None) -> Tuple[np.ndarray, np.ndarray]: """Computes coefficient matrix for basis states if a diabatic representation is chosen, no transformation takes place :param V: potential matrix @@ -609,19 +717,20 @@ def _compute_basis_states(self, V: ArrayLike, reference: Any = None) -> Tuple[Ar if reference is not None: try: for mo in range(self.nstates): - if (np.dot(coeff[:, mo], reference[:, mo]) < 0.0): + if np.dot(coeff[:, mo], reference[:, mo]) < 0.0: coeff[:, mo] *= -1.0 - except: - raise Exception("Failed to regularize new ElectronicStates from a reference object %s" % - (reference)) + except Exception as exc: + raise ComputeError( + f"Failed to regularize new ElectronicStates from a reference object {reference}" + ) from exc return coeff, np.diag(energies) elif self._representation == "diabatic": - raise Exception("Adiabatic models can only be run in adiabatic mode") - return None + raise ConfigurationError( + "Adiabatic models can only be run in adiabatic mode") else: - raise Exception("Unrecognized representation") + raise ConfigurationError("Unrecognized representation") - def _compute_force(self, dV: ArrayLike, coeff: ArrayLike) -> ArrayLike: + def _compute_force(self, dV: np.ndarray, coeff: np.ndarray) -> np.ndarray: r""":math:`-\langle \phi_{\mbox{state}} | \nabla H | \phi_{\mbox{state}} \rangle`""" nst = self.nstates ndof = self.ndof @@ -633,31 +742,45 @@ def _compute_force(self, dV: ArrayLike, coeff: ArrayLike) -> ArrayLike: out[ist, :] += -np.einsum("i,ix->x", coeff[:, ist], half[:, ist, :]) return out - def _compute_force_matrix(self, dV: ArrayLike, coeff: ArrayLike) -> ArrayLike: + def _compute_force_matrix(self, dV: np.ndarray, + coeff: np.ndarray) -> np.ndarray: r"""returns :math:`F^\xi{ij} = \langle \phi_i | -\nabla_\xi H | \phi_j\rangle`""" out = -np.einsum("ip,xij,jq->pqx", coeff, dV, coeff) return out - def _compute_derivative_coupling(self, coeff: ArrayLike, dV: ArrayLike, energies: ArrayLike) -> ArrayLike: - r"""returns :math:`\phi_{i} | \nabla_\alpha \phi_{j} = d^\alpha_{ij}`""" + def _compute_derivative_coupling(self, coeff: np.ndarray, dV: np.ndarray, + energies: np.ndarray) -> np.ndarray: + r"""Compute derivative couplings :math:`d^\alpha_{ij} = \langle \phi_i | \nabla_\alpha \phi_j \rangle`. + + Uses the Hellmann-Feynman relation to compute derivative couplings from + the energy gap. When the energy gap between two states is smaller than + :attr:`coupling_energy_threshold`, the gap is clamped to avoid numerical + instability. + """ if self._representation == "diabatic": - return np.zeros([self.nstates, self.nstates, self.ndof], dtype=np.float64) + return np.zeros([self.nstates, self.nstates, self.ndof], + dtype=np.float64) out = np.einsum("ip,xij,jq->pqx", coeff, dV, coeff) + thresh = self.coupling_energy_threshold for j in range(self.nstates): for i in range(j): dE = energies[j] - energies[i] - if abs(dE) < 1.0e-14: - dE = np.copysign(1.0e-14, dE) + if abs(dE) < thresh: + dE = np.copysign(thresh, dE) out[i, j, :] /= dE out[j, i, :] /= -dE return out - def V(self, X: ArrayLike) -> ArrayLike: - raise NotImplementedError("Diabatic models must implement the function V") + def V(self, X: np.ndarray) -> np.ndarray: + """Return the full electronic Hamiltonian matrix V(X).""" + raise NotImplementedError( + "Adiabatic models must implement the function V") - def dV(self, X: ArrayLike) -> ArrayLike: - raise NotImplementedError("Diabatic models must implement the function dV") + def dV(self, X: np.ndarray) -> np.ndarray: + """Return the gradient of the electronic Hamiltonian dV/dX.""" + raise NotImplementedError( + "Adiabatic models must implement the function dV") diff --git a/mudslide/models/harmonic_model.py b/mudslide/models/harmonic_model.py index 52b2f4d..62d122f 100644 --- a/mudslide/models/harmonic_model.py +++ b/mudslide/models/harmonic_model.py @@ -1,6 +1,8 @@ # -*- coding: utf-8 -*- """Harmonic model""" +from __future__ import annotations + from typing import Any, List import json import yaml @@ -8,15 +10,23 @@ import numpy as np from numpy.typing import ArrayLike +from ..exceptions import ConfigurationError from .electronics import ElectronicModel_ + class HarmonicModel(ElectronicModel_): r"""Adiabatic model for ground state dynamics """ - def __init__(self, x0: ArrayLike, E0: float, H0: ArrayLike, mass: ArrayLike, - atom_types: List[str] = None, ndims: int = 1, nparticles: int = 1): + def __init__(self, + x0: np.ndarray, + E0: float, + H0: np.ndarray, + mass: np.ndarray, + atom_types: List[str] | None = None, + ndims: int = 1, + nparticles: int = 1): """Constructor Args: @@ -28,22 +38,30 @@ def __init__(self, x0: ArrayLike, E0: float, H0: ArrayLike, mass: ArrayLike, ndims: Number of dimensions (e.g. 3 for 3D) nparticles: Number of particles (e.g. 1 for a single particle) """ - super().__init__(nstates=1, ndims=ndims, nparticles=nparticles, - atom_types=atom_types, representation="adiabatic") + super().__init__(nstates=1, + ndims=ndims, + nparticles=nparticles, + atom_types=atom_types, + representation="adiabatic") self.x0 = np.array(x0) self.E0 = E0 self.H0 = np.array(H0) self.mass = np.array(mass, dtype=np.float64).reshape(self._ndof) + self.energies = np.array([]) if self.H0.shape != (self._ndof, self._ndof): - raise ValueError("Incorrect shape of Hessian") + raise ConfigurationError("Incorrect shape of Hessian") if self.mass.shape != (self._ndof,): - raise ValueError("Incorrect shape of mass") + raise ConfigurationError("Incorrect shape of mass") - def compute(self, X: ArrayLike, gradients: Any = None, couplings: Any = None, reference: Any = None) -> None: + def compute(self, + X: np.ndarray, + couplings: Any = None, + gradients: Any = None, + reference: Any = None) -> None: """Compute and store the energies and gradients Args: @@ -65,7 +83,7 @@ def compute(self, X: ArrayLike, gradients: Any = None, couplings: Any = None, re self._forces_available = np.ones(self.nstates, dtype=bool) @classmethod - def from_dict(cls, model_dict: dict) -> "HarmonicModel": + def from_dict(cls, model_dict: dict) -> HarmonicModel: """Create a harmonic model from a dictionary Args: @@ -80,11 +98,16 @@ def from_dict(cls, model_dict: dict) -> "HarmonicModel": nparticles = len(atom_types) if atom_types is not None else 1 ndims = x0.size // nparticles - return cls(x0, E0, H0, mass, atom_types=atom_types, - ndims=ndims, nparticles=nparticles) + return cls(x0, + E0, + H0, + mass, + atom_types=atom_types, + ndims=ndims, + nparticles=nparticles) @classmethod - def from_file(cls, filename: str) -> "HarmonicModel": + def from_file(cls, filename: str) -> HarmonicModel: """Create a harmonic model from a file Args: @@ -98,7 +121,7 @@ def from_file(cls, filename: str) -> "HarmonicModel": elif filename.endswith(".yaml"): data = yaml.load(f, Loader=yaml.FullLoader) else: - raise ValueError("Unknown file format") + raise ConfigurationError("Unknown file format") return cls.from_dict(data) @@ -107,7 +130,12 @@ def to_file(self, filename: str) -> None: Use the ending on the filename to determine the format. """ - out = {"x0": self.x0.tolist(), "E0": self.E0, "H0": self.H0.tolist(), "mass": self.mass.tolist()} + out = { + "x0": self.x0.tolist(), + "E0": self.E0, + "H0": self.H0.tolist(), + "mass": self.mass.tolist() + } if self.atom_types is not None: out["atom_types"] = self.atom_types @@ -117,4 +145,4 @@ def to_file(self, filename: str) -> None: elif filename.endswith(".yaml"): yaml.dump(out, f) else: - raise ValueError("Unknown file format") + raise ConfigurationError("Unknown file format") diff --git a/mudslide/models/openmm_model.py b/mudslide/models/openmm_model.py index 11512dc..2a45312 100644 --- a/mudslide/models/openmm_model.py +++ b/mudslide/models/openmm_model.py @@ -1,11 +1,14 @@ # -*- coding: utf-8 -*- """ OpenMM interface for mudslide """ +from __future__ import annotations + from typing import Any import numpy as np from numpy.typing import ArrayLike +from ..exceptions import ConfigurationError from .electronics import ElectronicModel_ from ..constants import bohr_to_angstrom, amu_to_au, Hartree_to_kJmol from ..periodic_table import masses @@ -23,7 +26,7 @@ OPENMM_INSTALLED = False -def openmm_is_installed(): +def openmm_is_installed() -> bool: """Check if OpenMM is installed""" return OPENMM_INSTALLED @@ -31,8 +34,12 @@ def openmm_is_installed(): class OpenMM(ElectronicModel_): """OpenMM interface""" - def __init__(self, pdb, ff, system, platform_name: str = "Reference", - properties: dict = None): + def __init__(self, + pdb: Any, + ff: Any, + system: Any, + platform_name: str = "Reference", + properties: dict | None = None) -> None: """Initialize OpenMM interface""" super().__init__(representation="adiabatic", nstates=1, @@ -46,27 +53,36 @@ def __init__(self, pdb, ff, system, platform_name: str = "Reference", # check if there are constraints num_constraints = self._system.getNumConstraints() if num_constraints > 0: - raise ValueError( + raise ConfigurationError( "OpenMM system has constraints, which are not supported by mudslide. " + "Please remove constraints from the system and use the rigidWater=True option." ) # check if there are virtual sites - if any( self._system.isVirtualSite(i) for i in range(self._system.getNumParticles()) ): - raise ValueError( + if any( + self._system.isVirtualSite(i) + for i in range(self._system.getNumParticles())): + raise ConfigurationError( "OpenMM system has virtual sites, which are not supported by mudslide." ) # get charges try: - nonbonded = [ f for f in self._system.getForces() - if isinstance(f, openmm.NonbondedForce) ][0] - self._charges = np.array([ nonbonded.getParticleParameters(i)[0].value_in_unit( - openmm.unit.elementary_charge) for i in range(self.nparticles)]) + nonbonded = [ + f for f in self._system.getForces() + if isinstance(f, openmm.NonbondedForce) + ][0] + self._charges = np.array([ + nonbonded.getParticleParameters(i)[0].value_in_unit( + openmm.unit.elementary_charge) + for i in range(self.nparticles) + ]) except IndexError as exc: - raise ValueError("Can't find charges from OpenMM," - " probably because mudslide only understands Amber-like forces") from exc + raise ConfigurationError( + "Can't find charges from OpenMM," + " probably because mudslide only understands Amber-like forces" + ) from exc # make dummy integrator self._integrator = openmm.VerletIntegrator(0.001 * @@ -74,54 +90,47 @@ def __init__(self, pdb, ff, system, platform_name: str = "Reference", # create simulation platform = openmm.Platform.getPlatformByName(platform_name) - properties = properties if properties is not None else { } + properties = properties if properties is not None else {} self._simulation = openmm.app.Simulation(self._pdb.topology, self._system, self._integrator, platform, properties) - xyz = np.array( - self._convert_openmm_position_to_au( - pdb.getPositions())).reshape(-1) + xyz = np.array(self._convert_openmm_position_to_au( + pdb.getPositions())).reshape(-1) self._position = xyz self.atom_types = [ atom.element.symbol.lower() for atom in self._pdb.topology.atoms() ] - self.mass = np.array([ - masses[e] - for e in self.atom_types - for i in range(3) - ]) + self.mass = np.array( + [masses[e] for e in self.atom_types for i in range(3)]) self.mass *= amu_to_au self.energies = np.zeros([1], dtype=np.float64) - def _convert_au_position_to_openmm(self, xyz): + def _convert_au_position_to_openmm(self, xyz: np.ndarray) -> Any: """Convert position from bohr to nanometer using OpenMM units""" nm = openmm.unit.nanometer - return (xyz.reshape(-1, 3) * bohr_to_angstrom * - 0.1) * nm + return (xyz.reshape(-1, 3) * bohr_to_angstrom * 0.1) * nm - def _convert_openmm_position_to_au(self, xyz): + def _convert_openmm_position_to_au(self, xyz: Any) -> np.ndarray: """Convert position from nanometer to bohr using OpenMM units""" nm = openmm.unit.nanometer - return np.array( - xyz / nm).reshape(-1) * (10.0 / bohr_to_angstrom) + return np.array(xyz / nm).reshape(-1) * (10.0 / bohr_to_angstrom) - def _convert_openmm_force_to_au(self, force): + def _convert_openmm_force_to_au(self, force: Any) -> np.ndarray: kjmol = openmm.unit.kilojoules_per_mole nm = openmm.unit.nanometer return np.array(force * nm / kjmol).reshape( - 1, -1 - ) / Hartree_to_kJmol * 0.1 * bohr_to_angstrom + 1, -1) / Hartree_to_kJmol * 0.1 * bohr_to_angstrom - def _convert_energy_to_au(self, energy): + def _convert_energy_to_au(self, energy: Any) -> float: kjmol = openmm.unit.kilojoules_per_mole return energy / kjmol / Hartree_to_kJmol def compute(self, - X: ArrayLike, + X: np.ndarray, couplings: Any = None, gradients: Any = None, - reference: Any = None): + reference: Any = None) -> None: """Compute energy and forces""" self._position = X xyz = self._convert_au_position_to_openmm(X) @@ -131,9 +140,10 @@ def compute(self, getEnergy=True, getForces=True) - self.energies = np.array([self._convert_energy_to_au(state.getPotentialEnergy())]) + self.energies = np.array( + [self._convert_energy_to_au(state.getPotentialEnergy())]) self._hamiltonian = np.zeros([1, 1]) - self._hamiltonian[0,0] = self.energies[0] + self._hamiltonian[0, 0] = self.energies[0] self._force = self._convert_openmm_force_to_au( state.getForces(asNumpy=True)) diff --git a/mudslide/models/qmmm_model.py b/mudslide/models/qmmm_model.py index b1ef632..1d37a83 100644 --- a/mudslide/models/qmmm_model.py +++ b/mudslide/models/qmmm_model.py @@ -1,6 +1,8 @@ # -*- coding: utf-8 -*- """QM/MM model using turbomole and OpenMM""" +from __future__ import annotations + from typing import Any import numpy as np @@ -13,11 +15,15 @@ except ImportError: OPENMM_INSTALLED = False +from ..exceptions import ConfigurationError from .electronics import ElectronicModel_ + class QMMM(ElectronicModel_): """A QM/MM model""" - def __init__(self, qm_model, mm_model): + + def __init__(self, qm_model: ElectronicModel_, + mm_model: ElectronicModel_) -> None: self._qm_model = qm_model self._mm_model = mm_model @@ -26,7 +32,7 @@ def __init__(self, qm_model, mm_model): self._ndof_qm = qm_model.ndof self._ndof_mm = mm_model.ndof - self._ndof_qm - self._qm_atoms = list(range(self._ndof_qm//3)) + self._qm_atoms = list(range(self._ndof_qm // 3)) self._nqm = len(self._qm_atoms) # update position for qm atoms, just in case they are different @@ -35,19 +41,21 @@ def __init__(self, qm_model, mm_model): super().__init__(nstates=qm_model.nstates, ndof=mm_model.ndof) if not self.check_qm_and_mm_regions(self._qm_atoms): - raise ValueError("QM atoms must have the same elements in the QM and MM models.") + raise ConfigurationError( + "QM atoms must have the same elements in the QM and MM models.") self.remove_qm_interactions(self._qm_atoms) - def check_qm_and_mm_regions(self, qm_atoms): + def check_qm_and_mm_regions(self, qm_atoms: list[int]) -> bool: """Check to make sure that the atoms labelled QM have at least the same elements listed in the QM model and the MM model. """ qm_elements = self._qm_model.atom_types - qm_elements_in_mm = [ self._mm_model.atom_types[i] for i in qm_atoms ] + assert self._mm_model.atom_types is not None + qm_elements_in_mm = [self._mm_model.atom_types[i] for i in qm_atoms] return qm_elements == qm_elements_in_mm - def remove_qm_interactions(self, qm_atoms): + def remove_qm_interactions(self, qm_atoms: list[int]) -> None: """Remove bonded interactions of QM atoms from the MM model Args: @@ -56,41 +64,47 @@ def remove_qm_interactions(self, qm_atoms): num_bond_removed = 0 num_angl_removed = 0 num_tors_removed = 0 - for force in self._mm_model._system.getForces(): + for force in self._mm_model._system.getForces(): # type: ignore[attr-defined] if isinstance(force, openmm.HarmonicBondForce): for n in range(force.getNumBonds()): a, b, r, k = force.getBondParameters(n) if a in qm_atoms and b in qm_atoms: - force.setBondParameters(n, a, b, r, k*0.000) + force.setBondParameters(n, a, b, r, k * 0.000) num_bond_removed += 1 if (a in qm_atoms) != (b in qm_atoms): - raise ValueError("Bonded interactions between QM and MM regions not allowed." + raise ConfigurationError( + "Bonded interactions between QM and MM regions not allowed." f"Atoms {a:d} and {b:d} are bonded across regions.") elif isinstance(force, openmm.HarmonicAngleForce): for n in range(force.getNumAngles()): a, b, c, t, k = force.getAngleParameters(n) in_qm = [x in qm_atoms for x in [a, b, c]] if all(in_qm): - force.setAngleParameters(n, a, b, c, t, k*0.000) + force.setAngleParameters(n, a, b, c, t, k * 0.000) num_angl_removed += 1 elif any(in_qm): - raise ValueError("Bonded interactions between QM and MM regions not allowed." - f"Atoms {a:d}, {b:d}, and {c:d} are bonded across regions.") + raise ConfigurationError( + "Bonded interactions between QM and MM regions not allowed." + f"Atoms {a:d}, {b:d}, and {c:d} are bonded across regions." + ) elif isinstance(force, openmm.PeriodicTorsionForce): for n in range(force.getNumTorsions()): a, b, c, d, mult, phi, k = force.getTorsionParameters(n) in_qm = [x in qm_atoms for x in [a, b, c, d]] if all(in_qm): - force.setTorsionParameters(n, a, b, c, d, mult, phi, k*0.000) - num_tors_removed += 1 + force.setTorsionParameters(n, a, b, c, d, mult, phi, + k * 0.000) + num_tors_removed += 1 elif any(in_qm): - raise ValueError("Bonded interactions between QM and MM regions not allowed." - f"Atoms {a:d}, {b:d}, {c:d}, and {d:d} are bonded across regions.") + raise ConfigurationError( + "Bonded interactions between QM and MM regions not allowed." + f"Atoms {a:d}, {b:d}, {c:d}, and {d:d} are bonded across regions." + ) elif isinstance(force, openmm.NonbondedForce): for n in range(force.getNumParticles()): chg, sig, eps = force.getParticleParameters(n) if n in qm_atoms: - force.setParticleParameters(n, chg*0, sig, eps) + force.setParticleParameters(n, chg * 0, sig, eps) # add exceptions for qm-qm nonbonded interactions for n in range(force.getNumExceptions()): @@ -99,11 +113,19 @@ def remove_qm_interactions(self, qm_atoms): for j in range(i): force.addException(i, j, 0, 1, 0, replace=True) elif isinstance(force, openmm.CMMotionRemover): - raise ValueError("Cannot use CMMotionRemover in QM/MM model. Turn it off by setting removeCMMotion=False when preparing the System().") + raise ConfigurationError( + "Cannot use CMMotionRemover in QM/MM model. Turn it off by setting removeCMMotion=False when preparing the System()." + ) else: - raise ValueError(f"Force {force.__class__.__name__} not supported in QM/MM model.") - - def compute(self, X: ArrayLike, couplings: Any=None, gradients: Any=None, reference: Any=None) -> None: + raise ConfigurationError( + f"Force {force.__class__.__name__} not supported in QM/MM model." + ) + + def compute(self, + X: np.ndarray, + couplings: Any = None, + gradients: Any = None, + reference: Any = None) -> None: """Computes QM/MM energy by calling OpenMM and Turbomole and stitching together the results""" self._position = X qmxyz = X[:self._ndof_qm] @@ -111,8 +133,9 @@ def compute(self, X: ArrayLike, couplings: Any=None, gradients: Any=None, refere self._mm_model.compute(X) only_mm_xyz = X[self._ndof_qm:] - only_mm_charges = self._mm_model._charges[self._nqm:] - self._qm_model.control.add_point_charges(only_mm_xyz.reshape(-1,3), only_mm_charges) + only_mm_charges = self._mm_model._charges[self._nqm:] # type: ignore[attr-defined] + self._qm_model.control.add_point_charges( # type: ignore[attr-defined] + only_mm_xyz.reshape(-1, 3), only_mm_charges) self._qm_model.compute(qmxyz) self._hamiltonian = self._mm_model.hamiltonian + self._qm_model.hamiltonian @@ -121,14 +144,13 @@ def compute(self, X: ArrayLike, couplings: Any=None, gradients: Any=None, refere qmforce = self._qm_model._force self._force = np.zeros([self.nstates, self.ndof]) - self._force[:,:] = mmforce - self._force[:,:self._ndof_qm] += qmforce + self._force[:, :] = mmforce + self._force[:, :self._ndof_qm] += qmforce self._forces_available = self._qm_model._forces_available - a, b, qm_on_mm_force = self._qm_model.control.read_point_charge_gradients() - self._force[:,self._ndof_qm:] -= qm_on_mm_force.reshape(1,-1) - - def clone(self): - """Return a copy of the QMMM object""" - return QMMM(self._qm_model.clone(), self._mm_model.clone()) + a, b, qm_on_mm_force = self._qm_model.control.read_point_charge_gradients() # type: ignore[attr-defined] + self._force[:, self._ndof_qm:] -= qm_on_mm_force.reshape(1, -1) + def clone(self) -> QMMM: + """Return a copy of the QMMM object""" + return QMMM(self._qm_model.clone(), self._mm_model.clone()) diff --git a/mudslide/models/scattering_models.py b/mudslide/models/scattering_models.py index 27e1fa6..17a4220 100644 --- a/mudslide/models/scattering_models.py +++ b/mudslide/models/scattering_models.py @@ -1,6 +1,8 @@ # -*- coding: utf-8 -*- """Implementations of the one-dimensional two-state models Tully demonstrated FSSH on in Tully, J.C. J. Chem. Phys. 1990 93 1061.""" +from __future__ import annotations + import math from typing import Any @@ -67,6 +69,7 @@ class TullySimpleAvoidedCrossing(DiabaticModel_): V_{12} &= V_{21} = C e^{-D x^2} """ + def __init__(self, representation: str = "adiabatic", reference: Any = None, @@ -76,8 +79,11 @@ def __init__(self, d: float = 1.0, mass: float = 2000.0): """Constructor that defaults to the values reported in Tully's 1990 JCP""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=2, ndof=1) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=2, + ndof=1) self.A = a self.B = b @@ -85,18 +91,19 @@ def __init__(self, self.D = d self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, X: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" x = X[0] - v11 = float(np.copysign(self.A, x) * (1.0 - np.exp(-self.B * np.abs(x)))) + v11 = float( + np.copysign(self.A, x) * (1.0 - np.exp(-self.B * np.abs(x)))) v22 = -v11 v12 = float(self.C * np.exp(-self.D * x * x)) out = np.array([[v11, v12], [v12, v22]], dtype=np.float64) return out - def dV(self, x: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" - xx = np.array(x, dtype=np.float64) + xx = np.array(X, dtype=np.float64) v11 = self.A * self.B * np.exp(-self.B * abs(xx)) v22 = -v11 v12 = -2.0 * self.C * self.D * xx * np.exp(-self.D * xx * xx) @@ -112,6 +119,7 @@ class TullyDualAvoidedCrossing(DiabaticModel_): V_{22} &= -A e^{-Bx^2} + E_0 \\ V_{12} &= V_{21} = C e^{-D x^2} """ + def __init__(self, representation: str = "adiabatic", reference: Any = None, @@ -122,8 +130,11 @@ def __init__(self, e: float = 0.05, mass: float = 2000.0): """Constructor that defaults to the values reported in Tully's 1990 JCP""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=2, ndof=1) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=2, + ndof=1) self.A = a self.B = b self.C = c @@ -131,7 +142,7 @@ def __init__(self, self.E0 = e self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, X: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" x = X[0] v11 = 0.0 @@ -140,9 +151,9 @@ def V(self, X: ArrayLike) -> ArrayLike: out = np.array([[v11, v12], [v12, v22]], dtype=np.float64) return out - def dV(self, x: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" - xx = np.array(x, dtype=np.float64) + xx = np.array(X, dtype=np.float64) v11 = np.zeros_like(xx) v22 = 2.0 * self.A * self.B * xx * np.exp(-self.B * xx * xx) v12 = -2.0 * self.C * self.D * xx * np.exp(-self.D * xx * xx) @@ -171,14 +182,17 @@ def __init__(self, c: float = 0.90, mass: float = 2000.0): """Constructor that defaults to the values reported in Tully's 1990 JCP""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - ndof=1, nstates=2) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + ndof=1, + nstates=2) self.A = a self.B = b self.C = c self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, X: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" x = X[0] v11 = self.A @@ -191,9 +205,9 @@ def V(self, X: ArrayLike) -> ArrayLike: out = np.array([[v11, v12], [v12, v22]], dtype=np.float64) return out - def dV(self, x: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" - xx = np.array(x, dtype=np.float64) + xx = np.array(X, dtype=np.float64) v11 = np.zeros_like(xx) v22 = np.zeros_like(xx) v12 = self.B * self.C * np.exp(-self.C * np.abs(xx)) @@ -202,6 +216,8 @@ def dV(self, x: ArrayLike) -> ArrayLike: class SuperExchange(DiabaticModel_): + """Three-state super-exchange model from Prezhdo's GFSH paper.""" + def __init__(self, representation: str = "adiabatic", reference: Any = None, @@ -212,8 +228,11 @@ def __init__(self, v23: float = 0.01, mass: float = 2000.0): """Constructor defaults to Prezhdo paper on GFSH""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - ndof=1, nstates=3) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + ndof=1, + nstates=3) self.v11 = v11 self.v22 = v22 self.v33 = v33 @@ -221,27 +240,34 @@ def __init__(self, self.v23 = v23 self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, x: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" + x = X if np.ndim(x) != 0: x = x[0] v12 = self.v12 * np.exp(-0.5 * x * x) v23 = self.v23 * np.exp(-0.5 * x * x) - return np.array([[self.v11, v12, 0.0], [v12, self.v22, v23], [0.0, v23, self.v33]], dtype=np.float64) + return np.array( + [[self.v11, v12, 0.0], [v12, self.v22, v23], [0.0, v23, self.v33]], + dtype=np.float64) - def dV(self, x: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" + x = X if np.ndim(x) != 0: x = x[0] v12 = -x * self.v12 * np.exp(-0.5 * x * x) v23 = -x * self.v23 * np.exp(-0.5 * x * x) - out = np.array([[0.0, v12, 0.0], [v12, 0.0, v23], [0.0, v23, 0.0]], dtype=np.float64) + out = np.array([[0.0, v12, 0.0], [v12, 0.0, v23], [0.0, v23, 0.0]], + dtype=np.float64) return out.reshape([1, 3, 3]) class SubotnikModelX(DiabaticModel_): + """Three-state model X from Subotnik JPCA 2011 decoherence paper.""" + def __init__(self, representation: str = "adiabatic", reference: Any = None, @@ -251,16 +277,20 @@ def __init__(self, xp: float = 7.0, mass: float = 2000.0): """Constructor defaults to Subotnik JPCA 2011 paper on decoherence""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=3, ndof=1) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=3, + ndof=1) self.a = float(a) self.b = float(b) self.c = float(c) self.xp = float(xp) self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, x: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" + x = X if np.ndim(x) != 0: x = x[0] xx = np.array([x - self.xp, x, x + self.xp]) @@ -274,10 +304,12 @@ def V(self, x: ArrayLike) -> ArrayLike: v13 = ex[2] v23 = ex[0] - return np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], dtype=np.float64) + return np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], + dtype=np.float64) - def dV(self, x: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" + x = X if np.ndim(x) != 0: x = x[0] xx = np.array([x - self.xp, x, x + self.xp]) @@ -291,12 +323,14 @@ def dV(self, x: ArrayLike) -> ArrayLike: v13 = ex[2] v23 = ex[0] - out = np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], dtype=np.float64) + out = np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], + dtype=np.float64) return out.reshape([1, 3, 3]) class SubotnikModelS(DiabaticModel_): + """Three-state model S from Subotnik JPCA 2011 decoherence paper.""" def __init__(self, representation: str = "adiabatic", @@ -308,8 +342,11 @@ def __init__(self, xp: float = 7.0, mass: float = 2000.0): """Constructor defaults to Subotnik JPCA 2011 paper on decoherence""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=3, ndof=1) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=3, + ndof=1) self.a = float(a) self.b = float(b) self.c = float(c) @@ -317,8 +354,9 @@ def __init__(self, self.xp = float(xp) self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, x: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" + x = X if np.ndim(x) != 0: x = x[0] xx = np.array([x - self.xp, x, x + self.xp]) @@ -332,10 +370,12 @@ def V(self, x: ArrayLike) -> ArrayLike: v13 = ex[1] v23 = ex[1] - return np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], dtype=np.float64) + return np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], + dtype=np.float64) - def dV(self, x: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" + x = X if np.ndim(x) != 0: x = x[0] xx = np.array([x - self.xp, x, x + self.xp]) @@ -349,12 +389,14 @@ def dV(self, x: ArrayLike) -> ArrayLike: v13 = ex[1] v23 = ex[1] - out = np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], dtype=np.float64) + out = np.array([[v11, v12, v13], [v12, v22, v23], [v13, v23, v33]], + dtype=np.float64) return out.reshape([1, 3, 3]) class Subotnik2D(DiabaticModel_): + """Two-state, two-dimensional model from Subotnik JPCA 2011 decoherence paper.""" def __init__(self, representation: str = "adiabatic", @@ -368,8 +410,11 @@ def __init__(self, w: float = 2.0, mass: float = 2000.0): """Constructor defaults to Subotnik JPCA 2011 paper on decoherence""" - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=2, ndof=2) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=2, + ndof=2) self.a = float(a) self.b = float(b) self.c = float(c) @@ -379,9 +424,9 @@ def __init__(self, self.w = float(w) self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) - def V(self, r: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" - x, y = r[0], r[1] + x, y = X[0], X[1] v11 = -self.f * np.tanh(self.b * x) z = self.b * (x - 1.0) + self.w * np.cos(self.g * y + np.pi * 0.5) v22 = self.a * np.tanh(z) + 0.75 * self.a @@ -389,9 +434,9 @@ def V(self, r: ArrayLike) -> ArrayLike: return np.array([[v11, v12], [v12, v22]], dtype=np.float64) - def dV(self, r: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" - x, y = r[0], r[1] + x, y = X[0], X[1] z = self.b * (x - 1.0) + self.w * np.cos(self.g * y + np.pi * 0.5) @@ -414,6 +459,10 @@ def dV(self, r: ArrayLike) -> ArrayLike: class ShinMetiu(AdiabaticModel_): + """Shin-Metiu model for proton-coupled electron transfer.""" + + grid_boundary_offset: float = 1e-8 + def __init__(self, representation: str = "adiabatic", reference: Any = None, @@ -428,7 +477,11 @@ def __init__(self, box: Any = None): """Constructor defaults to classic Shin-Metiu as described in Gossel, Liacombe, Maitra JCP 2019""" - AdiabaticModel_.__init__(self, representation=representation, reference=reference, nstates=nstates, ndof=1) + AdiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=nstates, + ndof=1) self.L = L self.ion_left = -self.L * 0.5 @@ -442,24 +495,32 @@ def __init__(self, if box is None: box = L box_left, box_right = -0.5 * box, 0.5 * box - self.rr = np.linspace(box_left + 1e-12, box_right - 1e-12, nel, endpoint=True, dtype=np.float64) - - def soft_coulomb(self, r12: ArrayLike, gamma: np.floating) -> ArrayLike: + self.rr = np.linspace(box_left + self.grid_boundary_offset, + box_right - self.grid_boundary_offset, + nel, + endpoint=True, + dtype=np.float64) + + def soft_coulomb(self, r12: np.ndarray, gamma: float) -> np.ndarray: + """Soft Coulomb potential using the error function.""" abs_r12 = np.abs(r12) return erf(abs_r12 / gamma) / abs_r12 - def d_soft_coulomb(self, r12: ArrayLike, gamma: np.floating) -> ArrayLike: + def d_soft_coulomb(self, r12: np.ndarray, gamma: float) -> np.ndarray: + """Derivative of the soft Coulomb potential.""" abs_r12 = np.abs(r12) two_over_root_pi = 2.0 / np.sqrt(np.pi) out = r12 * erf(abs_r12/gamma) / (abs_r12**3) \ - two_over_root_pi * r12 * np.exp(-abs_r12**2/(gamma**2)) / (gamma * abs_r12 * abs_r12) return out - def V_nuc(self, R: ArrayLike) -> ArrayLike: + def V_nuc(self, R: np.ndarray) -> np.ndarray: + """Nuclear repulsion potential.""" v0 = 1.0 / np.abs(R - self.ion_left) + 1.0 / np.abs(R - self.ion_right) return v0 - def V_el(self, R: ArrayLike) -> ArrayLike: + def V_el(self, R: np.ndarray) -> np.ndarray: + """Electronic Hamiltonian matrix at nuclear position R.""" rr = self.rr v_en = -self.soft_coulomb(rr - R, self.Rf) @@ -470,52 +531,63 @@ def V_el(self, R: ArrayLike) -> ArrayLike: nr = len(rr) dr = rr[1] - rr[0] - T = (-0.5 / (self.m_el * dr * dr)) * (np.eye(nr, k=-1, dtype=np.float64) - 2.0 * np.eye(nr, dtype=np.float64) + - np.eye(nr, k=1, dtype=np.float64)) + T = (-0.5 / + (self.m_el * dr * dr)) * (np.eye(nr, k=-1, dtype=np.float64) - + 2.0 * np.eye(nr, dtype=np.float64) + + np.eye(nr, k=1, dtype=np.float64)) H = T + np.diag(vv) return H - def dV_nuc(self, R: ArrayLike) -> ArrayLike: + def dV_nuc(self, R: np.ndarray) -> np.ndarray: + """Derivative of nuclear repulsion potential.""" LmR = np.abs(0.5 * self.L - R) LpR = np.abs(0.5 * self.L + R) dv0 = LmR / np.abs(LmR**3) - LpR / np.abs(LpR**3) return dv0 - def dV_el(self, R: ArrayLike) -> ArrayLike: + def dV_el(self, R: np.ndarray) -> np.ndarray: + """Derivative of electronic Hamiltonian matrix.""" rr = self.rr rR = R - rr dvv = self.d_soft_coulomb(rR, self.Rf) return np.diag(dvv) - def V(self, R: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: """:math:`V(x)`""" - return self.V_el(R) + self.V_nuc(R) + return self.V_el(X) + self.V_nuc(X) - def dV(self, R: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: """:math:`\\nabla V(x)`""" - return (self.dV_el(R) + self.dV_nuc(R)).reshape([1, len(self.rr), len(self.rr)]) + return (self.dV_el(X) + self.dV_nuc(X)).reshape( + [1, len(self.rr), len(self.rr)]) class LinearVibronic(DiabaticModel_): + """Linear vibronic coupling model for conical intersections.""" def __init__( self, representation: str = "adiabatic", reference: Any = None, - mass: float = [243.6078782, 134.6412667, 99.93022402, 66.33593369, 3475.98736], + mass: np.ndarray | None = None, E1: float = 8.5037, E2: float = 9.4523, lamb: float = 0.3289, r0sqrtw5mh: float = 4.35, - om: float = np.array([0.1117, 0.2021, 0.2723, 0.4102]), - k1: float = np.array([-0.0456, 0.0399, -0.2139, -0.0864]), - k2: float = np.array([-0.0393, 0.0463, 0.2877, -0.1352]), - An: float = np.array([1.4823, -0.2191, 0.0525, -0.0118]), + om: np.ndarray = np.array([0.1117, 0.2021, 0.2723, 0.4102]), + k1: np.ndarray = np.array([-0.0456, 0.0399, -0.2139, -0.0864]), + k2: np.ndarray = np.array([-0.0393, 0.0463, 0.2877, -0.1352]), + An: np.ndarray = np.array([1.4823, -0.2191, 0.0525, -0.0118]), ): - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=2, ndof=5) + if mass is None: + mass = np.array([243.6078782, 134.6412667, 99.93022402, 66.33593369, 3475.98736]) + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=2, + ndof=5) self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) self.E1 = float(E1 / eVtoHartree) self.E2 = float(E2 / eVtoHartree) @@ -526,7 +598,7 @@ def __init__( self.k2 = k2 / eVtoHartree self.An = An / eVtoHartree - def V(self, X: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: w0 = 0 theta = X[4] q5 = np.zeros(4) @@ -550,9 +622,9 @@ def V(self, X: ArrayLike) -> ArrayLike: out = np.array([[w11, w12], [w21, w22]], dtype=np.float64) return out - def dV(self, X: ArrayLike) -> ArrayLike: - w0 = 0 - w12 = 0 + def dV(self, X: np.ndarray) -> np.ndarray: + w0: float = 0.0 + w12: float = 0.0 w21 = w12 theta = X[4] q5 = np.zeros(4) @@ -571,7 +643,8 @@ def dV(self, X: ArrayLike) -> ArrayLike: w22 = 0 for i in range(4): - q5[i] = self.An[i] * 2 * (i + 1) * (math.sin((i + 1) * theta) * (math.cos((i + 1) * theta))) + q5[i] = self.An[i] * 2 * (i + 1) * (math.sin( + (i + 1) * theta) * (math.cos((i + 1) * theta))) w11 = np.sum(q5) w22 = w11 @@ -582,21 +655,25 @@ def dV(self, X: ArrayLike) -> ArrayLike: return out.reshape([5, 2, 2]) + class SubotnikModelW(DiabaticModel_): - def __init__( - self, - representation: str = "adiabatic", - reference: Any = None, - mass: np.float64 = 2000.0, - nstates: int = 8, - eps: np.float64 = 0.1 - ): - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=nstates, ndof=1) + """N-state scattering model W from Subotnik and coworkers.""" + + def __init__(self, + representation: str = "adiabatic", + reference: Any = None, + mass: float = 2000.0, + nstates: int = 8, + eps: float = 0.1): + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=nstates, + ndof=1) self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) self.eps = eps - def V(self, X: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: N = self.nstates m = np.arange(0, N) + 1 @@ -609,7 +686,7 @@ def V(self, X: ArrayLike) -> ArrayLike: np.fill_diagonal(out, diag) return out - def dV(self, X: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: N = self.nstates m = np.arange(0, N) + 1 @@ -622,26 +699,30 @@ def dV(self, X: ArrayLike) -> ArrayLike: np.fill_diagonal(out[0], diag) return out + class SubotnikModelZ(DiabaticModel_): - def __init__( - self, - representation: str = "adiabatic", - reference: Any = None, - mass: np.float64 = 2000.0, - nstates: int = 8, - eps: np.float64 = 0.1 - ): - DiabaticModel_.__init__(self, representation=representation, reference=reference, - nstates=nstates, ndof=1) + """N-state scattering model Z from Subotnik and coworkers.""" + + def __init__(self, + representation: str = "adiabatic", + reference: Any = None, + mass: float = 2000.0, + nstates: int = 8, + eps: float = 0.1): + DiabaticModel_.__init__(self, + representation=representation, + reference=reference, + nstates=nstates, + ndof=1) self.mass = np.array(mass, dtype=np.float64).reshape(self.ndof) self.eps = eps - def V(self, X: ArrayLike) -> ArrayLike: + def V(self, X: np.ndarray) -> np.ndarray: N = self.nstates v = 0.1 / np.sqrt(N) - m1 = np.arange(0, N//2) + 1 - m2 = np.arange(N//2, N) + 1 + m1 = np.arange(0, N // 2) + 1 + m2 = np.arange(N // 2, N) + 1 diag = np.zeros(N, dtype=np.float64) d1 = X[0] + (m1 - 1) * self.eps @@ -654,12 +735,12 @@ def V(self, X: ArrayLike) -> ArrayLike: np.fill_diagonal(out, diag) return out - def dV(self, X: ArrayLike) -> ArrayLike: + def dV(self, X: np.ndarray) -> np.ndarray: N = self.nstates v = 0.1 / np.sqrt(N) - m1 = np.arange(0, N//2) + 1 - m2 = np.arange(N//2, N) + 1 + m1 = np.arange(0, N // 2) + 1 + m2 = np.arange(N // 2, N) + 1 diag = np.zeros(N, dtype=np.float64) d1 = X[0] + (m1 - 1) * self.eps diff --git a/mudslide/models/turbomole_model.py b/mudslide/models/turbomole_model.py index 4d08857..ce5c25d 100644 --- a/mudslide/models/turbomole_model.py +++ b/mudslide/models/turbomole_model.py @@ -1,10 +1,11 @@ # -*- coding: utf-8 -*- """Implementations of the interface between turbomole and mudslide.""" +from __future__ import annotations + import glob import io import os -import sys import shlex import shutil import re @@ -24,6 +25,7 @@ from ..constants import amu_to_au from ..periodic_table import masses from ..config import get_config +from ..exceptions import ConfigurationError, ConvergenceError, ExternalCodeError from .electronics import ElectronicModel_ @@ -37,7 +39,7 @@ def _resolve_command_prefix(explicit: Optional[str]) -> Optional[str]: return get_config("turbomole.command_prefix") -def turbomole_is_installed(): +def turbomole_is_installed() -> bool: """ Check if turbomole is installed by checking for environment variable TURBODIR and checking that the scripts and bin directories are available. @@ -52,7 +54,8 @@ def turbomole_is_installed(): return has_turbodir and has_scripts and has_bin -def turbomole_is_installed_or_prefixed(): + +def turbomole_is_installed_or_prefixed() -> bool: """ Check if turbomole is installed or a command prefix is set. :return: True if turbomole is installed or a command prefix is set, False otherwise @@ -85,7 +88,7 @@ def _verify_scf(module: str, data_dict: dict) -> None: warnings.warn(f"Convergence information not found for {module}") return if not converged: - raise RuntimeError(f"{module} SCF did not converge") + raise ConvergenceError(f"{module} SCF did not converge") def _verify_response(module: str, data_dict: dict) -> None: @@ -99,7 +102,7 @@ def _verify_response(module: str, data_dict: dict) -> None: ) continue if not converged: - raise RuntimeError(f"{module} {solver} did not converge") + raise ConvergenceError(f"{module} {solver} did not converge") class TurboControl: @@ -110,9 +113,9 @@ class TurboControl: control_file = "control" def __init__(self, - control_file=None, - workdir=None, - command_prefix: Optional[str] = None): + control_file: Optional[str] = None, + workdir: Optional[str] = None, + command_prefix: Optional[str] = None) -> None: self.command_prefix = _resolve_command_prefix(command_prefix) # workdir is directory of control file if control_file is not None: @@ -120,16 +123,16 @@ def __init__(self, elif workdir is not None: self.workdir = os.path.abspath(workdir) else: - raise ValueError("Must provide either control_file or workdir") + raise ConfigurationError("Must provide either control_file or workdir") self.control_file = control_file or "control" # make sure control file exists if not os.path.exists(os.path.join(self.workdir, self.control_file)): - raise RuntimeError( + raise ConfigurationError( f"control file not found in working directory {self.workdir:s}") # list of data groups and which file they are in, used to avoid rerunning sdg too much - self.dg_in_file = {} + self.dg_in_file: Dict[str, str] = {} def _build_command(self, cmd: list, cwd: Optional[str] = None) -> tuple: """Build command with prefix and working directory handling. @@ -150,12 +153,12 @@ def _build_command(self, cmd: list, cwd: Optional[str] = None) -> tuple: return ["sh", "-c", shell_cmd], None return cmd, cwd - def check_turbomole_is_installed(self): + def check_turbomole_is_installed(self) -> None: """Check that turbomole is installed, raise exception if not""" if not turbomole_is_installed(): - raise RuntimeError("Turbomole is not installed") + raise ExternalCodeError("Turbomole is not installed") - def where_is_dg(self, dg, absolute_path=False): + def where_is_dg(self, dg: str, absolute_path: bool = False) -> str: """Find which file a data group is in""" loc = self.dg_in_file[dg] if dg in self.dg_in_file \ else self.sdg(dg, show_filename_only=True) @@ -165,14 +168,14 @@ def where_is_dg(self, dg, absolute_path=False): def sdg( self, - dg, - file=None, - show_keyword=False, - show_body=False, - show_filename_only=False, - discard_comments=True, - quiet=False, - ): + dg: str, + file: Optional[str] = None, + show_keyword: bool = False, + show_body: bool = False, + show_filename_only: bool = False, + discard_comments: bool = True, + quiet: bool = False, + ) -> str: """Convenience function to run show data group (sdg) on a control""" sdg_command = "sdg" if file is not None: @@ -199,7 +202,7 @@ def sdg( check=False) return result.stdout.rstrip() - def adg(self, dg, data, newline=False): + def adg(self, dg: str, data: Union[str, list], newline: bool = False) -> None: """Convenience function to run add data group (adg) on a control""" if not isinstance(data, list): data = [data] @@ -216,9 +219,9 @@ def adg(self, dg, data, newline=False): check=True) # check that the command ran successfully if "abnormal" in result.stderr: - raise RuntimeError(f"Call to adg ended abnormally: {result.stderr}") + raise ExternalCodeError(f"Call to adg ended abnormally: {result.stderr}") - def cpc(self, dest): + def cpc(self, dest: str) -> None: """Copy the control file and other files to a new directory""" full_cmd, effective_cwd = self._build_command(["cpc", dest], cwd=self.workdir) @@ -237,7 +240,7 @@ def cpc(self, dest): shutil.copy(os.path.join(os.path.abspath(self.workdir), f), dest) - def use_weight_derivatives(self, use=True): + def use_weight_derivatives(self, use: bool = True) -> None: """Check if weight derivatives are used in the control file""" sdg_dft = self.sdg("dft", show_body=True) if use: # make sure weight derivatives turned on @@ -276,12 +279,12 @@ def run_single(self, else: print(output.stdout) if "abnormal" in output.stderr: - raise RuntimeError(f"Call to {module} ended abnormally") + raise ExternalCodeError(f"Call to {module} ended abnormally") data_dict = turboparse.parse_turbo(io.StringIO(output.stdout)) verify_module_output(module, data_dict, output.stderr) return data_dict - def read_coords(self): + def read_coords(self) -> tuple[list[str], np.ndarray]: """Read the coordinates from the control file :return: (symbols, X) where symbols is a list of element symbols @@ -300,14 +303,14 @@ def read_coords(self): X = np.array(coords, dtype=np.float64) return symbols, X - def get_masses(self, symbols): + def get_masses(self, symbols: list[str]) -> np.ndarray: """Get the masses of the atoms in the system""" atomic_masses = np.array([masses[s] for s in symbols for _ in range(3)], dtype=np.float64) atomic_masses *= amu_to_au return atomic_masses - def read_hessian(self): + def read_hessian(self) -> np.ndarray: """ Projected Hessian has a structure of $hessian (projected) @@ -327,7 +330,7 @@ def read_hessian(self): H = np.array(hessian, dtype=np.float64).reshape(ndof, ndof) return H - def add_point_charges(self, coords: ArrayLike, charges: ArrayLike): + def add_point_charges(self, coords: np.ndarray, charges: np.ndarray) -> None: """Add point charges to the control file point_charges data group has the structure: @@ -352,12 +355,12 @@ def add_point_charges(self, coords: ArrayLike, charges: ArrayLike): # make sure point charge gradients are requested drvopt = self.sdg("drvopt", show_body=True, show_keyword=False) if "point charges" not in drvopt: - drvopt = drvopt.rstrip().split("\n") - drvopt += [" point charges"] - self.adg("drvopt", drvopt, newline=True) + drvopt_lines = drvopt.rstrip().split("\n") + drvopt_lines += [" point charges"] + self.adg("drvopt", drvopt_lines, newline=True) self.adg("point_charge_gradients", [f"file={self.pcgrad_file}"]) - def read_point_charge_gradients(self): + def read_point_charge_gradients(self) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Read point charges and gradients from the control file point charges in dg $point_charges @@ -398,14 +401,14 @@ class TMModel(ElectronicModel_): def __init__( self, - states: ArrayLike, + states: np.ndarray, run_turbomole_dir: str = ".", workdir_stem: str = "run_turbomole", representation: str = "adiabatic", reference: Any = None, expert: bool = False, # when False, update turbomole parameters for NAMD - turbomole_modules: Dict = None, + turbomole_modules: Dict | None = None, command_prefix: Optional[str] = None, keep_output: int = 0, exopt_all: bool = False): @@ -454,7 +457,7 @@ def __init__( if not self.command_prefix: if not turbomole_is_installed(): - raise RuntimeError("Turbomole is not installed") + raise ExternalCodeError("Turbomole is not installed") if turbomole_modules is None: # always need energy and gradients @@ -468,14 +471,14 @@ def __init__( if not all( shutil.which(x) is not None for x in self.turbomole_modules.values()): - raise RuntimeError("Turbomole modules not found") + raise ExternalCodeError("Turbomole modules not found") if not self.expert: self.apply_suggested_parameters() self.exopt_all = exopt_all - def apply_suggested_parameters(self): + def apply_suggested_parameters(self) -> None: """ Apply suggested parameters for Turbomole to work well with NAMD This function will update the control file to ensure that Turbomole @@ -502,7 +505,7 @@ def apply_suggested_parameters(self): # probably force phaser on as well - def update_coords(self, X): + def update_coords(self, X: np.ndarray) -> None: """ Update the coordinates in the control file :param X: numpy array of shape (n_atoms * 3) with coordinates in Bohr @@ -523,11 +526,12 @@ def update_coords(self, X): # Reached end of file without finding $coord. if line == "": - raise ValueError(f"$coord entry not found in file: {coord_path}!") + raise ConfigurationError(f"$coord entry not found in file: {coord_path}!") coordline += 1 for i, coord_list in enumerate(X): x, y, z = coord_list[:3] + assert self.atom_types is not None s = self.atom_types[i] lines[coordline] = f"{x:26.16e} {y:28.16e} {z:28.16e} {s:>7}\n" coordline += 1 @@ -678,7 +682,7 @@ def _manage_output(self, outpath: Path) -> None: (parent / f"tm.{n}").unlink(missing_ok=True) def compute(self, - X, + X: np.ndarray, couplings: Any = None, gradients: Any = None, reference: Any = None) -> None: @@ -691,7 +695,7 @@ def compute(self, Parameters ---------- - X : ArrayLike + X : np.ndarray Position at which to compute properties couplings : list of tuple(int, int) or None, optional Which coupling pairs to compute. None means all. @@ -769,7 +773,7 @@ def compute_additional(self, for (i, j) in needed_c: self._derivative_couplings_available[i, j] = True - def clone(self): + def clone(self) -> TMModel: model_clone = cp.deepcopy(self) unique_workdir = find_unique_name(self.workdir_stem, self.run_turbomole_dir, diff --git a/mudslide/mud.py b/mudslide/mud.py index 41db312..94c4f70 100644 --- a/mudslide/mud.py +++ b/mudslide/mud.py @@ -4,25 +4,38 @@ This setup is temporary until I fully remove the __main__ and replace it with this. """ +from __future__ import annotations + import sys import argparse +from typing import Any import mudslide from .version import get_version_info -def mud_main(argv=None, file=sys.stdout) -> None: + +def mud_main(argv: list[str] | None = None, + file: Any = sys.stdout) -> None: """Mudslide CLI """ - parser = argparse.ArgumentParser(prog="mudslide", description="Mudslide CLI", - epilog=get_version_info(), - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('-v', '--version', action='version', + parser = argparse.ArgumentParser( + prog="mudslide", + description="Mudslide CLI", + epilog=get_version_info(), + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('-v', + '--version', + action='version', version=get_version_info()) - parser.add_argument('-d', '--debug', action='store_true', help="enable debug mode") + parser.add_argument('-d', + '--debug', + action='store_true', + help="enable debug mode") subparsers = parser.add_subparsers(title="commands", description='valid commands', - dest="subcommand", required=True) + dest="subcommand", + required=True) # subparser for the "collect" command mudslide.collect.add_collect_parser(subparsers) diff --git a/mudslide/propagation.py b/mudslide/propagation.py index 74924f4..ac4ef9b 100644 --- a/mudslide/propagation.py +++ b/mudslide/propagation.py @@ -1,12 +1,15 @@ # -*- coding: utf-8 -*- """Propagators for ODEs from quantum dynamics""" +from __future__ import annotations + from typing import Callable import numpy as np from numpy.typing import ArrayLike -def propagate_exponential(rho0: ArrayLike, H: ArrayLike, dt: np.floating) -> None: +def propagate_exponential(rho0: np.ndarray, H: np.ndarray, + dt: float) -> None: """Propagate density matrix in place using exponential of Hamiltonian. Propagates rho0 in place by exponentiating exp(-i H dt). @@ -14,11 +17,11 @@ def propagate_exponential(rho0: ArrayLike, H: ArrayLike, dt: np.floating) -> Non Parameters ---------- - rho0 : ArrayLike + rho0 : np.ndarray Initial density matrix - H : ArrayLike + H : np.ndarray Effective Hamiltonian for propagation - dt : np.floating + dt : float Time step for propagation Returns @@ -28,11 +31,18 @@ def propagate_exponential(rho0: ArrayLike, H: ArrayLike, dt: np.floating) -> Non """ diags, coeff = np.linalg.eigh(H) - U = np.linalg.multi_dot([ coeff, np.diag(np.exp(-1j * diags * dt)), coeff.T.conj() ]) + U = np.linalg.multi_dot( + [coeff, np.diag(np.exp(-1j * diags * dt)), + coeff.T.conj()]) np.dot(U, np.dot(rho0, U.T.conj()), out=rho0) -def propagate_interpolated_rk4(rho0: ArrayLike, h0: ArrayLike, tau0: ArrayLike, vel0: ArrayLike, - h1: ArrayLike, tau1: ArrayLike, vel1: ArrayLike, dt: np.floating, nsteps: int) -> None: + + +def propagate_interpolated_rk4(rho0: np.ndarray, h0: np.ndarray, + tau0: np.ndarray, vel0: np.ndarray, + h1: np.ndarray, tau1: np.ndarray, + vel1: np.ndarray, dt: float, + nsteps: int) -> None: """Propagate density matrix using linearly interpolated quantities and RK4. Propagate density matrix forward by linearly interpolating all quantities @@ -41,21 +51,21 @@ def propagate_interpolated_rk4(rho0: ArrayLike, h0: ArrayLike, tau0: ArrayLike, Parameters ---------- - rho0 : ArrayLike + rho0 : np.ndarray Input/output density matrix - h0 : ArrayLike + h0 : np.ndarray Hamiltonian from prior step - tau0 : ArrayLike + tau0 : np.ndarray Derivative coupling from prior step - vel0 : ArrayLike + vel0 : np.ndarray Velocity from prior step - h1 : ArrayLike + h1 : np.ndarray Hamiltonian from current step - tau1 : ArrayLike + tau1 : np.ndarray Derivative coupling from current step - vel1 : ArrayLike + vel1 : np.ndarray Velocity from current step - dt : np.floating + dt : float Time step nsteps : int Number of inner time steps @@ -67,24 +77,25 @@ def propagate_interpolated_rk4(rho0: ArrayLike, h0: ArrayLike, tau0: ArrayLike, """ TV00 = np.einsum("ijx,x->ij", tau0, vel0) TV11 = np.einsum("ijx,x->ij", tau1, vel1) - TV01 = np.einsum("ijx,x->ij", tau0, vel1) + np.einsum("ijx,x->ij", tau1, vel0) + TV01 = np.einsum("ijx,x->ij", tau0, vel1) + np.einsum( + "ijx,x->ij", tau1, vel0) eigs, vecs = np.linalg.eigh(h0) - H0 = np.linalg.multi_dot([vecs.T, h0, vecs]) - H1 = np.linalg.multi_dot([vecs.T, h1, vecs]) + H0 = np.linalg.multi_dot([vecs.T, h0, vecs]) + H1 = np.linalg.multi_dot([vecs.T, h1, vecs]) W00 = np.linalg.multi_dot([vecs.T, TV00, vecs]) W11 = np.linalg.multi_dot([vecs.T, TV11, vecs]) W01 = np.linalg.multi_dot([vecs.T, TV01, vecs]) - def ydot(rho: ArrayLike, t: np.floating) -> ArrayLike: + def ydot(rho: np.ndarray, t: float) -> np.ndarray: """Calculate time derivative of density matrix. Parameters ---------- - rho : ArrayLike + rho : np.ndarray Current density matrix - t : np.floating + t : float Current time point Returns @@ -93,17 +104,17 @@ def ydot(rho: ArrayLike, t: np.floating) -> ArrayLike: Time derivative of density matrix """ assert t >= 0.0 and t <= dt - w0 = 1.0 - t/dt - w1 = t/dt + w0 = 1.0 - t / dt + w1 = t / dt ergs = np.exp(1j * eigs * t).reshape([1, -1]) phases = np.dot(ergs.T, ergs.conj()) H = H0 * (w0 - 1.0) + H1 * w1 - Hbar = H - 1j * (w0*w0*W00 + w1*w1*W11 + w0*w1*W01) + Hbar = H - 1j * (w0 * w0 * W00 + w1 * w1 * W11 + w0 * w1 * W01) HI = Hbar * phases - out = -1j * ( np.dot(HI, rho) - np.dot(rho, HI) ) + out = -1j * (np.dot(HI, rho) - np.dot(rho, HI)) return out tmprho = np.linalg.multi_dot([vecs.T, rho0, vecs]) @@ -111,20 +122,23 @@ def ydot(rho: ArrayLike, t: np.floating) -> ArrayLike: ergs = np.exp(1j * eigs * dt).reshape([1, -1]) phases = np.dot(ergs.T.conj(), ergs) - rho0[:,:] = np.linalg.multi_dot([vecs, tmprho * phases, vecs.T]) + rho0[:, :] = np.linalg.multi_dot([vecs, tmprho * phases, vecs.T]) + + -def rk4(y0: ArrayLike, ydot: Callable, t0: np.floating, tf: np.floating, nsteps: int) -> ArrayLike: +def rk4(y0: np.ndarray, ydot: Callable, t0: float, tf: float, + nsteps: int) -> np.ndarray: """Propagate using 4th-order Runge-Kutta (RK4) method. Parameters ---------- - y0 : ArrayLike + y0 : np.ndarray Initial state vector ydot : Callable Function that computes the time derivative of y - t0 : np.floating + t0 : float Initial time - tf : np.floating + tf : float Final time nsteps : int Number of integration steps diff --git a/mudslide/propagator.py b/mudslide/propagator.py index 04c55b0..b5cbc3d 100644 --- a/mudslide/propagator.py +++ b/mudslide/propagator.py @@ -1,8 +1,10 @@ # -*- coding: utf-8 -*- """Propagators for ODEs from quantum dynamics""" +from __future__ import annotations + from typing import Any -import numpy as np + class Propagator_: """Base class for propagators. diff --git a/mudslide/surface.py b/mudslide/surface.py index 67e588c..4c1136f 100644 --- a/mudslide/surface.py +++ b/mudslide/surface.py @@ -1,7 +1,9 @@ # -*- coding: utf-8 -*- """Helper module for printing model surface""" -from typing import Any, List +from __future__ import annotations + +from typing import Any, List, TYPE_CHECKING import argparse import sys @@ -9,42 +11,64 @@ import numpy as np from numpy.typing import ArrayLike -from .models import scattering_models as models +from .exceptions import ConfigurationError +from .models.scattering_models import scattering_models as models from .version import get_version_info +if TYPE_CHECKING: + from .models.electronics import DiabaticModel_, ElectronicModel_ + + def add_surface_parser(subparsers: Any) -> None: """ Should accept a subparser object from argparse, add new subcommand, and then add arguments """ - parser = subparsers.add_parser('surface', help="Generate potential energy surface scans of two-state models") + parser = subparsers.add_parser( + 'surface', + help="Generate potential energy surface scans of two-state models") add_surface_arguments(parser) parser.set_defaults(func=surface_wrapper) + def add_surface_arguments(parser: Any) -> None: """ Add arguments to the parser object Note, this is spun out so that surface can act as a main function and a subcommand. The main function will eventually be deprecated and removed. """ - parser.add_argument('-m', '--model', default='simple', choices=[m for m in models], help="Tully model to plot") - parser.add_argument('-r', - '--range', - default=(-10.0, 10.0), - nargs=2, - type=float, - help="range over which to plot PES (default: %(default)s)") - parser.add_argument('-n', default=100, type=int, help="number of points to plot") - parser.add_argument('-s', - '--scan_dimension', - default=0, + parser.add_argument('-m', + '--model', + default='simple', + choices=[m for m in models], + help="Tully model to plot") + parser.add_argument( + '-r', + '--range', + default=(-10.0, 10.0), + nargs=2, + type=float, + help="range over which to plot PES (default: %(default)s)") + parser.add_argument('-n', + default=100, type=int, - help="which dimension to scan along for multi-dimensional models") + help="number of points to plot") + parser.add_argument( + '-s', + '--scan_dimension', + default=0, + type=int, + help="which dimension to scan along for multi-dimensional models") parser.add_argument('--x0', nargs='+', default=[0.0], type=float, help="reference point for multi-dimensional models") - parser.add_argument('-o', '--output', default=sys.stdout, type=argparse.FileType('w'), help="output file") + parser.add_argument('-o', + '--output', + default=sys.stdout, + type=argparse.FileType('w'), + help="output file") + def surface_wrapper(args: Any) -> None: """ Wrapper function for surface @@ -52,30 +76,46 @@ def surface_wrapper(args: Any) -> None: This function is called by the main function, and is used to wrap the surface function so that it can be called by the subcommand interface. """ - surface_main(args.model, args.range, args.n, args.scan_dimension, args.x0, output=args.output) + surface_main(args.model, + args.range, + args.n, + args.scan_dimension, + args.x0, + output=args.output) + -def main(argv=None) -> None: +def main(argv: List[str] | None = None) -> None: """ Main function for surface Deprecated """ - parser = argparse.ArgumentParser(description="Generate potential energy surface scans of two-state models", - epilog=get_version_info(), - formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('-v', '--version', action='version', version=get_version_info()) + parser = argparse.ArgumentParser( + description= + "Generate potential energy surface scans of two-state models", + epilog=get_version_info(), + formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument('-v', + '--version', + action='version', + version=get_version_info()) add_surface_arguments(parser) args = parser.parse_args(argv) surface_wrapper(args) -def surface_main(model: str, scan_range: List[float], n: int, scan_dimension: int, x0: List[float], + + +def surface_main(model_name: str, scan_range: List[float], n: int, + scan_dimension: int, x0: List[float], output: Any) -> None: """ Main function for surface""" - if model in models: - model = models[model]() + if model_name in models: + model = models[model_name]() else: - raise Exception("Unknown model chosen") # the argument parser should prevent this throw from being possible + raise ConfigurationError( + "Unknown model chosen" + ) # the argument parser should prevent this throw from being possible start, end = scan_range samples = n @@ -86,8 +126,11 @@ def surface_main(model: str, scan_range: List[float], n: int, scan_dimension: in ndof = model.ndof if len(x0) != ndof: - print("Must provide reference vector of same length as the model problem") - raise Exception("Expected reference vector of length {}, but received {}".format(ndof, len(x0))) + print( + "Must provide reference vector of same length as the model problem") + raise ConfigurationError( + f"Expected reference vector of length {ndof}, but received {len(x0)}" + ) xx = np.array(x0) xx[scan_dimension] = start @@ -96,28 +139,38 @@ def surface_main(model: str, scan_range: List[float], n: int, scan_dimension: in elec = model.update(xx) def headprinter() -> str: - xn = ["x{:d}".format(i) for i in range(ndof)] - diabats = ["V_%1d" % i for i in range(nstates)] - energies = ["E_%1d" % i for i in range(nstates)] - dc = ["d_%1d%1d" % (j, i) for i in range(nstates) for j in range(i)] - if model == "vibronic": - forces = ["dE_%1d" % i for i in range(ndof * 2)] + xn = [f"x{i}" for i in range(ndof)] + diabats = [f"V_{i}" for i in range(nstates)] + energies = [f"E_{i}" for i in range(nstates)] + dc = [f"d_{j}{i}" for i in range(nstates) for j in range(i)] + if model_name == "vibronic": + forces = [f"dE_{i}" for i in range(ndof * 2)] else: - forces = ["dE_%1d" % i for i in range(nstates)] + forces = [f"dE_{i}" for i in range(nstates)] plist = xn + diabats + energies + dc + forces - return "#" + " ".join(["%16s" % x for x in plist]) + return "#" + " ".join(f"{x:>16s}" for x in plist) - def lineprinter(x: ArrayLike, model: Any, estates: Any) -> str: + def lineprinter(x: np.ndarray, model: DiabaticModel_, estates: ElectronicModel_) -> str: V = model.V(x) ndof = estates.ndof diabats = [V[i, i] for i in range(nstates)] # type: List[float] - energies = [estates.hamiltonian[i, i] for i in range(nstates)] # type: List[float] - dc = [estates._derivative_coupling[j, i, 0] for i in range(nstates) for j in range(i)] # type: List[float] - forces = [float(-estates.force(i)[j]) for i in range(nstates) for j in range(ndof)] # type: List[float] - plist = list(x.flatten()) + diabats + energies + dc + forces # type: List[float] - - return " ".join(["{:16.10f}".format(x) for x in plist]) + energies = [estates.hamiltonian[i, i] for i in range(nstates) + ] # type: List[float] + dc = [ + estates._derivative_coupling[j, i, 0] + for i in range(nstates) + for j in range(i) + ] # type: List[float] + forces = [ + float(-estates.force(i)[j]) + for i in range(nstates) + for j in range(ndof) + ] # type: List[float] + plist = list( + x.flatten()) + diabats + energies + dc + forces # type: List[float] + + return " ".join(f"{x:16.10f}" for x in plist) #print("# scanning using model {}".format(model), file=output) #print("# reference point: {}".format(xx), file=output) diff --git a/mudslide/surface_hopping_md.py b/mudslide/surface_hopping_md.py index 7a669c1..ee57ee2 100644 --- a/mudslide/surface_hopping_md.py +++ b/mudslide/surface_hopping_md.py @@ -1,21 +1,25 @@ # -*- coding: utf-8 -*- """Propagate FSSH trajectory""" -import copy as cp -from typing import List, Dict, Union, Any +from __future__ import annotations +from typing import List, Dict, Union, Any, TYPE_CHECKING import numpy as np -from numpy.typing import ArrayLike -from .util import check_options -from .constants import boltzmann, fs_to_au +from .constants import boltzmann +from .exceptions import ConfigurationError from .propagation import propagate_exponential, propagate_interpolated_rk4 -from .tracer import Trace from .math import poisson_prob_scale +from .propagator import Propagator_ +from .trajectory_md import TrajectoryMD from .surface_hopping_propagator import SHPropagator -class SurfaceHoppingMD: +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + + +class SurfaceHoppingMD(TrajectoryMD): # pylint: disable=too-many-instance-attributes """Class to propagate a single FSSH trajectory. This class implements the Fewest Switches Surface Hopping (FSSH) algorithm @@ -23,26 +27,17 @@ class SurfaceHoppingMD: of both nuclear and electronic degrees of freedom, including surface hopping events between electronic states. """ - recognized_options = [ "propagator", "last_velocity", "bounds", - "dt", "t0", "previous_steps", - "duration", "max_steps", "max_time", - "seed_sequence", - "outcome_type", "trace_every", - "electronics", - "electronic_integration", "max_electronic_dt", "starting_electronic_intervals", - "weight", - "restarting", - "hopping_probability", "zeta_list", - "state0", - "hopping_method", - "forced_hop_threshold" - ] + recognized_options = TrajectoryMD.recognized_options + [ + "electronic_integration", "max_electronic_dt", + "starting_electronic_intervals", "hopping_probability", "zeta_list", + "state0", "hopping_method", "forced_hop_threshold" + ] def __init__(self, - model: Any, + model: ElectronicModel_, x0: np.ndarray, - v0: ArrayLike, - rho0: Union[ArrayLike, int, str], + v0: np.ndarray, + rho0: Union[np.ndarray, int, str], tracer: Any = "default", queue: Any = None, strict_option_check: bool = True, @@ -53,11 +48,11 @@ def __init__(self, ---------- model : Any Model object defining problem - x0 : ArrayLike + x0 : np.ndarray Initial position - v0 : ArrayLike + v0 : np.ndarray Initial velocity - rho0 : ArrayLike, int, or str + rho0 : np.ndarray, int, or str Initial density matrix or state index. If an integer, populates a single state. If a matrix, populates a density matrix (requires state0 for active state). tracer : Any, optional @@ -69,51 +64,21 @@ def __init__(self, Other Parameters ---------------- - propagator : str or dict, optional - The propagator to use for nuclear motion. Can be a string (e.g., 'vv', 'fssh') or a - dictionary with more options. Default is 'vv'. - last_velocity : array-like, optional - The velocity from the previous step, used for restarts. Default is zeros. - bounds : tuple or list, optional - Tuple or list of (lower, upper) bounds for the simulation box. Used to determine if - the trajectory is inside a region. Default is None. - duration : dict, optional - Dictionary controlling simulation duration (overrides max_steps, max_time, etc.). - Default is auto-generated. - dt : float, optional - Time step for nuclear propagation (in atomic units). Default is fs_to_au. - t0 : float, optional - Initial time. Default is 0.0. - previous_steps : int, optional - Number of previous steps (for restarts). Default is 0. - trace_every : int, optional - Interval (in steps) at which to record trajectory data. Default is 1. - max_steps : int, optional - Maximum number of simulation steps. Default is 1000000. - max_time : float, optional - Maximum simulation time. Default is 1e25. - seed_sequence : int or numpy.random.SeedSequence, optional - Seed or SeedSequence for random number generation. Default is None. - outcome_type : str, optional - Type of outcome to record (e.g., 'state'). Default is 'state'. - electronics : object, optional - Initial electronic state object. Default is None. electronic_integration : str, optional - Method for integrating electronic equations ('exp' or 'linear-rk4'). Default is 'exp'. + Method for integrating electronic equations ('exp' or 'linear-rk4'). + Default is 'exp'. max_electronic_dt : float, optional Maximum time step for electronic integration. Default is 0.1. starting_electronic_intervals : int, optional Initial number of intervals for electronic integration. Default is 4. - weight : float, optional - Statistical weight of the trajectory. Default is 1.0. - restarting : bool, optional - Whether this is a restarted trajectory. Default is False. hopping_probability : str, optional - Method for computing hopping probability ('tully' or 'poisson'). Default is 'tully'. + Method for computing hopping probability ('tully' or 'poisson'). + Default is 'tully'. zeta_list : list, optional List of pre-determined random numbers for hopping decisions. Default is []. state0 : int, optional - Initial electronic state (used if rho0 is a matrix). Required if rho0 is not scalar. + Initial electronic state (used if rho0 is a matrix). Required if rho0 + is not scalar. hopping_method : str, optional Hopping method: 'cumulative', 'cumulative_integrated', or 'instantaneous'. Default is 'cumulative'. @@ -122,77 +87,53 @@ def __init__(self, and the active state is not the lowest, force a hop to the lowest state. Default is None (off). """ - check_options(options, self.recognized_options, strict=strict_option_check) - - self.model = model - self.mass = model.mass - self.tracer = Trace(tracer) - self.queue: Any = queue - - # initial conditions - self.position = np.array(x0, dtype=np.float64).reshape(model.ndof) - self.last_position = np.zeros_like(self.position, dtype=np.float64) - self.velocity = np.array(v0, dtype=np.float64).reshape(model.ndof) - self.last_velocity = np.zeros_like(self.velocity, dtype=np.float64) - if "last_velocity" in options: - self.last_velocity[:] = options["last_velocity"] + super().__init__(model, + x0, + v0, + tracer=tracer, + queue=queue, + strict_option_check=strict_option_check, + **options) + + # Process initial electronic state if np.isscalar(rho0): try: - state = int(rho0) - self.rho = np.zeros([model.nstates, model.nstates], dtype=np.complex128) + state = int(rho0) # type: ignore[arg-type] + self.rho = np.zeros([model.nstates, model.nstates], + dtype=np.complex128) self.rho[state, state] = 1.0 self.state = state - except: - raise ValueError("Initial state rho0 must be convertible to an integer state " - "index") + except (TypeError, ValueError) as exc: + raise ConfigurationError( + "Initial state rho0 must be convertible to an integer state " + "index") from exc else: try: self.rho = np.copy(rho0) self.state = int(options["state0"]) - except KeyError: - raise KeyError("state0 option required when rho0 is a density matrix") - except (ValueError, TypeError): - raise ValueError("state0 option must be convertible to an integer state index") - - # function duration_initialize should get us ready to for future continue_simulating calls - # that decide whether the simulation has finished - if "duration" in options: - self.duration = options["duration"] - else: - self.duration_initialize(options) - - # fixed initial parameters - self.time = float(options.get("t0", 0.0)) - self.nsteps = int(options.get("previous_steps", 0)) - self.max_steps = int(options.get("max_steps", 1000000)) - self.max_time = float(options.get("max_time", 1e25)) - self.trace_every = int(options.get("trace_every", 1)) - self.dt = float(options.get("dt", fs_to_au)) - self.propagator = SHPropagator(self.model, options.get("propagator", "vv")) - - self.outcome_type = options.get("outcome_type", "state") - - ss = options.get("seed_sequence", None) - self.seed_sequence = ss if isinstance(ss, np.random.SeedSequence) \ - else np.random.SeedSequence(ss) - self.random_state = np.random.default_rng(self.seed_sequence) - - self.electronics = options.get("electronics", None) - self.last_electronics = options.get("last_electronics", None) + except KeyError as exc: + raise ConfigurationError( + "state0 option required when rho0 is a density matrix" + ) from exc + except (ValueError, TypeError) as exc: + raise ConfigurationError( + "state0 option must be convertible to an integer state index" + ) from exc + + # Surface hopping specific initialization self.hopping = 0.0 - self.electronic_integration = options.get("electronic_integration", "exp").lower() + self.electronic_integration = options.get("electronic_integration", + "exp").lower() self.max_electronic_dt = options.get("max_electronic_dt", 0.1) - self.starting_electronic_intervals = options.get("starting_electronic_intervals", 4) - - self.weight = float(options.get("weight", 1.0)) - - self.restarting = options.get("restarting", False) - self.force_quit = False + self.starting_electronic_intervals = options.get( + "starting_electronic_intervals", 4) self.hopping_probability = options.get("hopping_probability", "tully") if self.hopping_probability not in ["tully", "poisson"]: - raise ValueError("hopping_probability accepts only \"tully\" or \"poisson\" options") + raise ConfigurationError( + "hopping_probability accepts only \"tully\" or \"poisson\" options" + ) self.zeta_list = list(options.get("zeta_list", [])) self.zeta = 0.0 @@ -208,9 +149,12 @@ def __init__(self, self.hopping_method = hopping_method if self.hopping_method in aliases: self.hopping_method = aliases[hopping_method] - allowed_methods = ["instantaneous", "cumulative", "cumulative_integrated"] + allowed_methods = [ + "instantaneous", "cumulative", "cumulative_integrated" + ] if self.hopping_method not in allowed_methods: - raise ValueError(f"hopping_method should be one of {allowed_methods}") + raise ConfigurationError( + f"hopping_method should be one of {allowed_methods}") self.forced_hop_threshold = options.get("forced_hop_threshold", None) @@ -220,8 +164,27 @@ def __init__(self, if self.hopping_method == "cumulative_integrated": self.zeta = -np.log(1.0 - self.zeta) + def make_propagator(self, model: ElectronicModel_, + options: Dict[str, Any]) -> Propagator_: + """Create the surface hopping propagator. + + Parameters + ---------- + model : Any + Model object defining problem. + options : Dict[str, Any] + Options dictionary. + + Returns + ------- + Propagator_ + Surface hopping propagator instance. + """ + return SHPropagator(model, options.get("propagator", "vv")) # type: ignore[return-value] + @classmethod - def restart(cls, model, log, **options) -> 'SurfaceHoppingMD': + def restart(cls, model: ElectronicModel_, log: Any, + **options: Any) -> 'SurfaceHoppingMD': """Restart a simulation from a previous trajectory log. Parameters @@ -269,129 +232,11 @@ def restart(cls, model, log, **options) -> 'SurfaceHoppingMD': restarting=True, **options) - def update_weight(self, weight: float) -> None: - """Update weight held by trajectory and by trace. - - Parameters - ---------- - weight : float - New weight value - """ - self.weight = weight - self.tracer.weight = weight - - if self.weight == 0.0: - self.force_quit = True - - def __deepcopy__(self, memo: Any) -> 'SurfaceHoppingMD': - """Override deepcopy. - - Parameters - ---------- - memo : Any - Memo dictionary for deepcopy - - Returns - ------- - SurfaceHoppingMD - Deep copy of the instance - """ - cls = self.__class__ - result = cls.__new__(cls) - memo[id(self)] = result - shallow_only = ["queue"] - for k, v in self.__dict__.items(): - setattr(result, k, cp.deepcopy(v, memo) if v not in shallow_only else cp.copy(v)) - return result - - def clone(self) -> 'SurfaceHoppingMD': - """Clone existing trajectory for spawning. - - Returns - ------- - SurfaceHoppingMD - Copy of current object - """ - return cp.deepcopy(self) - - def random(self) -> np.float64: - """Get random number for hopping decisions. - - Returns - ------- - np.float64 - Uniform random number between 0 and 1 - """ - return self.random_state.uniform() - - def currently_interacting(self) -> bool: - """Determine whether trajectory is currently inside an interaction region. - - Returns - ------- - bool - True if trajectory is inside interaction region, False otherwise - """ - if self.duration["box_bounds"] is None: - return False - return np.all(self.duration["box_bounds"][0] < self.position) and np.all( - self.position < self.duration["box_bounds"][1]) - - def duration_initialize(self, options: Dict[str, Any]) -> None: - """Initialize variables related to continue_simulating. - - Parameters - ---------- - options : Dict[str, Any] - Dictionary with options for simulation duration - """ - duration = {} # type: Dict[str, Any] - duration['found_box'] = False - - bounds = options.get('bounds', None) - if bounds: - duration["box_bounds"] = (np.array(bounds[0], dtype=np.float64), - np.array(bounds[1], dtype=np.float64)) - else: - duration["box_bounds"] = None - - self.duration = duration - - def continue_simulating(self) -> bool: - """Decide whether trajectory should keep running. - - Returns - ------- - bool - True if trajectory should keep running, False if it should finish - """ - if self.force_quit: - return False - elif self.max_steps >= 0 and self.nsteps >= self.max_steps: - return False - elif self.time >= self.max_time or np.isclose( - self.time, self.max_time, atol=1e-8, rtol=0.0): - return False - elif self.duration["found_box"]: - return self.currently_interacting() - else: - if self.currently_interacting(): - self.duration["found_box"] = True - return True - - def trace(self, force: bool = False) -> None: - """Add results from current time point to tracing function. - - Only adds snapshot if nsteps%trace_every == 0, unless force=True. - - Parameters - ---------- - force : bool, optional - Force snapshot regardless of trace_every, by default False - """ - if force or (self.nsteps % self.trace_every) == 0: - self.tracer.collect(self.snapshot()) - #self.trouble_shooter() + def _report_columns(self) -> list[tuple[str, str]]: + """Append active state to report columns.""" + columns = super()._report_columns() + columns.append(("Active", f"{self.state:>6d}")) + return columns def snapshot(self) -> Dict[str, Any]: """Collect data from run for logging. @@ -401,141 +246,89 @@ def snapshot(self) -> Dict[str, Any]: Dict[str, Any] Dictionary with all data from current time step """ - out = { - "time": float(self.time), - "position": self.position.tolist(), - "velocity": self.velocity.tolist(), - "potential": float(self.potential_energy()), - "kinetic": float(self.kinetic_energy()), - "temperature": float(2 * self.kinetic_energy() / ( boltzmann * self.model.ndof)), - "energy": float(self.total_energy()), - "density_matrix": self.rho.view(dtype=np.float64).tolist(), - "active": int(self.state), - "electronics": self.electronics.as_dict(), - "hopping": float(self.hopping), - "zeta": float(self.zeta) - } + out = super().snapshot() + out["density_matrix"] = self.rho.view(dtype=np.float64).tolist() + out["active"] = int(self.state) + out["hopping"] = float(self.hopping) + out["zeta"] = float(self.zeta) if self.hopping_method in ["cumulative", "cumulative_integrated"]: out["prob_cum"] = float(self.prob_cum) return out - def kinetic_energy(self) -> np.float64: - """Calculate kinetic energy. - - Returns - ------- - np.float64 - Kinetic energy - """ - return 0.5 * np.sum(self.mass * self.velocity**2) - - def potential_energy(self, electronics: 'ElectronicModel_' = None) -> np.floating: + def potential_energy(self, + electronics: ElectronicModel_ | None = None + ) -> float: """Calculate potential energy. Parameters ---------- - electronics : ElectronicModel, optional + electronics : ElectronicModel_, optional electronic states from current step, by default None Returns ------- - np.floating + float Potential energy """ if electronics is None: electronics = self.electronics + assert electronics is not None return electronics.hamiltonian[self.state, self.state] - def total_energy(self, electronics: 'ElectronicModel_' = None) -> np.floating: - """Calculate total energy (kinetic + potential). - - Parameters - ---------- - electronics : ElectronicModel, optional - Electronic states from current step, by default None - - Returns - ------- - np.floating - Total energy - """ - potential = self.potential_energy(electronics) - kinetic = self.kinetic_energy() - return potential + kinetic - - def _force(self, electronics: 'ElectronicModel_' = None) -> ArrayLike: + def force(self, + electronics: ElectronicModel_ | None = None) -> np.ndarray: """Compute force on active state. Parameters ---------- - electronics : ElectronicModel, optional + electronics : ElectronicModel_, optional Electronic states from current step, by default None Returns ------- - ArrayLike + np.ndarray Force on active electronic state """ if electronics is None: electronics = self.electronics + assert electronics is not None return electronics.force(self.state) - def needed_gradients(self) -> List[int]: + def needed_gradients(self) -> list[int] | None: """States whose forces are needed during normal propagation. Returns ------- - List[int] + list[int] | None List of state indices for which gradients are needed. + None means all states are needed. + Standard FSSH only needs the active state gradient. """ return [self.state] - def needed_couplings(self): - """Coupling pairs needed during normal propagation. - - Returns None, meaning all couplings are needed. - """ - return None - - def NAC_matrix(self, electronics: 'ElectronicModel_' = None, - velocity: ArrayLike = None) -> ArrayLike: + def NAC_matrix(self, + electronics: ElectronicModel_ | None = None, + velocity: np.ndarray | None = None) -> np.ndarray: """Calculate nonadiabatic coupling matrix. Parameters ---------- - electronics : ElectronicModel, optional + electronics : ElectronicModel_, optional electronic states from current step, by default None - velocity : ArrayLike, optional + velocity : np.ndarray, optional Velocity used to compute NAC, by default None Returns ------- - ArrayLike + np.ndarray NAC matrix """ velo = velocity if velocity is not None else self.velocity if electronics is None: electronics = self.electronics + assert electronics is not None return electronics.NAC_matrix(velo) - def mode_kinetic_energy(self, direction: ArrayLike) -> np.float64: - """Calculate kinetic energy along given momentum mode. - - Parameters - ---------- - direction : ArrayLike - Numpy array defining direction - - Returns - ------- - np.float64 - Kinetic energy along specified direction - """ - u = direction / np.linalg.norm(direction) - momentum = self.velocity * self.mass - component = np.dot(u, momentum) * u - return 0.5 * np.einsum('m,m,m', 1.0 / self.mass, component, component) - def draw_new_zeta(self) -> float: """Get a new zeta value for hopping. @@ -549,15 +342,14 @@ def draw_new_zeta(self) -> float: """ if self.zeta_list: return self.zeta_list.pop(0) - else: - return self.random() + return self.random() - def hop_allowed(self, direction: ArrayLike, dE: float) -> bool: + def hop_allowed(self, direction: np.ndarray, dE: float) -> bool: """Determine if a hop with given rescale direction and energy change is allowed. Parameters ---------- - direction : ArrayLike + direction : np.ndarray Momentum unit vector dE : float Change in energy such that Enew = Eold + dE @@ -575,8 +367,11 @@ def hop_allowed(self, direction: ArrayLike, dE: float) -> bool: c = -2.0 * dE return b * b > 4.0 * a * c - def direction_of_rescale(self, source: int, target: int, - electronics: 'ElectronicModel_' = None) -> np.ndarray: + def direction_of_rescale( + self, + source: int, + target: int, + electronics: ElectronicModel_ | None = None) -> np.ndarray: """ Return direction in which to rescale momentum. @@ -586,8 +381,9 @@ def direction_of_rescale(self, source: int, target: int, Active state before hop target : int Active state after hop - electronics : ElectronicModel, optional - Electronic model information (used to pull derivative coupling), by default None + electronics : ElectronicModel_, optional + Electronic model information (used to pull derivative coupling), + by default None Returns ------- @@ -595,18 +391,20 @@ def direction_of_rescale(self, source: int, target: int, Unit vector pointing in direction of rescale """ elec_states = self.electronics if electronics is None else electronics + assert elec_states is not None out = elec_states.derivative_coupling(source, target) return np.copy(out) - def rescale_component(self, direction: ArrayLike, reduction: np.floating) -> None: + def rescale_component(self, direction: np.ndarray, + reduction: float) -> None: """ Update velocity by rescaling the *momentum* in the specified direction and amount. Parameters ---------- - direction : ArrayLike + direction : np.ndarray The direction of the *momentum* to rescale - reduction : np.floating + reduction : float How much kinetic energy should be damped """ # normalize @@ -620,40 +418,43 @@ def rescale_component(self, direction: ArrayLike, reduction: np.floating) -> Non self.velocity += scal * M_inv * u def hamiltonian_propagator(self, - last_electronics: 'ElectronicModel_', - this_electronics: 'ElectronicModel_', - velo: ArrayLike = None) -> np.ndarray: + last_electronics: ElectronicModel_, + this_electronics: ElectronicModel_, + velo: np.ndarray | None = None) -> np.ndarray: """ Compute the Hamiltonian used to propagate the electronic wavefunction. Parameters ---------- - last_electronics : ElectronicModel + last_electronics : ElectronicModel_ Electronic states at previous time step - this_electronics : ElectronicModel + this_electronics : ElectronicModel_ Electronic states at current time step - velo : ArrayLike, optional - Velocity at midpoint between current and previous time steps, by default None + velo : np.ndarray, optional + Velocity at midpoint between current and previous time steps, + by default None Returns ------- np.ndarray - Nonadiabatic coupling Hamiltonian at midpoint between current and previous time steps + Nonadiabatic coupling Hamiltonian at midpoint between current + and previous time steps """ if velo is None: velo = 0.5 * (self.velocity + self.last_velocity) if last_electronics is None: last_electronics = this_electronics - H = 0.5 * (this_electronics.hamiltonian + last_electronics.hamiltonian) # type: ignore + H = 0.5 * (this_electronics.hamiltonian + last_electronics.hamiltonian + ) # type: ignore this_tau = this_electronics.derivative_coupling_tensor last_tau = last_electronics.derivative_coupling_tensor TV = 0.5 * np.einsum("ijx,x->ij", this_tau + last_tau, velo) return H - 1j * TV - def propagate_electronics(self, last_electronics: 'ElectronicModel_', - this_electronics: 'ElectronicModel_', - dt: np.floating) -> None: + def propagate_electronics(self, last_electronics: ElectronicModel_, + this_electronics: ElectronicModel_, + dt: float) -> None: """ Propagate density matrix from t to t+dt. @@ -662,11 +463,11 @@ def propagate_electronics(self, last_electronics: 'ElectronicModel_', Parameters ---------- - last_electronics : ElectronicModel + last_electronics : ElectronicModel_ Electronic states at t - this_electronics : ElectronicModel + this_electronics : ElectronicModel_ Electronic states at t+dt - dt : np.floating + dt : float Time step """ if self.electronic_integration == "exp": @@ -681,26 +482,25 @@ def propagate_electronics(self, last_electronics: 'ElectronicModel_', this_tau = this_electronics.derivative_coupling_tensor last_tau = last_electronics.derivative_coupling_tensor - propagate_interpolated_rk4(self.rho, - last_electronics.hamiltonian, last_tau, self.last_velocity, - this_electronics.hamiltonian, this_tau, self.velocity, - self.dt, nsteps) + propagate_interpolated_rk4(self.rho, last_electronics.hamiltonian, + last_tau, self.last_velocity, + this_electronics.hamiltonian, this_tau, + self.velocity, self.dt, nsteps) else: - raise ValueError( + raise ConfigurationError( f"Unrecognized electronic integration option: {self.electronic_integration}. " - "Must be one of ['exp', 'linear-rk4']" - ) + "Must be one of ['exp', 'linear-rk4']") - def surface_hopping(self, last_electronics: 'ElectronicModel_', - this_electronics: 'ElectronicModel_'): + def surface_hopping(self, last_electronics: ElectronicModel_, + this_electronics: ElectronicModel_) -> None: """ Compute probability of hopping, generate random number, and perform hops. Parameters ---------- - last_electronics : ElectronicModel + last_electronics : ElectronicModel_ Electronic states at previous time step - this_electronics : ElectronicModel + this_electronics : ElectronicModel_ Electronic states at current time step """ H = self.hamiltonian_propagator(last_electronics, this_electronics) @@ -719,6 +519,7 @@ def surface_hopping(self, last_electronics: 'ElectronicModel_', old_state = self.state self.hop_to_it(hop_targets, this_electronics) if self.state != old_state: + assert self.electronics is not None self.electronics.compute_additional(gradients=[self.state]) def hopper(self, gkndt: np.ndarray) -> List[Dict[str, float]]: @@ -744,19 +545,27 @@ def hopper(self, gkndt: np.ndarray) -> List[Dict[str, float]]: # Forced hop check — short-circuit normal hopping logic if self.forced_hop_threshold is not None: + assert self.electronics is not None energies = np.diag(self.electronics.hamiltonian).real sorted_energies = np.sort(energies) gap = sorted_energies[1] - sorted_energies[0] - if gap < self.forced_hop_threshold and self.state != np.argmin(energies): + if gap < self.forced_hop_threshold and self.state != np.argmin( + energies): lowest_state = int(np.argmin(energies)) # Reset hopping state as a normal hop would - if self.hopping_method in ["cumulative", "cumulative_integrated"]: + if self.hopping_method in [ + "cumulative", "cumulative_integrated" + ]: self.prob_cum = np.longdouble(0.0) self.zeta = self.draw_new_zeta() if self.hopping_method == "cumulative_integrated": self.zeta = -np.log(1.0 - self.zeta) - return [{"target": lowest_state, "weight": 1.0, - "zeta": self.zeta, "prob": 1.0}] + return [{ + "target": lowest_state, + "weight": 1.0, + "zeta": self.zeta, + "prob": 1.0 + }] if self.hopping_method in ["cumulative", "cumulative_integrated"]: accumulated = np.longdouble(self.prob_cum) @@ -766,21 +575,29 @@ def hopper(self, gkndt: np.ndarray) -> List[Dict[str, float]]: elif self.hopping_method == "cumulative_integrated": accumulated += gkdt else: - raise ValueError(f"Unrecognized hopping method: {self.hopping_method}") + raise ConfigurationError( + f"Unrecognized hopping method: {self.hopping_method}") if accumulated > self.zeta: # then hop # where to hop hop_choice = gkndt / gkdt zeta = self.zeta - target = self.random_state.choice(list(range(self.model.nstates)), p=hop_choice) + target = self.random_state.choice(list(range( + self.model.nstates)), + p=hop_choice) # reset probabilities and random - self.prob_cum = 0.0 + self.prob_cum = np.longdouble(0.0) self.zeta = self.draw_new_zeta() if self.hopping_method == "cumulative_integrated": self.zeta = -np.log(1.0 - self.zeta) - return [{"target": target, "weight": 1.0, "zeta": zeta, "prob": accumulated}] + return [{ + "target": target, + "weight": 1.0, + "zeta": zeta, + "prob": accumulated + }] self.prob_cum = accumulated return [] @@ -804,7 +621,7 @@ def hopper(self, gkndt: np.ndarray) -> List[Dict[str, float]]: else: return [] - def hop_update(self, hop_from, hop_to): # pylint: disable=unused-argument + def hop_update(self, hop_from: int, hop_to: int) -> None: # pylint: disable=unused-argument """ Handle any extra operations that need to occur after a hop. @@ -817,9 +634,9 @@ def hop_update(self, hop_from, hop_to): # pylint: disable=unused-argument """ return - def hop_to_it(self, - hop_targets: List[Dict[str, Union[float,int]]], - electronics: 'ElectronicModel_' = None) -> None: + def hop_to_it(self, + hop_targets: List[Dict[str, Union[float, int]]], + electronics: ElectronicModel_ | None = None) -> None: """ Hop from the current active state to the given state, including rescaling the momentum. @@ -827,14 +644,16 @@ def hop_to_it(self, ---------- hop_targets : List[Dict[str, Union[float, int]]] List of (target, weight) pairs - electronics : ElectronicModel, optional + electronics : ElectronicModel_, optional Electronic states for current step, by default None """ hop_dict = hop_targets[0] hop_to = int(hop_dict["target"]) elec_states = electronics if electronics is not None else self.electronics + assert elec_states is not None H = elec_states.hamiltonian - new_potential, old_potential = H[hop_to, hop_to], H[self.state, self.state] + new_potential, old_potential = H[hop_to, hop_to], H[self.state, + self.state] delV = new_potential - old_potential rescale_vector = self.direction_of_rescale(self.state, hop_to) hop_from = self.state @@ -843,54 +662,18 @@ def hop_to_it(self, self.state = hop_to self.rescale_component(rescale_vector, -delV) self.hop_update(hop_from, hop_to) - self.tracer.record_event( - event_dict={ - "hop_from": int(hop_from), - "hop_to": int(hop_to), - "zeta": float(hop_dict["zeta"]), - "prob": float(hop_dict["prob"]) - }, - event_type="hop" - ) + self.tracer.record_event(event_dict={ + "hop_from": int(hop_from), + "hop_to": int(hop_to), + "zeta": float(hop_dict["zeta"]), + "prob": float(hop_dict["prob"]) + }, + event_type="hop") else: - self.tracer.record_event( - event_dict={ - "hop_from": int(hop_from), - "hop_to": int(hop_to), - "zeta": float(hop_dict["zeta"]), - "prob": float(hop_dict["prob"]) - }, - event_type="frustrated_hop" - ) - def simulate(self) -> 'Trace': - """ - Run the surface hopping molecular dynamics simulation. - - Returns - ------- - Trace - Trace of trajectory - """ - if not self.continue_simulating(): - return self.tracer - - if self.electronics is None: - self.electronics = self.model.update(self.position, - gradients=self.needed_gradients(), couplings=self.needed_couplings()) - - if not self.restarting: - self.trace() - - # propagation - while True: - self.propagator(self, 1) # pylint: disable=not-callable - - # ending condition - if not self.continue_simulating(): - break - - self.trace() - - self.trace(force=True) - - return self.tracer + self.tracer.record_event(event_dict={ + "hop_from": int(hop_from), + "hop_to": int(hop_to), + "zeta": float(hop_dict["zeta"]), + "prob": float(hop_dict["prob"]) + }, + event_type="frustrated_hop") diff --git a/mudslide/surface_hopping_propagator.py b/mudslide/surface_hopping_propagator.py index 123bbee..5fae659 100644 --- a/mudslide/surface_hopping_propagator.py +++ b/mudslide/surface_hopping_propagator.py @@ -1,11 +1,20 @@ # -*- coding: utf-8 -*- """Propagate FSSH trajectory""" +from __future__ import annotations + from typing import Any -from .util import is_string +from typing import TYPE_CHECKING + +from .exceptions import ConfigurationError +from .util import is_string, check_options from .propagator import Propagator_ +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + + class SHVVPropagator(Propagator_): """Velocity Verlet propagator for surface hopping dynamics. @@ -13,6 +22,8 @@ class SHVVPropagator(Propagator_): classical trajectories in surface hopping molecular dynamics simulations. """ + recognized_options: list[str] = ["type"] + def __init__(self, **options: Any) -> None: """Constructor @@ -21,6 +32,7 @@ def __init__(self, **options: Any) -> None: **options : Any Option dictionary for configuration. """ + check_options(options, self.recognized_options) super().__init__() def __call__(self, traj: Any, nsteps: int) -> None: @@ -37,28 +49,32 @@ def __call__(self, traj: Any, nsteps: int) -> None: # first update nuclear coordinates for _ in range(nsteps): # Advance position using Velocity Verlet - acceleration = traj._force(traj.electronics) / traj.mass + acceleration = traj.force(traj.electronics) / traj.mass traj.last_position = traj.position traj.position += traj.velocity * dt + 0.5 * acceleration * dt * dt # calculate electronics at new position traj.last_electronics, traj.electronics = traj.electronics, traj.model.update( - traj.position, electronics=traj.electronics, - gradients=traj.needed_gradients(), couplings=traj.needed_couplings()) + traj.position, + electronics=traj.electronics, + gradients=traj.needed_gradients(), + couplings=traj.needed_couplings()) # Update velocity using Velocity Verlet - last_acceleration = traj._force(traj.last_electronics) / traj.mass - this_acceleration = traj._force(traj.electronics) / traj.mass + last_acceleration = traj.force(traj.last_electronics) / traj.mass + this_acceleration = traj.force(traj.electronics) / traj.mass traj.last_velocity = traj.velocity traj.velocity += 0.5 * (last_acceleration + this_acceleration) * dt # now propagate the electronic wavefunction to the new time - traj.propagate_electronics(traj.last_electronics, traj.electronics, dt) + traj.propagate_electronics(traj.last_electronics, traj.electronics, + dt) traj.surface_hopping(traj.last_electronics, traj.electronics) traj.time += dt traj.nsteps += 1 + class SHPropagator: """Surface Hopping propagator factory. @@ -66,7 +82,7 @@ class SHPropagator: used in surface hopping molecular dynamics simulations. """ - def __new__(cls, model: Any, prop_options: Any = "vv") -> 'SHPropagator': + def __new__(cls, model: ElectronicModel_, prop_options: Any = "vv") -> Propagator_: # type: ignore[misc] """Create a new surface hopping propagator instance. Parameters @@ -89,10 +105,10 @@ def __new__(cls, model: Any, prop_options: Any = "vv") -> 'SHPropagator': if is_string(prop_options): prop_options = {"type": prop_options} elif not isinstance(prop_options, dict): - raise ValueError("prop_options must be a string or a dictionary") + raise ConfigurationError("prop_options must be a string or a dictionary") proptype = prop_options.get("type", "vv") if proptype.lower() == "vv": return SHVVPropagator(**prop_options) - else: - raise ValueError(f"Unrecognized surface hopping propagator type: {proptype}.") + raise ConfigurationError( + f"Unrecognized surface hopping propagator type: {proptype}.") diff --git a/mudslide/tracer.py b/mudslide/tracer.py index 7d37658..1767b4f 100644 --- a/mudslide/tracer.py +++ b/mudslide/tracer.py @@ -1,6 +1,8 @@ # -*- coding: utf-8 -*- """Collect results from single trajectories""" +from __future__ import annotations + import bz2 import copy as cp import gzip @@ -16,6 +18,7 @@ from numpy.typing import ArrayLike from .constants import fs_to_au +from .exceptions import ConfigurationError from .util import find_unique_name, is_string from .math import RollingAverage from .version import __version__ @@ -34,7 +37,7 @@ } -def _open_log(path: str, mode: str = "rt"): +def _open_log(path: str, mode: str = "rt") -> Any: """Open a log file, automatically handling compression based on extension.""" for ext, module in _COMPRESSION_EXTENSIONS.items(): if path.endswith(ext): @@ -132,7 +135,7 @@ def __len__(self) -> int: Number of snapshots """ - def form_data(self, snap_dict: Dict) -> Dict: + def form_data(self, snap_dict: Dict) -> Dict[str, Any]: """Convert snapshot dictionary to appropriate data types. Parameters @@ -145,7 +148,7 @@ def form_data(self, snap_dict: Dict) -> Dict: Dict Processed snapshot data """ - out = {} + out: Dict[str, Any] = {} for k, v in snap_dict.items(): if isinstance(v, list): o = np.array(v) @@ -238,7 +241,7 @@ def print_egylog(self, file: Any = sys.stdout, T_window: int = 50) -> None: print(line, file=file) - def outcome(self) -> ArrayLike: + def outcome(self) -> np.ndarray: """Classifies end of simulation: 2*state + [0 for left, 1 for right]""" last_snapshot = self[-1] ndof = len(last_snapshot["position"]) @@ -290,7 +293,7 @@ def collect(self, snapshot: Any) -> None: """collect and optionally process data""" self.data.append(snapshot) - def record_event(self, event_dict: Dict, event_type: str = "hop"): + def record_event(self, event_dict: Dict, event_type: str = "hop") -> None: if event_type == "hop": self.hops.append(event_dict) return @@ -363,9 +366,8 @@ def __init__(self, super().__init__(weight=weight) if compression is not None and compression not in _COMPRESSORS: - raise ValueError( - f"Unknown compression type: {compression}. " - f"Supported types: {list(_COMPRESSORS.keys())}") + raise ConfigurationError(f"Unknown compression type: {compression}. " + f"Supported types: {list(_COMPRESSORS.keys())}") self.weight: float = weight self.log_pitch = log_pitch @@ -421,7 +423,7 @@ def __init__(self, self.logfiles = logdata["logfiles"] self.main_log = self.unique_name + ".yaml" if self.main_log != os.path.basename(load_main_log): - raise RuntimeError( + raise ConfigurationError( f"It looks like the log file {load_main_log} was renamed. " "This is undefined behavior for now!") self.active_logfile = self.logfiles[-1] @@ -440,7 +442,7 @@ def __init__(self, self.logsize = self.log_pitch * (self.nlogs - 1) + self.active_logsize - def files(self, absolute_path=True): + def files(self, absolute_path: bool = True) -> List[str]: """returns a list of all files associated with this trace :param absolute_path: if True, returns the absolute path to the files, @@ -453,7 +455,7 @@ def files(self, absolute_path=True): return [os.path.join(self.location, x) for x in rel_files] return rel_files - def write_main_log(self): + def write_main_log(self) -> None: """Writes main log file, which points to other files for logging information""" out = { "name": self.unique_name, @@ -481,6 +483,7 @@ def _compress_log(self, log_index: int) -> None: log_index : int Index into self.logfiles of the log to compress """ + assert self.compression is not None module, ext = _COMPRESSORS[self.compression] old_name = self.logfiles[log_index] new_name = old_name + ext @@ -519,7 +522,7 @@ def collect(self, snapshot: Any) -> None: self.logsize += 1 - def record_event(self, event_dict: Dict, event_type: str = "hop"): + def record_event(self, event_dict: Dict, event_type: str = "hop") -> None: if event_type == "hop": log = self.hop_log else: @@ -541,7 +544,7 @@ def record_event(self, event_dict: Dict, event_type: str = "hop"): default_flow_style=False, explicit_start=False) - def clone(self): + def clone(self) -> YAMLTrace: """Create a deep copy of the trace. Returns @@ -614,9 +617,8 @@ def __getitem__(self, i: int) -> Any: target_log = i // self.log_pitch target_snap = i - target_log * self.log_pitch - with _open_log( - os.path.join(self.location, self.logfiles[target_log]), - "rt") as f: + with _open_log(os.path.join(self.location, self.logfiles[target_log]), + "rt") as f: chunk = yaml.safe_load(f) return self.form_data(chunk[target_snap]) @@ -637,14 +639,14 @@ def as_dict(self) -> Dict: } -def load_log(main_log_name): +def load_log(main_log_name: str) -> YAMLTrace: """prepare a trace object from a main log file""" # assuming online yaml logs for now out = YAMLTrace(load_main_log=main_log_name) return out -def trace_factory(trace_type: str = "yaml"): +def trace_factory(trace_type: str = "yaml") -> type: """returns the appropriate trace class based on the type specified Parameters @@ -656,10 +658,10 @@ def trace_factory(trace_type: str = "yaml"): return YAMLTrace if trace_type == "in_memory": return InMemoryTrace - raise ValueError(f"Invalid trace type specified: {trace_type}") + raise ConfigurationError(f"Invalid trace type specified: {trace_type}") -def Trace(trace_type, *args, **kwargs): +def Trace(trace_type: Any, *args: Any, **kwargs: Any) -> Trace_: """Create a trace object based on the type specified. Parameters @@ -695,23 +697,23 @@ def Trace(trace_type, *args, **kwargs): elif isinstance(trace_type, Trace_): return trace_type - raise ValueError("Unrecognized Trace option") + raise ConfigurationError("Unrecognized Trace option") class TraceManager: """Manage the collection of observables from a set of trajectories""" def __init__(self, - trace_type="default", - trace_args=None, - trace_kwargs=None) -> None: + trace_type: str = "default", + trace_args: Optional[List[Any]] = None, + trace_kwargs: Optional[Dict[str, Any]] = None) -> None: self.trace_type = trace_type self.trace_args = trace_args if trace_args is not None else [] self.trace_kwargs = trace_kwargs if trace_kwargs is not None else {} self.traces: List = [] - self.outcomes: ArrayLike + self.outcomes: np.ndarray def spawn_tracer(self) -> Trace_: """returns a Tracer object that collects all of the observables for a given trajectory""" @@ -731,14 +733,14 @@ def __iter__(self) -> Iterator: def __getitem__(self, i: int) -> Trace_: return self.traces[i] - def outcome(self) -> ArrayLike: + def outcome(self) -> np.ndarray: """summarize outcomes from entire set of traces""" weight_norm = sum((t.weight for t in self.traces)) outcome = sum( (t.weight * t.outcome() for t in self.traces)) / weight_norm return outcome - def counts(self) -> ArrayLike: + def counts(self) -> np.ndarray: """summarize outcomes from entire set of traces""" out = sum(t.outcome() for t in self.traces) return out @@ -804,7 +806,7 @@ def summarize(self, verbose: bool = False, file: Any = sys.stdout) -> None: def as_dict(self) -> Dict: """return the object as a dictionary""" - out = {"hops": [], "data": [], "weight": []} + out: Dict[str, list[Any]] = {"hops": [], "data": [], "weight": []} for x in self.traces: out["hops"].append(x.as_dict()["hops"]) out["data"].append(x.as_dict()["data"]) diff --git a/mudslide/trajectory_md.py b/mudslide/trajectory_md.py new file mode 100644 index 0000000..94933e7 --- /dev/null +++ b/mudslide/trajectory_md.py @@ -0,0 +1,554 @@ +# -*- coding: utf-8 -*- +"""Base class for MD trajectory propagation. + +This module provides the abstract base class TrajectoryMD, which contains +shared infrastructure for all molecular dynamics trajectory types, including +adiabatic, surface hopping, and Ehrenfest dynamics. +""" + +from __future__ import annotations + +import time +from abc import ABC, abstractmethod +from collections import deque +from typing import Dict, Any, TYPE_CHECKING +import copy as cp + +import numpy as np + +from .exceptions import ConfigurationError +from .util import check_options +from .constants import boltzmann, fs_to_au +from .tracer import Trace, Trace_ +from .propagator import Propagator_ + +if TYPE_CHECKING: + from .models.electronics import ElectronicModel_ + + +TIME_COMPARISON_ATOL: float = 1e-12 + + +class TrajectoryMD(ABC): # pylint: disable=too-many-instance-attributes + """Abstract base class for molecular dynamics trajectories. + + This class provides shared infrastructure for all trajectory types, + including common initialization, simulation loop, and analysis methods. + Subclasses must implement force and potential energy calculations, as + well as a factory method for creating the appropriate propagator. + """ + + recognized_options: list[str] = [ + "dt", "t0", "trace_every", "remove_com_every", + "remove_angular_momentum_every", "max_steps", "max_time", "bounds", + "propagator", "seed_sequence", "electronics", "outcome_type", + "weight", "last_velocity", "previous_steps", "restarting", "duration", + "report_every", "report_file", + ] + + def __init__(self, + model: ElectronicModel_, + x0: np.ndarray, + v0: np.ndarray, + tracer: Any = None, + queue: Any = None, + strict_option_check: bool = True, + **options: Any): + """Initialize common trajectory state. + + Parameters + ---------- + model : Any + Model object defining problem. + x0 : np.ndarray + Initial position. + v0 : np.ndarray + Initial velocity. + tracer : Any, optional + Spawn from TraceManager to collect results. + queue : Any, optional + Trajectory queue. + strict_option_check : bool, optional + Whether to strictly check options. + **options : Any + Additional options for the simulation. Recognized options are: + + dt : float + Time step for nuclear propagation (in atomic units). Required. + t0 : float, optional + Initial time. Default is 0.0. + trace_every : int, optional + Interval (in steps) at which to record trajectory data. Default is 1. + remove_com_every : int, optional + Interval for removing center-of-mass motion. Default is 0 (never). + remove_angular_momentum_every : int, optional + Interval for removing angular momentum. Default is 0 (never). + max_steps : int, optional + Maximum number of steps. Default depends on subclass. + max_time : float, optional + Maximum simulation time. Default is 1e25. + bounds : tuple or list, optional + Tuple or list of (lower, upper) bounds for the simulation box. + Default is None. + propagator : str or dict, optional + The propagator to use for nuclear motion. Default depends on subclass. + seed_sequence : int or numpy.random.SeedSequence, optional + Seed or SeedSequence for random number generation. Default is None. + electronics : object, optional + Initial electronic state object. Default is None. + outcome_type : str, optional + Type of outcome to record (e.g., 'state'). Default is 'state'. + weight : float, optional + Statistical weight of the trajectory. Default is 1.0. + restarting : bool, optional + Whether this is a restarted trajectory. Default is False. + """ + check_options(options, + self.recognized_options, + strict=strict_option_check) + + self.model = model + self.tracer = Trace(tracer) + self.queue: Any = queue + self.mass = model.mass + self.position = np.array(x0, dtype=np.float64).reshape(model.ndof) + self.velocity = np.array(v0, dtype=np.float64).reshape(model.ndof) + self.last_position = np.zeros_like(self.position, dtype=np.float64) + self.last_velocity = np.zeros_like(self.velocity, dtype=np.float64) + if "last_velocity" in options: + self.last_velocity[:] = options["last_velocity"] + + # function duration_initialize should get us ready for future + # continue_simulating calls that decide whether the simulation + # has finished + if "duration" in options: + self.duration = options["duration"] + else: + self.duration_initialize(options) + + # fixed initial parameters + self.time = float(options.get("t0", 0.0)) + self.nsteps = int(options.get("previous_steps", 0)) + self.max_steps = int( + options.get("max_steps", 100000)) + self.max_time = float(options.get("max_time", 1e25)) + self.trace_every = int(options.get("trace_every", 1)) + self.remove_com_every = int(options.get("remove_com_every", 0)) + self.remove_angular_momentum_every = int( + options.get("remove_angular_momentum_every", 0)) + if "dt" not in options: + raise ConfigurationError("dt option is required") + self.dt = float(options["dt"]) + + self.propagator: Propagator_ = self.make_propagator( + model, options) # type: ignore[assignment] + + self.outcome_type = options.get("outcome_type", "state") + + ss = options.get("seed_sequence", None) + self.seed_sequence = ss if isinstance(ss, np.random.SeedSequence) \ + else np.random.SeedSequence(ss) + self.random_state = np.random.default_rng(self.seed_sequence) + + self.electronics: ElectronicModel_ | None = options.get( + "electronics", None) + self.last_electronics: ElectronicModel_ | None = options.get( + "last_electronics", None) + + self.weight = float(options.get("weight", 1.0)) + + self.restarting = options.get("restarting", False) + self.force_quit = False + + # Progress reporting + self.report_file: str | None = options.get("report_file", + "mudslide.report") + report_every_opt = options.get("report_every", None) + if report_every_opt is not None: + self.report_every = int(report_every_opt) + else: + total_steps = self.max_steps + if total_steps <= 0 or total_steps > 1e15: + total_steps = int(self.max_time / self.dt) if self.max_time < 1e15 else 0 + self.report_every = max(1, total_steps // 100) if total_steps > 0 else 0 + self._temp_buffer: deque[float] = deque(maxlen=50) + self._last_report_time: float = 0.0 + + @abstractmethod + def make_propagator(self, model: ElectronicModel_, + options: Dict[str, Any]) -> Propagator_: + """Create the propagator for this trajectory type. + + Parameters + ---------- + model : Any + Model object defining problem. + options : Dict[str, Any] + Options dictionary. + + Returns + ------- + Propagator_ + Propagator instance for this trajectory. + """ + + def update_weight(self, weight: float) -> None: + """Update weight held by trajectory and by trace. + + Parameters + ---------- + weight : float + New weight value to set. + """ + self.weight = weight + self.tracer.weight = weight + + if self.weight == 0.0: + self.force_quit = True + + def __deepcopy__(self, memo: Any) -> 'TrajectoryMD': + """Override deepcopy. + + Parameters + ---------- + memo : Any + Memo dictionary for deepcopy. + + Returns + ------- + TrajectoryMD + Deep copy of the current object. + """ + cls = self.__class__ + result = cls.__new__(cls) + memo[id(self)] = result + shallow_only = ["queue"] + for k, v in self.__dict__.items(): + setattr( + result, k, + cp.deepcopy(v, memo) if k not in shallow_only else cp.copy(v)) + return result + + def clone(self) -> 'TrajectoryMD': + """Clone existing trajectory for spawning. + + Returns + ------- + TrajectoryMD + Copy of current object. + """ + return cp.deepcopy(self) + + def random(self) -> float: + """Get random number. + + Returns + ------- + float + Uniform random number between 0 and 1. + """ + return self.random_state.uniform() + + def currently_interacting(self) -> bool: + """Determine whether trajectory is currently inside an interaction region. + + Returns + ------- + bool + True if trajectory is inside interaction region, False otherwise. + """ + if self.duration["box_bounds"] is None: + return False + return bool( + np.all(self.duration["box_bounds"][0] < self.position) and np.all( + self.position < self.duration["box_bounds"][1])) + + def duration_initialize(self, options: Dict[str, Any]) -> None: + """Initialize variables related to continue_simulating. + + Parameters + ---------- + options : Dict[str, Any] + Dictionary with options. + """ + duration: Dict[str, Any] = {} + duration['found_box'] = False + + bounds = options.get('bounds', None) + if bounds: + b0 = np.array(bounds[0], dtype=np.float64) + b1 = np.array(bounds[1], dtype=np.float64) + duration["box_bounds"] = (b0, b1) + else: + duration["box_bounds"] = None + + self.duration = duration + + def continue_simulating(self) -> bool: + """Decide whether trajectory should keep running. + + Returns + ------- + bool + True if trajectory should keep running, False if it should finish. + """ + if self.force_quit: + return False + if self.max_steps >= 0 and self.nsteps >= self.max_steps: + return False + if self.time >= self.max_time or np.isclose( + self.time, self.max_time, atol=TIME_COMPARISON_ATOL, rtol=0.0): + return False + if self.duration["found_box"]: + return self.currently_interacting() + if self.currently_interacting(): + self.duration["found_box"] = True + return True + + def trace(self, force: bool = False) -> None: + """Add results from current time point to tracing function. + + Only adds snapshot if nsteps%trace_every == 0, unless force=True. + + Parameters + ---------- + force : bool, optional + Force snapshot regardless of trace_every interval. + """ + if force or (self.nsteps % self.trace_every) == 0: + self.tracer.collect(self.snapshot()) + + def snapshot(self) -> Dict[str, Any]: + """Collect data from run for logging. + + Subclasses may override to add additional fields. + + Returns + ------- + Dict[str, Any] + Dictionary with all data from current time step. + """ + assert self.electronics is not None + out = { + "time": + float(self.time), + "position": + self.position.tolist(), + "velocity": + self.velocity.tolist(), + "potential": + float(self.potential_energy()), + "kinetic": + float(self.kinetic_energy()), + "energy": + float(self.total_energy()), + "temperature": + float(2 * self.kinetic_energy() / + (boltzmann * self.model.ndof)), + "electronics": + self.electronics.as_dict() + } + return out + + def kinetic_energy(self) -> np.float64: + """Calculate kinetic energy. + + Returns + ------- + np.float64 + Kinetic energy. + """ + return 0.5 * np.einsum('m,m,m', self.mass, self.velocity, + self.velocity) + + @abstractmethod + def potential_energy(self, + electronics: ElectronicModel_ | None = None + ) -> float: + """Calculate potential energy. + + Parameters + ---------- + electronics : ElectronicModel_, optional + Electronic states from current step. + + Returns + ------- + float + Potential energy. + """ + + def total_energy(self, + electronics: ElectronicModel_ | None = None) -> float: + """Calculate total energy (kinetic + potential). + + Parameters + ---------- + electronics : ElectronicModel_, optional + Electronic states from current step. + + Returns + ------- + float + Total energy. + """ + potential = self.potential_energy(electronics) + kinetic = self.kinetic_energy() + return potential + kinetic + + @abstractmethod + def force(self, + electronics: ElectronicModel_ | None = None) -> np.ndarray: + """Compute force on active electronic state. + + Parameters + ---------- + electronics : ElectronicModel_, optional + ElectronicStates from current step. + + Returns + ------- + np.ndarray + Force on active electronic state. + """ + + def mode_kinetic_energy(self, direction: np.ndarray) -> np.float64: + """Calculate kinetic energy along given momentum mode. + + Parameters + ---------- + direction : np.ndarray + Array defining direction. + + Returns + ------- + np.float64 + Kinetic energy along specified direction. + """ + u = direction / np.linalg.norm(direction) + momentum = self.velocity * self.mass + component = np.dot(u, momentum) * u + return 0.5 * np.einsum('m,m,m', 1.0 / self.mass, component, component) + + def needed_gradients(self) -> list[int] | None: + """States whose forces are needed during propagation. + + Returns + ------- + list[int] | None + List of state indices for which gradients are needed. + None means all states are needed. + """ + return None + + def needed_couplings(self) -> list[tuple[int, int]] | None: + """Coupling pairs needed during propagation. + + Returns + ------- + list[tuple[int, int]] | None + List of (i, j) state pairs for which derivative couplings are needed. + None means all couplings are needed. + """ + return None + + def _reporting_enabled(self) -> bool: + """Check whether progress reporting is enabled.""" + return self.report_every > 0 and self.report_file is not None + + def _report_columns(self) -> list[tuple[str, str]]: + """Return (header, formatted_value) pairs for the current report line. + + Subclasses can override to append additional columns. + """ + assert self.electronics is not None + temp = float(2 * self.kinetic_energy() / + (boltzmann * self.model.ndof)) + self._temp_buffer.append(temp) + avg_temp = sum(self._temp_buffer) / len(self._temp_buffer) + + elapsed = time.monotonic() - self._last_report_time + self._last_report_time = time.monotonic() + + return [ + ("Step", f"{self.nsteps:>10d}"), + ("Time (fs)", f"{self.time / fs_to_au:>12.2f}"), + ("Total Energy (H)", f"{self.total_energy():>17.10f}"), + ("Avg Temp (K)", f"{avg_temp:>12.1f}"), + ("Wall (min)", f"{elapsed / 60.0:>10.2f}"), + ] + + def _print_report_header(self, fh: Any) -> None: + """Print column headers to stdout and report file.""" + columns = self._report_columns() + widths = [max(len(h), len(v)) for h, v in columns] + header = " ".join(f"{h:>{w}}" for (h, _), w in zip(columns, widths)) + sep = " ".join("-" * w for w in widths) + print(header) + print(sep) + fh.write(header + "\n") + fh.write(sep + "\n") + fh.flush() + + def _print_report(self, fh: Any) -> None: + """Print one report line to stdout and report file.""" + columns = self._report_columns() + line = " ".join(v for _, v in columns) + print(line) + fh.write(line + "\n") + fh.flush() + + def simulate(self) -> Trace_: + """Run the simulation. + + Returns + ------- + Trace_ + Trace of trajectory. + """ + if not self.continue_simulating(): + return self.tracer + + if self.electronics is None: + self.electronics = self.model.update( + self.position, + gradients=self.needed_gradients(), + couplings=self.needed_couplings()) + + if not self.restarting: + self.trace() + + # set up progress reporting + reporting = self._reporting_enabled() + report_fh = None + if reporting: + self._temp_buffer.clear() + self._last_report_time = time.monotonic() + assert self.report_file is not None + report_fh = open(self.report_file, "w") # noqa: SIM115 + self._print_report_header(report_fh) + + # propagation + try: + while True: + self.propagator(self, 1) # pylint: disable=not-callable + + if reporting and report_fh is not None \ + and self.nsteps % self.report_every == 0: + self._print_report(report_fh) + + # ending condition + if not self.continue_simulating(): + break + + self.trace() + + self.trace(force=True) + + if reporting and report_fh is not None: + self._print_report(report_fh) + finally: + if report_fh is not None: + report_fh.close() + + return self.tracer diff --git a/mudslide/turbo_make_harmonic.py b/mudslide/turbo_make_harmonic.py index c5fdf60..7f71f48 100644 --- a/mudslide/turbo_make_harmonic.py +++ b/mudslide/turbo_make_harmonic.py @@ -3,18 +3,23 @@ Extract harmonic parameters from a vibrational analysis """ +from __future__ import annotations + from typing import Any import argparse import sys import numpy as np +from .exceptions import ExternalCodeError from .models.turbomole_model import TurboControl, turbomole_is_installed_or_prefixed from .models.harmonic_model import HarmonicModel from .units import amu from .version import get_version_info -def add_make_harmonic_parser(subparsers): + + +def add_make_harmonic_parser(subparsers: Any) -> None: """Add make_harmonic subparser to an argument parser. Parameters @@ -29,12 +34,13 @@ def add_make_harmonic_parser(subparsers): """ parser = subparsers.add_parser( "make_harmonic", - help="Generate a harmonic model from a vibrational analysis" - ) + help="Generate a harmonic model from a vibrational analysis") add_make_harmonic_arguments(parser) parser.set_defaults(func=make_harmonic_main) -def add_make_harmonic_arguments(parser): + + +def add_make_harmonic_arguments(parser: Any) -> None: """Add command line arguments for make_harmonic command. Parameters @@ -47,13 +53,23 @@ def add_make_harmonic_arguments(parser): None Modifies parser in place by adding arguments """ - parser.add_argument("-c", "--control", default="control", help="Control file") - parser.add_argument("-d", "--model-dest", default="harmonic.json", + parser.add_argument("-c", + "--control", + default="control", + help="Control file") + parser.add_argument("-d", + "--model-dest", + default="harmonic.json", help="Where to write harmonic model as a json output") - parser.add_argument("-o", "--output", default=sys.stdout, type=argparse.FileType('w'), + parser.add_argument("-o", + "--output", + default=sys.stdout, + type=argparse.FileType('w'), help="Where to print output") -def main(argv=None): + + +def main(argv: list[str] | None = None) -> None: """Parse command line arguments and run make_harmonic command. Parameters @@ -70,13 +86,18 @@ def main(argv=None): description="Generate a harmonic model from a vibrational analysis", epilog=get_version_info(), formatter_class=argparse.RawDescriptionHelpFormatter) - parser.add_argument('-v', '--version', action='version', version=get_version_info()) + parser.add_argument('-v', + '--version', + action='version', + version=get_version_info()) add_make_harmonic_arguments(parser) args = parser.parse_args(argv) make_harmonic_wrapper(args) -def make_harmonic_wrapper(args): + + +def make_harmonic_wrapper(args: Any) -> None: """Wrapper function for make_harmonic command. Parameters @@ -91,7 +112,9 @@ def make_harmonic_wrapper(args): """ make_harmonic_main(args.control, args.model_dest, args.output) -def make_harmonic_main(control: str, model_dest: str, output: Any): + + +def make_harmonic_main(control: str, model_dest: str, output: Any) -> None: """Main function for make_harmonic command. Parameters @@ -109,7 +132,7 @@ def make_harmonic_main(control: str, model_dest: str, output: Any): Writes harmonic model to model_dest """ if not turbomole_is_installed_or_prefixed(): - raise RuntimeError("Turbomole is not available") + raise ExternalCodeError("Turbomole is not available") print(f"Reading Turbomole control file from {control}", file=output) print(file=output) @@ -122,15 +145,18 @@ def make_harmonic_main(control: str, model_dest: str, output: Any): masses = turbo.get_masses(symbols) print("Reference geometry:", file=output) - print(f"{'el':>3s} {'x (Å)':>20s} {'y (Å)':>20s} {'z (Å)':>20s} {'mass (amu)':>20s}", file=output) + print( + f"{'el':>3s} {'x (Å)':>20s} {'y (Å)':>20s} {'z (Å)':>20s} {'mass (amu)':>20s}", + file=output) print("-" * 100, file=output) ms = masses.reshape(-1, 3)[:, 0] for symbol, coord, mass in zip(symbols, coords.reshape(-1, 3), ms): - print(f"{symbol:3s} " - f"{coord[0]: 20.16g} {coord[1]: 20.16g} {coord[2]: 20.16g} " - f"{mass / amu: 20.16g}", - file=output) + print( + f"{symbol:3s} " + f"{coord[0]: 20.16g} {coord[1]: 20.16g} {coord[2]: 20.16g} " + f"{mass / amu: 20.16g}", + file=output) print(file=output) # read Hessian @@ -141,7 +167,12 @@ def make_harmonic_main(control: str, model_dest: str, output: Any): print(f" {i:3d} {val:20.10g}", file=output) print(file=output) - harmonic = HarmonicModel(coords, 0.0, hessian, masses, symbols, ndims=3, + harmonic = HarmonicModel(coords, + 0.0, + hessian, + masses, + symbols, + ndims=3, nparticles=len(symbols)) print(f"Writing harmonic model to {model_dest}", file=output) diff --git a/mudslide/turboparse/__init__.py b/mudslide/turboparse/__init__.py index a8657c9..a344e57 100644 --- a/mudslide/turboparse/__init__.py +++ b/mudslide/turboparse/__init__.py @@ -1,3 +1,5 @@ +"""Turboparse: a parser for Turbomole quantum chemistry output files.""" + from .parse_turbo import parse_turbo __version__ = "0.1.1" diff --git a/mudslide/turboparse/common_parser.py b/mudslide/turboparse/common_parser.py index fe09e8b..4c4bdb0 100644 --- a/mudslide/turboparse/common_parser.py +++ b/mudslide/turboparse/common_parser.py @@ -1,10 +1,21 @@ #!/usr/bin/env python +"""Parsers for data common across multiple Turbomole modules. -from .section_parser import ParseSection +Includes parsers for basis set information, DFT settings, ground state +properties, cartesian gradients, and nonadiabatic coupling vectors. +""" +from __future__ import annotations + +import re +from typing import Any, Callable + +from .section_parser import ParseSection, ParserProtocol +from .stack_iterator import StackIterator from .line_parser import LineParser, SimpleLineParser -def fortran_float(x): +def fortran_float(x: str) -> float: + """Convert a Fortran-formatted float string (using D/d exponent) to Python float.""" x = x.replace("D", "E") x = x.replace("d", "E") return float(x) @@ -13,28 +24,31 @@ def fortran_float(x): class CompParser(LineParser): """Parse separate components of quantities that go elec, nuc, total""" - def process(self, m, out): + def process(self, m: re.Match[str], out: dict[str, Any]) -> None: out["elec"].append(float(m.group(1))) out["nuc"].append(float(m.group(2))) out["total"].append(float(m.group(3))) class BasisParser(ParseSection): + """Parser for basis set information (atoms, primitives, contractions).""" name = "basis" - def __init__(self): + def __init__(self) -> None: super().__init__(r"basis set information", r"total:") self.parsers = [ - SimpleLineParser(r"([a-z]+)\s+([0-9]+)\s+([0-9]+)\s+([0-9]+)\s+(\S+)\s+(\S+)", - ["atom", "natom", "nprim", "ncont", "nick", "contraction"], - types=[str, int, int, int, str, str], - title="list", - multi=True), - SimpleLineParser(r"total:\s+([0-9]+)\s+([0-9]+)\s+([0-9]+)", ["natoms", "nprim", "ncont"], + SimpleLineParser( + r"([a-z]+)\s+([0-9]+)\s+([0-9]+)\s+([0-9]+)\s+(\S+)\s+(\S+)", + ["atom", "natom", "nprim", "ncont", "nick", "contraction"], + types=[str, int, int, int, str, str], + title="list", + multi=True), + SimpleLineParser(r"total:\s+([0-9]+)\s+([0-9]+)\s+([0-9]+)", + ["natoms", "nprim", "ncont"], types=[int, int, int]), ] - def clean(self, liter, out): + def clean(self, liter: StackIterator, out: dict[str, Any]) -> None: bases = {x["nick"] for x in out["list"]} if len(bases) == 1: out["nick"] = list(bases)[0] @@ -43,18 +57,24 @@ def clean(self, liter, out): class DFTParser(ParseSection): + """Parser for DFT settings (grid size, weight derivatives).""" name = "dft" - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*density functional\s*$", r"partition sharpness") self.parsers = [ - SimpleLineParser(r"spherical gridsize\s*:\s*(\S+)", ["gridsize"], types=[str]), - SimpleLineParser(r"iterations will be done with (small) grid", ["mgrid"], types=[str]), - SimpleLineParser(r"Derivatives of quadrature weights will be (included)", ["weightderivatives"], + SimpleLineParser(r"spherical gridsize\s*:\s*(\S+)", ["gridsize"], + types=[str]), + SimpleLineParser(r"iterations will be done with (small) grid", + ["mgrid"], types=[str]), + SimpleLineParser( + r"Derivatives of quadrature weights will be (included)", + ["weightderivatives"], + types=[str]), ] - def clean(self, liter, out): + def clean(self, liter: StackIterator, out: dict[str, Any]) -> None: if "mgrid" in out: out["gridsize"] = "m" + out["gridsize"] del out["mgrid"] @@ -64,15 +84,16 @@ class GroundDipole(ParseSection): """Parse ground dipole moment""" name = "dipole" - def __init__(self): - super().__init__(r"Electric dipole moment", r" z\s+\S+\s+\S+\s+\S+\s+Norm ") + def __init__(self) -> None: + super().__init__(r"Electric dipole moment", + r" z\s+\S+\s+\S+\s+\S+\s+Norm ") self.parsers = [ CompParser(r" x\s+(\S+)\s+(\S+)\s+(\S+)\s+Norm:"), CompParser(r" y\s+(\S+)\s+(\S+)\s+(\S+)"), CompParser(r" z\s+(\S+)\s+(\S+)\s+(\S+)\s+Norm"), ] - def prepare(self, out): + def prepare(self, out: dict[str, Any]) -> dict[str, Any]: out[self.name] = {} out[self.name]["elec"] = [] out[self.name]["nuc"] = [] @@ -85,7 +106,7 @@ class GroundParser(ParseSection): """Parser for ground state properties""" name = "ground" - def __init__(self): + def __init__(self) -> None: super().__init__(r"Ground state", r"\S*=====+") self.parsers = [GroundDipole()] @@ -93,47 +114,57 @@ def __init__(self): class VarLineParser(LineParser): """A line parser that can handle variable length lines""" - def __init__(self, reg, title, vars_type=str): + def __init__(self, reg: str, title: str, vars_type: Callable[..., Any] = str) -> None: super().__init__(reg) self.title = title self.vars_type = vars_type - def process(self, match, out): + def process(self, m: re.Match[str], out: dict[str, Any]) -> None: if self.title not in out: out[self.title] = [] - for group in match.groups(): + for group in m.groups(): if group is not None: out[self.title].append(self.vars_type(group)) class CoordParser(ParseSection): - """ - Parser for regex for parsing gradients and NAC coupling. + """Parser for atom-indexed Cartesian vector data (gradients and NAC couplings). + + Parses blocks of data formatted as atom labels followed by dx, dy, dz + components in Fortran notation. After parsing, the clean() method combines + the components into a single 'd/dR' list of [dx, dy, dz] per atom. """ name = "regex" - atom_reg = r"^\s*ATOM\s+(\d+\ \D+)" + 4 * r"(?:\s+(\d+\ \D+))?" + r"\s*$" + atom_reg: str = r"^\s*ATOM\s+(\d+\ \D+)" + 4 * r"(?:\s+(\d+\ \D+))?" + r"\s*$" - def __init__(self, head, tail): + def __init__(self, head: str, tail: str) -> None: super().__init__(head, tail, multi=True) - self.parsers = [ + parsers: list[ParserProtocol] = [ VarLineParser(reg=self.atom_reg, title="atom_list", vars_type=str), - ] + [ - VarLineParser( - reg=(rf"^\s*d\w*/d{coord}\s+(-*\d+.\d+D(?:\+|-)\d+)" + - 4 * r"(?:\s*(-*\d+.\d+D(?:\+|-)\d+))?" + r"\s*$"), - title=f"d_d{coord}", - vars_type=fortran_float) for coord in "xyz" ] - - def clean(self, liter, out): + for coord in "xyz": + parsers.append( + VarLineParser( + reg=(rf"^\s*d\w*/d{coord}\s+(-*\d+.\d+D(?:\+|-)\d+)" + + 4 * r"(?:\s*(-*\d+.\d+D(?:\+|-)\d+))?" + r"\s*$"), + title=f"d_d{coord}", + vars_type=fortran_float)) + self.parsers = parsers + + def clean(self, liter: StackIterator, out: dict[str, Any]) -> None: atom_list = [atom.rstrip() for atom in out["atom_list"]] out["atom_list"] = atom_list # not always sure the d_dx, d_dy, d_dz exist, but if they do, combine them # into (natoms, 3) array with xyz contiguous per atom - components = [ x for x in [ "dE_dx", "dE_dy", "dE_dz", "d_dx", "d_dy", "d_dz"] if x in out ] - out["d/dR"] = [list(vals) for vals in zip(*[out[c] for c in components])] + components = [ + x for x in ["dE_dx", "dE_dy", "dE_dz", "d_dx", "d_dy", "d_dz"] + if x in out + ] + out["d/dR"] = [ + list(vals) for vals in zip(*[out[c] for c in components]) + ] for component in components: del out[component] @@ -144,36 +175,41 @@ class NACParser(CoordParser): """ name = "coupling" - coupled_states_reg = r"^\s*<\s*(\d+)\s*\|\s*\w+/\w+\s*\|\s*(\d+)>\s*$" + coupled_states_reg: str = r"^\s*<\s*(\d+)\s*\|\s*\w+/\w+\s*\|\s*(\d+)>\s*$" - def __init__(self): + def __init__(self) -> None: # Header may or may not have (state/method) at end head = r"\s+cartesian\s+nonadiabatic\s+coupling\s+matrix\s+elements(?:\s+\((\d+)/(\w+)\))?" tail = r"maximum component of gradient" super().__init__(head, tail) self.parsers.insert( - 0, SimpleLineParser(self.coupled_states_reg, ["bra_state", "ket_state"], types=[int, int])) + 0, + SimpleLineParser(self.coupled_states_reg, + ["bra_state", "ket_state"], + types=[int, int])) # Constants for the two gradient types -EXCITED_STATE_GRADIENT_HEAD = (r"(?:cartesian\s+gradients\s+of\s+excited\s+state\s+(?P\d+)|" - r"cartesian\s+gradient\s+of\s+the\s+energy)\s+\((\w+)/(\w+)\)") -GROUND_STATE_GRADIENT_HEAD = r"cartesian\s+gradient\s+of\s+the\s+energy\s+\((\w+)/(\w+)\)" +EXCITED_STATE_GRADIENT_HEAD: str = ( + r"(?:cartesian\s+gradients\s+of\s+excited\s+state\s+(?P\d+)|" + r"cartesian\s+gradient\s+of\s+the\s+energy)\s+\((\w+)/(\w+)\)") +GROUND_STATE_GRADIENT_HEAD: str = r"cartesian\s+gradient\s+of\s+the\s+energy\s+\((\w+)/(\w+)\)" class GradientDataParser(CoordParser): """Parser for cartesian gradient data (ground or excited state).""" name = "gradient" - def __init__(self, head): + def __init__(self, head: str) -> None: # Tail matches end of gradient section: either "resulting FORCE" (when NAC follows), # "maximum component of gradient" (when no NAC), or start of NAC section tail = r"(?:resulting FORCE|maximum component of gradient|cartesian\s+nonadiabatic)" super().__init__(head, tail) - def prepare(self, out): + def prepare(self, out: dict[str, Any]) -> dict[str, Any]: dest = super().prepare(out) try: + assert self.lastsearch is not None index = self.lastsearch.group("index") if index is not None: dest["index"] = int(index) diff --git a/mudslide/turboparse/freeh_parser.py b/mudslide/turboparse/freeh_parser.py index 245006f..41ec8bc 100644 --- a/mudslide/turboparse/freeh_parser.py +++ b/mudslide/turboparse/freeh_parser.py @@ -1,40 +1,59 @@ #!/usr/bin/env python +"""Parser for Turbomole freeh (free enthalpy) module output. + +Extracts thermodynamic data including temperature, pressure, partition +functions (translational, rotational, vibrational), chemical potential, +energy, entropy, heat capacities (Cv, Cp), and enthalpy. +""" +from __future__ import annotations import re +from typing import Any from .section_parser import ParseSection +from .stack_iterator import StackIterator + +FLT: str = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?" -FLT = r"[-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?" -def try_to_fill_with_float(dict_obj, key, value): - """Try to fill a dict with a float, but do nothing if it fails.""" +def try_to_fill_with_float(dict_obj: dict[str, float], key: str, + value: str | None) -> None: + """Store value as float in dict_obj[key], silently ignoring conversion failures.""" try: - dict_obj[key] = float(value) - except: + dict_obj[key] = float(value) # type: ignore[arg-type] + except (ValueError, TypeError): pass + class FreeHData(ParseSection): + """Parser for the thermodynamic data table within a freeh section.""" name = '' - start1st_re = re.compile(r"T\s+p\s+ln\(qtrans\)\s+ln\(qrot\)\s+ln\(qvib\)\s+chem\.pot\.\s+energy\s+entropy") - data1st_re = re.compile(rf"({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})") - start2nd_re = re.compile(r"\(K\)\s+\(MPa\)\s+\(kJ/mol-K\)\s+\(kJ/mol-K\)(?:\s+\(kJ/mol\))?") - data2nd_re = re.compile(rf"({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s*({FLT})?") - end_re = re.compile(r"\*{50}") - - def __init__(self): + start1st_re: re.Pattern[str] = re.compile( + r"T\s+p\s+ln\(qtrans\)\s+ln\(qrot\)\s+ln\(qvib\)\s+chem\.pot\.\s+energy\s+entropy" + ) + data1st_re: re.Pattern[str] = re.compile( + rf"({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s+({FLT})" + ) + start2nd_re: re.Pattern[str] = re.compile( + r"\(K\)\s+\(MPa\)\s+\(kJ/mol-K\)\s+\(kJ/mol-K\)(?:\s+\(kJ/mol\))?") + data2nd_re: re.Pattern[str] = re.compile( + rf"({FLT})\s+({FLT})\s+({FLT})\s+({FLT})\s*({FLT})?") + end_re: re.Pattern[str] = re.compile(r"\*{50}") + + def __init__(self) -> None: super().__init__( r"T\s+p\s+ln\(qtrans\)\s+ln\(qrot\)\s+ln\(qvib\)\s+chem\.pot\.\s+energy\s+entropy", r"\*{50}", multi=False) - def parse_driver(self, liter, out): + def parse_driver(self, liter: StackIterator, out: dict[str, Any]) -> bool: """ Driver to FreeH section return: advanced (whether next has been called on liter) """ - data = [] + data: list[dict[str, float]] = [] done = False advanced = False while not done: @@ -47,15 +66,19 @@ def parse_driver(self, liter, out): while not done1st: m1st = self.data1st_re.search(liter.top()) if m1st: - new_data = {} + new_data: dict[str, float] = {} try_to_fill_with_float(new_data, 'T', m1st.group(1)) try_to_fill_with_float(new_data, 'P', m1st.group(2)) - try_to_fill_with_float(new_data, 'qtrans', m1st.group(3)) + try_to_fill_with_float(new_data, 'qtrans', + m1st.group(3)) try_to_fill_with_float(new_data, 'qrot', m1st.group(4)) try_to_fill_with_float(new_data, 'qvib', m1st.group(5)) - try_to_fill_with_float(new_data, 'chem.pot.', m1st.group(6)) - try_to_fill_with_float(new_data, 'energy', m1st.group(7)) - try_to_fill_with_float(new_data, 'entropy', m1st.group(8)) + try_to_fill_with_float(new_data, 'chem.pot.', + m1st.group(6)) + try_to_fill_with_float(new_data, 'energy', + m1st.group(7)) + try_to_fill_with_float(new_data, 'entropy', + m1st.group(8)) data.append(new_data) mend = self.start2nd_re.search(liter.top()) @@ -111,8 +134,9 @@ def parse_driver(self, liter, out): class FreeHParser(ParseSection): + """Parser for the freeh (free enthalpy) module output.""" name = "freeh" - def __init__(self): + def __init__(self) -> None: super().__init__(r"f r e e e n t h a l p y", r"freeh\s*:\s*all done") self.parsers = [FreeHData()] diff --git a/mudslide/turboparse/line_parser.py b/mudslide/turboparse/line_parser.py index 54ce86d..4adb906 100644 --- a/mudslide/turboparse/line_parser.py +++ b/mudslide/turboparse/line_parser.py @@ -1,15 +1,27 @@ #!/usr/bin/env python +"""Line-level parsers that match individual lines against regex patterns. + +LineParser is the base class; SimpleLineParser handles the common case +of extracting named groups with type conversion, and BooleanLineParser +stores True/False based on which of two patterns matches. +""" +from __future__ import annotations import re +from typing import Any, Callable + +from mudslide.exceptions import ConfigurationError + +from .stack_iterator import StackIterator class LineParser: """Base class to parse a single line and return results. Implementations require a process function""" - def __init__(self, reg): - self.reg = re.compile(reg) + def __init__(self, reg: str) -> None: + self.reg: re.Pattern[str] = re.compile(reg) - def parse(self, liter, out): + def parse(self, liter: StackIterator, out: dict[str, Any]) -> tuple[bool, bool]: """ Parse line found at liter.top() @@ -20,18 +32,45 @@ def parse(self, liter, out): self.process(result, out) return bool(result), False - def process(self, m, out): + def process(self, m: re.Match[str], out: dict[str, Any]) -> None: + """Process a regex match and store results in out. Must be overridden.""" raise NotImplementedError("Subclasses must implement process()") class SimpleLineParser(LineParser): - """Parse a single line and return a list of all matched groups""" - - def __init__(self, reg, names, converter=None, types=None, title="", multi=False, first_only=False): + """Line parser that extracts regex groups into named fields. + + Matches a line against a regex and stores the captured groups as named + entries in the output dict, with optional type conversion. + + Args: + reg: Regex pattern with capturing groups. + names: List of keys corresponding to each captured group. + converter: A single callable applied to all groups (e.g., float). + Mutually exclusive with types. + types: List of callables, one per group, for individual type + conversion. If neither converter nor types is given, all values + are stored as str. + title: If non-empty, results are stored as a sub-dict under this key + rather than merged directly into the output dict. + multi: If True, each match appends to a list under title (requires + title to be set). + first_only: If True, only the first match is stored; subsequent + matches for the same key are ignored. + """ + + def __init__(self, + reg: str, + names: list[str], + converter: Callable[..., Any] | None = None, + types: list[Callable[..., Any]] | None = None, + title: str = "", + multi: bool = False, + first_only: bool = False) -> None: super().__init__(reg) self.names = names if converter is None and types is None: - self.types = [str] * len(names) + self.types: list[Callable[..., Any]] = [str] * len(names) elif types is not None: self.types = types elif converter is not None: @@ -41,10 +80,12 @@ def __init__(self, reg, names, converter=None, types=None, title="", multi=False self.first_only = first_only if self.multi and self.title == "": - raise ValueError("SimpleLineParser in multi mode requires title") + raise ConfigurationError("SimpleLineParser in multi mode requires title") - def process(self, m, out): - data = {n: self.types[i](m.group(i + 1)) for i, n in enumerate(self.names)} + def process(self, m: re.Match[str], out: dict[str, Any]) -> None: + data = { + n: self.types[i](m.group(i + 1)) for i, n in enumerate(self.names) + } if not self.multi: if self.title != "": if not (self.first_only and self.title in out): @@ -61,15 +102,27 @@ def process(self, m, out): class BooleanLineParser(LineParser): - """Parse a line and store a boolean based on success/failure regex matches.""" - - def __init__(self, success_reg, failure_reg, key, first_only=False): - self.success_reg = re.compile(success_reg) - self.failure_reg = re.compile(failure_reg) + """Line parser that stores a boolean based on matching one of two patterns. + + Tests the current line against a success regex and a failure regex. + If the success regex matches, stores True; if the failure regex matches, + stores False. If neither matches, reports no match. + + Args: + success_reg: Regex pattern indicating a True result. + failure_reg: Regex pattern indicating a False result. + key: Dict key under which the boolean is stored. + first_only: If True, only the first match is stored. + """ + + def __init__(self, success_reg: str, failure_reg: str, key: str, + first_only: bool = False) -> None: + self.success_reg: re.Pattern[str] = re.compile(success_reg) + self.failure_reg: re.Pattern[str] = re.compile(failure_reg) self.key = key self.first_only = first_only - def parse(self, liter, out): + def parse(self, liter: StackIterator, out: dict[str, Any]) -> tuple[bool, bool]: """ Parse line found at liter.top() @@ -88,7 +141,6 @@ def parse(self, liter, out): return False, False - def process(self, value, out): + def process(self, value: bool, out: dict[str, Any]) -> None: # type: ignore[override] if not (self.first_only and self.key in out): out[self.key] = value - diff --git a/mudslide/turboparse/parse.py b/mudslide/turboparse/parse.py index 5734a4c..a4ffff2 100644 --- a/mudslide/turboparse/parse.py +++ b/mudslide/turboparse/parse.py @@ -1,4 +1,6 @@ #!/usr/bin/env python +"""CLI entry point for the turboparse command.""" +from __future__ import annotations import argparse import json @@ -10,7 +12,12 @@ from ..yaml_format import CompactDumper -def parse(argv=None): +def parse(argv: list[str] | None = None) -> None: + """CLI entry point for the turboparse command. + + Parses a Turbomole output file and prints the extracted data as + JSON or YAML (default: YAML). + """ ap = argparse.ArgumentParser( description= "Collects excited state information from an egrad run and prepares as JSON", diff --git a/mudslide/turboparse/parse_turbo.py b/mudslide/turboparse/parse_turbo.py index c651519..2b7539c 100644 --- a/mudslide/turboparse/parse_turbo.py +++ b/mudslide/turboparse/parse_turbo.py @@ -1,4 +1,13 @@ #!/usr/bin/env python +"""Top-level parser for Turbomole output files. + +Provides parse_turbo(), which reads a Turbomole output file (or any +iterable of lines) and returns a nested dict of parsed data organized +by Turbomole module (ridft, dscf, egrad, escf, freeh, etc.). +""" +from __future__ import annotations + +from typing import Any, Iterable from .stack_iterator import StackIterator from .section_parser import ParseSection @@ -9,9 +18,14 @@ class TurboParser(ParseSection): + """Root parser that aggregates all Turbomole module parsers. + + Uses head=r".*" to match any line and tail=r"a^" (unmatchable) so it + never terminates on its own -- parsing runs until the input is exhausted. + """ name = "" - def __init__(self): + def __init__(self) -> None: super().__init__(r".*", r"a^") self.parsers = [ RIDFTParser(), @@ -25,11 +39,23 @@ def __init__(self): ] -def parse_turbo(iterable): +def parse_turbo(iterable: Iterable[str]) -> dict[str, Any]: + """Parse a Turbomole output file and return structured data. + + Args: + iterable: An iterable of strings (lines), typically an open file + handle pointing to a Turbomole output file. + + Returns: + A nested dict keyed by module name (e.g., 'ridft', 'egrad', 'escf', + 'freeh'). Each value contains the parsed data for that module, + including energies, excited states, gradients, couplings, and + thermodynamic data as applicable. + """ liter = StackIterator(iterable) parser = TurboParser() - out = {} + out: dict[str, Any] = {} try: next(liter) while True: diff --git a/mudslide/turboparse/response_parser.py b/mudslide/turboparse/response_parser.py index 0aa507f..dba9b14 100644 --- a/mudslide/turboparse/response_parser.py +++ b/mudslide/turboparse/response_parser.py @@ -1,7 +1,19 @@ #!/usr/bin/env python +"""Parsers for Turbomole response property and excited state output. + +Handles output from egrad and escf modules, extracting excited state +energies, transition dipole moments, state-to-state properties, +two-photon absorption amplitudes, hyperpolarizabilities, Davidson +iteration convergence, CPKS iterations, and NAC couplings. +""" +from __future__ import annotations + +import re +from typing import Any from .line_parser import LineParser, SimpleLineParser, BooleanLineParser -from .section_parser import ParseSection +from .section_parser import ParseSection, ParserProtocol +from .stack_iterator import StackIterator from .common_parser import GroundParser, NACParser, GradientDataParser, EXCITED_STATE_GRADIENT_HEAD, fortran_float @@ -14,13 +26,13 @@ class ExoptLineParser(LineParser): - "3 excited states specified in $exopt: 1 2 3" -> [1, 2, 3] """ - def __init__(self): + def __init__(self) -> None: reg = (r"(?:Excited state no\.\s+(\d+)\s+chosen for optimization|" r"Default state chosen:\s+(\d+)|" r"\d+\s+excited states specified in \$exopt:\s+([\d\s]+))") super().__init__(reg) - def process(self, m, out): + def process(self, m: re.Match[str], out: dict[str, Any]) -> None: if m.group(1) is not None: out["exopt"] = [int(m.group(1))] elif m.group(2) is not None: @@ -32,7 +44,7 @@ def process(self, m, out): class ExcitedDipoleParser(ParseSection): """Parser for excited state dipole moments""" - def __init__(self, name, head, tail=r"z\s+(\S+)"): + def __init__(self, name: str, head: str, tail: str = r"z\s+(\S+)") -> None: super().__init__(head, tail) self.name = name self.parsers = [ @@ -41,7 +53,7 @@ def __init__(self, name, head, tail=r"z\s+(\S+)"): SimpleLineParser(r"z\s+(\S+)", ["z"], converter=float), ] - def prepare(self, out): + def prepare(self, out: dict[str, Any]) -> dict[str, Any]: out[self.name] = {"x": 0.0, "y": 0.0, "z": 0.0} return out[self.name] @@ -50,98 +62,126 @@ class ExcitedParser(ParseSection): """Parser for excited state properties""" name = "excited_state" - def __init__(self): - super().__init__(r"\d+ (singlet |triplet |)[abte1234567890]+ excitation", - r"Electric quadrupole transition moment", - multi=True) + def __init__(self) -> None: + super().__init__( + r"\d+ (singlet |triplet |)[abte1234567890]+ excitation", + r"Electric quadrupole transition moment", + multi=True) self.parsers = [ - SimpleLineParser(r"(\d+) (singlet |triplet |)([abte1234567890]*) excitation", - ["index", "spin", "irrep"], - types=[int, str, str]), - SimpleLineParser(r"Excitation energy:\s+(\S+)", ["energy"], converter=float), - ExcitedDipoleParser("diplen", r"Electric transition dipole moment \(length rep"), + SimpleLineParser( + r"(\d+) (singlet |triplet |)([abte1234567890]*) excitation", + ["index", "spin", "irrep"], + types=[int, str, str]), + SimpleLineParser(r"Excitation energy:\s+(\S+)", ["energy"], + converter=float), + ExcitedDipoleParser( + "diplen", r"Electric transition dipole moment \(length rep"), ] class MomentParser(LineParser): + """Parse transition dipole moment matrix elements.""" - def __init__(self, reg=r"<\s*(\d+)\|mu\|\s*(\d+)>:\s+(\S+)\s+(\S+)\s+(\S+)"): + def __init__(self, + reg: str = r"<\s*(\d+)\|mu\|\s*(\d+)>:\s+(\S+)\s+(\S+)\s+(\S+)" + ) -> None: super().__init__(reg) - def process(self, m, out): + def process(self, m: re.Match[str], out: dict[str, Any]) -> None: # type: ignore[override] i, j = int(m.group(1)), int(m.group(2)) dip = [float(m.group(x)) for x in range(3, 6)] if i not in out: - out[i] = {} - out[i][j] = {"diplen": dip} + out[i] = {} # type: ignore[index] + out[i][j] = {"diplen": dip} # type: ignore[index] if j not in out: - out[j] = {} - out[j][i] = {"diplen": dip} + out[j] = {} # type: ignore[index] + out[j][i] = {"diplen": dip} # type: ignore[index] class ExcitedMoments(ParseSection): + """Parser for a block of excited state dipole moment matrix elements.""" - def __init__(self, name, head, tail=r"^\s*$"): + def __init__(self, name: str, head: str, tail: str = r"^\s*$") -> None: super().__init__(head, tail) self.name = name self.parsers = [MomentParser()] class StateToStateParser(ParseSection): + """Parser for state-to-state transition and difference moments.""" name = "state-to-state" - def __init__(self): - super().__init__(r"<\s*\d+\s*\|\s*W\s*\|\s*\d+\s*>.*(?:transition|difference) moments", - r"(<\s*\d+\s*\|\s*W\s*\|\s*\d+\s*>.*(?:transition|difference) moments)|(S\+T\+V CONTRIBUTIONS TO)", - multi=True) + def __init__(self) -> None: + super().__init__( + r"<\s*\d+\s*\|\s*W\s*\|\s*\d+\s*>.*(?:transition|difference) moments", + r"(<\s*\d+\s*\|\s*W\s*\|\s*\d+\s*>.*(?:transition|difference) moments)|(S\+T\+V CONTRIBUTIONS TO)", + multi=True) self.parsers = [ - SimpleLineParser(r"<\s*(\d+)\s*\|\s*W\s*\|\s*(\d+)\s*>", ["bra", "ket"], + SimpleLineParser(r"<\s*(\d+)\s*\|\s*W\s*\|\s*(\d+)\s*>", + ["bra", "ket"], types=[int, int], first_only=True), - ExcitedDipoleParser("diplen", r"Relaxed electric transition dipole moment \(length rep"), + ExcitedDipoleParser( + "diplen", + r"Relaxed electric transition dipole moment \(length rep"), ] class TPAColParser(ParseSection): + """Parser for two-photon absorption tensor columns and cross sections.""" name = "columns" - def __init__(self): + def __init__(self) -> None: super().__init__(r"Column:s+(\S+)", r"sigma_0", multi=True) self.parsers = self._make_parsers() @staticmethod - def _make_parsers(): + def _make_parsers() -> list[ParserProtocol]: return [ SimpleLineParser(r"Column:\s+(\S+)", ["column"], converter=int), - SimpleLineParser(r"xx\s+(\S+)\s+xy\s+(\S+)\s+xz\s+(\S+)", ["xx", "xy", "xz"], converter=float), - SimpleLineParser(r"yx\s+(\S+)\s+yy\s+(\S+)\s+yz\s+(\S+)", ["yx", "yy", "yz"], converter=float), - SimpleLineParser(r"zx\s+(\S+)\s+zy\s+(\S+)\s+zz\s+(\S+)", ["zx", "zy", "zz"], converter=float), - SimpleLineParser(r"\(dF,dG,dH\):\s+(\S+)\s+(\S+)\s+(\S+)", ["df", "dg", "dh"], converter=float), - SimpleLineParser(r"transition strength \[a\.u\.\]:\s*(\S+)\s*(\S+)\s*(\S+)", - ["parallel_strength", "perp_strength", "circular_strength"], + SimpleLineParser(r"xx\s+(\S+)\s+xy\s+(\S+)\s+xz\s+(\S+)", + ["xx", "xy", "xz"], + converter=float), + SimpleLineParser(r"yx\s+(\S+)\s+yy\s+(\S+)\s+yz\s+(\S+)", + ["yx", "yy", "yz"], converter=float), - SimpleLineParser(r"sigma_0 \[1e-50 cm\^4 s\]:\s*(\S+)\s*(\S+)\s*(\S+)", - ["parallel_cross", "perp_cross", "circular_cross"], + SimpleLineParser(r"zx\s+(\S+)\s+zy\s+(\S+)\s+zz\s+(\S+)", + ["zx", "zy", "zz"], converter=float), + SimpleLineParser(r"\(dF,dG,dH\):\s+(\S+)\s+(\S+)\s+(\S+)", + ["df", "dg", "dh"], + converter=float), + SimpleLineParser( + r"transition strength \[a\.u\.\]:\s*(\S+)\s*(\S+)\s*(\S+)", + ["parallel_strength", "perp_strength", "circular_strength"], + converter=float), + SimpleLineParser( + r"sigma_0 \[1e-50 cm\^4 s\]:\s*(\S+)\s*(\S+)\s*(\S+)", + ["parallel_cross", "perp_cross", "circular_cross"], + converter=float), ] class TPAParser(ParseSection): + """Parser for two-photon absorption amplitudes per excited state.""" name = "tpa" - def __init__(self): - super().__init__(r"Two-photon absorption amplitudes for transition to " + - r"the\s+(\d+)\S+\s+electronic excitation in symmetry\s+(\S+)", - r"sigma_0", - multi=True) - self.parsers = [ - SimpleLineParser(r"Two-photon absorption amplitudes for transition to " + - r"the\s+(\d+)\S+\s+electronic excitation in symmetry\s+(\S+)", - ["state", "irrep"], - types=[int, str]), - SimpleLineParser(r"Exc\. energy:\s+(\S+)\s+Hartree,\s+(\S+)", ["E (H)", "E (eV)"], + def __init__(self) -> None: + super().__init__( + r"Two-photon absorption amplitudes for transition to " + + r"the\s+(\d+)\S+\s+electronic excitation in symmetry\s+(\S+)", + r"sigma_0", + multi=True) + parsers: list[ParserProtocol] = [ + SimpleLineParser( + r"Two-photon absorption amplitudes for transition to " + + r"the\s+(\d+)\S+\s+electronic excitation in symmetry\s+(\S+)", + ["state", "irrep"], + types=[int, str]), + SimpleLineParser(r"Exc\. energy:\s+(\S+)\s+Hartree,\s+(\S+)", + ["E (H)", "E (eV)"], converter=float), SimpleLineParser(r"omega_1\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)", ["w1 (H)", "w1 (eV)", "w1 (nm)", "w1 (rcm)"], @@ -150,48 +190,69 @@ def __init__(self): ["w2 (H)", "w2 (eV)", "w2 (nm)", "w2 (rcm)"], converter=float), TPAColParser(), - ] + TPAColParser._make_parsers() + ] + parsers.extend(TPAColParser._make_parsers()) + self.parsers = parsers class HyperParser(ParseSection): + """Parser for first hyperpolarizability tensor components.""" name = "hyper" - def __init__(self): - super().__init__(r"(\d+)\S+ pair of frequencies", r"septor norm", multi=True) - self.parsers = [ - SimpleLineParser(r"Frequencies:\s+(\S+)\s+(\S+)", ["omega_1", "omega_2"], converter=float), - SimpleLineParser(r"Frequencies / eV:\s+(\S+)\s+(\S+)", ["omega_1 (eV)", "omega_2 (eV)"], + def __init__(self) -> None: + super().__init__(r"(\d+)\S+ pair of frequencies", + r"septor norm", + multi=True) + parsers: list[ParserProtocol] = [ + SimpleLineParser(r"Frequencies:\s+(\S+)\s+(\S+)", + ["omega_1", "omega_2"], converter=float), - SimpleLineParser(r"Frequencies / nm:\s+(\S+)\s+(\S+)", ["omega_1 (nm)", "omega_2 (nm)"], + SimpleLineParser(r"Frequencies / eV:\s+(\S+)\s+(\S+)", + ["omega_1 (eV)", "omega_2 (eV)"], + converter=float), + SimpleLineParser(r"Frequencies / nm:\s+(\S+)\s+(\S+)", + ["omega_1 (nm)", "omega_2 (nm)"], converter=str), - SimpleLineParser(r"Frequencies / cm\^\(-1\):\s+(\S+)\s+(\S+)", ["omega_1 (rcm)", "omega_2 (rcm)"], + SimpleLineParser(r"Frequencies / cm\^\(-1\):\s+(\S+)\s+(\S+)", + ["omega_1 (rcm)", "omega_2 (rcm)"], converter=float), - SimpleLineParser(r"scalar norm:\s+(\S+)", ["scalar"], converter=float), - SimpleLineParser(r"vector norms:\s+(\S+)\s+(\S+)\s+(\S+)", ["v1", "v2", "v3"], converter=float), - SimpleLineParser(r"deviator 1 eigenvalues:\s+(\S+)\s+(\S+)\s+(\S+)\s+", ["d1_1", "d1_2", "d1_3"], + SimpleLineParser(r"scalar norm:\s+(\S+)", ["scalar"], converter=float), - SimpleLineParser(r"deviator 2 eigenvalues:\s+(\S+)\s+(\S+)\s+(\S+)\s+", ["d2_1", "d2_2", "d2_3"], + SimpleLineParser(r"vector norms:\s+(\S+)\s+(\S+)\s+(\S+)", + ["v1", "v2", "v3"], converter=float), + SimpleLineParser( + r"deviator 1 eigenvalues:\s+(\S+)\s+(\S+)\s+(\S+)\s+", + ["d1_1", "d1_2", "d1_3"], + converter=float), + SimpleLineParser( + r"deviator 2 eigenvalues:\s+(\S+)\s+(\S+)\s+(\S+)\s+", + ["d2_1", "d2_2", "d2_3"], + converter=float), SimpleLineParser(r"septor nom:\s+(\S+)", ["sep"], converter=float), - ] + [ - SimpleLineParser(rf"x{xy}\s+(\S+)\s+y{xy}\s+(\S+)\s+z{xy}\s+(\S+)", - [f"x{xy}", f"y{xy}", f"z{xy}"], - converter=float) - for xy in ["xx", "yx", "zx", "xy", "yy", "zy", "xz", "yz", "zz"] ] + for xy in ["xx", "yx", "zx", "xy", "yy", "zy", "xz", "yz", "zz"]: + parsers.append( + SimpleLineParser(rf"x{xy}\s+(\S+)\s+y{xy}\s+(\S+)\s+z{xy}\s+(\S+)", + [f"x{xy}", f"y{xy}", f"z{xy}"], + converter=float)) + self.parsers = parsers class _DavidsonIterationParsers: """Common parsers for Block Davidson iteration sections""" @staticmethod - def make_parsers(): + def make_parsers() -> list[ParserProtocol]: + """Create parsers for iteration step number, residual norm, and convergence.""" return [ - SimpleLineParser(r"^\s*(\d+)\s+\S+\s+\d+\s+(\S+)\s*$", ["step", "max_residual_norm"], + SimpleLineParser(r"^\s*(\d+)\s+\S+\s+\d+\s+(\S+)\s*$", + ["step", "max_residual_norm"], types=[int, fortran_float], title="iterations", multi=True), - BooleanLineParser(r"^\s*converged!", r"not converged!", "converged"), + BooleanLineParser(r"^\s*converged!", r"not converged!", + "converged"), ] @@ -199,7 +260,7 @@ class CPKSParser(ParseSection): """Parser for CPKS iterations""" name = "cpks" - def __init__(self): + def __init__(self) -> None: super().__init__(r"CPKS right-hand side", r"(not )?converged!") self.parsers = _DavidsonIterationParsers.make_parsers() @@ -208,14 +269,21 @@ class DavidsonParser(ParseSection): """Parser for excitation vector Davidson iterations""" name = "davidson" - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*excitation vector\s*$", r"(not )?converged!") self.parsers = _DavidsonIterationParsers.make_parsers() class EgradEscfParser(ParseSection): + """Base parser for egrad/escf output sections. + + Contains the full set of sub-parsers for excited state properties: + excitation energies, transition moments, gradients, NAC couplings, + TPA amplitudes, and hyperpolarizabilities. The clean() method assigns + gradient indices from exopt data when both are present. + """ - def __init__(self, head, tail): + def __init__(self, head: str, tail: str) -> None: super().__init__(head, tail) self.parsers = [ ExoptLineParser(), @@ -223,10 +291,14 @@ def __init__(self, head, tail): CPKSParser(), GroundParser(), ExcitedParser(), - ExcitedMoments("relaxed moments", r"Fully relaxed moments of the excited states"), - ExcitedMoments("relaxed transitions", r"Fully relaxed state-to-state transition moments"), - ExcitedMoments("unrelaxed moments", r"Unrelaxed moments of the excited states"), - ExcitedMoments("unrelaxed transitions", r"Unrelaxed state-to-state transition moments"), + ExcitedMoments("relaxed moments", + r"Fully relaxed moments of the excited states"), + ExcitedMoments("relaxed transitions", + r"Fully relaxed state-to-state transition moments"), + ExcitedMoments("unrelaxed moments", + r"Unrelaxed moments of the excited states"), + ExcitedMoments("unrelaxed transitions", + r"Unrelaxed state-to-state transition moments"), TPAParser(), HyperParser(), StateToStateParser(), @@ -234,7 +306,7 @@ def __init__(self, head, tail): GradientDataParser(EXCITED_STATE_GRADIENT_HEAD), ] - def clean(self, liter, out): + def clean(self, liter: StackIterator, out: dict[str, Any]) -> None: exopt = out.get("exopt") gradients = out.get("gradient") if not exopt or not gradients or len(exopt) != len(gradients): @@ -245,14 +317,16 @@ def clean(self, liter, out): class EgradParser(EgradEscfParser): + """Parser for the egrad (excited state gradient) module output.""" name = "egrad" - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*e g r a d", r"egrad\s*:\s*all done") class EscfParser(EgradEscfParser): + """Parser for the escf (excited state SCF) module output.""" name = "escf" - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*e s c f", r"escf\s*:\s*all done") diff --git a/mudslide/turboparse/scf_parser.py b/mudslide/turboparse/scf_parser.py index 2157804..9fb7cc0 100644 --- a/mudslide/turboparse/scf_parser.py +++ b/mudslide/turboparse/scf_parser.py @@ -1,4 +1,11 @@ #!/usr/bin/env python +"""Parsers for Turbomole SCF and gradient module output. + +Handles output from ridft, dscf (SCF calculations) and rdgrad, grad +(gradient calculations), extracting energies, convergence status, +basis set info, and gradient data. +""" +from __future__ import annotations from .section_parser import ParseSection from .line_parser import BooleanLineParser, SimpleLineParser @@ -6,8 +13,9 @@ class SCFParser(ParseSection): + """Base parser for SCF module output (energy, convergence, basis, DFT settings).""" - def __init__(self, head, tail): + def __init__(self, head: str, tail: str) -> None: super().__init__(head, tail) self.parsers = [ SimpleLineParser(r"\|\s*total energy\s*=\s*(\S+)\s*\|", @@ -27,23 +35,25 @@ def __init__(self, head, tail): class RIDFTParser(SCFParser): + """Parser for the ridft (RI-DFT SCF) module output.""" name = 'ridft' - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*r i d f t", r"ridft\s*:\s*all done") class DSCFParser(SCFParser): + """Parser for the dscf (conventional SCF) module output.""" name = 'dscf' - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*d s c f", r"dscf\s*:\s*all done") class GradientModuleParser(ParseSection): """Base parser for gradient module output (grad, rdgrad).""" - def __init__(self, head, tail): + def __init__(self, head: str, tail: str) -> None: super().__init__(head, tail) self.parsers = [ DFTParser(), @@ -55,7 +65,7 @@ class RdgradModuleParser(GradientModuleParser): """Parser for rdgrad module output.""" name = 'rdgrad' - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*r d g r a d", r"rdgrad\s+:\s*all done") @@ -63,5 +73,5 @@ class GradModuleParser(GradientModuleParser): """Parser for grad module output.""" name = 'grad' - def __init__(self): + def __init__(self) -> None: super().__init__(r"^\s*g r a d", r"grad\s+:\s*all done") diff --git a/mudslide/turboparse/section_parser.py b/mudslide/turboparse/section_parser.py index a211e87..eaaeeae 100644 --- a/mudslide/turboparse/section_parser.py +++ b/mudslide/turboparse/section_parser.py @@ -1,33 +1,67 @@ #!/usr/bin/env python +"""Section-based parser framework for structured text output. + +ParseSection matches a region of text delimited by head and tail regex +patterns, then delegates line-by-line parsing to child parsers. +""" +from __future__ import annotations import re +from typing import Any, Protocol + +from .stack_iterator import StackIterator + + +class ParserProtocol(Protocol): + """Protocol for objects that can parse lines from a StackIterator.""" + + def parse(self, liter: StackIterator, out: dict[str, Any]) -> tuple[bool, bool]: ... class ParseSection: - """Parse a section by looping over attached parsers until tail search is true""" - DEBUG = False - name = "" - parsers = [] - - def __init__(self, head, tail, multi=False): - self.head = re.compile(head) - self.tail = re.compile(tail) - self.lastsearch = None + """Parser for a delimited section of output. + + A section is defined by a head regex (start marker) and a tail regex + (end marker). When parse() detects the head pattern on the current line, + it enters parse_driver() which iterates through lines, calling each child + parser on every line until the tail pattern is matched. + + Child parsers (stored in self.parsers) can be LineParser instances for + single-line matches or nested ParseSection instances for subsections. + + Attributes: + name: Key under which parsed results are stored in the output dict. + Empty string means results merge into the parent dict. + multi: If True, each match appends a new dict to a list under + self.name, allowing multiple instances of the same section. + parsers: List of child parsers to apply within this section. + """ + DEBUG: bool = False + name: str = "" + parsers: list[ParserProtocol] = [] + + def __init__(self, head: str, tail: str, multi: bool = False) -> None: + self.head: re.Pattern[str] = re.compile(head) + self.tail: re.Pattern[str] = re.compile(tail) + self.lastsearch: re.Match[str] | None = None self.multi = multi - def test(self, reg, line): + def test(self, reg: re.Pattern[str], line: str) -> re.Match[str] | None: + """Test line against a regex and store the match result.""" self.lastsearch = reg.search(line) return self.lastsearch - def test_head(self, line): + def test_head(self, line: str) -> re.Match[str] | None: + """Test if line matches the head (section start) pattern.""" self.lastsearch = self.head.search(line) return self.lastsearch - def test_tail(self, line): + def test_tail(self, line: str) -> re.Match[str] | None: + """Test if line matches the tail (section end) pattern.""" self.lastsearch = self.tail.search(line) return self.lastsearch - def parse_driver(self, liter, out): + def parse_driver(self, liter: StackIterator, out: dict[str, Any]) -> bool: """ Driver to parse a section @@ -54,7 +88,7 @@ def parse_driver(self, liter, out): return advanced - def parse(self, liter, out): + def parse(self, liter: StackIterator, out: dict[str, Any]) -> tuple[bool, bool]: """ Parse line found at liter.top() @@ -74,7 +108,13 @@ def parse(self, liter, out): print(f"No match for {self.name} at {liter.top().strip()}") return found, advanced - def prepare(self, out): + def prepare(self, out: dict[str, Any]) -> dict[str, Any]: + """Set up the output dict entry for this section and return the destination. + + If name is empty, returns the parent dict directly (results merge in). + If multi is True, appends a new dict to a list under self.name. + Otherwise, creates a single dict under self.name. + """ if self.name == "": return out if self.multi: @@ -85,5 +125,10 @@ def prepare(self, out): out[self.name] = {} return out[self.name] - def clean(self, liter, out): + def clean(self, liter: StackIterator, out: dict[str, Any]) -> None: + """Post-processing hook called after a section is fully parsed. + + Override in subclasses to transform or validate parsed data. + Default implementation does nothing. + """ return diff --git a/mudslide/turboparse/stack_iterator.py b/mudslide/turboparse/stack_iterator.py index f8578f1..2393821 100644 --- a/mudslide/turboparse/stack_iterator.py +++ b/mudslide/turboparse/stack_iterator.py @@ -1,29 +1,51 @@ #!/usr/bin/env python +"""Iterator wrapper that maintains a lookback stack of recent lines. + +Used by the parsing framework to allow parsers to inspect the current +line (via top()) without consuming it, since multiple parsers may need +to test the same line. +""" +from __future__ import annotations + +from typing import Iterable, Iterator class StackIterator: - """FIFO stack used to iterate over file while holding onto most recent lines""" + """Iterator with a fixed-size lookback stack. + + Wraps any iterable and maintains a FIFO stack of the most recently + yielded items. The current item is always accessible via top() without + advancing the iterator. + + Args: + iterable: The underlying iterable to wrap. + stacksize: Maximum number of items to retain in the lookback stack. + current: Initial line counter value (defaults to -1 so first next() + sets it to 0). + """ - def __init__(self, iterable, stacksize=1, current=-1): + def __init__(self, iterable: Iterable[str], stacksize: int = 1, current: int = -1) -> None: self.stacksize = stacksize - self.stack = [] + self.stack: list[str] = [] self.current = current self.iterable = iterable - self.it = self.iterable.__iter__() + self.it: Iterator[str] = self.iterable.__iter__() - def __iter__(self): + def __iter__(self) -> StackIterator: return self - def __next__(self): + def __next__(self) -> str: nx = next(self.it) self.add_to_stack(nx) self.current += 1 return nx - def add_to_stack(self, item): + def add_to_stack(self, item: str) -> None: + """Add item to the stack, evicting the oldest if at capacity.""" self.stack.append(item) if len(self.stack) > self.stacksize: self.stack.pop(0) - def top(self): + def top(self) -> str: + """Return the most recently yielded item without advancing.""" return self.stack[-1] diff --git a/mudslide/turboparse/thermo_parser.py b/mudslide/turboparse/thermo_parser.py index 5271447..0ff878f 100644 --- a/mudslide/turboparse/thermo_parser.py +++ b/mudslide/turboparse/thermo_parser.py @@ -1,16 +1,25 @@ #!/usr/bin/env python +"""Parser for Turbomole thermo module output. + +Extracts thermochemical quantities: enthalpy H(T), entropy contribution +T*S, and Gibbs free energy G(T) in atomic units, kcal/mol, and kJ/mol. +""" +from __future__ import annotations from .section_parser import ParseSection from .line_parser import SimpleLineParser class ThermoParser(ParseSection): + """Parser for the thermo module output (H, T*S, G at given temperature).""" name = "thermo" - def __init__(self): + def __init__(self) -> None: super().__init__(r"T H E R M O", r"thermo\s*:\s*all done") self.parsers = [ - SimpleLineParser(r"(\S+)\s*VIB\.\s+\S+\s+\S+\s+\S+\s+\S+", names=['T'], types=[float]), + SimpleLineParser(r"(\S+)\s*VIB\.\s+\S+\s+\S+\s+\S+\s+\S+", + names=['T'], + types=[float]), SimpleLineParser(r"H\(T\)\s+(\S+)\s+(\S+)\s+(\S+)", names=["H(au)", "H(kcal/mol)", "H(kJ/mol)"], types=[float, float, float]), diff --git a/mudslide/util.py b/mudslide/util.py index 162b1b9..3c31d6f 100644 --- a/mudslide/util.py +++ b/mudslide/util.py @@ -1,12 +1,21 @@ # -*- coding: utf-8 -*- """Utility functions for the mudslide package.""" +from __future__ import annotations + import os import sys +from typing import Any import numpy as np -def find_unique_name(name: str, location="", always_enumerate: bool = False, +from .exceptions import ConfigurationError + + + +def find_unique_name(name: str, + location: str = "", + always_enumerate: bool = False, ending: str = "") -> str: """Generate a unique filename by adding a suffix if the file already exists. @@ -33,7 +42,8 @@ def find_unique_name(name: str, location="", always_enumerate: bool = False, """ name_yaml = f"{name}{ending}" name_yaml = f"{name}{ending}" - if not always_enumerate and not os.path.exists(os.path.join(location, name_yaml)): + if not always_enumerate and not os.path.exists( + os.path.join(location, name_yaml)): return name for i in range(sys.maxsize): out = f"{name}-{i:d}" @@ -44,7 +54,9 @@ def find_unique_name(name: str, location="", always_enumerate: bool = False, return out raise FileExistsError(f"No unique name could be made from base {name}.") -def is_string(x) -> bool: + + +def is_string(x: Any) -> bool: """Check if the input is a string type. Parameters @@ -59,7 +71,9 @@ def is_string(x) -> bool: """ return isinstance(x, str) -def remove_center_of_mass_motion(velocities: np.ndarray, masses: np.ndarray) -> np.ndarray: + +def remove_center_of_mass_motion(velocities: np.ndarray, + masses: np.ndarray) -> np.ndarray: """Remove the center of mass motion from a set of velocities. Parameters @@ -77,6 +91,7 @@ def remove_center_of_mass_motion(velocities: np.ndarray, masses: np.ndarray) -> com = np.sum(velocities * masses[:, np.newaxis], axis=0) / np.sum(masses) return velocities - com + def remove_angular_momentum(velocities: np.ndarray, masses: np.ndarray, coordinates: np.ndarray) -> np.ndarray: """Remove the angular momentum from a set of coordinates and velocities. @@ -102,7 +117,8 @@ def remove_angular_momentum(velocities: np.ndarray, masses: np.ndarray, momentum = np.einsum('ai,a->ai', velocities, masses, dtype=np.float64) # angular momentum - angular_momentum = np.cross(com_coord, momentum).sum(axis=0, dtype=np.float64) + angular_momentum = np.cross(com_coord, momentum).sum(axis=0, + dtype=np.float64) # inertial tensor inertia = np.zeros((3, 3), dtype=np.float64) @@ -120,6 +136,7 @@ def remove_angular_momentum(velocities: np.ndarray, masses: np.ndarray, return corrected_velocities + def check_options(options: dict, recognized: list, strict: bool = True) -> None: """Check whether the options dictionary contains only recognized options. @@ -137,10 +154,10 @@ def check_options(options: dict, recognized: list, strict: bool = True) -> None: ValueError If strict is True and unrecognized options are found. """ - problems = [ x for x in options if x not in recognized ] + problems = [x for x in options if x not in recognized] if problems: if strict: # pylint: disable=no-else-raise - raise ValueError(f"Unrecognized options found: {problems}.") + raise ConfigurationError(f"Unrecognized options found: {problems}.") else: print(f"WARNING: Ignoring unrecognized options: {problems}.") diff --git a/mudslide/yaml_format.py b/mudslide/yaml_format.py index d7ee020..99f4613 100644 --- a/mudslide/yaml_format.py +++ b/mudslide/yaml_format.py @@ -1,6 +1,10 @@ # -*- coding: utf-8 -*- """YAML formatting utilities for compact output.""" +from __future__ import annotations + +from typing import Any + import yaml @@ -24,7 +28,8 @@ def _is_scalar(value: object) -> bool: SHORT_LIST_THRESHOLD = 5 -def _compact_represent_list(orig_represent_list): # type: ignore[no-untyped-def] +def _compact_represent_list( + orig_represent_list: Any) -> Any: # type: ignore[no-untyped-def] """Create a compact list representer wrapping the given base representer.""" def representer(dumper, data): # type: ignore[no-untyped-def] diff --git a/pyproject.toml b/pyproject.toml index fbb1040..033f76c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -42,6 +42,7 @@ dev = [ "mypy", "pylint", "yapf", + "types-PyYAML", ] [project.urls] diff --git a/roadmap_1.0.md b/roadmap_1.0.md new file mode 100644 index 0000000..c1ccf40 --- /dev/null +++ b/roadmap_1.0.md @@ -0,0 +1,94 @@ +# Mudslide 1.0 Release: Suggested Improvements + +## Context + +Mudslide is a Python nonadiabatic molecular dynamics library at version 0.12.0. Before a 1.0 release, this audit identifies bugs, missing features, inconsistencies, and quality improvements across the entire codebase. Items are grouped by category and prioritized. + +--- + +## 1. Bugs and Correctness Issues + +**None left** +--- + +## 2. Missing Features + +### High Priority (expected for a serious NAMD code) + +- **Simple decoherence corrections for FSSH** (`surface_hopping_md.py`): Standard FSSH has *no* decoherence correction. The literature considers this a known deficiency. At minimum, add: + - **Energy-based decoherence (EDC)** from Granucci & Persico, *J. Chem. Phys.* 2007 — simple, widely used, only a few lines of code + - Optionally: Simplified Decay of Mixing (SDM, Truhlar) or Instantaneous Decoherence Correction (IDC) + - These should be opt-in via a `decoherence` keyword option. +- **AIMD → NAMD initialization pathway**: Currently no built-in way to take snapshots from an adiabatic MD (ground-state) trajectory and use them to initialize surface hopping trajectories. This is a very common workflow (run AIMD to equilibrate, then spawn NAMD from snapshots). Needs: + - A utility/trajectory generator that reads AIMD trace output (positions, velocities) and spawns FSSH initial conditions from them + - Could be a new `TrajGen` subclass like `TrajGenFromTrace` or `TrajGenFromAIMD` + +### Medium Priority (would strengthen the package) + +- **Trivial crossing / state reordering detection**: Near conical intersections, adiabatic state ordering can swap between timesteps. No detection or handling mechanism exists. At minimum, add a warning; ideally, implement overlap-based state tracking (some infrastructure exists in `DiabaticModel_` phase fixing but not for ab initio models). +- **Generalize AFSSH beyond 2 states**: The collapse equation is hardcoded for 2 states (now raises `NotImplementedError`). The Subotnik group has published multi-state generalizations. + +### Lower Priority (nice to have for 1.0) + +- **Configuration file support**: All options are CLI arguments or programmatic. A YAML config file would help usability for complex ab initio setups. + +--- + +## 3. Feature Inconsistencies + +### Medium + +- **Two CLI entry points**: `mudslide` (**main**.py) and `mud` (mud.py) overlap in functionality. For 1.0, consolidate into one or clearly document that `mud` is the successor and deprecate `mudslide`. + +--- + +## 4. Code Quality and Static Analysis + +### Medium + +- **No `conftest.py`**: No shared fixtures. Common setup (model creation, trajectory initialization) is duplicated across test files. +- **Under-tested areas**: AFSSH (6 tests), scattering models (2 tests), QM/MM (1 test), OpenMM (1 test), surface CLI (0 tests), quadrature (1 test). Increase coverage for core algorithms before 1.0. +- **Pickle for output** (`__main__.py:220`): `pickle.dump()` used for saving results. Consider safer alternatives (JSON, YAML which is already used elsewhere, or at minimum document the pickle security caveat). + +--- + +## 5. Packaging and Distribution + +### Medium + +- **Add MANIFEST.in**: Not present — could cause missing files in source distributions (e.g., example files, test data). +- **Add `py.typed` marker**: For PEP 561 compliance, add a `py.typed` file so downstream users get type checking benefits. +- **Consider dropping Python 3.9**: If 1.0 is coming soon, 3.9 reaches EOL October 2025. Supporting 3.10+ simplifies code (e.g., `match` statements, `X | Y` union types). + +--- + +## 6. Documentation + +### High + +- **Add a CHANGELOG**: Essential for 1.0. Document what changed from the pre-1.0 versions. + +### Medium + +- **Expand Sphinx docs**: The documentation structure is good but some pages may be thin. Ensure all public API classes have autodoc entries. +- **Add a "Getting Started" tutorial**: A Jupyter notebook or script walkthrough that goes from model → single trajectory → batch simulation → analysis. +- **Document decoherence options** (once implemented): Explain when and why to use each correction. + +--- + +## Summary: Suggested Priority Order for 1.0 + +**Do first (blockers for a credible 1.0):** + +1. Add simple decoherence correction (EDC at minimum) +2. Add CHANGELOG + +**Do next (significantly improves the package):** + +3. AIMD → NAMD initialization pathway +4. Increase test coverage for under-tested areas + +**If time permits:** + +5. Generalize AFSSH beyond 2 states +6. State reordering detection diff --git a/test/ref/tm-es-c2h4/es_turbo.py b/test/ref/tm-es-c2h4/es_turbo.py index 361419e..e165517 100644 --- a/test/ref/tm-es-c2h4/es_turbo.py +++ b/test/ref/tm-es-c2h4/es_turbo.py @@ -43,7 +43,6 @@ def run(): sample_stack = SpawnStack.from_quadrature(nsamples=[2, 2, 2]) sample_stack.sample_stack[0]["zeta"]=0.003 samples = 1 - nprocs = 1 trace_type = YAMLTrace trace_options = {} electronic_integration = 'exp' @@ -60,7 +59,6 @@ def run(): positions=positions, samples=samples, max_time = max_time, - nprocs=nprocs, dt=dt, t0=t0, tracemanager=TraceManager(TraceType=trace_type, trace_kwargs=trace_options), diff --git a/test/test_adiabatic_propagator.py b/test/test_adiabatic_propagator.py index 22ee9d7..2fcef66 100644 --- a/test/test_adiabatic_propagator.py +++ b/test/test_adiabatic_propagator.py @@ -7,6 +7,7 @@ import pytest import mudslide +from mudslide.exceptions import ConfigurationError from mudslide.adiabatic_propagator import (VVPropagator, NoseHooverChainPropagator, AdiabaticPropagator) @@ -169,5 +170,5 @@ def test_factory_nhc_from_dict(): def test_factory_unknown_type_raises(): """Factory raises ValueError for unknown propagator type""" - with pytest.raises(ValueError, match="Unknown propagator type"): + with pytest.raises(ConfigurationError, match="Unknown propagator type"): AdiabaticPropagator(water_model, {"type": "unknown"}) diff --git a/test/test_afssh.py b/test/test_afssh.py index cdfbbed..5a72f4c 100644 --- a/test/test_afssh.py +++ b/test/test_afssh.py @@ -10,18 +10,20 @@ import pytest from mudslide.afssh import AugmentedFSSH, AFSSHPropagator, AFSSHVVPropagator +from mudslide.constants import fs_to_au +from mudslide.exceptions import ConfigurationError from mudslide.models import scattering_models as models def make_afssh_traj(model, x0, v0, state0, **kwargs): """Helper to create an AugmentedFSSH trajectory with common defaults.""" - # Use strict_option_check=False since augmented_integration is AFSSH-specific + if "dt" not in kwargs: + kwargs["dt"] = fs_to_au return AugmentedFSSH( model, np.atleast_1d(x0), np.atleast_1d(v0), state0, - strict_option_check=False, **kwargs ) @@ -61,7 +63,7 @@ def test_dict_options_vv(self): def test_invalid_propagator_type_raises(self): """Test that invalid propagator type raises ValueError""" model = models["simple"](mass=2000.0) - with pytest.raises(ValueError, match="Unrecognized"): + with pytest.raises(ConfigurationError, match="Unrecognized"): AFSSHPropagator(model, "invalid_type") def test_invalid_options_type_raises(self): @@ -141,7 +143,7 @@ def test_invalid_integration_method_delR_raises(self): # Manually set invalid integration method traj.augmented_integration = "invalid" - with pytest.raises(Exception, match="Unrecognized propagate delR"): + with pytest.raises(ConfigurationError, match="Unrecognized augmented integration method"): step_trajectory(traj) def test_invalid_integration_method_delP_raises(self): @@ -157,7 +159,7 @@ def test_invalid_integration_method_delP_raises(self): # Then set invalid method traj.augmented_integration = "invalid" - with pytest.raises(Exception, match="Unrecognized propagate delP"): + with pytest.raises(ConfigurationError, match="Unrecognized augmented integration method"): # Call advance_delP directly traj.advance_delP(traj.last_electronics, traj.electronics) diff --git a/test/test_boltzmann_velocity.py b/test/test_boltzmann_velocity.py index 000b505..4da5d97 100644 --- a/test/test_boltzmann_velocity.py +++ b/test/test_boltzmann_velocity.py @@ -2,54 +2,46 @@ # -*- coding: utf-8 -*- """Unit testing for boltzmann function""" -import sys import numpy as np -import unittest from mudslide.math import boltzmann_velocities -class TrajectoryTest(unittest.TestCase): - - def test_boltzmann(self): - """test for boltzmann_velocity function""" - # yapf: disable - mass = np.array ([21874.6618344, 21874.6618344, 21874.6618344, - 21874.6618344, 21874.6618344, 21874.6618344, - 1837.15264736, 1837.15264736, 1837.15264736, - 1837.15264736, 1837.15264736, 1837.15264736, - 1837.15264736, 1837.15264736, 1837.15264736, - 1837.15264736, 1837.15264736, 1837.15264736]) - # yapf: enable - velocities = boltzmann_velocities(mass, temperature=300, remove_translation=False, - scale=True, seed=109) - - # yapf: disable - results = [ 2.56748743e-04, 3.44471310e-04, 2.46424017e-07, - -3.89747528e-04, 2.31250717e-04, 1.20055248e-04, - -2.49906132e-04, -5.58142454e-05, -2.39676921e-04, - 3.22950419e-04, 2.51235241e-04, -9.35539925e-04, - 7.00590108e-05, -6.98837331e-04, -1.31891980e-03, - -8.41937362e-05, 9.68859510e-04, 3.86297562e-04] - # yapf: enable - - np.testing.assert_almost_equal(results, velocities, decimal=8) - - # yapf: enable - velocities = boltzmann_velocities(mass, temperature=300, remove_translation=True, - scale=True, seed=109) - - # yapf: disable - results = np.array( - [ 3.52267636e-04, 9.18916486e-05, 2.77349560e-05, - -3.78682783e-04, -3.61193603e-05, 1.63194840e-04, - -2.20573399e-04, -3.60684658e-04, -2.43530444e-04, - 4.27117477e-04, -1.35241861e-05, -1.03029655e-03, - 1.41189948e-04, -1.08770818e-03, -1.46375873e-03, - -3.32133934e-05, 7.97845939e-04, 4.64217351e-04]) - # yapf: enable - - np.testing.assert_almost_equal(results, velocities, decimal=8) - - -if __name__ == '__main__': - unittest.main() +def test_boltzmann(): + """test for boltzmann_velocity function""" + # yapf: disable + mass = np.array ([21874.6618344, 21874.6618344, 21874.6618344, + 21874.6618344, 21874.6618344, 21874.6618344, + 1837.15264736, 1837.15264736, 1837.15264736, + 1837.15264736, 1837.15264736, 1837.15264736, + 1837.15264736, 1837.15264736, 1837.15264736, + 1837.15264736, 1837.15264736, 1837.15264736]) + # yapf: enable + velocities = boltzmann_velocities(mass, temperature=300, remove_translation=False, + scale=True, seed=109) + + # yapf: disable + results = [ 2.56748743e-04, 3.44471310e-04, 2.46424017e-07, + -3.89747528e-04, 2.31250717e-04, 1.20055248e-04, + -2.49906132e-04, -5.58142454e-05, -2.39676921e-04, + 3.22950419e-04, 2.51235241e-04, -9.35539925e-04, + 7.00590108e-05, -6.98837331e-04, -1.31891980e-03, + -8.41937362e-05, 9.68859510e-04, 3.86297562e-04] + # yapf: enable + + np.testing.assert_almost_equal(results, velocities, decimal=8) + + # yapf: enable + velocities = boltzmann_velocities(mass, temperature=300, remove_translation=True, + scale=True, seed=109) + + # yapf: disable + results = np.array( + [ 3.52267636e-04, 9.18916486e-05, 2.77349560e-05, + -3.78682783e-04, -3.61193603e-05, 1.63194840e-04, + -2.20573399e-04, -3.60684658e-04, -2.43530444e-04, + 4.27117477e-04, -1.35241861e-05, -1.03029655e-03, + 1.41189948e-04, -1.08770818e-03, -1.46375873e-03, + -3.32133934e-05, 7.97845939e-04, 4.64217351e-04]) + # yapf: enable + + np.testing.assert_almost_equal(results, velocities, decimal=8) diff --git a/test/test_collect.py b/test/test_collect.py index 0fdc31c..911df02 100644 --- a/test/test_collect.py +++ b/test/test_collect.py @@ -2,298 +2,246 @@ # -*- coding: utf-8 -*- """Tests for mudslide/collect.py""" -import unittest import os -import tempfile -import shutil import argparse +import pytest import yaml from mudslide.collect import add_collect_parser, collect, collect_wrapper, legend, legend_format -class TestAddCollectParser(unittest.TestCase): - """Tests for the add_collect_parser function""" - - def test_parser_added_to_subparsers(self): - """Test that collect parser is added correctly""" - parser = argparse.ArgumentParser() - subparsers = parser.add_subparsers() - - add_collect_parser(subparsers) - - # Parse a collect command - args = parser.parse_args(['collect', 'mylog']) - self.assertEqual(args.logname, 'mylog') - self.assertEqual(args.keys, 'tkpea') # default value - - def test_custom_keys(self): - """Test that custom keys can be specified""" - parser = argparse.ArgumentParser() - subparsers = parser.add_subparsers() - - add_collect_parser(subparsers) - - args = parser.parse_args(['collect', 'mylog', '-k', 'te']) - self.assertEqual(args.keys, 'te') - - def test_func_set_to_collect_wrapper(self): - """Test that func is set to collect_wrapper""" - parser = argparse.ArgumentParser() - subparsers = parser.add_subparsers() - - add_collect_parser(subparsers) - - args = parser.parse_args(['collect', 'mylog']) - self.assertEqual(args.func, collect_wrapper) - - -class TestCollect(unittest.TestCase): - """Tests for the collect function""" - - def setUp(self): - """Create temporary directory and YAML trace files for testing""" - self.tmpdir = tempfile.mkdtemp() - self.original_dir = os.getcwd() - os.chdir(self.tmpdir) - - def tearDown(self): - """Clean up temporary directory""" - os.chdir(self.original_dir) - shutil.rmtree(self.tmpdir) - - def _create_yaml_trace(self, basename, snapshots): - """Create a minimal YAML trace structure for testing. - - Creates the main log file and associated data log file. - """ - log_file = f"{basename}-log_0.yaml" - hop_file = f"{basename}-hops.yaml" - main_log = f"{basename}.yaml" - - # Write main log - main_data = { - "name": basename, - "logfiles": [log_file], - "nlogs": 1, - "log_pitch": 512, - "hop_log": hop_file, - "event_logs": {}, - "weight": 1.0 - } - with open(main_log, "w") as f: - yaml.dump(main_data, f) - - # Write data log with snapshots - with open(log_file, "w") as f: - yaml.dump(snapshots, f) - - # Write empty hop log - with open(hop_file, "w") as f: - pass - - return main_log - - def test_collect_basic(self): - """Test basic collect functionality with default keys""" - snapshots = [ - { - "time": 0.0, - "kinetic": 0.5, - "potential": -1.0, - "energy": -0.5, - "active": 0 - }, - { - "time": 1.0, - "kinetic": 0.6, - "potential": -1.1, - "energy": -0.5, - "active": 0 - }, - { - "time": 2.0, - "kinetic": 0.7, - "potential": -1.2, - "energy": -0.5, - "active": 1 - }, - ] - main_log = self._create_yaml_trace("test-traj-0", snapshots) - - collect(main_log) - - output_file = main_log + ".dat" - self.assertTrue(os.path.exists(output_file)) - - with open(output_file, "r") as f: - lines = f.readlines() - - # Check header - self.assertTrue(lines[0].startswith("#")) - self.assertIn("time", lines[0]) - self.assertIn("kinetic", lines[0]) - self.assertIn("potential", lines[0]) - self.assertIn("energy", lines[0]) - self.assertIn("active", lines[0]) - - # Check data lines (should be 3) - self.assertEqual(len(lines), 4) # 1 header + 3 data - - def test_collect_custom_keys(self): - """Test collect with subset of keys""" - snapshots = [ - { - "time": 0.0, - "kinetic": 0.5, - "potential": -1.0, - "energy": -0.5, - "active": 0 - }, - { - "time": 1.0, - "kinetic": 0.6, - "potential": -1.1, - "energy": -0.5, - "active": 1 - }, - ] - main_log = self._create_yaml_trace("test-traj-0", snapshots) - - collect(main_log, keys="te") - - output_file = main_log + ".dat" - with open(output_file, "r") as f: - lines = f.readlines() - - # Header should only have time and energy - header = lines[0] - self.assertIn("time", header) - self.assertIn("energy", header) - self.assertNotIn("kinetic", header) - self.assertNotIn("potential", header) - self.assertNotIn("active", header) - - def test_collect_output_values(self): - """Test that output values match input data""" - snapshots = [ - { - "time": 1.5, - "kinetic": 0.123, - "potential": -0.456, - "energy": -0.333, - "active": 2 - }, - ] - main_log = self._create_yaml_trace("test-traj-0", snapshots) - - collect(main_log) - - output_file = main_log + ".dat" - with open(output_file, "r") as f: - lines = f.readlines() - - # Parse data line - data_line = lines[1].strip().split() - self.assertAlmostEqual(float(data_line[0]), 1.5, places=6) - self.assertAlmostEqual(float(data_line[1]), 0.123, places=6) - self.assertAlmostEqual(float(data_line[2]), -0.456, places=6) - self.assertAlmostEqual(float(data_line[3]), -0.333, places=6) - self.assertEqual(int(data_line[4]), 2) - - -class TestCollectWrapper(unittest.TestCase): - """Tests for the collect_wrapper function""" - - def setUp(self): - """Create temporary directory for testing""" - self.tmpdir = tempfile.mkdtemp() - self.original_dir = os.getcwd() - os.chdir(self.tmpdir) - - def tearDown(self): - """Clean up temporary directory""" - os.chdir(self.original_dir) - shutil.rmtree(self.tmpdir) - - def _create_yaml_trace(self, basename, snapshots): - """Create a minimal YAML trace structure for testing.""" - log_file = f"{basename}-log_0.yaml" - hop_file = f"{basename}-hops.yaml" - main_log = f"{basename}.yaml" - - main_data = { - "name": basename, - "logfiles": [log_file], - "nlogs": 1, - "log_pitch": 512, - "hop_log": hop_file, - "event_logs": {}, - "weight": 1.0 - } - with open(main_log, "w") as f: - yaml.dump(main_data, f) - - with open(log_file, "w") as f: - yaml.dump(snapshots, f) - - with open(hop_file, "w") as f: - pass - - return main_log - - def test_wrapper_calls_collect(self): - """Test that wrapper correctly calls collect with args""" - snapshots = [ - { - "time": 0.0, - "kinetic": 0.5, - "potential": -1.0, - "energy": -0.5, - "active": 0 - }, - ] - main_log = self._create_yaml_trace("test-traj-0", snapshots) - - # Create mock args object - class MockArgs: - logname = main_log - keys = "te" - - collect_wrapper(MockArgs()) - - output_file = main_log + ".dat" - self.assertTrue(os.path.exists(output_file)) - - with open(output_file, "r") as f: - content = f.read() - - self.assertIn("time", content) - self.assertIn("energy", content) - - -class TestLegendMappings(unittest.TestCase): - """Tests for the legend and legend_format dictionaries""" - - def test_legend_keys_match_format_keys(self): - """Test that legend and legend_format have the same keys""" - self.assertEqual(set(legend.keys()), set(legend_format.keys())) - - def test_expected_keys_present(self): - """Test that expected keys are present in legend""" - expected_keys = ['t', 'k', 'p', 'e', 'a'] - for key in expected_keys: - self.assertIn(key, legend) - - def test_legend_values(self): - """Test legend mappings are correct""" - self.assertEqual(legend['t'], 'time') - self.assertEqual(legend['k'], 'kinetic') - self.assertEqual(legend['p'], 'potential') - self.assertEqual(legend['e'], 'energy') - self.assertEqual(legend['a'], 'active') - - -if __name__ == '__main__': - unittest.main() +def test_parser_added_to_subparsers(): + """Test that collect parser is added correctly""" + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + + add_collect_parser(subparsers) + + # Parse a collect command + args = parser.parse_args(['collect', 'mylog']) + assert args.logname == 'mylog' + assert args.keys == 'tkpea' # default value + + +def test_custom_keys(): + """Test that custom keys can be specified""" + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + + add_collect_parser(subparsers) + + args = parser.parse_args(['collect', 'mylog', '-k', 'te']) + assert args.keys == 'te' + + +def test_func_set_to_collect_wrapper(): + """Test that func is set to collect_wrapper""" + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + + add_collect_parser(subparsers) + + args = parser.parse_args(['collect', 'mylog']) + assert args.func == collect_wrapper + + +def _create_yaml_trace(basename, snapshots): + """Create a minimal YAML trace structure for testing. + + Creates the main log file and associated data log file. + """ + log_file = f"{basename}-log_0.yaml" + hop_file = f"{basename}-hops.yaml" + main_log = f"{basename}.yaml" + + # Write main log + main_data = { + "name": basename, + "logfiles": [log_file], + "nlogs": 1, + "log_pitch": 512, + "hop_log": hop_file, + "event_logs": {}, + "weight": 1.0 + } + with open(main_log, "w") as f: + yaml.dump(main_data, f) + + # Write data log with snapshots + with open(log_file, "w") as f: + yaml.dump(snapshots, f) + + # Write empty hop log + with open(hop_file, "w") as f: + pass + + return main_log + + +@pytest.fixture() +def working_dir(tmp_path, monkeypatch): + """Change to a temporary directory for tests that need file I/O""" + monkeypatch.chdir(tmp_path) + return tmp_path + + +def test_collect_basic(working_dir): + """Test basic collect functionality with default keys""" + snapshots = [ + { + "time": 0.0, + "kinetic": 0.5, + "potential": -1.0, + "energy": -0.5, + "active": 0 + }, + { + "time": 1.0, + "kinetic": 0.6, + "potential": -1.1, + "energy": -0.5, + "active": 0 + }, + { + "time": 2.0, + "kinetic": 0.7, + "potential": -1.2, + "energy": -0.5, + "active": 1 + }, + ] + main_log = _create_yaml_trace("test-traj-0", snapshots) + + collect(main_log) + + output_file = main_log + ".dat" + assert os.path.exists(output_file) + + with open(output_file, "r") as f: + lines = f.readlines() + + # Check header + assert lines[0].startswith("#") + assert "time" in lines[0] + assert "kinetic" in lines[0] + assert "potential" in lines[0] + assert "energy" in lines[0] + assert "active" in lines[0] + + # Check data lines (should be 3) + assert len(lines) == 4 # 1 header + 3 data + + +def test_collect_custom_keys(working_dir): + """Test collect with subset of keys""" + snapshots = [ + { + "time": 0.0, + "kinetic": 0.5, + "potential": -1.0, + "energy": -0.5, + "active": 0 + }, + { + "time": 1.0, + "kinetic": 0.6, + "potential": -1.1, + "energy": -0.5, + "active": 1 + }, + ] + main_log = _create_yaml_trace("test-traj-0", snapshots) + + collect(main_log, keys="te") + + output_file = main_log + ".dat" + with open(output_file, "r") as f: + lines = f.readlines() + + # Header should only have time and energy + header = lines[0] + assert "time" in header + assert "energy" in header + assert "kinetic" not in header + assert "potential" not in header + assert "active" not in header + + +def test_collect_output_values(working_dir): + """Test that output values match input data""" + snapshots = [ + { + "time": 1.5, + "kinetic": 0.123, + "potential": -0.456, + "energy": -0.333, + "active": 2 + }, + ] + main_log = _create_yaml_trace("test-traj-0", snapshots) + + collect(main_log) + + output_file = main_log + ".dat" + with open(output_file, "r") as f: + lines = f.readlines() + + # Parse data line + data_line = lines[1].strip().split() + assert float(data_line[0]) == pytest.approx(1.5, abs=1e-6) + assert float(data_line[1]) == pytest.approx(0.123, abs=1e-6) + assert float(data_line[2]) == pytest.approx(-0.456, abs=1e-6) + assert float(data_line[3]) == pytest.approx(-0.333, abs=1e-6) + assert int(data_line[4]) == 2 + + +def test_wrapper_calls_collect(working_dir): + """Test that wrapper correctly calls collect with args""" + snapshots = [ + { + "time": 0.0, + "kinetic": 0.5, + "potential": -1.0, + "energy": -0.5, + "active": 0 + }, + ] + main_log = _create_yaml_trace("test-traj-0", snapshots) + + # Create mock args object + class MockArgs: + logname = main_log + keys = "te" + + collect_wrapper(MockArgs()) + + output_file = main_log + ".dat" + assert os.path.exists(output_file) + + with open(output_file, "r") as f: + content = f.read() + + assert "time" in content + assert "energy" in content + + +def test_legend_keys_match_format_keys(): + """Test that legend and legend_format have the same keys""" + assert set(legend.keys()) == set(legend_format.keys()) + + +def test_expected_keys_present(): + """Test that expected keys are present in legend""" + expected_keys = ['t', 'k', 'p', 'e', 'a'] + for key in expected_keys: + assert key in legend + + +def test_legend_values(): + """Test legend mappings are correct""" + assert legend['t'] == 'time' + assert legend['k'] == 'kinetic' + assert legend['p'] == 'potential' + assert legend['e'] == 'energy' + assert legend['a'] == 'active' diff --git a/test/test_harmonic.py b/test/test_harmonic.py index 3aad8b2..1162f34 100644 --- a/test/test_harmonic.py +++ b/test/test_harmonic.py @@ -7,6 +7,7 @@ import numpy as np import pytest import mudslide +from mudslide.exceptions import ConfigurationError water_json = """ {"x0": [0.0, 0.0, -0.12178983933899, 1.41713420892173, 0.0, 0.96657854674257, -1.41713420892173, 0.0, 0.96657854674257], "E0": 0.0, "H0": [[0.783125763, 0.0, 5e-10, -0.3915628813, 0.0, -0.3007228667, -0.3915628817, 0.0, 0.3007228662], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [5e-10, 0.0, 0.482147457, -0.245482284, 0.0, -0.2410737287, 0.2454822835, 0.0, -0.2410737283], [-0.3915628813, 0.0, -0.245482284, 0.4189939059, 0.0, 0.2731025751, -0.0274310246, 0.0, -0.0276202911], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [-0.3007228667, 0.0, -0.2410737287, 0.2731025751, 0.0, 0.2360154334, 0.0276202915, 0.0, 0.0050582953], [-0.3915628817, 0.0, 0.2454822835, -0.0274310246, 0.0, 0.0276202915, 0.4189939063, 0.0, -0.273102575], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.3007228662, 0.0, -0.2410737283, -0.0276202911, 0.0, 0.0050582953, -0.273102575, 0.0, 0.236015433]], "mass": [29156.945697766205, 29156.945697766205, 29156.945697766205, 1837.1526473562108, 1837.1526473562108, 1837.1526473562108, 1837.1526473562108, 1837.1526473562108, 1837.1526473562108], "atom_types": ["O", "H", "H"]}""" @@ -75,7 +76,7 @@ def test_from_file_unknown_format(tmp_path): with open(filepath, "w") as f: f.write("hello") - with pytest.raises(ValueError, match="Unknown file format"): + with pytest.raises(ConfigurationError, match="Unknown file format"): mudslide.models.HarmonicModel.from_file(filepath) @@ -107,7 +108,7 @@ def test_to_file_unknown_format(tmp_path): model = mudslide.models.HarmonicModel.from_dict(water_model) filepath = str(tmp_path / "out.txt") - with pytest.raises(ValueError, match="Unknown file format"): + with pytest.raises(ConfigurationError, match="Unknown file format"): model.to_file(filepath) diff --git a/test/test_header.py b/test/test_header.py index 31ba9ec..ab05637 100644 --- a/test/test_header.py +++ b/test/test_header.py @@ -2,7 +2,6 @@ # -*- coding: utf-8 -*- """Unit testing for header""" -import unittest import re from io import StringIO from unittest.mock import patch @@ -11,68 +10,71 @@ from mudslide.header import print_header, BANNER -class TestHeader(unittest.TestCase): - """Test Suite for header functions""" - - def _capture_header(self) -> str: - with patch('sys.stdout', new_callable=StringIO) as mock_stdout: - print_header() - return mock_stdout.getvalue() - - def test_print_header_runs(self) -> None: - """print_header should run without error""" - output = self._capture_header() - self.assertIn("MUDSLIDE", output) - - def test_header_contains_version(self) -> None: - """print_header should contain the mudslide version""" - output = self._capture_header() - self.assertIn(mudslide.__version__, output) - - def test_header_contains_python_version(self) -> None: - """print_header should contain the Python version""" - import sys - output = self._capture_header() - self.assertIn(sys.version.split()[0], output) - - def test_header_contains_numpy_version(self) -> None: - """print_header should contain the NumPy version""" - import numpy - output = self._capture_header() - self.assertIn(numpy.__version__, output) - - def test_header_contains_scipy_version(self) -> None: - """print_header should contain the SciPy version""" - import scipy - output = self._capture_header() - self.assertIn(scipy.__version__, output) - - def test_header_contains_platform(self) -> None: - """print_header should contain the platform""" - import platform - output = self._capture_header() - self.assertIn(platform.platform(), output) - - def test_header_contains_path(self) -> None: - """print_header should contain the install path""" - from mudslide.version import get_install_path - output = self._capture_header() - self.assertIn(get_install_path(), output) - - def test_header_contains_date(self) -> None: - """print_header should contain a date string""" - output = self._capture_header() - self.assertRegex(output, r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}") - - def test_header_contains_banner(self) -> None: - """print_header should contain the ASCII banner""" - output = self._capture_header() - self.assertIn(".-'~~~`-.", output) - - def test_accessible_from_mudslide(self) -> None: - """print_header should be accessible as mudslide.print_header""" - self.assertIs(mudslide.print_header, print_header) - - -if __name__ == '__main__': - unittest.main() +def _capture_header() -> str: + with patch('sys.stdout', new_callable=StringIO) as mock_stdout: + print_header() + return mock_stdout.getvalue() + + +def test_print_header_runs() -> None: + """print_header should run without error""" + output = _capture_header() + assert "MUDSLIDE" in output + + +def test_header_contains_version() -> None: + """print_header should contain the mudslide version""" + output = _capture_header() + assert mudslide.__version__ in output + + +def test_header_contains_python_version() -> None: + """print_header should contain the Python version""" + import sys + output = _capture_header() + assert sys.version.split()[0] in output + + +def test_header_contains_numpy_version() -> None: + """print_header should contain the NumPy version""" + import numpy + output = _capture_header() + assert numpy.__version__ in output + + +def test_header_contains_scipy_version() -> None: + """print_header should contain the SciPy version""" + import scipy + output = _capture_header() + assert scipy.__version__ in output + + +def test_header_contains_platform() -> None: + """print_header should contain the platform""" + import platform + output = _capture_header() + assert platform.platform() in output + + +def test_header_contains_path() -> None: + """print_header should contain the install path""" + from mudslide.version import get_install_path + output = _capture_header() + assert get_install_path() in output + + +def test_header_contains_date() -> None: + """print_header should contain a date string""" + output = _capture_header() + assert re.search(r"\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}", output) + + +def test_header_contains_banner() -> None: + """print_header should contain the ASCII banner""" + output = _capture_header() + assert ".-'~~~`-." in output + + +def test_accessible_from_mudslide() -> None: + """print_header should be accessible as mudslide.print_header""" + assert mudslide.print_header is print_header diff --git a/test/test_io.py b/test/test_io.py index d290aa7..b4cd8fc 100644 --- a/test/test_io.py +++ b/test/test_io.py @@ -2,244 +2,222 @@ # -*- coding: utf-8 -*- """Tests for mudslide/io.py""" -import unittest import os -import tempfile import io import numpy as np +import pytest from mudslide.io import write_xyz, write_trajectory_xyz from mudslide.constants import bohr_to_angstrom -class TestWriteXYZ(unittest.TestCase): - """Tests for the write_xyz function""" +def test_basic_xyz_output(): + """Test basic XYZ format output with simple coordinates""" + coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) + atom_types = ["H", "H", "O"] + output = io.StringIO() - def test_basic_xyz_output(self): - """Test basic XYZ format output with simple coordinates""" - coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]) - atom_types = ["H", "H", "O"] - output = io.StringIO() + write_xyz(coords, atom_types, output, comment="test molecule") - write_xyz(coords, atom_types, output, comment="test molecule") + result = output.getvalue() + lines = result.strip().split("\n") - result = output.getvalue() - lines = result.strip().split("\n") + assert lines[0] == "3" + assert lines[1] == "test molecule" + assert len(lines) == 5 # natom + comment + 3 atoms - self.assertEqual(lines[0], "3") - self.assertEqual(lines[1], "test molecule") - self.assertEqual(len(lines), 5) # natom + comment + 3 atoms - def test_bohr_to_angstrom_conversion(self): - """Test that coordinates are converted from bohr to angstrom""" - # 1 bohr should become ~0.529 angstrom - coords = np.array([[1.0, 0.0, 0.0]]) - atom_types = ["H"] - output = io.StringIO() +def test_bohr_to_angstrom_conversion(): + """Test that coordinates are converted from bohr to angstrom""" + # 1 bohr should become ~0.529 angstrom + coords = np.array([[1.0, 0.0, 0.0]]) + atom_types = ["H"] + output = io.StringIO() - write_xyz(coords, atom_types, output) + write_xyz(coords, atom_types, output) - result = output.getvalue() - lines = result.strip().split("\n") - # Parse the coordinate line - parts = lines[2].split() - x_coord = float(parts[1]) + result = output.getvalue() + lines = result.strip().split("\n") + # Parse the coordinate line + parts = lines[2].split() + x_coord = float(parts[1]) - self.assertAlmostEqual(x_coord, bohr_to_angstrom, places=10) + assert x_coord == pytest.approx(bohr_to_angstrom, abs=1e-10) - def test_atom_type_capitalization(self): - """Test that atom types are capitalized correctly""" - coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) - atom_types = ["h", "he"] # lowercase input - output = io.StringIO() - write_xyz(coords, atom_types, output) +def test_atom_type_capitalization(): + """Test that atom types are capitalized correctly""" + coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) + atom_types = ["h", "he"] # lowercase input + output = io.StringIO() - result = output.getvalue() - lines = result.strip().split("\n") + write_xyz(coords, atom_types, output) - # Check atom labels are capitalized - self.assertTrue(lines[2].startswith("H ")) - self.assertTrue(lines[3].startswith("He ")) + result = output.getvalue() + lines = result.strip().split("\n") - def test_empty_comment(self): - """Test with empty comment string""" - coords = np.array([[0.0, 0.0, 0.0]]) - atom_types = ["C"] - output = io.StringIO() + # Check atom labels are capitalized + assert lines[2].startswith("H ") + assert lines[3].startswith("He ") - write_xyz(coords, atom_types, output, comment="") - result = output.getvalue() - lines = result.split("\n") +def test_empty_comment(): + """Test with empty comment string""" + coords = np.array([[0.0, 0.0, 0.0]]) + atom_types = ["C"] + output = io.StringIO() - self.assertEqual(lines[0], "1") - self.assertEqual(lines[1], "") # empty comment line + write_xyz(coords, atom_types, output, comment="") - def test_coordinate_precision(self): - """Test that coordinates are written with proper precision""" - coords = np.array([[1.123456789012, 2.234567890123, 3.345678901234]]) - atom_types = ["N"] - output = io.StringIO() + result = output.getvalue() + lines = result.split("\n") - write_xyz(coords, atom_types, output) + assert lines[0] == "1" + assert lines[1] == "" # empty comment line - result = output.getvalue() - lines = result.strip().split("\n") - parts = lines[2].split() - # Check we have 12 decimal places (20.12f format) - # The actual values are converted to angstrom - expected_x = 1.123456789012 * bohr_to_angstrom - self.assertAlmostEqual(float(parts[1]), expected_x, places=10) +def test_coordinate_precision(): + """Test that coordinates are written with proper precision""" + coords = np.array([[1.123456789012, 2.234567890123, 3.345678901234]]) + atom_types = ["N"] + output = io.StringIO() + write_xyz(coords, atom_types, output) -class TestWriteTrajectoryXYZ(unittest.TestCase): - """Tests for the write_trajectory_xyz function""" + result = output.getvalue() + lines = result.strip().split("\n") + parts = lines[2].split() - def setUp(self): - """Create a mock model and trace for testing""" - self.tmpdir = tempfile.mkdtemp() + # Check we have 12 decimal places (20.12f format) + # The actual values are converted to angstrom + expected_x = 1.123456789012 * bohr_to_angstrom + assert float(parts[1]) == pytest.approx(expected_x, abs=1e-10) - def tearDown(self): - """Clean up temporary directory""" - import shutil - shutil.rmtree(self.tmpdir) - def _make_mock_model(self, natom=3, ndim=3, atom_types=None): - """Create a mock model object""" +def _make_mock_model(natom=3, ndim=3, atom_types=None): + """Create a mock model object""" - class MockModel: + class MockModel: - def __init__(self, natom, ndim, atom_types): - self.dimensionality = (natom, ndim) - self.atom_types = atom_types + def __init__(self, natom, ndim, atom_types): + self.dimensionality = (natom, ndim) + self.atom_types = atom_types - return MockModel(natom, ndim, atom_types) - - def _make_trace(self, nframes=5, natom=3, ndim=3): - """Create a mock trace (list of frame dictionaries)""" - trace = [] - for i in range(nframes): - frame = { - "energy": -100.0 + i * 0.1, - "time": i * 10.0, - "position": np.random.randn(natom * ndim) - } - trace.append(frame) - return trace - - def test_basic_trajectory_output(self): - """Test basic trajectory XYZ file output""" - model = self._make_mock_model(natom=2, ndim=3, atom_types=["H", "O"]) - trace = self._make_trace(nframes=3, natom=2, ndim=3) - filename = os.path.join(self.tmpdir, "test_traj.xyz") - - write_trajectory_xyz(model, trace, filename) - - self.assertTrue(os.path.exists(filename)) - - with open(filename, "r") as f: - content = f.read() - - lines = content.strip().split("\n") - # 3 frames * (1 natom + 1 comment + 2 atoms) = 12 lines - self.assertEqual(len(lines), 12) - - def test_every_parameter(self): - """Test that 'every' parameter skips frames correctly""" - model = self._make_mock_model(natom=1, ndim=3, atom_types=["C"]) - trace = self._make_trace(nframes=10, natom=1, ndim=3) - filename = os.path.join(self.tmpdir, "test_every.xyz") - - write_trajectory_xyz(model, trace, filename, every=3) - - with open(filename, "r") as f: - content = f.read() - - lines = content.strip().split("\n") - # Frames 0, 3, 6, 9 should be written (4 frames) - # Each frame: 1 natom + 1 comment + 1 atom = 3 lines - self.assertEqual(len(lines), 12) - - def test_fallback_to_x_atom_types(self): - """Test fallback to 'X' when model.atom_types is None""" - model = self._make_mock_model(natom=2, ndim=3, atom_types=None) - trace = self._make_trace(nframes=1, natom=2, ndim=3) - filename = os.path.join(self.tmpdir, "test_fallback.xyz") - - write_trajectory_xyz(model, trace, filename) - - with open(filename, "r") as f: - content = f.read() - - lines = content.strip().split("\n") - # Both atoms should be labeled "X" - self.assertTrue(lines[2].startswith("X ")) - self.assertTrue(lines[3].startswith("X ")) - - def test_comment_line_format(self): - """Test that comment line contains energy and time""" - model = self._make_mock_model(natom=1, ndim=3, atom_types=["H"]) - trace = [{ - "energy": -123.456, - "time": 42.0, - "position": np.array([1.0, 2.0, 3.0]) - }] - filename = os.path.join(self.tmpdir, "test_comment.xyz") - - write_trajectory_xyz(model, trace, filename) - - with open(filename, "r") as f: - lines = f.readlines() - - comment = lines[1].strip() - self.assertIn("E=", comment) - self.assertIn("t=", comment) - self.assertIn("-123.456", comment) - self.assertIn("42", comment) - - def test_position_reshaping(self): - """Test that flat position array is reshaped correctly""" - model = self._make_mock_model(natom=2, ndim=3, atom_types=["H", "O"]) - # Position as flat array: [x1, y1, z1, x2, y2, z2] - trace = [{ - "energy": 0.0, - "time": 0.0, - "position": np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) - }] - filename = os.path.join(self.tmpdir, "test_reshape.xyz") - - write_trajectory_xyz(model, trace, filename) - - with open(filename, "r") as f: - lines = f.readlines() - - # First atom should have coords (1, 2, 3) * bohr_to_angstrom - parts1 = lines[2].split() - self.assertAlmostEqual(float(parts1[1]), - 1.0 * bohr_to_angstrom, - places=8) - self.assertAlmostEqual(float(parts1[2]), - 2.0 * bohr_to_angstrom, - places=8) - self.assertAlmostEqual(float(parts1[3]), - 3.0 * bohr_to_angstrom, - places=8) - - # Second atom should have coords (4, 5, 6) * bohr_to_angstrom - parts2 = lines[3].split() - self.assertAlmostEqual(float(parts2[1]), - 4.0 * bohr_to_angstrom, - places=8) - self.assertAlmostEqual(float(parts2[2]), - 5.0 * bohr_to_angstrom, - places=8) - self.assertAlmostEqual(float(parts2[3]), - 6.0 * bohr_to_angstrom, - places=8) - - -if __name__ == '__main__': - unittest.main() + return MockModel(natom, ndim, atom_types) + + +def _make_trace(nframes=5, natom=3, ndim=3): + """Create a mock trace (list of frame dictionaries)""" + trace = [] + for i in range(nframes): + frame = { + "energy": -100.0 + i * 0.1, + "time": i * 10.0, + "position": np.random.randn(natom * ndim) + } + trace.append(frame) + return trace + + +def test_basic_trajectory_output(tmp_path): + """Test basic trajectory XYZ file output""" + model = _make_mock_model(natom=2, ndim=3, atom_types=["H", "O"]) + trace = _make_trace(nframes=3, natom=2, ndim=3) + filename = os.path.join(str(tmp_path), "test_traj.xyz") + + write_trajectory_xyz(model, trace, filename) + + assert os.path.exists(filename) + + with open(filename, "r") as f: + content = f.read() + + lines = content.strip().split("\n") + # 3 frames * (1 natom + 1 comment + 2 atoms) = 12 lines + assert len(lines) == 12 + + +def test_every_parameter(tmp_path): + """Test that 'every' parameter skips frames correctly""" + model = _make_mock_model(natom=1, ndim=3, atom_types=["C"]) + trace = _make_trace(nframes=10, natom=1, ndim=3) + filename = os.path.join(str(tmp_path), "test_every.xyz") + + write_trajectory_xyz(model, trace, filename, every=3) + + with open(filename, "r") as f: + content = f.read() + + lines = content.strip().split("\n") + # Frames 0, 3, 6, 9 should be written (4 frames) + # Each frame: 1 natom + 1 comment + 1 atom = 3 lines + assert len(lines) == 12 + + +def test_fallback_to_x_atom_types(tmp_path): + """Test fallback to 'X' when model.atom_types is None""" + model = _make_mock_model(natom=2, ndim=3, atom_types=None) + trace = _make_trace(nframes=1, natom=2, ndim=3) + filename = os.path.join(str(tmp_path), "test_fallback.xyz") + + write_trajectory_xyz(model, trace, filename) + + with open(filename, "r") as f: + content = f.read() + + lines = content.strip().split("\n") + # Both atoms should be labeled "X" + assert lines[2].startswith("X ") + assert lines[3].startswith("X ") + + +def test_comment_line_format(tmp_path): + """Test that comment line contains energy and time""" + model = _make_mock_model(natom=1, ndim=3, atom_types=["H"]) + trace = [{ + "energy": -123.456, + "time": 42.0, + "position": np.array([1.0, 2.0, 3.0]) + }] + filename = os.path.join(str(tmp_path), "test_comment.xyz") + + write_trajectory_xyz(model, trace, filename) + + with open(filename, "r") as f: + lines = f.readlines() + + comment = lines[1].strip() + assert "E=" in comment + assert "t=" in comment + assert "-123.456" in comment + assert "42" in comment + + +def test_position_reshaping(tmp_path): + """Test that flat position array is reshaped correctly""" + model = _make_mock_model(natom=2, ndim=3, atom_types=["H", "O"]) + # Position as flat array: [x1, y1, z1, x2, y2, z2] + trace = [{ + "energy": 0.0, + "time": 0.0, + "position": np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + }] + filename = os.path.join(str(tmp_path), "test_reshape.xyz") + + write_trajectory_xyz(model, trace, filename) + + with open(filename, "r") as f: + lines = f.readlines() + + # First atom should have coords (1, 2, 3) * bohr_to_angstrom + parts1 = lines[2].split() + assert float(parts1[1]) == pytest.approx(1.0 * bohr_to_angstrom, abs=1e-8) + assert float(parts1[2]) == pytest.approx(2.0 * bohr_to_angstrom, abs=1e-8) + assert float(parts1[3]) == pytest.approx(3.0 * bohr_to_angstrom, abs=1e-8) + + # Second atom should have coords (4, 5, 6) * bohr_to_angstrom + parts2 = lines[3].split() + assert float(parts2[1]) == pytest.approx(4.0 * bohr_to_angstrom, abs=1e-8) + assert float(parts2[2]) == pytest.approx(5.0 * bohr_to_angstrom, abs=1e-8) + assert float(parts2[3]) == pytest.approx(6.0 * bohr_to_angstrom, abs=1e-8) diff --git a/test/test_mains.py b/test/test_mains.py index ed30b88..fa5c390 100644 --- a/test/test_mains.py +++ b/test/test_mains.py @@ -1,10 +1,11 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -import unittest import sys import os, shutil +import pytest + import mudslide import mudslide.__main__ import mudslide.surface @@ -54,8 +55,6 @@ def compare(x, y, typekey): types = {"f": float, "d": int, "s": str} typelist = [types[x] for x in typespec] - failed = False - problems = [] for l1, l2 in zip(f1, f2): @@ -87,175 +86,128 @@ def compare(x, y, typekey): return problems -class TrajectoryTest(object): - """Test base class""" - samples = 1 - method = "fssh" - x = -10 - dt = 5 - n = 1 - seed = 200 - o = "single" - j = 1 - electronic = "exp" - log = "memory" - - def capture_traj_problems(self, k, tol, extra_options=[]): - options = "-s {0:d} -m {1:s} -k {2:f} {2:f} -x {3:f} --dt {4:f} -n {5:d} -z {6:d} -o {7:s} -j {8:d} -a {9:s} --electronic {10:s} --log {11:s}".format( - self.samples, self.model, k, self.x, self.dt, self.n, self.seed, self.o, self.j, self.method, - self.electronic, self.log).split() - options += extra_options - - checkdir = os.path.join(testdir, "checks", self.method) - options += "--logdir {}".format(checkdir).split() - - # clean_directory(checkdir) - os.makedirs(checkdir, exist_ok=True) - outfile = os.path.join(checkdir, f"{self.model:s}_k{k:d}.out") - with open(outfile, "w") as f: - mudslide.__main__.main(options, f) - - if self.o == "single": - form = "f" * (6 + 2 * self.nstate) + "df" - elif self.o == "averaged": - form = "ffff" - reffile = os.path.join(testdir, "ref", self.method, "{:s}_k{:d}.ref".format(self.model, k)) - with open(reffile) as ref, open(outfile) as out: - problems = compare_line_by_line(ref, out, form, tol) - for p in problems: - print_problem(p) - - return problems - - -class TestTSAC(unittest.TestCase, TrajectoryTest): - """Test Suite for tully simple avoided crossing""" - model = "simple" - nstate = 2 - - def test_tsac(self): - """Tully Simple Avoided Crossing""" - for k in [8, 14, 20]: - with self.subTest(k=k): - probs = self.capture_traj_problems(k, 1e-3) - self.assertEqual(len(probs), 0) - - -class TestDual(unittest.TestCase, TrajectoryTest): - """Test Suite for tully dual avoided crossing""" - model = "dual" - nstate = 2 - - def test_dual(self): - """Tully Dual Avoided Crossing""" - for k in [20, 50, 100]: - with self.subTest(k=k): - probs = self.capture_traj_problems(k, 1e-3) - self.assertEqual(len(probs), 0) - - -class TestExtended(unittest.TestCase, TrajectoryTest): - """Test Suite for tully extended coupling""" - model = "extended" - nstate = 2 - - def test_extended(self): - """Tully Extended Coupling""" - for k in [10, 15, 20]: - with self.subTest(k=k): - probs = self.capture_traj_problems(k, 1e-3) - self.assertEqual(len(probs), 0) - - -class TestTSACc(unittest.TestCase, TrajectoryTest): - """Test Suite for tully simple avoided crossing with cumulative hopping""" - model = "simple" - nstate = 2 - seed = 756396545 - method = "cumulative-sh" - electronic = "linear-rk4" - - def test_tsac_c(self): - """Tully Simple Avoided Crossing (FSSH-c)""" - for k in [10, 20]: - with self.subTest(k=k): - probs = self.capture_traj_problems(k, 1e-3) - self.assertEqual(len(probs), 0) - - -class TestEhrenfest(unittest.TestCase, TrajectoryTest): - """Test suite for ehrenfest trajectory""" - model = "simple" - nstate = 2 - method = "ehrenfest" - - def test_ehrenfest(self): - """Tully Simple Avoided Crossing (Ehrenfest)""" - k = 15 - probs = self.capture_traj_problems(k, 1e-3) - self.assertEqual(len(probs), 0) - -class TestAFSSH(unittest.TestCase, TrajectoryTest): - """Test suite for AFSSH trajectory""" - model = "dual" - nstate = 2 - method = "afssh" - seed = 78341 - - def test_afssh(self): - """Tully Dual Avoided Crossing (A-FSSH)""" - k = 14 - probs = self.capture_traj_problems(k, 1e-3) - self.assertEqual(len(probs), 0) - print(probs) - -class TestES(unittest.TestCase, TrajectoryTest): - """Test Suite for tully simple avoided crossing with cumulative hopping""" - model = "simple" - nstate = 2 - dt = 20 - seed = 84329 - method = "even-sampling" - o = "averaged" - log = "yaml" - - def test_es_tsac(self): - """Even Sampling""" - for k in [10, 20]: - with self.subTest(k=k): - probs = self.capture_traj_problems(k, 1e-3, extra_options=["--sample-stack", "5"]) - self.assertEqual(len(probs), 0) - - -class TestSurface(unittest.TestCase): - """Test Suite for surface writer""" - - def test_surface(self): - """Surface Writer""" - tol = 1e-3 - for m in ["simple", "extended", "dual", "super", "shin-metiu", "modelx", "models", "vibronic"]: - with self.subTest(m=m): - if m in ["vibronic"]: - options = "-m {:s} --x0 0 0 0 0 0 -s 2 -r -5 5".format(m).split() - else: - options = "-m {:s} -r -11 11 -n 200".format(m).split() - checkdir = os.path.join(testdir, "checks", "surface") - os.makedirs(checkdir, exist_ok=True) - outfile = os.path.join(checkdir, f"{m:s}.out") - options.append(f"--output={outfile}") - with open(outfile, "w") as f: - mudslide.surface.main(options) - - form = "f" * (8 if m in ["simple", "extended", "dual"] else 13) - if m in ["vibronic"]: - form = "f" * 20 - reffile = os.path.join(testdir, "ref", "surface", "{:s}.ref".format(m)) - with open(reffile) as ref, open(outfile) as out: - problems = compare_line_by_line(ref, out, form, tol) - for p in problems: - print_problem(p) - self.assertEqual(len(problems), 0) - - -if __name__ == '__main__': - unittest.main() +def capture_traj_problems(model, nstate, k, tol, method="fssh", x=-10, dt=5, + n=1, seed=200, o="single", electronic="exp", + log="memory", extra_options=None): + if extra_options is None: + extra_options = [] + options = ("-s {0:d} -m {1:s} -k {2:f} {2:f} -x {3:f} --dt {4:f} -n {5:d} " + "-z {6:d} -o {7:s} -a {8:s} --electronic {9:s} " + "--log {10:s}").format( + 1, model, k, x, dt, n, seed, o, method, electronic, log).split() + options += extra_options + + checkdir = os.path.join(testdir, "checks", method) + options += "--logdir {}".format(checkdir).split() + + os.makedirs(checkdir, exist_ok=True) + outfile = os.path.join(checkdir, f"{model:s}_k{k:d}.out") + with open(outfile, "w") as f: + mudslide.__main__.main(options, f) + + if o == "single": + form = "f" * (6 + 2 * nstate) + "df" + elif o == "averaged": + form = "ffff" + reffile = os.path.join(testdir, "ref", method, "{:s}_k{:d}.ref".format(model, k)) + with open(reffile) as ref, open(outfile) as out: + problems = compare_line_by_line(ref, out, form, tol) + for p in problems: + print_problem(p) + + return problems + + +# -- Tully Simple Avoided Crossing (FSSH) -- + +@pytest.mark.parametrize("k", [8, 14, 20]) +def test_tsac(k): + """Tully Simple Avoided Crossing""" + probs = capture_traj_problems("simple", 2, k, 1e-3) + assert len(probs) == 0 + + +# -- Tully Dual Avoided Crossing (FSSH) -- + +@pytest.mark.parametrize("k", [20, 50, 100]) +def test_dual(k): + """Tully Dual Avoided Crossing""" + probs = capture_traj_problems("dual", 2, k, 1e-3) + assert len(probs) == 0 + + +# -- Tully Extended Coupling (FSSH) -- + +@pytest.mark.parametrize("k", [10, 15, 20]) +def test_extended(k): + """Tully Extended Coupling""" + probs = capture_traj_problems("extended", 2, k, 1e-3) + assert len(probs) == 0 + + +# -- Tully Simple Avoided Crossing (cumulative-sh with linear-rk4) -- + +@pytest.mark.parametrize("k", [10, 20]) +def test_tsac_cumulative(k): + """Tully Simple Avoided Crossing (FSSH-c)""" + probs = capture_traj_problems("simple", 2, k, 1e-3, + method="cumulative-sh", seed=756396545, + electronic="linear-rk4") + assert len(probs) == 0 + + +# -- Ehrenfest -- + +def test_ehrenfest(): + """Tully Simple Avoided Crossing (Ehrenfest)""" + probs = capture_traj_problems("simple", 2, 15, 1e-3, method="ehrenfest") + assert len(probs) == 0 + + +# -- A-FSSH -- + +def test_afssh(): + """Tully Dual Avoided Crossing (A-FSSH)""" + probs = capture_traj_problems("dual", 2, 14, 1e-3, method="afssh", seed=78341) + assert len(probs) == 0 + + +# -- Even Sampling -- + +@pytest.mark.parametrize("k", [10, 20]) +def test_even_sampling(k): + """Even Sampling""" + probs = capture_traj_problems("simple", 2, k, 1e-3, + method="even-sampling", dt=20, seed=84329, + o="averaged", log="yaml", + extra_options=["--sample-stack", "5"]) + assert len(probs) == 0 + + +# -- Surface Writer -- + +@pytest.mark.parametrize("m", ["simple", "extended", "dual", "super", + "shin-metiu", "modelx", "models", "vibronic"]) +def test_surface(m): + """Surface Writer""" + tol = 1e-3 + if m in ["vibronic"]: + options = "-m {:s} --x0 0 0 0 0 0 -s 2 -r -5 5".format(m).split() + else: + options = "-m {:s} -r -11 11 -n 200".format(m).split() + checkdir = os.path.join(testdir, "checks", "surface") + os.makedirs(checkdir, exist_ok=True) + outfile = os.path.join(checkdir, f"{m:s}.out") + options.append(f"--output={outfile}") + with open(outfile, "w") as f: + mudslide.surface.main(options) + + form = "f" * (8 if m in ["simple", "extended", "dual"] else 13) + if m in ["vibronic"]: + form = "f" * 20 + reffile = os.path.join(testdir, "ref", "surface", "{:s}.ref".format(m)) + with open(reffile) as ref, open(outfile) as out: + problems = compare_line_by_line(ref, out, form, tol) + for p in problems: + print_problem(p) + assert len(problems) == 0 diff --git a/test/test_math.py b/test/test_math.py index 78ae1cd..5976ffe 100644 --- a/test/test_math.py +++ b/test/test_math.py @@ -2,16 +2,17 @@ # -*- coding: utf-8 -*- """Unit testing for mudslide""" -import unittest +import pytest import sys from mudslide import poisson_prob_scale +from mudslide.exceptions import ConfigurationError from mudslide.math import boltzmann_velocities from mudslide.constants import boltzmann import numpy as np -class TestMath(unittest.TestCase): +class TestMath: """Test Suite for math functions""" def test_poisson_scale(self): @@ -20,97 +21,103 @@ def test_poisson_scale(self): fx = poisson_prob_scale(args) refs = np.array([1.0, 0.7869386805747332, 0.6321205588285577]) for i, x in enumerate(args): - self.assertAlmostEqual(fx[i], refs[i], places=10) - self.assertAlmostEqual(poisson_prob_scale(args[i]), refs[i], places=10) + assert fx[i] == pytest.approx(refs[i], abs=1e-10) + assert poisson_prob_scale(args[i]) == pytest.approx(refs[i], abs=1e-10) -class TestBoltzmannVelocities(unittest.TestCase): +class TestBoltzmannVelocities: """Test Suite for boltzmann_velocities function""" - def setUp(self): - self.natom = 4 - self.temperature = 300.0 - self.seed = 42 + def _setup(self): + natom = 4 + temperature = 300.0 + seed = 42 # masses for 4 atoms, repeated for each DOF atom_masses = np.array([16.0, 1.0, 1.0, 12.0]) - self.mass_flat = np.repeat(atom_masses, 3) # shape (12,) - self.mass_2d = np.column_stack([atom_masses] * 3) # shape (4, 3) + mass_flat = np.repeat(atom_masses, 3) # shape (12,) + mass_2d = np.column_stack([atom_masses] * 3) # shape (4, 3) # coordinates for angular momentum removal - self.coords_flat = np.array([ + coords_flat = np.array([ 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, ]) - self.coords_2d = self.coords_flat.reshape((4, 3)) + coords_2d = coords_flat.reshape((4, 3)) + return (natom, temperature, seed, atom_masses, + mass_flat, mass_2d, coords_flat, coords_2d) def test_flat_input_returns_flat_output(self): """Flat mass input (ndof,) should return flat velocities (ndof,)""" - v = boltzmann_velocities(self.mass_flat, self.temperature, seed=self.seed) - self.assertEqual(v.shape, self.mass_flat.shape) + _, temperature, seed, _, mass_flat, _, _, _ = self._setup() + v = boltzmann_velocities(mass_flat, temperature, seed=seed) + assert v.shape == mass_flat.shape def test_2d_input_returns_2d_output(self): """2D mass input (natom, 3) should return 2D velocities (natom, 3)""" - v = boltzmann_velocities(self.mass_2d, self.temperature, seed=self.seed) - self.assertEqual(v.shape, self.mass_2d.shape) + _, temperature, seed, _, _, mass_2d, _, _ = self._setup() + v = boltzmann_velocities(mass_2d, temperature, seed=seed) + assert v.shape == mass_2d.shape def test_flat_and_2d_give_same_velocities(self): """Flat and 2D inputs with same seed should produce equivalent velocities""" - v_flat = boltzmann_velocities(self.mass_flat, self.temperature, seed=self.seed) - v_2d = boltzmann_velocities(self.mass_2d, self.temperature, seed=self.seed) + _, temperature, seed, _, mass_flat, mass_2d, _, _ = self._setup() + v_flat = boltzmann_velocities(mass_flat, temperature, seed=seed) + v_2d = boltzmann_velocities(mass_2d, temperature, seed=seed) np.testing.assert_allclose(v_flat, v_2d.reshape(-1), atol=1e-12) def test_com_removed(self): """Center of mass momentum should be zero after removal""" - v = boltzmann_velocities(self.mass_flat, self.temperature, - remove_translation=True, seed=self.seed) + _, temperature, seed, _, mass_flat, _, _, _ = self._setup() + v = boltzmann_velocities(mass_flat, temperature, + remove_translation=True, seed=seed) v3 = v.reshape((-1, 3)) - atom_masses = self.mass_flat.reshape((-1, 3))[:, 0] + atom_masses = mass_flat.reshape((-1, 3))[:, 0] com_momentum = np.sum(v3 * atom_masses[:, np.newaxis], axis=0) np.testing.assert_allclose(com_momentum, 0.0, atol=1e-10) def test_angular_momentum_removed(self): """Angular momentum should be approximately zero after removal""" - v = boltzmann_velocities(self.mass_flat, self.temperature, + _, temperature, seed, _, mass_flat, _, coords_flat, coords_2d = self._setup() + v = boltzmann_velocities(mass_flat, temperature, remove_translation=True, - coords=self.coords_flat, + coords=coords_flat, remove_rotation=True, - seed=self.seed) + seed=seed) v3 = v.reshape((-1, 3)) - atom_masses = self.mass_flat.reshape((-1, 3))[:, 0] + atom_masses = mass_flat.reshape((-1, 3))[:, 0] momentum = v3 * atom_masses[:, np.newaxis] - angular_momentum = np.cross(self.coords_2d, momentum).sum(axis=0) + angular_momentum = np.cross(coords_2d, momentum).sum(axis=0) np.testing.assert_allclose(angular_momentum, 0.0, atol=1e-10) def test_angular_momentum_removed_2d(self): """Angular momentum removal should work with 2D mass and coords""" - v = boltzmann_velocities(self.mass_2d, self.temperature, + _, temperature, seed, _, _, mass_2d, _, coords_2d = self._setup() + v = boltzmann_velocities(mass_2d, temperature, remove_translation=True, - coords=self.coords_2d, + coords=coords_2d, remove_rotation=True, - seed=self.seed) - self.assertEqual(v.shape, self.mass_2d.shape) - atom_masses = self.mass_2d[:, 0] + seed=seed) + assert v.shape == mass_2d.shape + atom_masses = mass_2d[:, 0] momentum = v * atom_masses[:, np.newaxis] - angular_momentum = np.cross(self.coords_2d, momentum).sum(axis=0) + angular_momentum = np.cross(coords_2d, momentum).sum(axis=0) np.testing.assert_allclose(angular_momentum, 0.0, atol=1e-10) def test_angular_momentum_auto_enabled_with_coords(self): """Rotation removal should be auto-enabled when coords are provided""" - v = boltzmann_velocities(self.mass_flat, self.temperature, - coords=self.coords_flat, seed=self.seed) + _, temperature, seed, _, mass_flat, _, coords_flat, coords_2d = self._setup() + v = boltzmann_velocities(mass_flat, temperature, + coords=coords_flat, seed=seed) v3 = v.reshape((-1, 3)) - atom_masses = self.mass_flat.reshape((-1, 3))[:, 0] + atom_masses = mass_flat.reshape((-1, 3))[:, 0] momentum = v3 * atom_masses[:, np.newaxis] - angular_momentum = np.cross(self.coords_2d, momentum).sum(axis=0) + angular_momentum = np.cross(coords_2d, momentum).sum(axis=0) np.testing.assert_allclose(angular_momentum, 0.0, atol=1e-10) def test_remove_rotation_without_coords_raises(self): """Requesting rotation removal without coords should raise ValueError""" - with self.assertRaises(ValueError): - boltzmann_velocities(self.mass_flat, self.temperature, - remove_rotation=True, seed=self.seed) - - -if __name__ == '__main__': - unittest.main() + _, temperature, seed, _, mass_flat, _, _, _ = self._setup() + with pytest.raises(ConfigurationError): + boltzmann_velocities(mass_flat, temperature, + remove_rotation=True, seed=seed) diff --git a/test/test_openmm.py b/test/test_openmm.py index 510e89f..87dcb19 100644 --- a/test/test_openmm.py +++ b/test/test_openmm.py @@ -5,7 +5,6 @@ import numpy as np import os import shutil -import unittest import pytest import mudslide diff --git a/test/test_qmmm.py b/test/test_qmmm.py index f1aa8df..8df7bbd 100644 --- a/test/test_qmmm.py +++ b/test/test_qmmm.py @@ -5,7 +5,6 @@ import numpy as np import os import shutil -import unittest import pytest import mudslide diff --git a/test/test_quadrature.py b/test/test_quadrature.py index 267004e..3dcf5f1 100644 --- a/test/test_quadrature.py +++ b/test/test_quadrature.py @@ -2,107 +2,107 @@ # -*- coding: utf-8 -*- """Unit testing for mudslide""" -import unittest -import sys -import re - import mudslide import numpy as np -class TestQuadratures(unittest.TestCase): - """Test Suite for integration quadratures""" - - def setUp(self): - """Setup function""" - self.a = 0 - self.b = 1 - self.n = 11 - - self.evenx = [i / (self.n - 1) for i in range(self.n)] - self.flatw = [ - 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, - 0.09090909, 0.09090909 - ] - - def test_midpoint_quadrature(self): - """Midpoint quadrature""" - refx = [ - 0.04545455, 0.13636364, 0.22727273, 0.31818182, 0.40909091, 0.5, 0.59090909, 0.68181818, 0.77272727, - 0.86363636, 0.95454545 - ] - refw = self.flatw - x, w = mudslide.integration.quadrature(self.n, self.a, self.b, "midpoint") - self.assertEqual(len(x), self.n) - self.assertTrue(np.all(np.isclose(x, refx))) - - self.assertEqual(len(w), self.n) - self.assertTrue(np.all(np.isclose(w, refw))) - - def test_trapezoid_quadrature(self): - """Trapezoid quadrature""" - refx = self.evenx - refw = [0.1 for i in range(self.n)] - refw[0] *= 0.5 - refw[-1] *= 0.5 - - x, w = mudslide.integration.quadrature(self.n, self.a, self.b, "trapezoid") - self.assertEqual(len(x), self.n) - self.assertTrue(np.all(np.isclose(x, refx))) - - self.assertEqual(len(w), self.n) - self.assertTrue(np.all(np.isclose(w, refw))) - - def test_simpson_quadrature(self): - """Simpson quadrature""" - refx = self.evenx - refw = [ - 0.03333333, 0.13333333, 0.06666667, 0.13333333, 0.06666667, 0.13333333, 0.06666667, 0.13333333, 0.06666667, - 0.13333333, 0.03333333 - ] - - x, w = mudslide.integration.quadrature(self.n, self.a, self.b, "simpson") - self.assertEqual(len(x), self.n) - self.assertTrue(np.all(np.isclose(x, refx))) - - self.assertEqual(len(w), self.n) - self.assertTrue(np.all(np.isclose(w, refw))) - - def test_clenshawcurtis_quadrature(self): - """Clenshaw-Curtis quadrature""" - refx = [ - 0., 0.02447174, 0.0954915, 0.20610737, 0.3454915, 0.5, 0.6545085, 0.79389263, 0.9045085, 0.97552826, 1. - ] - refw = [ - 0.00505051, 0.04728953, 0.09281761, 0.12679417, 0.14960664, 0.15688312, 0.14960664, 0.12679417, 0.09281761, - 0.04728953, 0.00505051 - ] - - x, w = mudslide.integration.quadrature(self.n, self.a, self.b, "cc") - self.assertEqual(len(x), self.n) - self.assertTrue(np.all(np.isclose(x, refx))) - - self.assertEqual(len(w), self.n) - self.assertTrue(np.all(np.isclose(w, refw))) - - def test_gasslegendre_quadrature(self): - """Gauss-Legendre quadrature""" - refx = [ - 0.01088567, 0.0564687, 0.134924, 0.24045194, 0.36522842, 0.5, 0.63477158, 0.75954806, 0.865076, 0.9435313, - 0.98911433 - ] - refw = [ - 0.02783428, 0.06279018, 0.09314511, 0.11659688, 0.13140227, 0.13646254, 0.13140227, 0.11659688, 0.09314511, - 0.06279018, 0.02783428 - ] - - x, w = mudslide.integration.quadrature(self.n, self.a, self.b, "gl") - self.assertEqual(len(x), self.n) - self.assertTrue(np.all(np.isclose(x, refx))) - - self.assertEqual(len(w), self.n) - self.assertTrue(np.all(np.isclose(w, refw))) - - -if __name__ == '__main__': - unittest.main() +def _setup_quadrature(): + """Setup function""" + a = 0 + b = 1 + n = 11 + + evenx = [i / (n - 1) for i in range(n)] + flatw = [ + 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, 0.09090909, + 0.09090909, 0.09090909 + ] + return a, b, n, evenx, flatw + + +def test_midpoint_quadrature(): + """Midpoint quadrature""" + a, b, n, evenx, flatw = _setup_quadrature() + refx = [ + 0.04545455, 0.13636364, 0.22727273, 0.31818182, 0.40909091, 0.5, 0.59090909, 0.68181818, 0.77272727, + 0.86363636, 0.95454545 + ] + refw = flatw + x, w = mudslide.integration.quadrature(n, a, b, "midpoint") + assert len(x) == n + assert np.all(np.isclose(x, refx)) + + assert len(w) == n + assert np.all(np.isclose(w, refw)) + + +def test_trapezoid_quadrature(): + """Trapezoid quadrature""" + a, b, n, evenx, flatw = _setup_quadrature() + refx = evenx + refw = [0.1 for i in range(n)] + refw[0] *= 0.5 + refw[-1] *= 0.5 + + x, w = mudslide.integration.quadrature(n, a, b, "trapezoid") + assert len(x) == n + assert np.all(np.isclose(x, refx)) + + assert len(w) == n + assert np.all(np.isclose(w, refw)) + + +def test_simpson_quadrature(): + """Simpson quadrature""" + a, b, n, evenx, flatw = _setup_quadrature() + refx = evenx + refw = [ + 0.03333333, 0.13333333, 0.06666667, 0.13333333, 0.06666667, 0.13333333, 0.06666667, 0.13333333, 0.06666667, + 0.13333333, 0.03333333 + ] + + x, w = mudslide.integration.quadrature(n, a, b, "simpson") + assert len(x) == n + assert np.all(np.isclose(x, refx)) + + assert len(w) == n + assert np.all(np.isclose(w, refw)) + + +def test_clenshawcurtis_quadrature(): + """Clenshaw-Curtis quadrature""" + a, b, n, evenx, flatw = _setup_quadrature() + refx = [ + 0., 0.02447174, 0.0954915, 0.20610737, 0.3454915, 0.5, 0.6545085, 0.79389263, 0.9045085, 0.97552826, 1. + ] + refw = [ + 0.00505051, 0.04728953, 0.09281761, 0.12679417, 0.14960664, 0.15688312, 0.14960664, 0.12679417, 0.09281761, + 0.04728953, 0.00505051 + ] + + x, w = mudslide.integration.quadrature(n, a, b, "cc") + assert len(x) == n + assert np.all(np.isclose(x, refx)) + + assert len(w) == n + assert np.all(np.isclose(w, refw)) + + +def test_gasslegendre_quadrature(): + """Gauss-Legendre quadrature""" + a, b, n, evenx, flatw = _setup_quadrature() + refx = [ + 0.01088567, 0.0564687, 0.134924, 0.24045194, 0.36522842, 0.5, 0.63477158, 0.75954806, 0.865076, 0.9435313, + 0.98911433 + ] + refw = [ + 0.02783428, 0.06279018, 0.09314511, 0.11659688, 0.13140227, 0.13646254, 0.13140227, 0.11659688, 0.09314511, + 0.06279018, 0.02783428 + ] + + x, w = mudslide.integration.quadrature(n, a, b, "gl") + assert len(x) == n + assert np.all(np.isclose(x, refx)) + + assert len(w) == n + assert np.all(np.isclose(w, refw)) diff --git a/test/test_scattering_models.py b/test/test_scattering_models.py index 2371418..889059f 100644 --- a/test/test_scattering_models.py +++ b/test/test_scattering_models.py @@ -5,7 +5,6 @@ import numpy as np import os import shutil -import unittest import pytest import mudslide diff --git a/test/test_spawnstack.py b/test/test_spawnstack.py index 15c45bc..163c26a 100644 --- a/test/test_spawnstack.py +++ b/test/test_spawnstack.py @@ -3,170 +3,167 @@ """Unit testing for larger order integrants""" import numpy as np -import unittest +import pytest +from mudslide.exceptions import ConfigurationError from mudslide.even_sampling import SpawnStack import itertools -class Test_ES(unittest.TestCase): - def test_2D_integral(self): - def two_d_poly_np(x, y): - c = np.array( - [ - [-12, 0, -1], - [0, 3, 0], - [2, 0, 0], - ] - ) - return np.polynomial.polynomial.polyval2d(x, y, c) - - - ss = SpawnStack.from_quadrature(nsamples=[5,5]) - pnts_wghts = ss.unravel() - - results_2Dpoly = sum ( - [ - two_d_poly_np(x,y)* w for ((x, y), w) in pnts_wghts - ] +def test_2D_integral(): + def two_d_poly_np(x, y): + c = np.array( + [ + [-12, 0, -1], + [0, 3, 0], + [2, 0, 0], + ] ) + return np.polynomial.polynomial.polyval2d(x, y, c) + + + ss = SpawnStack.from_quadrature(nsamples=[5,5]) + pnts_wghts = ss.unravel() - results_analytical_2D = -131/12 + results_2Dpoly = sum ( + [ + two_d_poly_np(x,y)* w for ((x, y), w) in pnts_wghts + ] + ) - self.assertAlmostEqual(results_2Dpoly, results_analytical_2D) + results_analytical_2D = -131/12 - def test_3D_integral(self): - def three_d_poly_np(x, y, z): - c = np.array( + assert results_2Dpoly == pytest.approx(results_analytical_2D) + +def test_3D_integral(): + def three_d_poly_np(x, y, z): + c = np.array( + [ + [ + [1, 2, 1], + [2, 4, 2], + [1, 2, 1], + ], [ - [ - [1, 2, 1], - [2, 4, 2], - [1, 2, 1], - ], - [ - [2, 4, 2], - [4, 8, 4], - [2, 4, 2], - ], - [ - [1, 2, 1], - [2, 4, 2], - [1, 2, 1], - ] ] - - ) - return np.polynomial.polynomial.polyval3d(x, y, z, c) - ss = SpawnStack.from_quadrature(nsamples=[7,7,7]) - pnts_wghts = ss.unravel() - - results_3Dpoly = sum ( - [ - three_d_poly_np(x,y,z)* w for ((x, y, z), w) in pnts_wghts - ] + [2, 4, 2], + [4, 8, 4], + [2, 4, 2], + ], + [ + [1, 2, 1], + [2, 4, 2], + [1, 2, 1], + ] ] + ) + return np.polynomial.polynomial.polyval3d(x, y, z, c) + ss = SpawnStack.from_quadrature(nsamples=[7,7,7]) + pnts_wghts = ss.unravel() - results_analytical_3D = 343/27 - results_analytical_3D = 343/27 - self.assertAlmostEqual(results_3Dpoly, results_analytical_3D) + results_3Dpoly = sum ( + [ + three_d_poly_np(x,y,z)* w for ((x, y, z), w) in pnts_wghts + ] + ) - def test_dimension(self): - sample_stack = [] - ss = SpawnStack.from_quadrature (nsamples = [2,2,2]) + results_analytical_3D = 343/27 + results_analytical_3D = 343/27 + assert results_3Dpoly == pytest.approx(results_analytical_3D) - ss_2D = SpawnStack.from_quadrature (nsamples = [2,2]) +def test_dimension(): + sample_stack = [] + ss = SpawnStack.from_quadrature (nsamples = [2,2,2]) - ss_1D = SpawnStack.from_quadrature (nsamples = [2]) + ss_2D = SpawnStack.from_quadrature (nsamples = [2,2]) - results_empty = [{'zeta': 1.0, 'dw': 1.0, 'children': [], 'spawn_size': 1}] - results_append = [{'zeta': 1.0, 'dw': 1.0, 'children': [{'zeta': 0.1, 'dw': 0.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}] + ss_1D = SpawnStack.from_quadrature (nsamples = [2]) - zetas = [ss_1D.sample_stack[i]["zeta"] for i in range (2)] - dws = [ss_1D.sample_stack[i]["dw"] for i in range (2)] + results_empty = [{'zeta': 1.0, 'dw': 1.0, 'children': [], 'spawn_size': 1}] + results_append = [{'zeta': 1.0, 'dw': 1.0, 'children': [{'zeta': 0.1, 'dw': 0.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}] - ss_2D.append_layer(zetas = zetas, dws = dws, stack = ss_2D.sample_stack) + zetas = [ss_1D.sample_stack[i]["zeta"] for i in range (2)] + dws = [ss_1D.sample_stack[i]["dw"] for i in range (2)] - zetas = [0.1] - dws = [0.1] - self.assertDictEqual(ss_2D.sample_stack[0], ss.sample_stack[0]) - self.assertDictEqual(ss_2D.sample_stack[1], ss.sample_stack[1]) + ss_2D.append_layer(zetas = zetas, dws = dws, stack = ss_2D.sample_stack) - merged = list(itertools.chain(*ss.sample_stack)) - merged = list(itertools.chain(*ss.sample_stack)) + zetas = [0.1] + dws = [0.1] + assert ss_2D.sample_stack[0] == ss.sample_stack[0] + assert ss_2D.sample_stack[1] == ss.sample_stack[1] + merged = list(itertools.chain(*ss.sample_stack)) + merged = list(itertools.chain(*ss.sample_stack)) - def test_unravel_2D(self): - ss = SpawnStack.from_quadrature(nsamples=[2,2]) - ss.unravel() - pnts_wghts = ss.unravel() - results_unravel_2D = [((0.21132486540518713, 0.21132486540518713), 0.25), - ((0.21132486540518713, 0.7886751345948129), 0.25), - ((0.7886751345948129, 0.21132486540518713), 0.25), - ((0.7886751345948129, 0.7886751345948129), 0.25)] +def test_unravel_2D(): + ss = SpawnStack.from_quadrature(nsamples=[2,2]) + ss.unravel() + pnts_wghts = ss.unravel() - self.assertEqual(results_unravel_2D, pnts_wghts) + results_unravel_2D = [((0.21132486540518713, 0.21132486540518713), 0.25), + ((0.21132486540518713, 0.7886751345948129), 0.25), + ((0.7886751345948129, 0.21132486540518713), 0.25), + ((0.7886751345948129, 0.7886751345948129), 0.25)] + assert results_unravel_2D == pnts_wghts - def test_unravel_3D(self): - ss = SpawnStack.from_quadrature(nsamples=[2,2,2]) - ss.unravel() - pnts_wghts = ss.unravel() - results_unravel_3D = [((0.21132486540518713, 0.21132486540518713, 0.21132486540518713), 0.125), - ((0.21132486540518713, 0.21132486540518713, 0.7886751345948129), 0.125), - ((0.21132486540518713, 0.7886751345948129, 0.21132486540518713), 0.125), - ((0.21132486540518713, 0.7886751345948129, 0.7886751345948129), 0.125), - ((0.7886751345948129, 0.21132486540518713, 0.21132486540518713), 0.125), - ((0.7886751345948129, 0.21132486540518713, 0.7886751345948129), 0.125), - ((0.7886751345948129, 0.7886751345948129, 0.21132486540518713), 0.125), - ((0.7886751345948129, 0.7886751345948129, 0.7886751345948129), 0.125)] +def test_unravel_3D(): + ss = SpawnStack.from_quadrature(nsamples=[2,2,2]) + ss.unravel() + pnts_wghts = ss.unravel() - self.assertEqual(results_unravel_3D, pnts_wghts) + results_unravel_3D = [((0.21132486540518713, 0.21132486540518713, 0.21132486540518713), 0.125), + ((0.21132486540518713, 0.21132486540518713, 0.7886751345948129), 0.125), + ((0.21132486540518713, 0.7886751345948129, 0.21132486540518713), 0.125), + ((0.21132486540518713, 0.7886751345948129, 0.7886751345948129), 0.125), + ((0.7886751345948129, 0.21132486540518713, 0.21132486540518713), 0.125), + ((0.7886751345948129, 0.21132486540518713, 0.7886751345948129), 0.125), + ((0.7886751345948129, 0.7886751345948129, 0.21132486540518713), 0.125), + ((0.7886751345948129, 0.7886751345948129, 0.7886751345948129), 0.125)] + assert results_unravel_3D == pnts_wghts - def test_append_layer(self): - sample_stack = [] - ss = SpawnStack(sample_stack = sample_stack) - results_empty = [{'zeta': 1.0, 'dw': 1.0, 'children': [], 'spawn_size': 1}] - results_append = [{'zeta': 1.0, 'dw': 1.0, 'children': [{'zeta': 0.1, 'dw': 0.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}] +def test_append_layer(): + sample_stack = [] + ss = SpawnStack(sample_stack = sample_stack) - zetas = [1.0] - dws = [1.0] - ss.append_layer(zetas=zetas, dws=dws, stack=ss.sample_stack) + results_empty = [{'zeta': 1.0, 'dw': 1.0, 'children': [], 'spawn_size': 1}] + results_append = [{'zeta': 1.0, 'dw': 1.0, 'children': [{'zeta': 0.1, 'dw': 0.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}] - self.assertDictEqual(results_empty[0], ss.sample_stack[0]) + zetas = [1.0] + dws = [1.0] + ss.append_layer(zetas=zetas, dws=dws, stack=ss.sample_stack) + assert results_empty[0] == ss.sample_stack[0] - zetas = [0.1] - dws = [0.1] - ss.append_layer(zetas=zetas, dws=dws, stack=ss.sample_stack) - self.assertDictEqual(results_append[0], ss.sample_stack[0]) + zetas = [0.1] + dws = [0.1] + ss.append_layer(zetas=zetas, dws=dws, stack=ss.sample_stack) - def test_ss_dimesion(self): - ss = SpawnStack.from_quadrature(nsamples=[2, 2], mcsamples=3) - sample_stack = ss.sample_stack + assert results_append[0] == ss.sample_stack[0] - zetas=[1.0, 1.1] - dws=[2.0, 2.1] - for layers in range(5): - ss.append_layer(zetas, dws) +def test_ss_dimesion(): + ss = SpawnStack.from_quadrature(nsamples=[2, 2], mcsamples=3) + sample_stack = ss.sample_stack + zetas=[1.0, 1.1] + dws=[2.0, 2.1] + for layers in range(5): + ss.append_layer(zetas, dws) - results_ss_dimension= [{'zeta': 0.21132486540518713, 'dw': 0.5, 'children': [{'zeta': 0.21132486540518713, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}, {'zeta': 0.7886751345948129, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}], 'spawn_size': 1}, {'zeta': 0.7886751345948129, 'dw': 0.5, 'children': [{'zeta': 0.21132486540518713, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}, {'zeta': 0.7886751345948129, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}], 'spawn_size': 1}] - self.assertEqual(results_ss_dimension, sample_stack) + results_ss_dimension= [{'zeta': 0.21132486540518713, 'dw': 0.5, 'children': [{'zeta': 0.21132486540518713, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}, {'zeta': 0.7886751345948129, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}], 'spawn_size': 1}, {'zeta': 0.7886751345948129, 'dw': 0.5, 'children': [{'zeta': 0.21132486540518713, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}, {'zeta': 0.7886751345948129, 'dw': 0.5, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [{'zeta': 1.0, 'dw': 2.0, 'children': [], 'spawn_size': 1}, {'zeta': 1.1, 'dw': 2.1, 'children': [], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 1}], 'spawn_size': 3}], 'spawn_size': 1}] + assert results_ss_dimension == sample_stack - def test_input(self): - zetas = [0.1] - dws = [0.1, 0.0] - sample_stack = [] - ss = SpawnStack(sample_stack = sample_stack) +def test_input(): + zetas = [0.1] + dws = [0.1, 0.0] - with self.assertRaises(ValueError): - ss.append_layer(zetas=zetas, dws=dws, stack=ss.sample_stack) + sample_stack = [] + ss = SpawnStack(sample_stack = sample_stack) -if __name__ == '__main__': - unittest.main() + with pytest.raises(ConfigurationError): + ss.append_layer(zetas=zetas, dws=dws, stack=ss.sample_stack) diff --git a/test/test_tracer.py b/test/test_tracer.py index a1d404f..6bb779a 100644 --- a/test/test_tracer.py +++ b/test/test_tracer.py @@ -1,10 +1,11 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -import unittest -import os, shutil +import os +import shutil import numpy as np +import pytest import mudslide testdir = os.path.dirname(__file__) @@ -18,76 +19,69 @@ def clean_directory(dirname): shutil.rmtree(dirname) -class TrajectoryTest(unittest.TestCase): +def test_log_yaml(): + refdir = os.path.join(_refdir, "tracer") + rundir = os.path.join(_checkdir, "tracer") + clean_directory(rundir) - def test_log_yaml(self): - refdir = os.path.join(_refdir, "tracer") - rundir = os.path.join(_checkdir, "tracer") - clean_directory(rundir) + model = mudslide.models.TullySimpleAvoidedCrossing() + log = mudslide.YAMLTrace(base_name="test-traj", location=rundir, log_pitch=8) + traj = mudslide.SurfaceHoppingMD(model, [-3.0], [10.0 / model.mass], + 0, + dt=4, + tracer=log, + max_time=80, + zeta_list=[0.2, 0.2, 0.9], + hopping_method="instantaneous") + results = traj.simulate() - model = mudslide.models.TullySimpleAvoidedCrossing() - log = mudslide.YAMLTrace(base_name="test-traj", location=rundir, log_pitch=8) - traj = mudslide.SurfaceHoppingMD(model, [-3.0], [10.0 / model.mass], - 0, - dt=4, - tracer=log, - max_time=80, - zeta_list=[0.2, 0.2, 0.9], - hopping_method="instantaneous") - results = traj.simulate() + main_log = results.main_log - main_log = results.main_log + assert main_log == "test-traj-0.yaml" - assert main_log == "test-traj-0.yaml" + snap_t16 = results[16] - snap_t16 = results[16] + refs = mudslide.load_log(os.path.join(refdir, "test-traj-0.yaml")) - refs = mudslide.load_log(os.path.join(refdir, "test-traj-0.yaml")) + ref_t16 = refs[16] - ref_t16 = refs[16] + for prop in ["position", "velocity", "density_matrix"]: + np.testing.assert_almost_equal(snap_t16[prop], ref_t16[prop], decimal=8) - for prop in ["position", "velocity", "density_matrix"]: - np.testing.assert_almost_equal(snap_t16[prop], ref_t16[prop], decimal=8) + for prop in ["potential", "kinetic", "hopping"]: + assert snap_t16[prop] == pytest.approx(ref_t16[prop], abs=1e-8) - for prop in ["potential", "kinetic", "hopping"]: - self.assertAlmostEqual(snap_t16[prop], ref_t16[prop], places=8) - def test_restart_from_trace(self): - refdir = os.path.join(_refdir, "trace_restart") - rundir = os.path.join(_checkdir, "trace_restart") - clean_directory(rundir) +def test_restart_from_trace(): + refdir = os.path.join(_refdir, "trace_restart") + rundir = os.path.join(_checkdir, "trace_restart") + clean_directory(rundir) - model = mudslide.models.TullySimpleAvoidedCrossing() - log = mudslide.YAMLTrace(base_name="test-traj", location=rundir, log_pitch=8) - traj = mudslide.SurfaceHoppingMD(model, [-3.0], [10.0 / model.mass], - 0, - dt=4, - tracer=log, - max_time=40, - zeta_list=[0.2, 0.2, 0.9], - hopping_method="instantaneous") - results = traj.simulate() + model = mudslide.models.TullySimpleAvoidedCrossing() + log = mudslide.YAMLTrace(base_name="test-traj", location=rundir, log_pitch=8) + traj = mudslide.SurfaceHoppingMD(model, [-3.0], [10.0 / model.mass], + 0, + dt=4, + tracer=log, + max_time=40, + zeta_list=[0.2, 0.2, 0.9], + hopping_method="instantaneous") + results = traj.simulate() - main_log = results.main_log + main_log = results.main_log - assert main_log == "test-traj-0.yaml" + assert main_log == "test-traj-0.yaml" - yaml_trace = mudslide.load_log(os.path.join(results.location, main_log)) - traj2 = mudslide.SurfaceHoppingMD.restart(model, yaml_trace, max_time=80) - results2 = traj2.simulate() - snap_t16 = results2[16] + yaml_trace = mudslide.load_log(os.path.join(results.location, main_log)) + traj2 = mudslide.SurfaceHoppingMD.restart(model, yaml_trace, max_time=80) + results2 = traj2.simulate() + snap_t16 = results2[16] - refs = mudslide.load_log(os.path.join(refdir, "test-traj-0.yaml")) - ref_t16 = refs[16] + refs = mudslide.load_log(os.path.join(refdir, "test-traj-0.yaml")) + ref_t16 = refs[16] - for prop in ["position", "velocity", "density_matrix"]: - with self.subTest(property=prop): - np.testing.assert_almost_equal(snap_t16[prop], ref_t16[prop], decimal=8) + for prop in ["position", "velocity", "density_matrix"]: + np.testing.assert_almost_equal(snap_t16[prop], ref_t16[prop], decimal=8) - for prop in ["potential", "kinetic", "hopping"]: - with self.subTest(property=prop): - self.assertAlmostEqual(snap_t16[prop], ref_t16[prop], places=8) - - -if __name__ == '__main__': - unittest.main() + for prop in ["potential", "kinetic", "hopping"]: + assert snap_t16[prop] == pytest.approx(ref_t16[prop], abs=1e-8) diff --git a/test/test_tracer_extended.py b/test/test_tracer_extended.py index c54546b..c37657f 100644 --- a/test/test_tracer_extended.py +++ b/test/test_tracer_extended.py @@ -8,6 +8,7 @@ import pytest import mudslide +from mudslide.exceptions import ConfigurationError from mudslide.tracer import (_sanitize_for_yaml, InMemoryTrace, YAMLTrace, TraceManager, trace_factory, Trace, load_log, _COMPRESSORS, _COMPRESSION_EXTENSIONS) @@ -230,7 +231,7 @@ def test_trace_factory_yaml(): def test_trace_factory_invalid(): - with pytest.raises(ValueError, match="Invalid trace type"): + with pytest.raises(ConfigurationError, match="Invalid trace type"): trace_factory("unknown") @@ -259,7 +260,7 @@ def test_trace_function_passthrough(): def test_trace_function_invalid(): - with pytest.raises(ValueError, match="Unrecognized Trace option"): + with pytest.raises(ConfigurationError, match="Unrecognized Trace option"): Trace(12345) @@ -459,7 +460,7 @@ def test_yaml_trace_load_and_reload(tmp_path): def test_yaml_trace_invalid_compression(): - with pytest.raises(ValueError, match="Unknown compression type"): + with pytest.raises(ConfigurationError, match="Unknown compression type"): YAMLTrace(base_name="test", compression="lz4") diff --git a/test/test_traj_gen.py b/test/test_traj_gen.py index be33eac..711c658 100644 --- a/test/test_traj_gen.py +++ b/test/test_traj_gen.py @@ -2,9 +2,6 @@ # -*- coding: utf-8 -*- """Unit testing for mudslide""" -import unittest -import re - import mudslide import numpy as np @@ -14,57 +11,57 @@ def get_initial_from_gen(g): return l[0] -class TestTrajGen(unittest.TestCase): - """Test Suite for trajectory initial state generators""" +def _setup_traj_gen(): + """Setup function""" + n = 4 + rng = np.random.default_rng(7) + x = np.array([1, 2, 3, 4]) + v = np.array([5, 6, 7, 8]) + i = "ground" + masses = np.abs(rng.normal(0.0, 1e4, size=n)) + print(masses) + seed = 9 + seed2 = 11 + return n, x, v, i, masses, seed, seed2 - def setUp(self): - """Setup function""" - self.n = 4 - rng = np.random.default_rng(7) - self.x = np.array([1, 2, 3, 4]) - self.v = np.array([5, 6, 7, 8]) - self.i = "ground" - self.masses = np.abs(rng.normal(0.0, 1e4, size=self.n)) - print(self.masses) - self.seed = 9 - self.seed2 = 11 - def test_const_gen(self): - """Test constant generator""" - g = mudslide.TrajGenConst(self.x, self.v, self.i, seed=self.seed) +def test_const_gen(): + """Test constant generator""" + n, x, v, i, masses, seed, seed2 = _setup_traj_gen() + g = mudslide.TrajGenConst(x, v, i, seed=seed) - x, v, i, o = get_initial_from_gen(g) + xo, vo, io, o = get_initial_from_gen(g) - self.assertTrue(np.all(np.isclose(x, self.x))) - self.assertTrue(np.all(np.isclose(v, self.v))) - self.assertEqual(i, "ground") + assert np.all(np.isclose(xo, x)) + assert np.all(np.isclose(vo, v)) + assert io == "ground" - def test_normal_gen(self): - """Test normal generator""" - refx = [1.17096384, 8.7987377, 9.12360539, 1.44846462] - refv = [4.97020305, 5.94726158, 7.05697264, 7.99439356] - g = mudslide.TrajGenNormal(self.x, self.v, self.i, 10.0, seed=self.seed, seed_traj=self.seed2) +def test_normal_gen(): + """Test normal generator""" + n, x, v, i, masses, seed, seed2 = _setup_traj_gen() + refx = [1.17096384, 8.7987377, 9.12360539, 1.44846462] + refv = [4.97020305, 5.94726158, 7.05697264, 7.99439356] - x, v, i, o = get_initial_from_gen(g) + g = mudslide.TrajGenNormal(x, v, i, 10.0, seed=seed, seed_traj=seed2) - self.assertTrue(np.all(np.isclose(x, refx))) - self.assertTrue(np.all(np.isclose(v, refv))) - self.assertEqual(i, "ground") + xo, vo, io, o = get_initial_from_gen(g) - def test_boltzmann_gen(self): - """Test Boltzmann generator""" - refv = [0.00389077, 2.41118654, 2.08038404, -1.56240191] / self.masses + assert np.all(np.isclose(xo, refx)) + assert np.all(np.isclose(vo, refv)) + assert io == "ground" - g = mudslide.TrajGenBoltzmann(self.x, self.masses, 300, self.i, seed=self.seed, velocity_seed=self.seed2) - x, v, i, o = get_initial_from_gen(g) - print(v) - print(refv) - self.assertTrue(np.all(np.isclose(x, self.x))) - self.assertTrue(np.all(np.isclose(v, refv))) - self.assertEqual(i, "ground") +def test_boltzmann_gen(): + """Test Boltzmann generator""" + n, x, v, i, masses, seed, seed2 = _setup_traj_gen() + refv = [0.00389077, 2.41118654, 2.08038404, -1.56240191] / masses + g = mudslide.TrajGenBoltzmann(x, masses, 300, i, seed=seed, velocity_seed=seed2) -if __name__ == '__main__': - unittest.main() + xo, vo, io, o = get_initial_from_gen(g) + print(vo) + print(refv) + assert np.all(np.isclose(xo, x)) + assert np.all(np.isclose(vo, refv)) + assert io == "ground" diff --git a/test/test_turbomole-es.py b/test/test_turbomole-es.py index e504ff7..51650de 100644 --- a/test/test_turbomole-es.py +++ b/test/test_turbomole-es.py @@ -1,11 +1,11 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -"""Unit testing for Turbomole_class""" +"""Unit testing for Turbomole_class with Even Sampling""" import numpy as np import os import shutil -import unittest +import pytest from pathlib import Path import mudslide import yaml @@ -15,6 +15,7 @@ from mudslide.tracer import TraceManager from mudslide.batch import TrajGenConst, TrajGenNormal, BatchedTraj +from mudslide.exceptions import ConvergenceError, ExternalCodeError from mudslide.models import TMModel, turbomole_is_installed from mudslide.config import get_config @@ -31,124 +32,118 @@ def _turbomole_available(): or "MUDSLIDE_TURBOMOLE_PREFIX" in os.environ or get_config("turbomole.command_prefix") is not None) -@unittest.skipUnless(_turbomole_available(), "Turbomole must be installed") -class TestTMModel(unittest.TestCase): - """Test Suite for TMModel class""" - - def setUp(self): - self.refdir = os.path.join(_refdir, "tm-es-c2h4") - - self.rundir = os.path.join(_checkdir, "tm-es-c2h4") - - clean_directory(self.rundir) - os.makedirs(self.rundir, exist_ok=True) - - self.origin = os.getcwd() - - os.chdir(self.rundir) - with os.scandir(self.refdir) as it: - for fil in it: - if fil.name.endswith(".input") and fil.is_file(): - filename = fil.name[:-6] - shutil.copy(os.path.join(self.refdir, fil.name), filename) - - def test_get_gs_ex_properties(self): - """test for gs_ex_properties function""" - tm_model = TMModel(states=[0, 1, 2, 3], expert=True) - - # yapf: disable - mom = [ 5.583286976987380000, -2.713959745507320000, 0.392059702162967000, - -0.832994241764031000, -0.600752326053757000, -0.384006560250834000, - -1.656414687719690000, 1.062437820195600000, -1.786171104341720000, - -2.969087779972610000, 1.161804203506510000, -0.785009852486148000, - 2.145175145340160000, 0.594918215579156000, 1.075977514428970000, - -2.269965412856570000, 0.495551832268249000, 1.487150300486560000] - # yapf: enable - velocities = mom / tm_model.mass - positions = tm_model._position - mass = tm_model.mass - q = queue.Queue() - dt = 20 - max_time=41 - t0 = 1 - sample_stack = SpawnStack.from_quadrature(nsamples=[2, 2, 2]) - sample_stack.sample_stack[0]["zeta"]=0.003 - samples = 1 - nprocs = 1 - trace_type = "yaml" - electronic_integration = 'exp' - trace_options = {"location": "", "base_name": "TMtrace"} - every = 1 - model=tm_model - traj_gen = TrajGenConst(positions, velocities, 3, dt) - - fssh = mudslide.BatchedTraj(model, - traj_gen, - trajectory_type=EvenSamplingTrajectory, - samples=samples, - max_time=max_time, - nprocs=nprocs, - dt=dt, - t0=t0, - tracemanager=TraceManager(trace_type=trace_type), - trace_every=every, - spawn_stack=sample_stack, - electronic_integration=electronic_integration) - results = fssh.compute() - outcomes = results.outcomes - - refs = mudslide.load_log(os.path.join(self.refdir, "traj-1.yaml")) - - ref_times = [0, 1, 2] - states = [0, 1, 2, 3] - - for t in ref_times: - available = results[1][t]["electronics"]["forces_available"] - for s in states: - self.assertAlmostEqual(results[1][t]["electronics"]["hamiltonian"][s][s], - refs[t]["electronics"]["hamiltonian"][s][s], - places=8) - - if available[s]: - np.testing.assert_almost_equal(results[1][t]["electronics"]["force"][s], - refs[t]["electronics"]["force"][s], - decimal=8) - - for t in ref_times: - np.testing.assert_almost_equal(results[1][t]["density_matrix"], refs[t]["density_matrix"], decimal=8) - np.testing.assert_almost_equal(results[1][t]["position"], refs[t]["position"], decimal=8) - np.testing.assert_almost_equal(results[1][t]["velocity"], refs[t]["velocity"], decimal=8) - - for t in ref_times: - for s1 in states: - for s2 in range(s1, 3): - np.testing.assert_almost_equal(results[1][t]["electronics"]["derivative_coupling"][s1][s2], - refs[t]["electronics"]["derivative_coupling"][s1][s2], - decimal=6) - - def test_ridft_convergence_failure(self): - """Test that non-convergence of ridft raises RuntimeError""" - tm_model = TMModel(states=[0, 1, 2, 3], expert=True) - tm_model.control.adg("scfiterlimit", 2) - # perturb geometry so the converged MOs are no longer a good guess - positions = tm_model._position.copy() - positions += 0.5 * np.random.default_rng(42).standard_normal(positions.shape) - with self.assertRaises(RuntimeError): - tm_model.compute(positions) - - def test_egrad_convergence_failure(self): - """Test that non-convergence of egrad raises RuntimeError""" - tm_model = TMModel(states=[0, 1, 2, 3], expert=True) - tm_model.control.adg("escfiterlimit", 2) - # perturb geometry so the excited state solver needs more iterations - positions = tm_model._position.copy() - #positions += 0.5 * np.random.default_rng(42).standard_normal(positions.shape) - with self.assertRaises(RuntimeError): - tm_model.compute(positions) - - def tearDown(self): - os.chdir(self.origin) - - -if __name__ == '__main__': - unittest.main() +pytestmark = pytest.mark.skipif(not _turbomole_available(), + reason="Turbomole must be installed") + + +@pytest.fixture +def tm_es_setup(): + """Set up turbomole even-sampling test directory.""" + refdir = os.path.join(_refdir, "tm-es-c2h4") + rundir = os.path.join(_checkdir, "tm-es-c2h4") + + clean_directory(rundir) + os.makedirs(rundir, exist_ok=True) + + origin = os.getcwd() + os.chdir(rundir) + with os.scandir(refdir) as it: + for fil in it: + if fil.name.endswith(".input") and fil.is_file(): + filename = fil.name[:-6] + shutil.copy(os.path.join(refdir, fil.name), filename) + + yield {"refdir": refdir, "rundir": rundir} + + os.chdir(origin) + + +def test_get_gs_ex_properties(tm_es_setup): + """test for gs_ex_properties function""" + refdir = tm_es_setup["refdir"] + + tm_model = TMModel(states=[0, 1, 2, 3], expert=True) + + # yapf: disable + mom = [ 5.583286976987380000, -2.713959745507320000, 0.392059702162967000, + -0.832994241764031000, -0.600752326053757000, -0.384006560250834000, + -1.656414687719690000, 1.062437820195600000, -1.786171104341720000, + -2.969087779972610000, 1.161804203506510000, -0.785009852486148000, + 2.145175145340160000, 0.594918215579156000, 1.075977514428970000, + -2.269965412856570000, 0.495551832268249000, 1.487150300486560000] + # yapf: enable + velocities = mom / tm_model.mass + positions = tm_model._position + dt = 20 + max_time = 41 + t0 = 1 + sample_stack = SpawnStack.from_quadrature(nsamples=[2, 2, 2]) + sample_stack.sample_stack[0]["zeta"] = 0.003 + samples = 1 + trace_type = "yaml" + electronic_integration = 'exp' + every = 1 + model = tm_model + traj_gen = TrajGenConst(positions, velocities, 3, dt) + + fssh = mudslide.BatchedTraj(model, + traj_gen, + trajectory_type=EvenSamplingTrajectory, + samples=samples, + max_time=max_time, + dt=dt, + t0=t0, + tracemanager=TraceManager(trace_type=trace_type), + trace_every=every, + spawn_stack=sample_stack, + electronic_integration=electronic_integration) + results = fssh.compute() + outcomes = results.outcomes + + refs = mudslide.load_log(os.path.join(refdir, "traj-1.yaml")) + + ref_times = [0, 1, 2] + states = [0, 1, 2, 3] + + for t in ref_times: + available = results[1][t]["electronics"]["forces_available"] + for s in states: + assert results[1][t]["electronics"]["hamiltonian"][s][s] == pytest.approx( + refs[t]["electronics"]["hamiltonian"][s][s], abs=1e-8) + + if available[s]: + np.testing.assert_almost_equal(results[1][t]["electronics"]["force"][s], + refs[t]["electronics"]["force"][s], + decimal=8) + + for t in ref_times: + np.testing.assert_almost_equal(results[1][t]["density_matrix"], refs[t]["density_matrix"], decimal=8) + np.testing.assert_almost_equal(results[1][t]["position"], refs[t]["position"], decimal=8) + np.testing.assert_almost_equal(results[1][t]["velocity"], refs[t]["velocity"], decimal=8) + + for t in ref_times: + for s1 in states: + for s2 in range(s1, 3): + np.testing.assert_almost_equal(results[1][t]["electronics"]["derivative_coupling"][s1][s2], + refs[t]["electronics"]["derivative_coupling"][s1][s2], + decimal=6) + + +def test_ridft_convergence_failure(tm_es_setup): + """Test that non-convergence of ridft raises RuntimeError""" + tm_model = TMModel(states=[0, 1, 2, 3], expert=True) + tm_model.control.adg("scfiterlimit", 2) + # perturb geometry so the converged MOs are no longer a good guess + positions = tm_model._position.copy() + positions += 0.5 * np.random.default_rng(42).standard_normal(positions.shape) + with pytest.raises(ConvergenceError): + tm_model.compute(positions) + + +def test_egrad_convergence_failure(tm_es_setup): + """Test that non-convergence of egrad raises ExternalCodeError""" + tm_model = TMModel(states=[0, 1, 2, 3], expert=True) + tm_model.control.adg("escfiterlimit", 2) + positions = tm_model._position.copy() + with pytest.raises(ExternalCodeError): + tm_model.compute(positions) diff --git a/test/test_turbomole.py b/test/test_turbomole.py index d7e67c8..2250df0 100644 --- a/test/test_turbomole.py +++ b/test/test_turbomole.py @@ -5,12 +5,12 @@ import numpy as np import os import shutil -import unittest import pytest from pathlib import Path import mudslide import yaml +from mudslide.exceptions import ConfigurationError from mudslide.models import TMModel, turbomole_is_installed from mudslide.config import get_config from mudslide.tracer import YAMLTrace @@ -32,157 +32,152 @@ def _turbomole_available(): def test_raise_on_missing_control(): """Test if an exception is raised if no control file is found""" - with pytest.raises(RuntimeError): + with pytest.raises(ConfigurationError): model = TMModel(states=[0]) -class _TestTM(unittest.TestCase): - """Base class for TMModel class""" - testname = None +@pytest.fixture +def tm_setup(request, tmp_path): + """Set up a Turbomole test directory from reference inputs.""" + testname = request.param + refdir = os.path.join(_refdir, testname) + rundir = os.path.join(_checkdir, testname) - def setUp(self): - self.refdir = os.path.join(_refdir, self.testname) - self.rundir = os.path.join(_checkdir, self.testname) + if os.path.isdir(rundir): + shutil.rmtree(rundir) + os.makedirs(rundir, exist_ok=True) - if os.path.isdir(self.rundir): - shutil.rmtree(self.rundir) - os.makedirs(self.rundir, exist_ok=True) + origin = os.getcwd() + os.chdir(rundir) + with os.scandir(refdir) as it: + for fil in it: + if fil.name.endswith(".input") and fil.is_file(): + filename = fil.name[:-6] + shutil.copy(os.path.join(refdir, fil.name), filename) - self.origin = os.getcwd() + yield {"refdir": refdir, "rundir": rundir} - os.chdir(self.rundir) - with os.scandir(self.refdir) as it: - for fil in it: - if fil.name.endswith(".input") and fil.is_file(): - filename = fil.name[:-6] - shutil.copy(os.path.join(self.refdir, fil.name), filename) + os.chdir(origin) - def tearDown(self): - os.chdir(self.origin) - -class TestTMGround(_TestTM): +@pytest.mark.parametrize("tm_setup", ["tm-c2h4-ground"], indirect=True) +def test_ridft_rdgrad(tm_setup): """Test ground state calculation""" - testname = "tm-c2h4-ground" - - def test_ridft_rdgrad(self): - model = TMModel(states=[0]) - xyz = model._position + model = TMModel(states=[0]) + xyz = model._position - model.compute(xyz) + model.compute(xyz) - Eref = -78.40037210973 - Fref = np.loadtxt("force.ref.txt") + Eref = -78.40037210973 + Fref = np.loadtxt("force.ref.txt") - assert np.isclose(model.hamiltonian[0, 0], Eref) - assert np.allclose(model.force(0), Fref) + assert np.isclose(model.hamiltonian[0, 0], Eref) + assert np.allclose(model.force(0), Fref) -class TestTMGroundPC(_TestTM): +@pytest.mark.parametrize("tm_setup", ["tm-c2h4-ground-pc"], indirect=True) +def test_ridft_rdgrad_w_pc(tm_setup): """Test ground state calculation with point charges""" - testname = "tm-c2h4-ground-pc" - - def test_ridft_rdgrad_w_pc(self): - model = TMModel(states=[0]) - xyzpc = np.array([[3.0, 3.0, 3.0], [-3.0, -3.0, -3.0]]) - pcharges = np.array([2, -2]) - model.control.add_point_charges(xyzpc, pcharges) - - xyz = model._position - model.compute(xyz) - - Eref = -78.63405047062 - Fref = np.loadtxt("force.ref.txt") - - assert np.isclose(model.hamiltonian[0, 0], Eref) - assert np.allclose(model.force(0), Fref) - - xyzpc1, q1, dpc = model.control.read_point_charge_gradients() - - forcepcref = np.loadtxt("forcepc.ref.txt") - - assert np.allclose(xyzpc, xyzpc1) - assert np.allclose(pcharges, q1) - assert np.allclose(dpc, forcepcref) - - -class TestTMExDynamics(_TestTM): - """Test short excited state dynamics""" - testname = "tm-c2h4" - - def test_get_gs_ex_properties(self): - """test for gs_ex_properties function""" - model = TMModel(states=[0, 1, 2, 3], expert=True) - positions = model._position - - # yapf: disable - mom = [ 5.583286976987380000, -2.713959745507320000, 0.392059702162967000, - -0.832994241764031000, -0.600752326053757000, -0.384006560250834000, - -1.656414687719690000, 1.062437820195600000, -1.786171104341720000, - -2.969087779972610000, 1.161804203506510000, -0.785009852486148000, - 2.145175145340160000, 0.594918215579156000, 1.075977514428970000, - -2.269965412856570000, 0.495551832268249000, 1.487150300486560000] - # yapf: enable - - velocities = np.array(mom) / model.mass - - log = mudslide.YAMLTrace(base_name="TMtrace", - location=self.rundir, - log_pitch=8) - traj = mudslide.SurfaceHoppingMD(model, - positions, - velocities, - 3, - tracer=log, - dt=20, - max_time=81, - t0=1, - seed_sequence=57892, - hopping_method="instantaneous") - results = traj.simulate() - - main_log = results.main_log - - assert main_log == "TMtrace-0.yaml" - - refs = mudslide.load_log(os.path.join(self.refdir, "traj-0.yaml")) - - ref_times = [0, 1, 2, 3] - states = [0, 1, 2, 3] - - for t in ref_times: - available = results[t]["electronics"]["forces_available"] - act_ham = results[t]["electronics"]["hamiltonian"] - ref_ham = refs[t]["electronics"]["hamiltonian"] - act_force = results[t]["electronics"]["force"] - ref_force = refs[t]["electronics"]["force"] - for s in states: - np.testing.assert_almost_equal(act_ham[s][s], - ref_ham[s][s], - decimal=8) - if available[s]: - np.testing.assert_almost_equal(act_force[s], - ref_force[s], - decimal=8) - - for t in ref_times: - act = results[t] - ref = refs[t] - np.testing.assert_almost_equal(act["density_matrix"], - ref["density_matrix"], - decimal=8) - np.testing.assert_almost_equal(act["position"], - ref["position"], - decimal=8) - np.testing.assert_almost_equal(act["velocity"], - ref["velocity"], + model = TMModel(states=[0]) + xyzpc = np.array([[3.0, 3.0, 3.0], [-3.0, -3.0, -3.0]]) + pcharges = np.array([2, -2]) + model.control.add_point_charges(xyzpc, pcharges) + + xyz = model._position + model.compute(xyz) + + Eref = -78.63405047062 + Fref = np.loadtxt("force.ref.txt") + + assert np.isclose(model.hamiltonian[0, 0], Eref) + assert np.allclose(model.force(0), Fref) + + xyzpc1, q1, dpc = model.control.read_point_charge_gradients() + + forcepcref = np.loadtxt("forcepc.ref.txt") + + assert np.allclose(xyzpc, xyzpc1) + assert np.allclose(pcharges, q1) + assert np.allclose(dpc, forcepcref) + + +@pytest.mark.parametrize("tm_setup", ["tm-c2h4"], indirect=True) +def test_get_gs_ex_properties(tm_setup): + """test for gs_ex_properties function""" + refdir = tm_setup["refdir"] + rundir = tm_setup["rundir"] + + model = TMModel(states=[0, 1, 2, 3], expert=True) + positions = model._position + + # yapf: disable + mom = [ 5.583286976987380000, -2.713959745507320000, 0.392059702162967000, + -0.832994241764031000, -0.600752326053757000, -0.384006560250834000, + -1.656414687719690000, 1.062437820195600000, -1.786171104341720000, + -2.969087779972610000, 1.161804203506510000, -0.785009852486148000, + 2.145175145340160000, 0.594918215579156000, 1.075977514428970000, + -2.269965412856570000, 0.495551832268249000, 1.487150300486560000] + # yapf: enable + + velocities = np.array(mom) / model.mass + + log = mudslide.YAMLTrace(base_name="TMtrace", + location=rundir, + log_pitch=8) + traj = mudslide.SurfaceHoppingMD(model, + positions, + velocities, + 3, + tracer=log, + dt=20, + max_time=81, + t0=1, + seed_sequence=57892, + hopping_method="instantaneous") + results = traj.simulate() + + main_log = results.main_log + + assert main_log == "TMtrace-0.yaml" + + refs = mudslide.load_log(os.path.join(refdir, "traj-0.yaml")) + + ref_times = [0, 1, 2, 3] + states = [0, 1, 2, 3] + + for t in ref_times: + available = results[t]["electronics"]["forces_available"] + act_ham = results[t]["electronics"]["hamiltonian"] + ref_ham = refs[t]["electronics"]["hamiltonian"] + act_force = results[t]["electronics"]["force"] + ref_force = refs[t]["electronics"]["force"] + for s in states: + np.testing.assert_almost_equal(act_ham[s][s], + ref_ham[s][s], decimal=8) + if available[s]: + np.testing.assert_almost_equal(act_force[s], + ref_force[s], + decimal=8) - for t in ref_times: - act_tau = results[t]["electronics"]["derivative_coupling"] - ref_tau = refs[t]["electronics"]["derivative_coupling"] - for s1 in states: - for s2 in range(s1, 3): - np.testing.assert_almost_equal(act_tau[s1][s2], - ref_tau[s1][s2], - decimal=6) + for t in ref_times: + act = results[t] + ref = refs[t] + np.testing.assert_almost_equal(act["density_matrix"], + ref["density_matrix"], + decimal=8) + np.testing.assert_almost_equal(act["position"], + ref["position"], + decimal=8) + np.testing.assert_almost_equal(act["velocity"], + ref["velocity"], + decimal=8) + + for t in ref_times: + act_tau = results[t]["electronics"]["derivative_coupling"] + ref_tau = refs[t]["electronics"]["derivative_coupling"] + for s1 in states: + for s2 in range(s1, 3): + np.testing.assert_almost_equal(act_tau[s1][s2], + ref_tau[s1][s2], + decimal=6)