diff --git a/.claude/CLAUDE.md b/.claude/CLAUDE.md new file mode 100644 index 0000000..e7933bd --- /dev/null +++ b/.claude/CLAUDE.md @@ -0,0 +1,111 @@ +# CLAUDE.md + +This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository. + +## What this is + +autostorage is a SQLModel/SQLAlchemy persistence layer for computational chemistry workflow +data, built on top of `automol`. It stores molecular geometries, identities, trajectories, +stationary points, calculation results, and the calculations/steps that connect them, as a +graph of related database rows. + +## Commands + +All tasks run through Pixi (`pixi run `), defined in `pixi.toml` under `[feature.dev.tasks]`: + +- `pixi run fmt` — format with Ruff +- `pixi run lint` — lint with Ruff (`--fix`) +- `pixi run types` — static type-check with `ty` +- `pixi run imports` — enforce module layering with `lint-imports` (import-linter) +- `pixi run test` — run the full pytest suite +- `pixi run pre-commit` — run all of the above via lefthook, in order (fmt → lint → types → imports → test), then check the tree is clean +- `pixi run cov-view` — open the HTML coverage report +- `pixi run local` — toggle `pixi.toml`'s `automol` dependency between the pinned release and a + local `../automol` checkout, via the `# local:true`/`# local:false` comment markers already in + the file (see "Relationship to automol" below) +- `pixi run local-pre-commit` — a second lefthook target, distinct from `pre-commit` +- `pixi run docs-build` / `pixi run docs-view` — build/view the Sphinx docs (`feature.docs` env) + +Single test: invoke `pytest` directly inside the pixi env, e.g. +`pixi run -e dev pytest tests/test_models.py::TestGeometryRow::test_name` (tests are organized +into `Test*` classes grouped by row/function under test). + +Note: pytest is configured with `--doctest-modules`, so doctests in `src/` docstrings are +collected and run as part of the suite. Coverage must stay ≥80% (`fail_under = 80` in +`pyproject.toml`), with branch coverage enabled. + +Don't invoke bare `python`/`python3` — `automol` and other deps aren't on the system interpreter, +only inside the pixi env. Always go through `pixi run` (e.g. `pixi run -e dev pytest ...`). + +## Architecture + +### Module layering (enforced by import-linter) + +`pyproject.toml` defines a strict layer contract ("Autostorage Layering") — higher layers may +depend on lower ones, never the reverse: + +``` +autostorage.database (highest) +autostorage.events +autostorage.models +autostorage.types (lowest) +``` + +Adding an import that violates this order will fail `pixi run imports`. `autostorage` is a flat +module structure — all modules live directly in `src/autostorage/`. + +### Relationship to automol + +Row models extend automol's core data models directly rather than wrapping them: `GeometryRow` +extends `automol.Geometry`, `IdentityRow` extends `automol.Identity`. Any conversion to/from +other external formats is delegated to automol's own conversion functions rather than +reimplemented here. + +### Current module map + +- `models.py` — SQLModel row definitions, organized in sections: + - Link tables (named alphabetically by the entities they connect): `CalculationGeometryLink`, + `CalculationTrajectoryLink`, `GeometryTrajectoryLink`, `IdentityStationaryLink`, + `StageStationaryLink`, `StepValidationLink` + - Existential data rows: `GeometryRow` (extends `automol.Geometry`), `TrajectoryRow`, + `ModelRow`, `CalculationRow`, result rows (`EnergyRow`, `GradientRow`, `HessianRow`), + `ValidationRow` + - Stationary point rows: `StationaryPointRow` + - Reaction network rows: `StageRow`, `StepRow` (a step between two stages, with a barrierless + flag) + - Identity rows: `IdentityRow` (extends `automol.Identity`), `IdentityExtraRow` + +- `events.py` — SQLAlchemy ORM event listeners, by concern: + - Shape validation: `verify_gradient_shapes_before_flush`, `verify_hessian_shapes_before_flush` + - Trajectory validation: `verify_trajectory_geometry_ndim_insert` (ensures geometry index + length matches trajectory ndim) + - Auto-managed identities: `add_inchi_identities_before_flush` (attaches an InChI `IdentityRow` + to newly inserted stationary points), `add_smiles_extras_before_flush` / + `add_hill_extras_before_flush` (attach SMILES / Hill formula as `IdentityExtraRow`s once an + InChI identity is present). Private `_find_or_create_identity`/`_find_or_create_identity_extra` + helpers dedup these against existing rows and pending session inserts. + - Step validation: `sort_step_stage_ids` (auto-sorts stage_id1 < stage_id2), + `verify_step_barrierless_consistency` (verifies is_barrierless matches stage_id_ts state) + +- `database.py` — `Database`: SQLite engine/session manager. `__init__` creates the engine + (with `PRAGMA foreign_keys=ON` and a sort-keys JSON serializer) and the schema via + `SQLModel.metadata.create_all`; `session()` returns a fresh `Session` bound to that engine + (use as a context manager; nothing auto-commits); `close()` disposes the engine. + +- `types.py` — Type definitions and utilities: + - `Role` (StrEnum: INPUT/OUTPUT) — relationship between calculations and geometries/trajectories + - `CompressedArrayTypeDecorator` — SQLAlchemy `TypeDecorator` storing NumPy arrays as + zlib-compressed binary data in SQLite + - `_fk_field()` — helper for building foreign-key fields with ON DELETE CASCADE + +### Docstrings + +NumPy docstring convention (`tool.ruff.lint.pydocstyle` = `"numpy"`), and doctest examples in +docstrings are executed as tests — keep them runnable and accurate. + +### Notes + +- Minimize chat/response verbosity when performing work to reduce unnecessary token costs. +- Keep docstrings and comments minimal: one-line NumPy-style summaries where the convention + allows, no restating what a name/type hint already conveys. Reserve comments for genuinely + non-obvious invariants — most docstrings in this repo don't need that much. \ No newline at end of file diff --git a/.claude/agents/autostorage-explorer.md b/.claude/agents/autostorage-explorer.md new file mode 100644 index 0000000..0e61f9a --- /dev/null +++ b/.claude/agents/autostorage-explorer.md @@ -0,0 +1,57 @@ +--- +name: autostorage-explorer +description: Use to explore/investigate the autostorage codebase before planning a feature or bugfix — pre-loaded with the module map, layering rules, and known subtleties so it doesn't need to rediscover them from scratch. Read-only; reports file:line references, does not propose implementations. +tools: Read, Grep, Glob, Bash +model: haiku +--- + +You are a read-only research agent for the `autostorage` repo (a SQLModel/SQLAlchemy persistence +layer for computational chemistry workflow data, built on `automol`). Your job is to locate the +exact rows, functions, event listeners, and tests relevant to a given feature/bug description, and +report `file:line` references — not to design or write the implementation. + +## Layout + +Flat module structure under `src/autostorage/`, layered (higher depends on lower, never reverse, +enforced by import-linter): `utils` > `database` > `merge` > `events` > `models` > `types`/`exc`. + +- `models.py` — SQLModel row definitions (`GeometryRow`, `EnergyRow`/`GradientRow`/`HessianRow`, + `TrajectoryRow`, `StationaryPointRow`, `IdentityRow`/`IdentityExtraRow`, `StageRow`, `StepRow`, + `ModelRow`, `CalculationRow`, `ValidationRow`, plus link tables). Base classes: `TimestampMixin`, + `BaseRow`, `BaseResultRow`, `BaseLink`. Several rows expose a shared `find_or_create` classmethod + (get-or-insert pattern) — check there first for any "does X already exist" question. +- `events.py` — SQLAlchemy ORM event listeners: shape validation for Gradient/Hessian; geometry + order-consensus recompute (`revalidate_geometry_orders_on_insert_update`/`_on_hessian_delete` — + session-level `before_flush` listeners, not mapper events, because they mutate sibling rows that + may already be clean going into the flush); `verify_geometry_immutable_fields`; + `compute_geometry_hash`; auto-managed identity attachment (`add_inchi_identities`, + `assign_conformer_ids`); `StepRow` stage-order/TS-consistency checks. +- `database.py` — `Database`: SQLite engine/session manager. +- `merge.py` — `merge_databases`: copies one database's rows into another, deduplicating + `ModelRow`, `GeometryRow`, non-auto-managed `IdentityRow`s, `CalculationRow`, and + `StationaryPointRow` via their `find_or_create` methods. +- `types.py` — `CalcType`, `CalcStatus`, `Role`, `IndexType`, `CompressedArrayTypeDecorator`. +- `exc.py` — `ResultShapeError`, `MissingPrimaryKeyError`. +- `utils.py` — MESS input export and PES plotting. + +## Known gotchas (check these before assuming a bug is novel) + +1. **`compute_geometry_hash`** writes `target.__dict__["geometry_hash"] = ...` + + `flag_modified(...)` instead of `target.geometry_hash = ...`. Plain attribute assignment inside + a mapper event breaks under `Geometry`'s `validate_assignment=True` pydantic config — it + corrupts SQLAlchemy's flush-time identity-key bookkeeping. +2. **`before_flush` vs mapper events**: anything that needs to mutate a *different* row than the + one that triggered the change (e.g. recomputing `StationaryPointRow.is_valid` when a sibling + `HessianRow` changes) must be a session-level `before_flush` listener. A per-instance + `before_insert`/`before_update` mapper event fires too late for such a mutation to be included + in the same flush — SQLAlchemy silently drops it instead of writing it. +3. **No migrations currently**: `migrations/` and `alembic.ini` were removed; `alembic` remains + a dev dependency for when migrations are reintroduced, but there is no active migration path — + schema changes only need to work with `SQLModel.metadata.create_all`. + +## What to report + +For a given feature/bug description: the specific row classes, event listeners, and existing +tests involved, with `file:line` references, plus which layering tier(s) a change would touch (to +flag likely `pixi run imports` fallout early). Do not propose an implementation — that's a +separate planning step. diff --git a/.claude/agents/autostorage-release.md b/.claude/agents/autostorage-release.md new file mode 100644 index 0000000..9afb0f4 --- /dev/null +++ b/.claude/agents/autostorage-release.md @@ -0,0 +1,31 @@ +--- +name: autostorage-release +description: Use when preparing an autostorage version release — walks the CHANGELOG/README/docstring/pre-commit/version-bump checklist. Can edit files and run pixi commands, but must stop and confirm before any publish/push/tag step. +tools: Read, Edit, Bash, Grep, Glob +model: haiku +--- + +You are the release-prep checklist runner for the `autostorage` repo. Work through these steps in +order, reporting status after each one. Stop and ask before any step marked irreversible. + +1. Run `pixi run pre-commit` first. Fix any failures (fmt/lint/types/imports/test) before + proceeding — don't paper over a failing step. +2. Review docstrings on code touched since the last tag (`git log ..HEAD --stat` or + similar) for NumPy-convention compliance and terseness: one-line summaries where possible, no + restating what a name/type hint already conveys. +3. Update `CHANGELOG.md` via `pixi run changelog ` (wraps `keepachangelog`) so entries + match the commits since the last release. Check `git log ..HEAD --oneline` against + what's already there. +4. Update `README.md` if the public API surface changed (new/removed exports in + `src/autostorage/__init__.py`, new Pixi tasks relevant to users, etc.). +5. Bump the version via `pixi run version` (check current) and `pixi run release` (backed by + `tbump`) — confirm `pyproject.toml`, `pixi.toml`, and `src/autostorage/__init__.py`'s + `__version__` all move together after the bump. +6. Re-run `pixi run pre-commit` to confirm a clean tree. + +## Do not run without explicit user confirmation first + +`pixi run build-conda`, `pixi run publish-conda`, `pixi run publish-pypi`, +`pixi run publish-test-pypi`, or any `git push`/tag push. These are external, hard-to-reverse +actions (publishing a package, pushing to a shared branch) — surface that the checklist is ready +for them and wait for the user to say go. diff --git a/.gitignore b/.gitignore index e196d81..9672a62 100644 --- a/.gitignore +++ b/.gitignore @@ -200,7 +200,7 @@ cython_debug/ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore # and can be added to the global gitignore or merged into this file. However, if you prefer, # you could uncomment the following to ignore the entire vscode folder -# .vscode/ +.vscode/ # Ruff stuff: .ruff_cache/ @@ -229,6 +229,6 @@ __marimo__/ # pixi build *.conda -# database schema for publication -schema/ -CLAUDE.md \ No newline at end of file +# Scratch files +.scratch/ +example*.db \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 143bd51..f8307b7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,31 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/). ## [Unreleased] +### Added + +- **`.claude/CLAUDE.md`**: Project-level guidance for Claude Code, documenting module architecture, layering rules, Pixi task commands, and autostorage-specific conventions. +- **`docs/source/models.md`**: Dedicated documentation page for the row models and their relationships. +- **Schema diagrams** (`schema/`): Added Pintora source files and rendered SVG/PNG diagrams for both simplified and full database schemas. +- **`examples/stationary.py`**: Example demonstrating stationary points, stages, and steps with H2O geometries. +- **`examples/scan.py`**: Example showing PES scan workflow with trajectory data. +- **`examples/transition.py`**: Example illustrating transition state and reaction pathway modeling. +- **`tests/test_events.py`**: Comprehensive test suite for SQLAlchemy ORM event listeners. +- **Top-level exports**: Added `CalculationTrajectoryLink`, `GeometryTrajectoryLink`, and `events` module to `autostorage.__all__`. + +### Changed + +- **Module consolidation**: Collapsed `src/autostorage/models/` directory (with separate `calc.py`, `data.py`, `geom.py`, `link.py`, `rxn.py`, `traj.py` files) into a single `src/autostorage/models.py` file — all row models, link tables, and helpers now live in one module, eliminating circular import complexity while maintaining logical organization via internal sections (link tables, existential data, stationary points, reaction network, identities). +- **Documentation reorganization**: Streamlined Sphinx docs — consolidated data model, development, and events pages into focused `models.md` and updated `database.md`; refreshed `index.md` and `quickstart.md` to reflect current architecture. +- **Examples refactor**: Replaced single `stationary_point.py` example with three focused examples (`stationary.py`, `scan.py`, `transition.py`), each demonstrating a distinct workflow pattern. +- **Test simplification**: Removed `tests/conftest.py` and test data files (`tests/data/*.xyz`, `*.gz`) — tests now use inline synthetic data instead of external fixtures. + +### Removed + +- **Alembic migrations** (`migrations/`, `alembic.ini`, `tests/test_migrations.py`, `pixi run migrate`): Removed for now; `alembic` remains a dev dependency for when migrations are reintroduced. Schema for fresh/in-memory `Database` instances is unaffected, still built via `create_all()`. +- **`src/autostorage/exc.py`**: Custom exception classes (`DataIntegrityError`, `ResultShapeError`) removed — validation errors now use standard Python exceptions or SQLAlchemy's built-in constraint violations. +- **`CalcType` and `CalcStatus` enums**: Removed from `types.py` and top-level exports — calculation type/status classification deferred to workflow layer or removed entirely. +- **Documentation files**: Deleted `docs/source/data-model.md`, `docs/source/development.md`, and `docs/source/events.md` (content consolidated into `models.md` and other updated docs). +- **`examples/stationary_point.py`**: Superseded by the three new focused examples. ## [0.0.12] - 2026-07-23 ### Added diff --git a/README.md b/README.md index 4115fbb..d3ad696 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ Install as a [Pixi](https://pixi.sh) dependency: ```toml [dependencies] -autostorage = ">=0.0.10" +autostorage = ">=0.0.12" ``` Or with `uv`/`pip` from PyPI: @@ -31,7 +31,6 @@ Requires Python ≥3.12. ```python import numpy as np from autostorage import ( - CalcType, CalculationGeometryLink, CalculationRow, Database, @@ -44,51 +43,59 @@ from autostorage import ( # Open (or create) a SQLite database; ":memory:" also works for scratch use. db = Database("workflow.db") -# `find_or_create` dedups on (program, program_version, method, basis). -model = ModelRow.find_or_create(db, program="orca", method="b3lyp", basis="def2-svp") - -calc = CalculationRow(model=model, calc_type=CalcType.ENERGY) -geo = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.8], [0, 0, 0], [0.8, 0, 0]]), - charge=0, - spin=0, -) -link = CalculationGeometryLink.create(calc, geo, role=Role.INPUT) -db.add_all([model, calc, geo, link]) -db.commit() - -# Attach a result to the geometry/calculation pair. -energy = EnergyRow(geometry=geo, calculation=calc, value=-76.02) -db.add(energy) -db.commit() - -# Look the result back up by geometry, model, and input provenance. -found = EnergyRow.query(db, geo=geo, model=model) -assert found is not None -print(found.value) +# Work within a session context. +with db.session() as session: + # Create a model specifying the calculation type, program, method, and basis. + model = ModelRow( + calc_type="energy", + program="orca", + method="b3lyp", + basis="def2-svp", + ) + + # Create a calculation using this model. + calc = CalculationRow(model=model) + + # Create a geometry. + geo = GeometryRow( + symbols=["H", "O", "H"], + coordinates=np.array([[0, 0, 0.8], [0, 0, 0], [0.8, 0, 0]]), + charge=0, + spin=0, + ) + + # Link the geometry to the calculation as an input. + link = CalculationGeometryLink( + calculation=calc, + geometry=geo, + role=Role.INPUT, + ) + + # Add all objects to the session and commit. + session.add_all([model, calc, geo, link]) + session.commit() + + # Attach an energy result to the geometry/calculation pair. + energy = EnergyRow(geometry=geo, calculation=calc, value=-76.02) + session.add(energy) + session.commit() + + # Query the result back by filtering on geometry and calculation. + found = session.query(EnergyRow).filter_by( + geometry_id=geo.id, + calculation_id=calc.id, + ).first() + assert found is not None + print(found.value) db.close() ``` -`Database` also supports the `with` statement, which rolls back on an unhandled exception and closes the connection on exit: - -```python -with Database("workflow.db") as db: - ... -``` - -Two databases can be combined with `Database.merge_from()`, which copies every row from one into the other, remapping ids/foreign keys and deduplicating content-unique rows (models, non-auto-managed identities) against the target's existing data: - -```python -with Database("combined.db") as target, Database("other.db") as other: - report = target.merge_from(other) - print(report.copied, report.reused) # per-table row counts -``` +`Database.session()` returns a standard SQLAlchemy `Session` that supports the context manager protocol. Sessions don't commit automatically — call `session.commit()` explicitly to persist changes. -For a full worked example covering geometries, trajectories, results, and identities, see [`examples/stationary_point.py`](examples/stationary_point.py). Reaction networks (`StageRow`/`StepRow`) can be exported as [MESS](https://tcg.cse.anl.gov/papr/codes/mess.html) input via `autostorage.utils.export_mess_input()`, or rendered as a potential energy surface diagram via `autostorage.utils.plot_pes()`. +For a full worked example covering geometries, trajectories, results, stationary points, stages, and steps, see [`examples/stationary.py`](examples/stationary.py). -See `CLAUDE.md` for the full module map and architecture notes. +See [CLAUDE.md](.claude/CLAUDE.md) for the full module map and architecture notes, or the [Sphinx docs](docs/source) for a rendered quickstart and API reference. ## Contributing diff --git a/alembic.ini b/alembic.ini deleted file mode 100644 index f72d5af..0000000 --- a/alembic.ini +++ /dev/null @@ -1,154 +0,0 @@ -# A generic, single database configuration. - -[alembic] -# path to migration scripts. -# this is typically a path given in POSIX (e.g. forward slashes) -# format, relative to the token %(here)s which refers to the location of this -# ini file -script_location = %(here)s/migrations - -# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s -# Uncomment the line below if you want the files to be prepended with date and time -# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file -# for all available tokens -# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s -# Or organize into date-based subdirectories (requires recursive_version_locations = true) -# file_template = %%(year)d/%%(month).2d/%%(day).2d_%%(hour).2d%%(minute).2d_%%(second).2d_%%(rev)s_%%(slug)s - -# sys.path path, will be prepended to sys.path if present. -# defaults to the current working directory. for multiple paths, the path separator -# is defined by "path_separator" below. -prepend_sys_path = . - - -# timezone to use when rendering the date within the migration file -# as well as the filename. -# If specified, requires the tzdata library which can be installed by adding -# `alembic[tz]` to the pip requirements. -# string value is passed to ZoneInfo() -# leave blank for localtime -# timezone = - -# max length of characters to apply to the "slug" field -# truncate_slug_length = 40 - -# set to 'true' to run the environment during -# the 'revision' command, regardless of autogenerate -# revision_environment = false - -# set to 'true' to allow .pyc and .pyo files without -# a source .py file to be detected as revisions in the -# versions/ directory -# sourceless = false - -# version location specification; This defaults -# to /versions. When using multiple version -# directories, initial revisions must be specified with --version-path. -# The path separator used here should be the separator specified by "path_separator" -# below. -# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions - -# path_separator; This indicates what character is used to split lists of file -# paths, including version_locations and prepend_sys_path within configparser -# files such as alembic.ini. -# The default rendered in new alembic.ini files is "os", which uses os.pathsep -# to provide os-dependent path splitting. -# -# Note that in order to support legacy alembic.ini files, this default does NOT -# take place if path_separator is not present in alembic.ini. If this -# option is omitted entirely, fallback logic is as follows: -# -# 1. Parsing of the version_locations option falls back to using the legacy -# "version_path_separator" key, which if absent then falls back to the legacy -# behavior of splitting on spaces and/or commas. -# 2. Parsing of the prepend_sys_path option falls back to the legacy -# behavior of splitting on spaces, commas, or colons. -# -# Valid values for path_separator are: -# -# path_separator = : -# path_separator = ; -# path_separator = space -# path_separator = newline -# -# Use os.pathsep. Default configuration used for new projects. -path_separator = os - -# set to 'true' to search source files recursively -# in each "version_locations" directory -# new in Alembic version 1.10 -# recursive_version_locations = false - -# the output encoding used when revision files -# are written from script.py.mako -# output_encoding = utf-8 - -# database URL. This is consumed by the user-maintained env.py script only. -# other means of configuring database URLs may be customized within the env.py -# file. -# -# This is a placeholder; the real target is normally supplied at runtime via the -# AUTOSTORAGE_DATABASE_URL environment variable (see migrations/env.py), since -# autostorage.Database accepts an arbitrary SQLite file path rather than having -# one fixed database. -sqlalchemy.url = sqlite:///autostorage.db - - -[post_write_hooks] -# post_write_hooks defines scripts or Python functions that are run -# on newly generated revision scripts. See the documentation for further -# detail and examples - -# format using "black" - use the console_scripts runner, against the "black" entrypoint -# hooks = black -# black.type = console_scripts -# black.entrypoint = black -# black.options = -l 79 REVISION_SCRIPT_FILENAME - -# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module -# hooks = ruff -# ruff.type = module -# ruff.module = ruff -# ruff.options = check --fix REVISION_SCRIPT_FILENAME - -# Alternatively, use the exec runner to execute a binary found on your PATH -# hooks = ruff -# ruff.type = exec -# ruff.executable = ruff -# ruff.options = check --fix REVISION_SCRIPT_FILENAME - -# Logging configuration. This is also consumed by the user-maintained -# env.py script only. -[loggers] -keys = root,sqlalchemy,alembic - -[handlers] -keys = console - -[formatters] -keys = generic - -[logger_root] -level = WARNING -handlers = console -qualname = - -[logger_sqlalchemy] -level = WARNING -handlers = -qualname = sqlalchemy.engine - -[logger_alembic] -level = INFO -handlers = -qualname = alembic - -[handler_console] -class = StreamHandler -args = (sys.stderr,) -level = NOTSET -formatter = generic - -[formatter_generic] -format = %(levelname)-5.5s [%(name)s] %(message)s -datefmt = %H:%M:%S diff --git a/docs/source/data-model.md b/docs/source/data-model.md deleted file mode 100644 index a9ccf86..0000000 --- a/docs/source/data-model.md +++ /dev/null @@ -1,114 +0,0 @@ -# Data model - -autostorage stores a workflow as a graph of related rows rather than a single flat table. -This page describes the module layout, the row/link types that make up the schema, and the -automatic behaviors that keep the graph consistent. - -## Module layering - -`autostorage` is a flat module structure (no sub-packages), with a strict dependency order -enforced by import-linter — higher layers may depend on lower ones, never the reverse: - -``` -autostorage.utils (highest) -autostorage.database -autostorage.events -autostorage.models -autostorage.types | autostorage.exc (lowest) -``` - -- {py:mod}`autostorage.types` — enums (`CalcType`, `CalcStatus`, `Role`) and - `CompressedArrayTypeDecorator`, which stores NumPy arrays as zlib-compressed `.npy` bytes - so gradients, Hessians, and coordinates round-trip through SQLite without precision loss. -- {py:mod}`autostorage.exc` — `ResultShapeError`, `MissingPrimaryKeyError`. -- {py:mod}`autostorage.models` — the SQLModel row/link definitions described below. -- {py:mod}`autostorage.events` — SQLAlchemy ORM event listeners that validate and enrich rows - as they flow through a session. -- {py:mod}`autostorage.database` — `Database`, the SQLite engine/session manager applications - interact with directly. - -## Relationship to automol - -Row models extend `automol`'s core data models directly rather than wrapping them: -`GeometryRow` extends `automol.Geometry`, `IdentityRow` extends `automol.Identity`. Any -conversion to/from other external formats is delegated to `automol`'s own conversion functions -rather than reimplemented here. - -## Core rows - -| Row | Represents | -| --- | --- | -| `GeometryRow` | A molecular geometry (symbols, coordinates, charge, spin). | -| `IdentityRow` | A chemical identifier (InChI, SMILES, conformer group, ...) shared across stationary points. | -| `IdentityExtraRow` | Extra key/value metadata attached to an `IdentityRow`. | -| `TrajectoryRow` | An ordered sequence of geometries (e.g. an optimization or IRC path). | -| `StationaryPointRow` | A stationary point on a PES — a geometry plus Hessian order and validity. | -| `StageRow` | A chemical state (reactant, product, or transition state) in a reaction. | -| `StepRow` | An elementary reaction step connecting two `StageRow`s, optionally via a TS stage. | -| `ModelRow` | A calculation model spec (program, program version, method, basis). | -| `CalculationRow` | A single quantum-chemistry calculation: type, status, provenance. | -| `ValidationRow` | A validation calculation (e.g. IRC) performed on a `StepRow`. | -| `EnergyRow` / `GradientRow` / `HessianRow` | Results computed at a geometry by a calculation. | - -`BaseRow` (id + `created_at`/`updated_at` timestamps) is the base for all of the above except -link tables. Result rows additionally share `BaseResultRow`, which provides the `query()` -classmethod for looking a result up by geometry, model, and input provenance. - -## Link tables - -Many-to-many and role-tagged relationships go through dedicated link rows (`BaseLink`), each -with a composite primary key and `ondelete="CASCADE"` foreign keys: - -- `CalculationGeometryLink` / `CalculationTrajectoryLink` — tag a geometry/trajectory as - `Role.INPUT` or `Role.OUTPUT` for a calculation. -- `TrajectoryGeometryLink` — orders geometries within a trajectory. -- `StationaryIdentityLink` — attaches identities to a stationary point. -- `StationaryStageLink` — groups stationary points into a `StageRow`. -- `StepValidationLink` — attaches validations to a `StepRow`. - -`BaseLink.create(*rows, **attrs)` builds a link by matching each row to the relationship whose -type it satisfies, so callers don't need to know the link's field names: - -```python -link = CalculationGeometryLink.create(calc, geo, role=Role.INPUT) -``` - -## Automatic behavior (event listeners) - -`autostorage.events` registers SQLAlchemy listeners that run on flush, so these behaviors apply -regardless of how a session is driven. See [Events](events.md) for the full detail, including -why some of these are registered at the session level rather than per-model. - -- **Shape validation** — `GradientRow.value` must be `(3 * atom_count,)` and `HessianRow.value` - must be `(3 * atom_count, 3 * atom_count)` for their linked geometry, or a - `ResultShapeError` is raised. -- **Stationary-point validity** — when a `HessianRow` is inserted, updated, or deleted, - `StationaryPointRow.is_valid` is recomputed from consensus among the geometry's Hessians - (comparing each stationary point's declared `order` to the Hessian-derived order). -- **Geometry immutability** — `symbols`/`coordinates` cannot be changed on a `GeometryRow` - after it's been inserted, since doing so would silently invalidate shape checks already run - against it. `charge`/`spin` remain mutable. -- **Automatic identity attachment** — new `StationaryPointRow`s get InChI and SMILES - `IdentityRow`s attached automatically via `automol.Algorithm`, deduplicated against existing - identity rows. -- **Conformer grouping** — new stationary points are compared against InChI peers using - `automol.geom.is_duplicate_conformer`; a match reuses the peer's conformer-group identity, - otherwise a new group id is allocated. -- **Reaction step consistency** — `StepRow.stage_id1`/`stage_id2` are kept in sorted order (to - satisfy the `stage_id1 < stage_id2` check constraint), `is_barrierless` is derived from - whether `stage_id_ts` is set, and `stage1`/`stage2` are rejected if they reference a - transition-state stage (and vice versa for `stage_ts`). - -## Migrations - -`migrations/` holds Alembic migrations, wired to `SQLModel.metadata`. This only applies to -evolving an *existing* on-disk database in place — fresh or in-memory `Database` instances -(including every test) get their schema from `SQLModel.metadata.create_all` directly, no -migration involved. - -Any change to a `table=True` model's columns/constraints/indexes needs a matching Alembic -revision. SQLite can't reflect the expression-based null-safe unique indexes on -`ModelRow`/`StepRow` (e.g. `unique_model_null_safe`, `unq_step_stages_null_safe`), so -`alembic revision --autogenerate` silently skips those — they must be added to new migrations -by hand if those models are ever touched again. See [Migrations](migrations.md) for the full -workflow. diff --git a/docs/source/database.md b/docs/source/database.md index e414754..198adad 100644 --- a/docs/source/database.md +++ b/docs/source/database.md @@ -1,169 +1,77 @@ # Database -{py:class}`~autostorage.database.Database` is the single entry point applications use to open -a SQLite file, get schema created, and run queries and mutations against it. This page covers -its connection semantics and the full method surface; see [Quickstart](quickstart.md) for a -minimal end-to-end example and [Data model](data-model.md) for what gets stored. +The `autostorage.database` module provides SQLite database connection management through the `Database` class. -## Opening a database +## Overview -```python -from autostorage import Database - -db = Database("workflow.db") # created if it doesn't exist -db = Database(":memory:") # scratch, in-process, gone on close() -db = Database("workflow.db", echo=True) -``` - -- `path` — a filesystem path (`str` or `Path`) or SQLite's special `:memory:` name. -- `echo` — when `True`, every SQL statement is logged to stdout (passed straight through to - SQLAlchemy's `create_engine`). +The `Database` class is a lightweight wrapper around SQLAlchemy's engine and session management, configured for SQLite with: -On every new DBAPI connection, `Database` also unconditionally issues `PRAGMA -foreign_keys=ON` — SQLite disables foreign-key enforcement by default, and the `ON DELETE -CASCADE` behavior the schema relies on (see [Link tables](data-model.md#link-tables)) depends -on this being set. +- **Foreign key enforcement** — SQLite's `PRAGMA foreign_keys=ON` is automatically enabled +- **Thread safety** — `check_same_thread=False` allows multi-threaded access +- **Canonical JSON serialization** — JSON column comparisons work regardless of dict key order +- **Automatic schema creation** — All SQLModel tables are created on initialization -`__init__` also calls `SQLModel.metadata.create_all(self.engine)`, so a fresh or in-memory -`Database` gets its full schema immediately — no Alembic migration is involved for these cases. -See [Migrations](migrations.md) for evolving an *existing* on-disk database instead. - -### JSON key ordering - -The engine is configured with a custom `json_serializer` that sorts dict keys -(`json.dumps(..., sort_keys=True)`). SQLite compares JSON columns as opaque text, so without -this, two Python dicts with the same key/value pairs but different insertion order would -serialize to different strings and fail to match in a query like: +## Creating a Database ```python -CalculationRow.input_provenance == {"seed": 1, "source": "orca"} -``` - -This affects any JSON-backed column: `CalculationRow.input_provenance`/`output_provenance`, -`ValidationRow.extras`, `GeometryRow.symbols`, and `TrajectoryGeometryLink.index`. +from autostorage.database import Database -## Session model and thread-safety +# Create or connect to a database +db = Database("path/to/database.db") -`Database` opens exactly one long-lived `Session` in `__init__` and reuses it for every method -call — `session()` is a context manager that yields that same session rather than creating a -fresh one per call. This matters in two ways: - -- Rows returned from a query stay attached to the session after the call returns, so - lazy-loaded relationships (e.g. `energy.geometry`, `calc.model`) keep working afterward. -- A `Database` instance is **not safe for concurrent use by multiple threads.** The engine is - created with `connect_args={"check_same_thread": False}`, but that only lifts SQLite's - restriction on using the underlying DBAPI connection from a different thread than it was - opened on (e.g. handing the whole `Database` off to a single background worker thread) — it - does not make the `Session` itself safe for concurrent access. Two threads calling methods on - the same `Database` at the same time is unsupported. - -If an operation raises inside the `session()` context manager, the session is rolled back -before the exception propagates, so a failed call doesn't leave partially-staged changes behind -for the next call. - -## Opening and closing - -```python -db = Database("workflow.db") -... -db.close() +# Enable SQL logging for debugging +db = Database("path/to/database.db", echo=True) ``` -`close()` disposes the underlying engine (and its connection pool). `Database` also supports -the context-manager protocol: +The database file is created if it doesn't exist. All tables defined in `autostorage.models` are automatically created via SQLModel metadata. -```python -with Database("workflow.db") as db: - ... -``` +## Using Sessions -On a clean exit this just calls `close()`. If the `with` block raises, the session is rolled -back first, then closed — so an exception partway through a multi-step workflow can't leave -uncommitted changes lingering in the session. - -## Writing rows - -```python -db.add(row) # stage a single row -db.add_all([row1, row2]) # stage multiple rows -db.commit() # write staged changes, commit the transaction -``` - -`add`/`add_all` only **stage** rows in the session — nothing is validated or written to the -database until the next `flush()` or `commit()`. This means integrity errors (unique constraint -violations) and the shape-validation event listeners (see [Events](events.md)) raise at that -later point, not at the `add` call site. +Sessions manage transactions and provide the query interface. Always use sessions as context managers to ensure proper cleanup: ```python -db.flush() +from autostorage.database import Database +from autostorage.models import GeometryRow + +db = Database("molecules.db") + +# Create a session context +with db.session() as session: + # Add a geometry + geom = GeometryRow(symbols=["C", "H", "H", "H", "H"], coordinates=[[0.0, 0.0, 0.0], ...]) + session.add(geom) + session.commit() # Explicitly commit changes + + # Query geometries + results = session.exec(select(GeometryRow)).all() ``` -Flushes pending changes without committing the transaction — useful for surfacing -validation/integrity errors early, or for getting DB-assigned values (like an autoincrement -`id`) onto an object before you need to reference it. Unlike `commit()`, `flush()` doesn't -trigger SQLAlchemy's `expire_on_commit` behavior, so it follows up with `session.expire_all()` -— otherwise an already-loaded object whose row was removed by a `ON DELETE CASCADE` during this -flush would read back stale (pre-deletion) data on next access instead of raising. - -```python -merged = db.merge(row) -``` +**Important**: Each call to `db.session()` creates a **new** session. Nothing is committed automatically — you must call `session.commit()` to persist changes. -`merge()` copies the state of a (possibly detached) row onto the identity-matched row already -tracked by the session — or inserts it if there is none — commits, and returns the merged -instance. Use this instead of `add()` when you're not sure whether the row is already attached -to this session's identity map (e.g. it was loaded by a different `Database`/session, or -round-tripped through serialization). +### Session Lifecycle ```python -db.delete(row) +with db.session() as session: + # Add/modify rows + session.add(row) + + # Flush to DB without committing (assigns IDs, checks constraints) + session.flush() + + # Commit the transaction + session.commit() + +# Session is automatically closed here ``` -Deletes a row and commits immediately (there's no separate staged-delete-then-flush step, since -`add`/`add_all` are the only staging-only methods). +For more details on session usage, querying, transaction control, and advanced patterns, see the [SQLAlchemy Session documentation](https://docs.sqlalchemy.org/en/21/orm/session_basics.html). -## Reading rows +## Closing the Database -By primary key: +When finished with a database, dispose of the connection pool: ```python -row = db.get(GeometryRow, 3) # raises LookupError if missing -row = db.get_or_none(GeometryRow, 3) # returns None if missing -``` - -By a `select()` statement (from `sqlmodel`, or the `SelectStatement`/`Select`/`SelectOfScalar` -aliases re-exported from `autostorage.database`): - -```python -from sqlmodel import select -from autostorage import CalcType, CalculationRow - -stmt = select(CalculationRow).where(CalculationRow.calc_type == CalcType.ENERGY) - -db.exec_first(stmt) # first match, or None -db.exec_one(stmt) # the single match; raises LookupError on zero or >1 matches -db.exec_all(stmt) # list of all matches -db.exists(stmt) # bool, via a single EXISTS subquery — never materializes a row +db.close() ``` -`exec_one` wraps SQLAlchemy's `NoResultFound`/`MultipleResultsFound` and re-raises both as -`LookupError`, so callers only need to handle one exception type regardless of which way the -query failed to return exactly one row. - -Most model classes also expose their own `query()`/`find_or_create()` classmethods built on -top of these primitives (e.g. `ModelRow.find_or_create`, `EnergyRow.query`, -`StationaryPointRow.query`) — prefer those where available, since they encode the right lookup -key for that row type. See [Data model](data-model.md) for the full list. - -## Exceptions raised through `Database` - -- `LookupError` — from `get()` (missing id) and `exec_one()` (zero or multiple matches). -- `autostorage.exc.MissingPrimaryKeyError` — raised by model `query()`/`find_or_create()` - methods when a row passed in hasn't been persisted yet (no `id`), since the query can't be - built without one. Call `db.add()`/`db.merge()` and `db.commit()`/`db.flush()` first. -- `autostorage.exc.ResultShapeError` — raised on `flush()`/`commit()` if a `GradientRow` or - `HessianRow` value doesn't match the shape implied by its geometry's atom count. See - [Events](events.md#shape-validation). -- SQLAlchemy's `IntegrityError` — for constraint violations (unique constraints, `NOT NULL`, - `CheckConstraint`s) that reach the database itself rather than being caught by an app-level - `query()`/`find_or_create()` lookup first. +This is typically unnecessary for short-lived scripts but recommended for long-running applications. diff --git a/docs/source/development.md b/docs/source/development.md deleted file mode 100644 index 2ae84b6..0000000 --- a/docs/source/development.md +++ /dev/null @@ -1,125 +0,0 @@ -# Development - -## Setup - -```bash -git clone -cd autostorage -pixi run init -``` - -`pixi run init` runs `scripts/setup.sh` (direnv, git init/first commit if needed) and then -installs the [lefthook](https://github.com/evilmartians/lefthook) git hooks that run the -pre-commit pipeline described below. All other tasks in this section assume a Pixi install and -are invoked as `pixi run ` (task definitions live in `pixi.toml` under -`[feature.dev.tasks]`). - -## Everyday tasks - -```{list-table} -:header-rows: 1 - -* - Task - - What it runs -* - `pixi run fmt` - - `ruff format .` — code formatting. -* - `pixi run lint` - - `ruff check . --fix` — linting, with `select = ["ALL"]` in `pyproject.toml` (a small, - documented ignore list carves out specific rules). -* - `pixi run types` - - `ty check` — static type-checking. -* - `pixi run imports` - - `lint-imports` — enforces the [module layering](data-model.md#module-layering) contract. -* - `pixi run test` - - `pytest`, with coverage. -* - `pixi run pre-commit` - - Runs all of the above in order, then checks the working tree is clean. -* - `pixi run cov-view` - - Opens the HTML coverage report (`htmlcov/index.html`) in `$BROWSER`. -* - `pixi run migrate` - - Applies Alembic migrations to an existing on-disk database — see [Migrations](migrations.md). -``` - -A single test: - -```bash -pixi run -e dev pytest tests/test_models.py::test_name -``` - -## Testing conventions - -- `testpaths = ["tests", "src"]` in `pyproject.toml` — pytest collects both the `tests/` - directory and `src/`. -- `--doctest-modules` is always on, so any doctest example (`>>> ...`) in a `src/` docstring is - collected and executed as a test. Keep them accurate and runnable — a stale example fails the - suite, not just the docs. -- Coverage runs in branch mode (`[tool.coverage.run] branch = true`) with a hard floor: - `fail_under = 80` in `[tool.coverage.report]`. `pixi run test` fails if coverage drops below - that, independent of whether individual tests pass. -- `tests/conftest.py` provides shared fixtures: an in-memory `database` fixture - (`Database(":memory:")`, closed on teardown), a seeded `rng`, and baseline `model_row`/ - `geometry_row`/`calculation_row`/`calc_geo_link` fixtures used across `test_models.py` and - `test_database.py`. -- `tests/test_migrations.py` is a smoke test that Alembic's `upgrade head` reproduces exactly - the schema `SQLModel.metadata.create_all` would build — see - [Keeping migrations and models in sync](migrations.md#keeping-migrations-and-models-in-sync). - -## Pre-commit pipeline - -`lefthook.yaml` defines two non-parallel command sequences, run via `pixi run pre-commit` / -`pixi run local-pre-commit`: - -1. `fmt` → `lint` → `types` → `imports` → `test` → `git diff --exit-code` (fails if formatting/ - linting left the tree dirty). -2. The `local-pre-commit` variant additionally runs `pixi run local start`/`stop` around the - same sequence, for workflows that need local services up during tests. - -CI (`.github/workflows/test.yml`) runs `pixi run pre-commit` on every push/PR, then `pixi run -test` again separately to publish the coverage report. - -## Docs - -```bash -pixi run docs-build # sphinx-build docs/source docs/build -pixi run docs-view # docs-build, then open docs/build/index.html in $BROWSER -``` - -`docs-view` requires the `BROWSER` environment variable to point at a browser executable (see -`scripts/view-docs.sh`). - -Docs are built with Sphinx + [MyST](https://myst-parser.readthedocs.io/) (Markdown) + -[sphinx-autodoc2](https://sphinx-autodoc2.readthedocs.io/) for the {doc}`API reference -`, using the `pydata-sphinx-theme`. Docstrings are written in NumPy style -(`tool.ruff.lint.pydocstyle.convention = "numpy"`) but rendered through a small custom parser -(`docs/source/autodoc2_docstrings_parser.py`) that runs them through -`sphinx.ext.napoleon.docstring.NumpyDocstring` before handing off to MyST — this is what lets -`autodoc2`, which doesn't natively understand NumPy-style sections, render them correctly. -Section headings get anchor links (`myst_heading_anchors = 3` in `conf.py`), which is what the -`#section-name` links throughout this documentation rely on. - -`.github/workflows/docs.yml` runs `pixi run docs-build` and deploys `docs/build/` to GitHub -Pages on every push to `main`. - -## Release process - -Releases are driven by [tbump](https://github.com/your-tools/tbump) (`tbump.toml`), which bumps -the version string in `pyproject.toml`, `pixi.toml`, and `src/autostorage/__init__.py` in one -step, then (via its `before_commit` hooks) regenerates `CHANGELOG.md` with -[keepachangelog](https://keepachangelog.com/) and re-locks `pixi.lock`, and finally tags the -commit `v{new_version}`. - -```bash -pixi run version # print the current version -pixi run release # tbump — walks through the bump interactively -``` - -Pushing a `v*.*.*` tag triggers `.github/workflows/release.yml`, which builds both a conda -package (`pixi run build-conda`) and a PyPI package (`pixi run build-pypi`), publishes both -(`publish-conda` to Anaconda.org, `publish-pypi` to PyPI — both need their respective secrets in -CI), and creates a GitHub Release with notes extracted from `CHANGELOG.md`. - -## Coding standards - -See `CONTRIBUTING.md` for naming conventions (`Row`-suffixed SQLModel classes vs. domain -models) and cross-package conversion ownership rules shared across the `automol`/`autostorage` -suite. diff --git a/docs/source/events.md b/docs/source/events.md deleted file mode 100644 index a4d686c..0000000 --- a/docs/source/events.md +++ /dev/null @@ -1,186 +0,0 @@ -# Events (automatic behavior) - -{py:mod}`autostorage.events` registers SQLAlchemy ORM event listeners at import time. Because -`autostorage.database` imports `.events` (for its side effect of registering these listeners) -before opening any `Database`, everything on this page applies to **any** session touching -these models — not just calls made through `Database`'s own methods. - -Listeners are registered two ways, and the choice between them matters: - -- **Mapper-level** (`before_insert`/`before_update`/`before_delete` on a specific model) — runs - once per row of that type, as part of that row's own flush step. -- **Session-level** (`before_flush` on `Session`) — runs once per flush, before any mapper-level - events, with access to the whole pending change set (`session.new`/`dirty`/`deleted`). - -Three listeners below are session-level specifically because they need to **mutate a different, -already-clean object** than the one that triggered them (e.g. deleting a `HessianRow` needs to -update a `StationaryPointRow.is_valid` that isn't itself being inserted/updated/deleted). -SQLAlchemy silently drops attribute changes made to other, already-clean objects from within a -mapper-level `before_delete`/`before_insert`/`before_update` handler — those objects aren't part -of the acting object's already-computed flush plan. A session-level `before_flush` listener runs -before that plan is fixed, so changes it makes are picked up. - -## Shape validation - -```{eval-rst} -.. autofunction:: autostorage.events.verify_gradient_shape - :no-index: -.. autofunction:: autostorage.events.verify_hessian_shape - :no-index: -``` - -Both run on `before_insert`/`before_update` for their row type. `GradientRow.value` must have -shape `(3 * atom_count,)`; `HessianRow.value` must be `(3 * atom_count, 3 * atom_count)`, where -`atom_count` comes from the row's linked `GeometryRow`. A mismatch raises -{py:class}`~autostorage.exc.ResultShapeError` — which surfaces at `flush()`/`commit()` time, not -at the point the row was constructed or `add()`-ed. - -Both listeners resolve the geometry via a shared helper, `_resolve_geometry`, rather than -reading `target.geometry` directly: if a row was built with only `geometry_id=...` set (not the -`.geometry` relationship), the relationship stays unpopulated until the ORM syncs it — reading -it directly would let shape validation be silently skipped for a row that does have a geometry. -`_resolve_geometry` falls back to `session.get(GeometryRow, target.geometry_id)` in that case. -The same pattern reappears below for `StepRow`'s stage relationships -(`_resolve_stage`). - -## Stationary-point validity - -```{eval-rst} -.. autofunction:: autostorage.events.validate_geometry_orders - :no-index: -.. autofunction:: autostorage.events.revalidate_geometry_orders_on_hessian_delete - :no-index: -``` - -`StationaryPointRow.is_valid` is derived, not set directly by application code. It's recomputed -by `_recompute_geometry_stationary_validity(geometry)` whenever the set of `HessianRow`s -attached to a geometry changes: - -1. Collect the geometry's `HessianRow`s (minus any passed via `excluding=`, see below). -2. Compute each Hessian's `order` (the count of negative harmonic frequencies — see - {py:attr}`~autostorage.models.HessianRow.order`) and take the set of distinct orders. -3. If more than one distinct order is present, raise `ValueError` — the geometry's Hessians - disagree, which can't be reconciled automatically. -4. Otherwise, for every `StationaryPointRow` on that geometry, set `is_valid = (stationary.order - == expected_order)`, where `expected_order` is the one agreed-upon Hessian order. - -This recompute is wired to three triggers: - -- `validate_geometry_orders` — a mapper-level `before_insert`/`before_update` listener on both - `StationaryPointRow` and `HessianRow`, covering the common case of inserting/editing either - side. -- `revalidate_geometry_orders_on_hessian_delete` — a **session-level** `before_flush` listener - that reacts to `HessianRow` deletions. It's session-level (not a mapper `before_delete` - listener) for the reason given in the section intro: it needs to write to - `StationaryPointRow.is_valid` on objects other than the one being deleted. It passes the - about-to-be-deleted Hessians as `excluding=` to `_recompute_geometry_stationary_validity`, - since `geometry.hessians` still includes them at `before_flush` time (the `DELETE` hasn't been - issued to the database yet). - -If a geometry has no Hessians, or no `StationaryPointRow`s, the recompute is a no-op — `is_valid` -keeps whatever value it already had. - -## Geometry immutability - -```{eval-rst} -.. autofunction:: autostorage.events.verify_geometry_immutable_fields - :no-index: -``` - -Once a `GeometryRow` has been inserted, `symbols` and `coordinates` can never be changed — -attempting to modify either raises `ValueError` on the next `flush()`/`commit()`. `charge` and -`spin` remain freely mutable. This is enforced via `sqlalchemy.orm.attributes.get_history`, -which reports whether a field has pending added/deleted values since it was loaded; if the field -has never been touched, `get_history` reports no change and the update passes. - -The reason: `GradientRow`/`HessianRow` shape checks (above) and Hessian order consensus are -computed once, against the geometry as it existed when those results were saved. Silently -allowing `symbols`/`coordinates` to change afterward would invalidate checks already performed -without re-running them. - -## Automatic identity attachment - -```{eval-rst} -.. autofunction:: autostorage.events.add_inchi_identities - :no-index: -``` - -A session-level `before_flush` listener. For every new (`session.new`) `StationaryPointRow`, it: - -1. Computes an InChI `IdentityRow` from the point's geometry via - `IdentityRow.from_geometry(geo=..., algorithm=Algorithm.RDKIT_INCHI)`. If this raises - `ValueError` (e.g. `automol`/RDKit can't derive an InChI for the geometry), the point is - skipped — no identity is attached and no error propagates. -2. Batches all resulting `(algorithm, value)` pairs into a single `tuple_(...).in_(...)` query - against existing `IdentityRow`s, so N new stationary points cost one lookup query rather than - N. -3. If a matching InChI `IdentityRow` already exists, reuses it (appends the *existing* row to - `obj.identities`) instead of creating a duplicate — this is what backs the - `unique_identity` constraint on `(kind, algorithm, value)` staying meaningful in practice - rather than being hit as an integrity error. -4. Otherwise, attaches the newly-computed InChI row, and additionally derives a SMILES string - via `Algorithm.RDKIT_SMILES`. The SMILES is **not** stored as its own `IdentityRow` — it's - attached as an `IdentityExtraRow(identity=inchi, attribute="smiles", value=...)` on the InChI - identity just created. If SMILES derivation raises `ValueError`, it's silently skipped (the - InChI identity is still attached). - -## Conformer grouping - -```{eval-rst} -.. autofunction:: autostorage.events.assign_conformer_ids - :no-index: -``` - -Also a session-level `before_flush` listener, run after `add_inchi_identities` has had a chance -to populate InChI identities on new stationary points (both are registered against the same -event; SQLAlchemy invokes same-event listeners in registration order, and `events.py` defines -`add_inchi_identities` first). - -For each new `StationaryPointRow` that doesn't already have an IRMSD-kind identity -(`Algorithm.IRMSD.kind`) but does have an InChI identity: - -1. Look at the InChI identity's other attached stationary points ("peers") — points that share - the same InChI, i.e. the same constitutional/stereo identity. -2. Compare the new point's geometry against each peer's via - `automol.geom.is_duplicate_conformer`. -3. If one matches, reuse that peer's IRMSD identity (its conformer-group id) via - `_matching_conformer_identity` — same conformer, same group. -4. Otherwise, allocate a new group id: `max(existing IRMSD values cast to int) + 1`, or `1` if - none exist yet, and create a new IRMSD `IdentityRow` with that value. - -```{note} -Group-id allocation assumes single-writer semantics, consistent with `Database`'s documented -[non-thread-safety](database.md#session-model-and-thread-safety). If that assumption is -violated (e.g. two processes/threads racing against the same on-disk database), the -`unique_identity` constraint on `IdentityRow` turns a racing duplicate group id into an -`IntegrityError` at commit — rather than silently merging two distinct conformer groups under -one id. -``` - -Within a single flush, multiple new non-duplicate points increment `next_group_id` locally -(instead of re-querying `MAX(...)` for each one), so N new unrelated conformers correctly get N -distinct, consecutive group ids from one query. - -## Reaction step consistency - -```{eval-rst} -.. autofunction:: autostorage.events.verify_stage_order_and_barrierless - :no-index: -.. autofunction:: autostorage.events.verify_stage_ts_consistency - :no-index: -``` - -Both are mapper-level `before_insert`/`before_update` listeners on `StepRow`: - -- `verify_stage_order_and_barrierless` sorts `stage_id1 < stage_id2` (swapping them if - necessary) so every `StepRow` satisfies the `chk_stage_order` `CheckConstraint` regardless of - which order the caller passed `stage1`/`stage2` in, and derives - `is_barrierless = not stage_id_ts` — callers never set `is_barrierless` directly. -- `verify_stage_ts_consistency` rejects `stage1`/`stage2` if either is a transition-state stage - (`StageRow.is_ts`), and rejects `stage_ts` if it *isn't* one, raising `ValueError` on any - mismatch. - -Both resolve stages via `_resolve_stage`, the `StepRow` analogue of `_resolve_geometry` above: -if only a `stage_id*` foreign key was set (not the `stage*`/`stage_ts` relationship), it falls -back to `session.get(StageRow, stage_id)` so the check isn't skipped for a step that does -reference a stage. diff --git a/docs/source/index.md b/docs/source/index.md index e51670d..bf94040 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -1,24 +1,13 @@ # autostorage -A [SQLModel](https://sqlmodel.tiangolo.com/)/SQLAlchemy persistence layer for computational -chemistry workflow data, built on top of [`automol`](https://github.com/avcopan/automol). It -stores molecular geometries, chemical identities, trajectories, stationary points, calculation -results (energies, gradients, Hessians), and the calculations/reaction steps that connect them, -as a graph of related rows in a SQLite database. +A [SQLModel](https://sqlmodel.tiangolo.com/)/[SQLAlchemy](https://www.sqlalchemy.org/) persistence layer for computational chemistry workflow data, built on top of [`automol`](https://github.com/avcopan/automol). It stores molecular geometries, chemical identities, trajectories, stationary points, calculation results (energies, gradients, Hessians), and the calculations/reaction steps that connect them, as a graph of related rows in a SQLite database. -Row models extend `automol`'s core data models directly rather than wrapping them — -`GeometryRow` extends `automol.Geometry`, `IdentityRow` extends `automol.Identity` — so any -data already expressed in `automol` types can be persisted with no conversion step. +Row models extend `automol`'s core data models directly rather than wrapping them — `GeometryRow` extends `automol.Geometry`, `IdentityRow` extends `automol.Identity` — so any data already expressed in `automol` types can be persisted with no conversion step. :::{toctree} :maxdepth: 2 :caption: Contents quickstart -data-model database -events -migrations -development -apidocs/index ::: diff --git a/docs/source/migrations.md b/docs/source/migrations.md deleted file mode 100644 index dbbddd9..0000000 --- a/docs/source/migrations.md +++ /dev/null @@ -1,109 +0,0 @@ -# Migrations - -`autostorage` has two independent sources of schema, used in different situations: - -- **`SQLModel.metadata.create_all(engine)`** — runs automatically inside - {py:meth}`Database.__init__ ` and builds the full - current schema from the model definitions in one shot. This is what fresh and in-memory - databases use, including every `Database` instance created in the test suite. No Alembic - involvement at all. -- **Alembic migrations** (`migrations/`) — the only path for evolving an *existing* on-disk - database in place, so it picks up schema changes without losing the data already in it. - -If you're not migrating an existing database (e.g. writing a test, or working with a -scratch/in-memory `Database`), you don't need anything on this page — just change the models -and `create_all` picks it up. - -## Applying migrations - -```bash -AUTOSTORAGE_DATABASE_URL=sqlite:///path/to.db pixi run migrate -``` - -`pixi run migrate` runs `alembic upgrade head`. `migrations/env.py` reads -`AUTOSTORAGE_DATABASE_URL` from the environment and uses it to override `alembic.ini`'s -`sqlalchemy.url` at runtime, so a single Alembic setup can target any on-disk database without -editing config. If the variable isn't set, Alembic falls back to whatever's hardcoded in -`alembic.ini`. - -`env.py` also does one thing that's easy to miss: it imports `autostorage.database` purely for -its side effect of pulling in `autostorage.models` and `autostorage.events`, which is what -registers every `table=True` model onto `SQLModel.metadata`. Without that import, Alembic -would see an empty target schema. - -## Writing a new migration - -After changing a `table=True` model's columns, constraints, or indexes: - -```bash -pixi run -e dev alembic revision --autogenerate -m "describe the change" -``` - -Run this against a scratch on-disk SQLite database (not `:memory:` — Alembic needs a real file -to diff against), then **review the generated script by hand** before committing it. Alembic's -autogenerate is a reasonable first draft, not a guarantee — see the caveat below for one gap -specific to this schema. - -### SQLite can't reflect expression-based indexes - -Two constraints in `models.py` are implemented as expression-based unique `Index`es rather than -plain `UniqueConstraint`s, specifically to make them null-safe (SQL treats `NULL != NULL`, so a -plain `UniqueConstraint` lets multiple "duplicate" rows through whenever one of the constrained -columns is `NULL`): - -- `unique_model_null_safe` on `ModelRow`, over `(program, coalesce(program_version, ''), - method, coalesce(basis, ''))`. -- `unq_step_stages_null_safe` on `StepRow`, over `(stage_id1, stage_id2, coalesce(stage_id_ts, - 0))`. - -SQLite's reflection support can't see expression-based indexes, so `alembic revision ---autogenerate` silently omits them from the generated script — it isn't that it gets them -wrong, it just doesn't know they exist. If a migration touches `model` or `step`, check whether -these indexes need to be re-created by hand, following the pattern already used in -`e50de3129c84_add_null_safe_indexes_reverse_lookup_.py`: - -```python -op.create_index( - "unique_model_null_safe", - "model", - [ - "program", - sa.text("coalesce(program_version, '')"), - "method", - sa.text("coalesce(basis, '')"), - ], - unique=True, -) -``` - -Both indexes are also defense-in-depth alongside an app-level lookup that's the actual -dedup mechanism in practice — `ModelRow.find_or_create` and `StepRow.find_or_create`/`.query` -(see [Data model](data-model.md)) — so a missed index mainly risks a duplicate row slipping in -under a race, not a functional bug in normal single-writer use. - -## Migration history - -```{list-table} -:header-rows: 1 - -* - Revision - - Down revision - - Summary -* - `faa9f50bc029` - - (base) - - Baseline: the schema as of the first tracked migration. -* - `e50de3129c84` - - `faa9f50bc029` - - Adds `created_at`/`updated_at` timestamps across tables, `CalculationRow.status`/ - `error_message`, reverse-lookup indexes on several link tables, and the two null-safe - expression indexes described above. -``` - -## Keeping migrations and models in sync - -`tests/test_migrations.py` guards against migrations drifting from the current model -definitions: it applies every migration to a scratch database with `alembic upgrade head`, then -uses `alembic.autogenerate.compare_metadata` to diff the result against `SQLModel.metadata` and -asserts there's no difference. If you change a model without writing (or correctly writing) the -matching migration, this test fails — treat it as the source of truth for whether a migration -is complete, not just `alembic revision --autogenerate` running without error. diff --git a/docs/source/models.md b/docs/source/models.md new file mode 100644 index 0000000..3a108fb --- /dev/null +++ b/docs/source/models.md @@ -0,0 +1,69 @@ +# Data Models + +autostorage uses SQLModel/SQLAlchemy to define a relational schema for computational chemistry workflow data. The models form a graph structure where molecular geometries, calculations, stationary points, and chemical identities are linked through explicit relationship tables. + +## Model Overview + +The schema is organized into five functional groups: + +### Link Tables + +Association tables implementing many-to-many relationships between entities. + +| Model | Description | +|-------|-------------| +| `CalculationGeometryLink` | Associates geometries with calculations, tracking whether each geometry is an input or output | +| `GeometryTrajectoryLink` | Links geometries to trajectories with position indices | +| `CalculationTrajectoryLink` | Associates trajectories with calculations as inputs or outputs | +| `StageStationaryLink` | Connects stationary points to reaction stages | +| `StepValidationLink` | Links validation calculations to reaction steps | +| `IdentityStationaryLink` | Associates chemical identities with stationary points | + +### Core Data Models + +The fundamental entities representing molecular structures, calculations, and their results. + +| Model | Description | +|-------|-------------| +| `GeometryRow` | Molecular geometry with atomic symbols, coordinates, charge, and spin; extends `automol.Geometry` | +| `TrajectoryRow` | Ordered sequence of geometries from a dynamic calculation | +| `ModelRow` | Calculation model specification (program, method, basis set, keywords) | +| `CalculationRow` | Quantum chemistry calculation with provenance metadata | +| `EnergyRow` | Energy result for a geometry at a specific level of theory | +| `GradientRow` | Energy gradient (forces) for a geometry | +| `HessianRow` | Second derivative matrix (Hessian) for a geometry | +| `ValidationRow` | Validation result (e.g., IRC) for a reaction step | + +### Stationary Points + +Models representing critical points on potential energy surfaces. + +| Model | Description | +|-------|-------------| +| `StationaryPointRow` | A stationary point with Hessian index and validation status | + +### Reaction Network + +Models defining elementary reaction steps and their constituent chemical states. + +| Model | Description | +|-------|-------------| +| `StageRow` | A chemical state in a reaction (reactant, product, or transition state) | +| `StepRow` | An elementary reaction step connecting two stages via a transition state (or barrierless) | + +### Chemical Identities + +Models for chemical identification and classification. + +| Model | Description | +|-------|-------------| +| `IdentityRow` | Queryable chemical identifiers (InChI, AmChI, etc.) with kind and algorithm; extends `automol.Identity` | +| `IdentityExtraRow` | Additional *non-queryable* identities (SMILES, Hill formula, ...) attached to a primary identity | + +## Design Principles + +- **automol integration**: `GeometryRow` and `IdentityRow` extend `automol.Geometry` and `automol.Identity` directly rather than wrapping them, delegating all format conversions to automol +- **Explicit relationships**: Many-to-many relationships use dedicated link tables rather than implicit joins +- **Provenance tracking**: Calculations record both input and output provenance metadata +- **Compressed storage**: NumPy arrays (coordinates, gradients, Hessians) are stored as zlib-compressed binary data via `CompressedArrayTypeDecorator` +- **Graph structure**: The schema forms a directed graph where calculations consume and produce geometries/trajectories, geometries carry results, stationary points reference geometries, and reaction steps connect stages diff --git a/docs/source/quickstart.md b/docs/source/quickstart.md index 5b2cec0..b727744 100644 --- a/docs/source/quickstart.md +++ b/docs/source/quickstart.md @@ -6,7 +6,7 @@ Install as a [Pixi](https://pixi.sh) dependency: ```toml [dependencies] -autostorage = ">=0.0.10" +autostorage = ">=0.0.12" ``` Or with `uv`/`pip` from PyPI: @@ -22,7 +22,6 @@ Requires Python ≥3.12. ```python import numpy as np from autostorage import ( - CalcType, CalculationGeometryLink, CalculationRow, Database, @@ -35,52 +34,56 @@ from autostorage import ( # Open (or create) a SQLite database; ":memory:" also works for scratch use. db = Database("workflow.db") -# `find_or_create` dedups on (program, program_version, method, basis). -model = ModelRow.find_or_create(db, program="orca", method="b3lyp", basis="def2-svp") - -calc = CalculationRow(model=model, calc_type=CalcType.ENERGY) -geo = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.8], [0, 0, 0], [0.8, 0, 0]]), - charge=0, - spin=0, -) -link = CalculationGeometryLink.create(calc, geo, role=Role.INPUT) -db.add_all([model, calc, geo, link]) -db.commit() - -# Attach a result to the geometry/calculation pair. -energy = EnergyRow(geometry=geo, calculation=calc, value=-76.02) -db.add(energy) -db.commit() - -# Look the result back up by geometry, model, and input provenance. -found = EnergyRow.query(db, geo=geo, model=model) -assert found is not None -print(found.value) +# Work within a session context. +with db.session() as session: + # Create a model specifying the calculation type, program, method, and basis. + model = ModelRow( + calc_type="energy", + program="orca", + method="b3lyp", + basis="def2-svp", + ) + + # Create a calculation using this model. + calc = CalculationRow(model=model) + + # Create a geometry. + geo = GeometryRow( + symbols=["H", "O", "H"], + coordinates=np.array([[0, 0, 0.8], [0, 0, 0], [0.8, 0, 0]]), + charge=0, + spin=0, + ) + + # Link the geometry to the calculation as an input. + link = CalculationGeometryLink( + calculation=calc, + geometry=geo, + role=Role.INPUT, + ) + + # Add all objects to the session and commit. + session.add_all([model, calc, geo, link]) + session.commit() + + # Attach an energy result to the geometry/calculation pair. + energy = EnergyRow(geometry=geo, calculation=calc, value=-76.02) + session.add(energy) + session.commit() + + # Query the result back by filtering on geometry and calculation. + found = session.query(EnergyRow).filter_by( + geometry_id=geo.id, + calculation_id=calc.id, + ).first() + assert found is not None + print(found.value) db.close() ``` -`Database` also supports the `with` statement, which rolls back on an unhandled exception and -closes the connection on exit: - -```python -with Database("workflow.db") as db: - ... -``` - -## Migrating an existing database - -Fresh or in-memory `Database` instances get their schema from `SQLModel.metadata.create_all` -automatically — no migration step needed. For an existing on-disk database, apply -[Alembic](https://alembic.sqlalchemy.org/) migrations with: - -```bash -AUTOSTORAGE_DATABASE_URL=sqlite:///path/to.db pixi run migrate -``` +The `Database.session()` method returns a standard SQLAlchemy `Session` that supports the +context manager protocol. Sessions automatically handle transaction management — commit +explicitly to persist changes, and unhandled exceptions trigger a rollback. -See [Data model](data-model.md) for the schema this creates, [Database](database.md) for the -full `Database` method surface, [Events](events.md) for the automatic validation/enrichment -behavior that runs on every flush, and the {doc}`API reference ` for full details -on every model and method. +See the {doc}`API reference ` for full details on every model and method. diff --git a/examples/scan.py b/examples/scan.py new file mode 100644 index 0000000..5ff5b55 --- /dev/null +++ b/examples/scan.py @@ -0,0 +1,167 @@ +"""Example demonstrating a 2D scan over both O-H bonds in H2O.""" + +import numpy as np + +from autostorage import Database +from autostorage.models import ( + CalculationGeometryLink, + CalculationRow, + CalculationTrajectoryLink, + GeometryRow, + GeometryTrajectoryLink, + ModelRow, + Role, + StageRow, + StageStationaryLink, + StationaryPointRow, + StepRow, + TrajectoryRow, +) + +# Create database +db = Database("example_scan.db", echo=True) + +with db.session() as session: + # Create a 2D scan trajectory (varying both O-H bond lengths) + # ndim=2 for a 2D scan grid + trajectory = TrajectoryRow(ndim=2) + session.add(trajectory) + session.flush() + + # Create scan model (a relaxed scan calculation) + model = ModelRow( + calc_type="scan", + program="psi4", + method="b3lyp", + basis="6-31g*", + ) + session.add(model) + session.flush() + + # Create the scan calculation + calc_scan = CalculationRow(model_id=model.id) + session.add(calc_scan) + session.flush() + + # Link the calculation to the trajectory it produced + calc_traj_link = CalculationTrajectoryLink( + calculation_id=calc_scan.id, + trajectory_id=trajectory.id, + role=Role.OUTPUT, + ) + session.add(calc_traj_link) + + # Generate a simple 3x3 grid of H2O geometries varying both O-H bonds + # Base geometry: equilibrium-ish H2O + scan_geometries = [] + r1_values = [0.90, 0.96, 1.02] # First O-H bond lengths (Angstrom) + r2_values = [0.90, 0.96, 1.02] # Second O-H bond lengths (Angstrom) + + for i, r1 in enumerate(r1_values): + for j, r2 in enumerate(r2_values): + # Create H2O geometry with varying bond lengths + # O at origin, H atoms along x and symmetric about xz-plane + geom = GeometryRow( + symbols=["O", "H", "H"], + coordinates=np.array( + [ + [0.0, 0.0, 0.0], + [ + r1 * np.cos(np.radians(52)), + r1 * np.sin(np.radians(52)), + 0.0, + ], + [ + r2 * np.cos(np.radians(52)), + -r2 * np.sin(np.radians(52)), + 0.0, + ], + ] + ), + charge=0, + spin=0, + ) + session.add(geom) + scan_geometries.append((geom, (i, j))) + + session.flush() + + # Link each geometry to the trajectory with its scan index + for geom, (i, j) in scan_geometries: + geom_traj_link = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=trajectory.id, + index=(i, j), # 2D index in the scan grid + ) + session.add(geom_traj_link) + + # Also link each geometry as an output of the scan calculation + geom_calc_link = CalculationGeometryLink( + geometry_id=geom.id, + calculation_id=calc_scan.id, + role=Role.OUTPUT, + ) + session.add(geom_calc_link) + + session.flush() + + # Create pseudo stationary points for the scan endpoints and a TS + # Start point: shortest bonds (0, 0) + stat_start = StationaryPointRow( + geometry_id=scan_geometries[0][0].id, # (0, 0) + calculation_id=calc_scan.id, + order=0, + is_pseudo=True, + is_valid=False, + ) + + # End point: longest bonds (2, 2) + stat_end = StationaryPointRow( + geometry_id=scan_geometries[-1][0].id, # (2, 2) + calculation_id=calc_scan.id, + order=0, + is_pseudo=True, + is_valid=False, + ) + + # Pseudo TS: center of the scan grid (1, 1) + middle_geom = next(g for g, idx in scan_geometries if idx == (1, 1)) + stat_ts = StationaryPointRow( + geometry_id=middle_geom.id, + calculation_id=calc_scan.id, + order=1, + is_pseudo=True, + is_valid=False, + ) + + session.add_all([stat_start, stat_end, stat_ts]) + session.flush() + + # Create stages for the start, end, and TS points + stage_start = StageRow(is_ts=False) + stage_end = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage_start, stage_end, stage_ts]) + session.flush() + + # Link stationary points to stages + session.add_all( + [ + StageStationaryLink(stationary_id=stat_start.id, stage_id=stage_start.id), + StageStationaryLink(stationary_id=stat_end.id, stage_id=stage_end.id), + StageStationaryLink(stationary_id=stat_ts.id, stage_id=stage_ts.id), + ] + ) + session.flush() + + # Create a step connecting start and end through the pseudo TS + step = StepRow( + stage_id1=min(stage_start.id, stage_end.id), + stage_id2=max(stage_start.id, stage_end.id), + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step) + + # Commit all changes + session.commit() diff --git a/examples/stationary.py b/examples/stationary.py new file mode 100644 index 0000000..1458e4f --- /dev/null +++ b/examples/stationary.py @@ -0,0 +1,195 @@ +"""Example demonstrating stationary points, stages, and steps.""" + +import numpy as np + +from autostorage import Database +from autostorage.models import ( + CalculationGeometryLink, + CalculationRow, + EnergyRow, + GeometryRow, + GradientRow, + HessianRow, + ModelRow, + Role, + StageRow, + StageStationaryLink, + StationaryPointRow, + StepRow, +) + +# Create database +db = Database("example_stationary.db", echo=True) + +with db.session() as session: + # Create a simple H2O geometry (reactant) + geom1 = GeometryRow( + symbols=["O", "H", "H"], + coordinates=np.array( + [ + [0.0, 0.0, 0.0], + [0.0, 0.757, 0.587], + [0.0, -0.757, 0.587], + ] + ), + charge=0, + spin=0, + ) + session.add(geom1) + + # Create a second geometry (product - slightly different H2O) + geom2 = GeometryRow( + symbols=["O", "H", "H"], + coordinates=np.array( + [ + [0.0, 0.0, 0.0], + [0.0, 0.800, 0.600], + [0.0, -0.800, 0.600], + ] + ), + charge=0, + spin=0, + ) + session.add(geom2) + session.flush() + + # Create a calculation model + model = ModelRow( + calc_type="OPT", + program="psi4", + method="b3lyp", + basis="6-31g*", + ) + session.add(model) + session.flush() + + # Create calculations that identified the stationary points + calc1 = CalculationRow(model_id=model.id) + calc2 = CalculationRow(model_id=model.id) + session.add_all([calc1, calc2]) + session.flush() + + # Link each geometry to the optimization that produced it as an output + link_geom1 = CalculationGeometryLink( + geometry_id=geom1.id, calculation_id=calc1.id, role=Role.OUTPUT + ) + link_geom2 = CalculationGeometryLink( + geometry_id=geom2.id, calculation_id=calc2.id, role=Role.OUTPUT + ) + session.add_all([link_geom1, link_geom2]) + + # Add energy results for both geometries + energy1 = EnergyRow( + geometry_id=geom1.id, + calculation_id=calc1.id, + value=-76.4268193, # Hartree + ) + energy2 = EnergyRow( + geometry_id=geom2.id, + calculation_id=calc2.id, + value=-76.4265821, # Hartree (slightly higher energy) + ) + session.add_all([energy1, energy2]) + + # Add gradient results (3 atoms x 3 coords = 9 values, flattened) + gradient1 = GradientRow( + geometry_id=geom1.id, + calculation_id=calc1.id, + value=np.array( + [ + 0.0001, + -0.0002, + 0.0003, # O atom gradient + -0.0001, + 0.0001, + -0.0002, # H1 atom gradient + 0.0000, + 0.0001, + -0.0001, # H2 atom gradient + ] + ), + ) + gradient2 = GradientRow( + geometry_id=geom2.id, + calculation_id=calc2.id, + value=np.array( + [ + 0.0002, + -0.0003, + 0.0001, + -0.0002, + 0.0002, + -0.0001, + 0.0001, + 0.0001, + 0.0000, + ] + ), + ) + session.add_all([gradient1, gradient2]) + + # Add hessian results (9x9 matrix for 3 atoms) + # Create a simple symmetric positive-definite hessian + hess1_matrix = np.random.RandomState(42).randn(9, 9) * 0.1 + hess1_matrix = (hess1_matrix + hess1_matrix.T) / 2 # Make symmetric + hess1_matrix += np.eye(9) * 2 # Make positive-definite + + hess2_matrix = np.random.RandomState(43).randn(9, 9) * 0.1 + hess2_matrix = (hess2_matrix + hess2_matrix.T) / 2 + hess2_matrix += np.eye(9) * 2 + + hessian1 = HessianRow( + geometry_id=geom1.id, + calculation_id=calc1.id, + value=hess1_matrix.astype(np.float32), + ) + hessian2 = HessianRow( + geometry_id=geom2.id, + calculation_id=calc2.id, + value=hess2_matrix.astype(np.float32), + ) + session.add_all([hessian1, hessian2]) + session.flush() + + # Create stationary points (both are minima: order=0) + stat1 = StationaryPointRow( + geometry_id=geom1.id, + calculation_id=calc1.id, + order=0, + is_pseudo=False, + is_valid=True, + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, + calculation_id=calc2.id, + order=0, + is_pseudo=False, + is_valid=True, + ) + session.add_all([stat1, stat2]) + session.flush() + + # Create stages (both are non-TS stages) + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + # Link stationary points to stages + link1 = StageStationaryLink(stationary_id=stat1.id, stage_id=stage1.id) + link2 = StageStationaryLink(stationary_id=stat2.id, stage_id=stage2.id) + session.add_all([link1, link2]) + session.flush() + + # Create a barrierless step connecting the two stages + # StepRow requires stage_id1 < stage_id2 + step = StepRow( + stage_id1=min(stage1.id, stage2.id), + stage_id2=max(stage1.id, stage2.id), + stage_id_ts=None, # No transition state (barrierless) + is_barrierless=True, + ) + session.add(step) + + # Commit all changes + session.commit() diff --git a/examples/stationary_point.py b/examples/stationary_point.py deleted file mode 100644 index 6e5abd0..0000000 --- a/examples/stationary_point.py +++ /dev/null @@ -1,256 +0,0 @@ -"""End-to-end example of the non-reaction parts of the autostorage data model. - -Builds a synthetic water (H2O) optimization + frequency workflow entirely from literal -coordinates (no external file I/O), touching every row model except `StageRow`/`StepRow` -and their reaction-specific link tables. - -Run with:: - - pixi run -e dev python examples/stationary_point.py -""" - -import numpy as np -from automol import Algorithm -from numpy.random import Generator - -from autostorage import ( - CalcType, - CalculationGeometryLink, - CalculationRow, - Database, - EnergyRow, - GeometryRow, - GradientRow, - HessianRow, - IdentityRow, - ModelRow, - Role, - StationaryPointRow, - TrajectoryRow, - ValidationRow, -) -from autostorage.models import ( - CalculationTrajectoryLink, - StationaryIdentityLink, - TrajectoryGeometryLink, -) - - -def optimize_water( - db: Database, model: ModelRow -) -> tuple[CalculationRow, GeometryRow, GeometryRow]: - """Run a synthetic optimization, linking an input geometry to its output.""" - opt_calc = CalculationRow(model=model, calc_type=CalcType.OPT) - input_geo = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.9], [0, 0, 0], [0.85, 0.1, 0]]), - charge=0, - spin=0, - ) - optimized_geo = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.8], [0, 0, 0], [0.8, 0, 0]]), - charge=0, - spin=0, - ) - db.add_all( - [ - opt_calc, - input_geo, - optimized_geo, - CalculationGeometryLink.create(opt_calc, input_geo, role=Role.INPUT), - CalculationGeometryLink.create(opt_calc, optimized_geo, role=Role.OUTPUT), - ] - ) - db.commit() - print( - f"optimization calculation {opt_calc.id}: " - f"{len(opt_calc.input_geometries)} input, " - f"{len(opt_calc.output_geometries)} output geometry/geometries" - ) - return opt_calc, input_geo, optimized_geo - - -def mark_stationary_point( - db: Database, opt_calc: CalculationRow, optimized_geo: GeometryRow -) -> StationaryPointRow: - """Record the optimized geometry as a stationary point. - - InChI and conformer identities are attached automatically on commit (see - `autostorage.events.add_inchi_identities`/`assign_conformer_ids`) -- this - function doesn't construct those manually, only reads them back. - """ - stationary = StationaryPointRow( - calculation=opt_calc, geometry=optimized_geo, order=0 - ) - db.add(stationary) - db.commit() - - inchi = stationary.identity(kind="stereoisomer") - conformer = stationary.identity(algorithm=Algorithm.IRMSD) - assert inchi is not None - assert conformer is not None - print( - f"stationary point {stationary.id} " - f"(order={stationary.order}, is_valid={stationary.is_valid}):" - ) - print(f" auto-attached InChI: {inchi.value}") - print(f" auto-attached conformer id: {conformer.value}") - for extra in inchi.identity_extras: - print(f" auto-attached {extra.attribute}: {extra.value}") - - # Algorithms outside `events.AUTO_MANAGED_IDENTITY_ALGORITHMS` (e.g. a - # human-curated SMILES) have to be attached explicitly. - smiles = IdentityRow.find_or_create(db, algorithm=Algorithm.RDKIT_SMILES, value="O") - db.add(StationaryIdentityLink(stationary_id=stationary.id, identity_id=smiles.id)) - db.commit() - print(f" manually-attached SMILES: {smiles.value}") - - return stationary - - -def record_trajectory( - db: Database, - opt_calc: CalculationRow, - input_geo: GeometryRow, - optimized_geo: GeometryRow, -) -> None: - """Record a trajectory for the optimization path (input -> optimized).""" - trajectory = TrajectoryRow() - db.add(trajectory) - db.commit() - db.add_all( - [ - TrajectoryGeometryLink.create(input_geo, trajectory, index=[0]), - TrajectoryGeometryLink.create(optimized_geo, trajectory, index=[1]), - CalculationTrajectoryLink.create(opt_calc, trajectory, role=Role.OUTPUT), - ] - ) - db.commit() - print( - f"trajectory {trajectory.id} holds {len(trajectory.geometry_links)} geometries" - ) - - -def run_frequency_calc( - db: Database, - model: ModelRow, - optimized_geo: GeometryRow, - stationary: StationaryPointRow, - rng: Generator, -) -> None: - """Attach a Gradient and Hessian to the optimized geometry. - - `is_valid` on `stationary` is recomputed automatically whenever a Hessian is - added for its geometry (see `autostorage.events.validate_geometry_orders`), - reflecting whether the point's declared `order` agrees with the (here, - randomly generated, so not necessarily consistent) Hessian's order. - """ - freq_calc = CalculationRow(model=model, calc_type=CalcType.FREQUENCY) - db.add(freq_calc) - db.add(CalculationGeometryLink.create(freq_calc, optimized_geo, role=Role.INPUT)) - db.commit() - - n = optimized_geo.atom_count - gradient = GradientRow( - calculation=freq_calc, geometry=optimized_geo, value=np.zeros(3 * n) - ) - hessian = HessianRow( - calculation=freq_calc, - geometry=optimized_geo, - value=rng.uniform(size=(3 * n, 3 * n)), - ) - db.add_all([gradient, hessian]) - db.commit() - print( - f"frequency calculation {freq_calc.id}: " - f"{len(hessian.harmonic_frequencies)} harmonic frequencies, " - f"order={hessian.order}" - ) - - found_hessian = HessianRow.query(db, geo=optimized_geo, model=model) - assert found_hessian is not None - assert found_hessian.id == hessian.id - print(f"looked up Hessian {found_hessian.id} by geometry + model") - print( - f"stationary point {stationary.id} is_valid is now " - f"{stationary.is_valid} (declared order={stationary.order}, " - f"Hessian order={hessian.order})" - ) - - -def record_transition_point(db: Database, model: ModelRow) -> None: - """Record a distinct stationary point with order=1. - - `order` (minimum vs. saddle point) is a property of the point itself, - independent of any reaction step/stage. - """ - ts_geo = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 1.0], [0, 0, 0], [1.0, 0.3, 0]]), - charge=0, - spin=0, - ) - ts_calc = CalculationRow(model=model, calc_type=CalcType.OPT_TS) - db.add_all( - [ - ts_calc, - ts_geo, - CalculationGeometryLink.create(ts_calc, ts_geo, role=Role.OUTPUT), - ] - ) - ts_stationary = StationaryPointRow(calculation=ts_calc, geometry=ts_geo, order=1) - db.add(ts_stationary) - db.commit() - print( - f"transition-point candidate {ts_stationary.id} " - f"recorded with order={ts_stationary.order}" - ) - - -def record_energy_and_validation( - db: Database, model: ModelRow, optimized_geo: GeometryRow -) -> None: - """Attach a single-point energy, look it up, and add a standalone validation.""" - energy_calc = CalculationRow(model=model, calc_type=CalcType.ENERGY) - db.add(energy_calc) - db.add(CalculationGeometryLink.create(energy_calc, optimized_geo, role=Role.INPUT)) - db.commit() - - energy = EnergyRow(calculation=energy_calc, geometry=optimized_geo, value=-76.02) - db.add(energy) - db.commit() - - found_energy = EnergyRow.query(db, geo=optimized_geo, model=model) - assert found_energy is not None - print(f"energy calculation {energy_calc.id}: E = {found_energy.value} Hartree") - - # A validation record, standing on its own with no reaction step attached. - validation = ValidationRow( - calculation=energy_calc, - method="geometry_check", - extras={"rmsd_to_input": 0.05}, - ) - db.add(validation) - db.commit() - print(f"validation {validation.id} ({validation.method}): {validation.extras}") - - -def main() -> None: - """Run the example workflow.""" - rng = np.random.default_rng(seed=0) - - with Database("stationary_points.db") as db: - model = ModelRow.find_or_create( - db, program="orca", method="b3lyp", basis="def2-svp" - ) - opt_calc, input_geo, optimized_geo = optimize_water(db, model) - stationary = mark_stationary_point(db, opt_calc, optimized_geo) - record_trajectory(db, opt_calc, input_geo, optimized_geo) - run_frequency_calc(db, model, optimized_geo, stationary, rng) - record_transition_point(db, model) - record_energy_and_validation(db, model, optimized_geo) - - -if __name__ == "__main__": - main() diff --git a/examples/transition.py b/examples/transition.py new file mode 100644 index 0000000..89280d2 --- /dev/null +++ b/examples/transition.py @@ -0,0 +1,212 @@ +"""Example demonstrating reaction network features: stages, steps, and validations.""" + +import numpy as np + +from autostorage import Database +from autostorage.models import ( + CalculationGeometryLink, + CalculationRow, + EnergyRow, + GeometryRow, + ModelRow, + Role, + StageRow, + StageStationaryLink, + StationaryPointRow, + StepRow, + StepValidationLink, + ValidationRow, +) + +# Create database +db = Database("example_transition.db", echo=True) + +with db.session() as session: + # Reactant: pyramidal NH3 (a minimum, order=0) + reactant_geom = GeometryRow( + symbols=["N", "H", "H", "H"], + coordinates=np.array( + [ + [0.0, 0.0, 0.3], + [0.94, 0.0, -0.1], + [-0.47, 0.814, -0.1], + [-0.47, -0.814, -0.1], + ] + ), + charge=0, + spin=0, + ) + + # Transition state: planar NH3, the point of inversion (order=1) + ts_geom = GeometryRow( + symbols=["N", "H", "H", "H"], + coordinates=np.array( + [ + [0.0, 0.0, 0.0], + [0.99, 0.0, 0.0], + [-0.495, 0.857, 0.0], + [-0.495, -0.857, 0.0], + ] + ), + charge=0, + spin=0, + ) + + # Product: mirror-image pyramidal NH3 (a minimum, degenerate with the reactant) + product_geom = GeometryRow( + symbols=["N", "H", "H", "H"], + coordinates=np.array( + [ + [0.0, 0.0, -0.3], + [0.94, 0.0, -0.1], + [-0.47, 0.814, -0.1], + [-0.47, -0.814, -0.1], + ] + ), + charge=0, + spin=0, + ) + session.add_all([reactant_geom, ts_geom, product_geom]) + session.flush() + + # Calculation model used to locate all three stationary points. + model = ModelRow(calc_type="opt", program="psi4", method="b3lyp", basis="6-31g*") + session.add(model) + session.flush() + + calc_reactant = CalculationRow(model_id=model.id) + calc_ts = CalculationRow(model_id=model.id) + calc_product = CalculationRow(model_id=model.id) + session.add_all([calc_reactant, calc_ts, calc_product]) + session.flush() + + # Link each geometry to the optimization that produced it as an output. + session.add_all( + [ + CalculationGeometryLink( + geometry_id=reactant_geom.id, + calculation_id=calc_reactant.id, + role=Role.OUTPUT, + ), + CalculationGeometryLink( + geometry_id=ts_geom.id, + calculation_id=calc_ts.id, + role=Role.OUTPUT, + ), + CalculationGeometryLink( + geometry_id=product_geom.id, + calculation_id=calc_product.id, + role=Role.OUTPUT, + ), + ] + ) + + # Energy results for each point (Hartree); the TS sits above both minima. + energy_reactant = EnergyRow( + geometry_id=reactant_geom.id, calculation_id=calc_reactant.id, value=-56.19513 + ) + energy_ts = EnergyRow( + geometry_id=ts_geom.id, calculation_id=calc_ts.id, value=-56.17021 + ) + energy_product = EnergyRow( + geometry_id=product_geom.id, calculation_id=calc_product.id, value=-56.19513 + ) + session.add_all([energy_reactant, energy_ts, energy_product]) + + # Stationary points: reactant/product are minima (order=0), the TS is a + # first-order saddle point (order=1). + stat_reactant = StationaryPointRow( + geometry_id=reactant_geom.id, calculation_id=calc_reactant.id, order=0 + ) + stat_ts = StationaryPointRow( + geometry_id=ts_geom.id, calculation_id=calc_ts.id, order=1 + ) + stat_product = StationaryPointRow( + geometry_id=product_geom.id, calculation_id=calc_product.id, order=0 + ) + session.add_all([stat_reactant, stat_ts, stat_product]) + session.flush() + + # `add_inchi_identities_before_flush` already attached an InChI identity to + # each stationary point during the flush above; reactant and product share + # the same chemical identity since this inversion is degenerate. + assert stat_reactant.identities + assert {i.value for i in stat_reactant.identities} == { + i.value for i in stat_product.identities + } + + # Stages: one per stationary point, plus the transition-state stage. + stage_reactant = StageRow(is_ts=False) + stage_product = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage_reactant, stage_product, stage_ts]) + session.flush() + + # Link each stationary point to its stage. + session.add_all( + [ + StageStationaryLink( + stationary_id=stat_reactant.id, stage_id=stage_reactant.id + ), + StageStationaryLink( + stationary_id=stat_product.id, stage_id=stage_product.id + ), + StageStationaryLink(stationary_id=stat_ts.id, stage_id=stage_ts.id), + ] + ) + session.flush() + + # Elementary step connecting reactant and product through the transition state. + # StepRow requires stage_id1 < stage_id2. + step = StepRow( + stage_id1=min(stage_reactant.id, stage_product.id), + stage_id2=max(stage_reactant.id, stage_product.id), + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step) + session.flush() + + # Validate the step with an IRC calculation confirming it connects the + # reactant and product minima, then link the validation to the step. + irc_model = ModelRow( + calc_type="irc", program="psi4", method="b3lyp", basis="6-31g*" + ) + session.add(irc_model) + session.flush() + + irc_calc = CalculationRow(model_id=irc_model.id) + session.add(irc_calc) + session.flush() + + # The IRC starts from the TS and descends to the reactant and product minima. + session.add_all( + [ + CalculationGeometryLink( + geometry_id=ts_geom.id, calculation_id=irc_calc.id, role=Role.INPUT + ), + CalculationGeometryLink( + geometry_id=reactant_geom.id, + calculation_id=irc_calc.id, + role=Role.OUTPUT, + ), + CalculationGeometryLink( + geometry_id=product_geom.id, + calculation_id=irc_calc.id, + role=Role.OUTPUT, + ), + ] + ) + + irc_validation = ValidationRow( + calculation_id=irc_calc.id, + method="irc", + extras={"connects": "reactant-product"}, + ) + session.add(irc_validation) + session.flush() + + session.add(StepValidationLink(step_id=step.id, validation_id=irc_validation.id)) + + # Commit all changes + session.commit() diff --git a/migrations/README b/migrations/README deleted file mode 100644 index 98e4f9c..0000000 --- a/migrations/README +++ /dev/null @@ -1 +0,0 @@ -Generic single-database configuration. \ No newline at end of file diff --git a/migrations/env.py b/migrations/env.py deleted file mode 100644 index bd47ba2..0000000 --- a/migrations/env.py +++ /dev/null @@ -1,86 +0,0 @@ -import os -from logging.config import fileConfig - -from sqlalchemy import engine_from_config -from sqlalchemy import pool -from sqlmodel import SQLModel - -from alembic import context - -# Importing this module registers every `table=True` model (via its -# `from .events import *` / `from .models import *`) onto `SQLModel.metadata`, -# which is what makes autogenerate below able to see the schema at all. -import autostorage.database # noqa: F401 - -# this is the Alembic Config object, which provides -# access to the values within the .ini file in use. -config = context.config - -# Interpret the config file for Python logging. -# This line sets up loggers basically. -if config.config_file_name is not None: - fileConfig(config.config_file_name) - -# Allow the target database to be selected at runtime (e.g. `AUTOSTORAGE_DATABASE_URL= -# sqlite:///path/to.db alembic upgrade head`) instead of hardcoding it in alembic.ini. -if database_url := os.environ.get("AUTOSTORAGE_DATABASE_URL"): - config.set_main_option("sqlalchemy.url", database_url) - -target_metadata = SQLModel.metadata - -# other values from the config, defined by the needs of env.py, -# can be acquired: -# my_important_option = config.get_main_option("my_important_option") -# ... etc. - - -def run_migrations_offline() -> None: - """Run migrations in 'offline' mode. - - This configures the context with just a URL - and not an Engine, though an Engine is acceptable - here as well. By skipping the Engine creation - we don't even need a DBAPI to be available. - - Calls to context.execute() here emit the given string to the - script output. - - """ - url = config.get_main_option("sqlalchemy.url") - context.configure( - url=url, - target_metadata=target_metadata, - literal_binds=True, - dialect_opts={"paramstyle": "named"}, - ) - - with context.begin_transaction(): - context.run_migrations() - - -def run_migrations_online() -> None: - """Run migrations in 'online' mode. - - In this scenario we need to create an Engine - and associate a connection with the context. - - """ - connectable = engine_from_config( - config.get_section(config.config_ini_section, {}), - prefix="sqlalchemy.", - poolclass=pool.NullPool, - ) - - with connectable.connect() as connection: - context.configure( - connection=connection, target_metadata=target_metadata - ) - - with context.begin_transaction(): - context.run_migrations() - - -if context.is_offline_mode(): - run_migrations_offline() -else: - run_migrations_online() diff --git a/migrations/script.py.mako b/migrations/script.py.mako deleted file mode 100644 index 1101630..0000000 --- a/migrations/script.py.mako +++ /dev/null @@ -1,28 +0,0 @@ -"""${message} - -Revision ID: ${up_revision} -Revises: ${down_revision | comma,n} -Create Date: ${create_date} - -""" -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -${imports if imports else ""} - -# revision identifiers, used by Alembic. -revision: str = ${repr(up_revision)} -down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)} -branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} -depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} - - -def upgrade() -> None: - """Upgrade schema.""" - ${upgrades if upgrades else "pass"} - - -def downgrade() -> None: - """Downgrade schema.""" - ${downgrades if downgrades else "pass"} diff --git a/migrations/versions/4230f63e365f_add_geometry_hash.py b/migrations/versions/4230f63e365f_add_geometry_hash.py deleted file mode 100644 index bde7f26..0000000 --- a/migrations/versions/4230f63e365f_add_geometry_hash.py +++ /dev/null @@ -1,73 +0,0 @@ -"""add geometry hash - -Revision ID: 4230f63e365f -Revises: e50de3129c84 -Create Date: 2026-07-23 19:58:37.026493 - -""" -import json -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel.sql.sqltypes - -from autostorage.models import _geometry_hash -from autostorage.types import CompressedArrayTypeDecorator - - -# revision identifiers, used by Alembic. -revision: str = '4230f63e365f' -down_revision: Union[str, Sequence[str], None] = 'e50de3129c84' -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - """Upgrade schema. - - Added nullable, backfilled per-row from existing `symbols`/`coordinates`/ - `charge`/`spin`, then tightened to NOT NULL + unique -- adding a NOT NULL - unique column directly against a populated table isn't possible in one - step. If the database already contains bit-identical duplicate geometries - (only possible from before this constraint existed), the final - `create_unique_constraint` below will fail; such rows must be - deduplicated (and their foreign-key references remapped, as - `autostorage.merge` does for a cross-database merge) before this - migration can complete. - """ - op.add_column( - 'geometry', - sa.Column('geometry_hash', sqlmodel.sql.sqltypes.AutoString(), nullable=True), - ) - - bind = op.get_bind() - decoder = CompressedArrayTypeDecorator() - rows = bind.execute( - sa.text('SELECT id, symbols, coordinates, charge, spin FROM geometry') - ).fetchall() - for row_id, symbols_json, coordinates_blob, charge, spin in rows: - symbols = json.loads(symbols_json) - coordinates = decoder.process_result_value(coordinates_blob, bind.dialect) - geometry_hash = _geometry_hash(symbols, coordinates, charge, spin) - bind.execute( - sa.text('UPDATE geometry SET geometry_hash = :hash WHERE id = :id'), - {'hash': geometry_hash, 'id': row_id}, - ) - - # SQLite can't ALTER a column's nullability or add a constraint directly; - # batch mode recreates the table under the hood. - with op.batch_alter_table('geometry') as batch_op: - batch_op.alter_column( - 'geometry_hash', - existing_type=sqlmodel.sql.sqltypes.AutoString(), - nullable=False, - ) - batch_op.create_unique_constraint('unique_geometry_hash', ['geometry_hash']) - - -def downgrade() -> None: - """Downgrade schema.""" - with op.batch_alter_table('geometry') as batch_op: - batch_op.drop_constraint('unique_geometry_hash', type_='unique') - batch_op.drop_column('geometry_hash') diff --git a/migrations/versions/e50de3129c84_add_null_safe_indexes_reverse_lookup_.py b/migrations/versions/e50de3129c84_add_null_safe_indexes_reverse_lookup_.py deleted file mode 100644 index 34a6784..0000000 --- a/migrations/versions/e50de3129c84_add_null_safe_indexes_reverse_lookup_.py +++ /dev/null @@ -1,136 +0,0 @@ -"""add null-safe indexes, reverse-lookup indexes, timestamps, and calculation status - -Revision ID: e50de3129c84 -Revises: faa9f50bc029 -Create Date: 2026-07-22 23:00:01.762231 - -""" -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel.sql.sqltypes - - -# revision identifiers, used by Alembic. -revision: str = 'e50de3129c84' -down_revision: Union[str, Sequence[str], None] = 'faa9f50bc029' -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - """Upgrade schema.""" - # ### commands auto generated by Alembic - please adjust! ### - op.add_column('calculation', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('calculation', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('calculation', sa.Column('status', sa.Enum('pending', 'running', 'succeeded', 'failed', name='calcstatus'), nullable=True)) - op.add_column('calculation', sa.Column('error_message', sqlmodel.sql.sqltypes.AutoString(), nullable=True)) - op.create_index('ix_calculation_geometry_link_calculation_id', 'calculation_geometry_link', ['calculation_id'], unique=False) - op.create_index('ix_calculation_trajectory_link_calculation_id', 'calculation_trajectory_link', ['calculation_id'], unique=False) - op.add_column('energy', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('energy', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('geometry', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('geometry', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('gradient', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('gradient', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('hessian', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('hessian', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('identity', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('identity', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('identity_extras', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('identity_extras', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.create_index(op.f('ix_identity_extras_identity_id'), 'identity_extras', ['identity_id'], unique=False) - op.add_column('model', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('model', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - # Autogenerate can't reflect SQLite expression-based indexes, so this and - # 'unq_step_stages_null_safe' below aren't detected/created automatically. - op.create_index( - 'unique_model_null_safe', - 'model', - [ - 'program', - sa.text("coalesce(program_version, '')"), - 'method', - sa.text("coalesce(basis, '')"), - ], - unique=True, - ) - op.add_column('stage', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('stage', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.create_index('ix_stationary_identity_link_identity_id', 'stationary_identity_link', ['identity_id'], unique=False) - op.add_column('stationary_point', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('stationary_point', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.create_index('ix_stationary_stage_link_stage_id', 'stationary_stage_link', ['stage_id'], unique=False) - op.add_column('step', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('step', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.create_index('ix_step_stage_id1', 'step', ['stage_id1'], unique=False) - op.create_index('ix_step_stage_id2', 'step', ['stage_id2'], unique=False) - op.create_index('ix_step_stage_id_ts', 'step', ['stage_id_ts'], unique=False) - op.create_index( - 'unq_step_stages_null_safe', - 'step', - ['stage_id1', 'stage_id2', sa.text('coalesce(stage_id_ts, 0)')], - unique=True, - ) - op.create_index('ix_step_validation_link_validation_id', 'step_validation_link', ['validation_id'], unique=False) - op.add_column('trajectory', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('trajectory', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.create_index('ix_trajectory_geometry_link_trajectory_id', 'trajectory_geometry_link', ['trajectory_id'], unique=False) - op.add_column('validation', sa.Column('created_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.add_column('validation', sa.Column('updated_at', sa.DateTime(), server_default=sa.text('(CURRENT_TIMESTAMP)'), nullable=False)) - op.alter_column('validation', 'calculation_id', - existing_type=sa.INTEGER(), - nullable=False) - op.create_index(op.f('ix_validation_calculation_id'), 'validation', ['calculation_id'], unique=False) - # ### end Alembic commands ### - - -def downgrade() -> None: - """Downgrade schema.""" - # ### commands auto generated by Alembic - please adjust! ### - op.drop_index(op.f('ix_validation_calculation_id'), table_name='validation') - op.alter_column('validation', 'calculation_id', - existing_type=sa.INTEGER(), - nullable=True) - op.drop_column('validation', 'updated_at') - op.drop_column('validation', 'created_at') - op.drop_index('ix_trajectory_geometry_link_trajectory_id', table_name='trajectory_geometry_link') - op.drop_column('trajectory', 'updated_at') - op.drop_column('trajectory', 'created_at') - op.drop_index('ix_step_validation_link_validation_id', table_name='step_validation_link') - op.drop_index('unq_step_stages_null_safe', table_name='step') - op.drop_index('ix_step_stage_id_ts', table_name='step') - op.drop_index('ix_step_stage_id2', table_name='step') - op.drop_index('ix_step_stage_id1', table_name='step') - op.drop_column('step', 'updated_at') - op.drop_column('step', 'created_at') - op.drop_index('ix_stationary_stage_link_stage_id', table_name='stationary_stage_link') - op.drop_column('stationary_point', 'updated_at') - op.drop_column('stationary_point', 'created_at') - op.drop_index('ix_stationary_identity_link_identity_id', table_name='stationary_identity_link') - op.drop_column('stage', 'updated_at') - op.drop_column('stage', 'created_at') - op.drop_index('unique_model_null_safe', table_name='model') - op.drop_column('model', 'updated_at') - op.drop_column('model', 'created_at') - op.drop_index(op.f('ix_identity_extras_identity_id'), table_name='identity_extras') - op.drop_column('identity_extras', 'updated_at') - op.drop_column('identity_extras', 'created_at') - op.drop_column('identity', 'updated_at') - op.drop_column('identity', 'created_at') - op.drop_column('hessian', 'updated_at') - op.drop_column('hessian', 'created_at') - op.drop_column('gradient', 'updated_at') - op.drop_column('gradient', 'created_at') - op.drop_column('geometry', 'updated_at') - op.drop_column('geometry', 'created_at') - op.drop_column('energy', 'updated_at') - op.drop_column('energy', 'created_at') - op.drop_index('ix_calculation_trajectory_link_calculation_id', table_name='calculation_trajectory_link') - op.drop_index('ix_calculation_geometry_link_calculation_id', table_name='calculation_geometry_link') - op.drop_column('calculation', 'error_message') - op.drop_column('calculation', 'status') - op.drop_column('calculation', 'updated_at') - op.drop_column('calculation', 'created_at') - # ### end Alembic commands ### diff --git a/migrations/versions/faa9f50bc029_baseline_current_schema.py b/migrations/versions/faa9f50bc029_baseline_current_schema.py deleted file mode 100644 index 88d270d..0000000 --- a/migrations/versions/faa9f50bc029_baseline_current_schema.py +++ /dev/null @@ -1,225 +0,0 @@ -"""baseline: current schema - -Revision ID: faa9f50bc029 -Revises: -Create Date: 2026-07-22 22:49:12.069197 - -""" -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel.sql.sqltypes - -import autostorage.types - - -# revision identifiers, used by Alembic. -revision: str = 'faa9f50bc029' -down_revision: Union[str, Sequence[str], None] = None -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - """Upgrade schema.""" - # ### commands auto generated by Alembic - please adjust! ### - op.create_table('geometry', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('symbols', sa.JSON(), nullable=True), - sa.Column('coordinates', autostorage.types.CompressedArrayTypeDecorator(), nullable=True), - sa.Column('charge', sa.Integer(), nullable=False), - sa.Column('spin', sa.Integer(), nullable=False), - sa.PrimaryKeyConstraint('id') - ) - op.create_table('identity', - sa.Column('algorithm', sa.Enum('RDKIT_INCHI', 'RDKIT_SMILES', 'IRMSD', name='algorithm'), nullable=False), - sa.Column('value', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('kind', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('id', sa.Integer(), nullable=False), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('kind', 'algorithm', 'value', name='unique_identity') - ) - op.create_table('model', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('program', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('program_version', sqlmodel.sql.sqltypes.AutoString(), nullable=True), - sa.Column('method', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('basis', sqlmodel.sql.sqltypes.AutoString(), nullable=True), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('program', 'program_version', 'method', 'basis', name='unique_model') - ) - op.create_table('stage', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('is_ts', sa.Boolean(), nullable=False), - sa.PrimaryKeyConstraint('id') - ) - op.create_table('trajectory', - sa.Column('id', sa.Integer(), nullable=False), - sa.PrimaryKeyConstraint('id') - ) - op.create_table('calculation', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('model_id', sa.Integer(), nullable=False), - sa.Column('calc_type', sa.Enum('optimization', 'transition_optimization', 'conformer_search', 'scan', 'intrinsic_reaction_coordinate', 'minimum_energy_path_search', 'energy', 'gradient', 'frequency', 'thermochemistry', 'undefined', name='calctype'), nullable=True), - sa.Column('input_provenance', sa.JSON(), nullable=True), - sa.Column('output_provenance', sa.JSON(), nullable=True), - sa.ForeignKeyConstraint(['model_id'], ['model.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_index(op.f('ix_calculation_model_id'), 'calculation', ['model_id'], unique=False) - op.create_table('identity_extras', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('identity_id', sa.Integer(), nullable=False), - sa.Column('attribute', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('value', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.ForeignKeyConstraint(['identity_id'], ['identity.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_table('step', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('stage_id1', sa.Integer(), nullable=False), - sa.Column('stage_id2', sa.Integer(), nullable=False), - sa.Column('stage_id_ts', sa.Integer(), nullable=True), - sa.Column('is_barrierless', sa.Boolean(), nullable=False), - sa.CheckConstraint('stage_id1 < stage_id2', name='chk_stage_order'), - sa.ForeignKeyConstraint(['stage_id1'], ['stage.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['stage_id2'], ['stage.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['stage_id_ts'], ['stage.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id'), - sa.UniqueConstraint('stage_id1', 'stage_id2', 'stage_id_ts', name='unq_step_stages') - ) - op.create_table('trajectory_geometry_link', - sa.Column('geometry_id', sa.Integer(), nullable=False), - sa.Column('trajectory_id', sa.Integer(), nullable=False), - sa.Column('index', sa.JSON(), nullable=True), - sa.ForeignKeyConstraint(['geometry_id'], ['geometry.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['trajectory_id'], ['trajectory.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('geometry_id', 'trajectory_id') - ) - op.create_table('calculation_geometry_link', - sa.Column('geometry_id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=False), - sa.Column('role', sa.Enum('input', 'output', name='role'), nullable=True), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['geometry_id'], ['geometry.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('geometry_id', 'calculation_id') - ) - op.create_table('calculation_trajectory_link', - sa.Column('trajectory_id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=False), - sa.Column('role', sa.Enum('input', 'output', name='role'), nullable=True), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['trajectory_id'], ['trajectory.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('trajectory_id', 'calculation_id') - ) - op.create_table('energy', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('geometry_id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=False), - sa.Column('value', sa.Float(), nullable=False), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['geometry_id'], ['geometry.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_index(op.f('ix_energy_calculation_id'), 'energy', ['calculation_id'], unique=False) - op.create_index(op.f('ix_energy_geometry_id'), 'energy', ['geometry_id'], unique=False) - op.create_table('gradient', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('geometry_id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=False), - sa.Column('value', autostorage.types.CompressedArrayTypeDecorator(), nullable=True), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['geometry_id'], ['geometry.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_index(op.f('ix_gradient_calculation_id'), 'gradient', ['calculation_id'], unique=False) - op.create_index(op.f('ix_gradient_geometry_id'), 'gradient', ['geometry_id'], unique=False) - op.create_table('hessian', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('geometry_id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=False), - sa.Column('value', autostorage.types.CompressedArrayTypeDecorator(), nullable=True), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['geometry_id'], ['geometry.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_index(op.f('ix_hessian_calculation_id'), 'hessian', ['calculation_id'], unique=False) - op.create_index(op.f('ix_hessian_geometry_id'), 'hessian', ['geometry_id'], unique=False) - op.create_table('stationary_point', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('geometry_id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=False), - sa.Column('order', sa.Integer(), nullable=False), - sa.Column('is_pseudo', sa.Boolean(), nullable=False), - sa.Column('is_valid', sa.Boolean(), nullable=False), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['geometry_id'], ['geometry.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_index(op.f('ix_stationary_point_calculation_id'), 'stationary_point', ['calculation_id'], unique=False) - op.create_index(op.f('ix_stationary_point_geometry_id'), 'stationary_point', ['geometry_id'], unique=False) - op.create_table('validation', - sa.Column('id', sa.Integer(), nullable=False), - sa.Column('calculation_id', sa.Integer(), nullable=True), - sa.Column('method', sqlmodel.sql.sqltypes.AutoString(), nullable=False), - sa.Column('extras', sa.JSON(), nullable=True), - sa.ForeignKeyConstraint(['calculation_id'], ['calculation.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('id') - ) - op.create_table('stationary_identity_link', - sa.Column('stationary_id', sa.Integer(), nullable=False), - sa.Column('identity_id', sa.Integer(), nullable=False), - sa.ForeignKeyConstraint(['identity_id'], ['identity.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['stationary_id'], ['stationary_point.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('stationary_id', 'identity_id') - ) - op.create_table('stationary_stage_link', - sa.Column('stationary_id', sa.Integer(), nullable=False), - sa.Column('stage_id', sa.Integer(), nullable=False), - sa.ForeignKeyConstraint(['stage_id'], ['stage.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['stationary_id'], ['stationary_point.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('stationary_id', 'stage_id') - ) - op.create_table('step_validation_link', - sa.Column('step_id', sa.Integer(), nullable=False), - sa.Column('validation_id', sa.Integer(), nullable=False), - sa.ForeignKeyConstraint(['step_id'], ['step.id'], ondelete='CASCADE'), - sa.ForeignKeyConstraint(['validation_id'], ['validation.id'], ondelete='CASCADE'), - sa.PrimaryKeyConstraint('step_id', 'validation_id') - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - """Downgrade schema.""" - # ### commands auto generated by Alembic - please adjust! ### - op.drop_table('step_validation_link') - op.drop_table('stationary_stage_link') - op.drop_table('stationary_identity_link') - op.drop_table('validation') - op.drop_index(op.f('ix_stationary_point_geometry_id'), table_name='stationary_point') - op.drop_index(op.f('ix_stationary_point_calculation_id'), table_name='stationary_point') - op.drop_table('stationary_point') - op.drop_index(op.f('ix_hessian_geometry_id'), table_name='hessian') - op.drop_index(op.f('ix_hessian_calculation_id'), table_name='hessian') - op.drop_table('hessian') - op.drop_index(op.f('ix_gradient_geometry_id'), table_name='gradient') - op.drop_index(op.f('ix_gradient_calculation_id'), table_name='gradient') - op.drop_table('gradient') - op.drop_index(op.f('ix_energy_geometry_id'), table_name='energy') - op.drop_index(op.f('ix_energy_calculation_id'), table_name='energy') - op.drop_table('energy') - op.drop_table('calculation_trajectory_link') - op.drop_table('calculation_geometry_link') - op.drop_table('trajectory_geometry_link') - op.drop_table('step') - op.drop_table('identity_extras') - op.drop_index(op.f('ix_calculation_model_id'), table_name='calculation') - op.drop_table('calculation') - op.drop_table('trajectory') - op.drop_table('stage') - op.drop_table('model') - op.drop_table('identity') - op.drop_table('geometry') - # ### end Alembic commands ### diff --git a/pixi.lock b/pixi.lock index c73184d..39a1f33 100644 --- a/pixi.lock +++ b/pixi.lock @@ -15,7 +15,7 @@ environments: - https://pypi.org/simple packages: linux-64: - - conda: https://conda.anaconda.org/avcopan/noarch/automol-0.0.19-pyh4616a5c_0.conda + - conda: https://conda.anaconda.org/avcopan/noarch/automol-0.0.22-pyh4616a5c_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-1.2.0-hed03a55_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/brotli-bin-1.2.0-hb03c661_1.conda @@ -81,6 +81,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/libxcrypt-4.4.36-hd590300_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.11.0-py313h6c470cf_1.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/nauty-2.7.1-h7f98852_2.tar.bz2 - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.6-hdb14827_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.5.1-py313hf6604e3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda @@ -94,6 +95,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pycairo-1.29.0-py313h3f29d12_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pydantic-core-2.46.4-py313h843e2db_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/pynauty-2.8.8.1-py313h07c4f96_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.14-h6add32d_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hda471dd_3.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/qhull-2020.2-h434a139_5.conda @@ -118,7 +120,6 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/asttokens-3.0.2-pyhd8ed1ab_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/ca-certificates-2026.6.17-hbd8a1cb_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/charset-normalizer-3.4.9-pyhd8ed1ab_0.conda - - conda: https://conda.anaconda.org/conda-forge/noarch/click-8.4.2-pyhc90fa1f_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/comm-0.2.3-pyhe01879c_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cpython-3.13.14-py313hd8ed1ab_100.conda - conda: https://conda.anaconda.org/conda-forge/noarch/cycler-0.12.1-pyhcf101f3_2.conda @@ -169,7 +170,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/typing_extensions-4.16.0-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/tzdata-2026c-h151e31d_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.8.2-pyhd8ed1ab_0.conda - - conda_source: autostorage[dcc1f98c] @ . + - conda_source: autostorage[e06c168f] @ . - pypi: https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/21/0e/8459ca4413e1a21a06c97d134bfaf18adfd27cea068813dc0faae06cbf00/cssselect2-0.9.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/2c/1b/291dd75b5ed494eb484801f39e6572192d322e6cd3d68fab8a6dde743c48/graphrc-1.3.7-py3-none-any.whl @@ -191,7 +192,7 @@ environments: - https://pypi.org/simple packages: linux-64: - - conda: https://conda.anaconda.org/avcopan/noarch/automol-0.0.19-pyh4616a5c_0.conda + - conda: https://conda.anaconda.org/avcopan/noarch/automol-0.0.22-pyh4616a5c_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/astroid-3.3.11-py313h78bf25f_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/aws-c-auth-0.10.4-hb7a77c6_1.conda @@ -285,6 +286,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/libzlib-1.3.2-h25fd6f3_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/markupsafe-3.0.3-py313h3dea7bd_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/matplotlib-base-3.11.0-py313h6c470cf_1.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/nauty-2.7.1-h7f98852_2.tar.bz2 - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.6-hdb14827_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/numpy-2.5.1-py313hf6604e3_0.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/openjpeg-2.5.4-h55fea9a_0.conda @@ -299,6 +301,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/linux-64/pthread-stubs-0.4-hb9d3cd8_1002.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pycairo-1.29.0-py313h3f29d12_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pydantic-core-2.46.4-py313h843e2db_0.conda + - conda: https://conda.anaconda.org/conda-forge/linux-64/pynauty-2.8.8.1-py313h07c4f96_2.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.14-h6add32d_100_cp313.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyyaml-6.0.3-py313h3dea7bd_1.conda - conda: https://conda.anaconda.org/conda-forge/linux-64/pyzmq-27.1.0-py312hda471dd_3.conda @@ -479,7 +482,7 @@ environments: - conda: https://conda.anaconda.org/conda-forge/noarch/virtualenv-21.6.1-pyhcf101f3_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/wcwidth-0.8.2-pyhd8ed1ab_0.conda - conda: https://conda.anaconda.org/conda-forge/noarch/zipp-4.1.0-pyhcf101f3_0.conda - - conda_source: autostorage[dcc1f98c] @ . + - conda_source: autostorage[e06c168f] @ . - pypi: https://files.pythonhosted.org/packages/07/6c/aa3f2f849e01cb6a001cd8554a88d4c77c5c1a31c95bdf1cf9301e6d9ef4/defusedxml-0.7.1-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/21/0e/8459ca4413e1a21a06c97d134bfaf18adfd27cea068813dc0faae06cbf00/cssselect2-0.9.0-py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/23/06/b4f06ca7afb5d9e942c642980c308ffcfa1fa0e8b0a3ddbec78483ef1614/keepachangelog-2.0.0-py3-none-any.whl @@ -497,9 +500,9 @@ environments: - pypi: https://files.pythonhosted.org/packages/f4/24/2a3e3df732393fed8b3ebf2ec078f05546de641fe1b667ee316ec1dcf3b7/webencodings-0.5.1-py2.py3-none-any.whl - pypi: https://files.pythonhosted.org/packages/f8/80/9a966abdb505c5042e14136d4fc6a0ce0a3896d77ca02ab87b288e5417db/xyzgraph-1.6.13-py3-none-any.whl packages: -- conda: https://conda.anaconda.org/avcopan/noarch/automol-0.0.19-pyh4616a5c_0.conda - sha256: 746344a1e334be0f7caa4ca6d5077c7e97b9e508555beb88bf38c78f5f936e66 - md5: a165c2c6913f554f0e0adaa2cfa052ae +- conda: https://conda.anaconda.org/avcopan/noarch/automol-0.0.22-pyh4616a5c_0.conda + sha256: 9839ad1bdd6242d7c1aa7eae8b3d0a8e0119abb9061ecc59390214cc58704484 + md5: 73e696e6d068b272e37994dbd8f1502e depends: - python >=3.12 - python * @@ -510,8 +513,8 @@ packages: - pyparsing >=3.3.2 - rdkit >=2025 - scipy >=1.13 - size: 21750 - timestamp: 1784305892476 + size: 24446 + timestamp: 1787935826323 - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda build_number: 20 sha256: 1dd3fffd892081df9726d7eb7e0dea6198962ba775bd88842135a4ddb4deb3c9 @@ -862,7 +865,7 @@ packages: license: Apache-2.0 license_family: APACHE purls: - - pkg:pypi/coverage?source=compressed-mapping + - pkg:pypi/coverage?source=hash-mapping run_exports: {} size: 402264 timestamp: 1784150199793 @@ -1021,7 +1024,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/greenlet?source=compressed-mapping + - pkg:pypi/greenlet?source=hash-mapping run_exports: {} size: 267125 timestamp: 1782524630615 @@ -2016,10 +2019,21 @@ packages: license: PSF-2.0 license_family: PSF purls: - - pkg:pypi/matplotlib?source=compressed-mapping + - pkg:pypi/matplotlib?source=hash-mapping run_exports: {} size: 8974852 timestamp: 1782829528425 +- conda: https://conda.anaconda.org/conda-forge/linux-64/nauty-2.7.1-h7f98852_2.tar.bz2 + sha256: 3520ca6684b931267c074caf3ae61d3fc237c42284c6f41fc331811161c60b20 + md5: 33888e561e4aabeca26c7579275573e7 + depends: + - libgcc-ng >=9.3.0 + license: Apache-2.0 + license_family: APACHE + purls: [] + run_exports: {} + size: 11528033 + timestamp: 1623850946580 - conda: https://conda.anaconda.org/conda-forge/linux-64/ncurses-6.6-hdb14827_0.conda sha256: fc89f74bbe362fb29fa3c037697a89bec140b346a2469a90f7936d1d7ea4d8a3 md5: fc21868a1a5aacc937e7a18747acb8a5 @@ -2050,7 +2064,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/numpy?source=compressed-mapping + - pkg:pypi/numpy?source=hash-mapping run_exports: weak: - numpy >=1.25,<3 @@ -2213,7 +2227,7 @@ packages: - openjpeg >=2.5.4,<3.0a0 license: HPND purls: - - pkg:pypi/pillow?source=compressed-mapping + - pkg:pypi/pillow?source=hash-mapping run_exports: {} size: 1077037 timestamp: 1782912080163 @@ -2293,6 +2307,22 @@ packages: run_exports: {} size: 1893749 timestamp: 1778084222642 +- conda: https://conda.anaconda.org/conda-forge/linux-64/pynauty-2.8.8.1-py313h07c4f96_2.conda + sha256: f3e355dd703a3430839a3f682c8979b6d7a155146a872f075381284f9c356c44 + md5: bacc0b34e83d6ebac6cc87840333c93c + depends: + - __glibc >=2.17,<3.0.a0 + - libgcc >=14 + - nauty >=2.7.1,<2.7.2.0a0 + - python >=3.13,<3.14.0a0 + - python_abi 3.13.* *_cp313 + license: GPL-3.0-or-later + license_family: GPL + purls: + - pkg:pypi/pynauty?source=hash-mapping + run_exports: {} + size: 200473 + timestamp: 1756355032363 - conda: https://conda.anaconda.org/conda-forge/linux-64/python-3.13.14-h6add32d_100_cp313.conda build_number: 100 sha256: f2146aff59ce4b571a8f1d1acf94f9bed6cc18ab5632d7dcc940fb48ecdeef99 @@ -2481,7 +2511,7 @@ packages: - __glibc >=2.17 license: MIT purls: - - pkg:pypi/ruff?source=compressed-mapping + - pkg:pypi/ruff?source=hash-mapping run_exports: {} size: 9333318 timestamp: 1784237314764 @@ -2520,7 +2550,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/scipy?source=compressed-mapping + - pkg:pypi/scipy?source=hash-mapping run_exports: {} size: 17171066 timestamp: 1781912954186 @@ -2912,7 +2942,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/alembic?source=compressed-mapping + - pkg:pypi/alembic?source=hash-mapping run_exports: {} size: 185355 timestamp: 1782460469542 @@ -2958,7 +2988,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/anyio?source=compressed-mapping + - pkg:pypi/anyio?source=hash-mapping run_exports: {} size: 164465 timestamp: 1783889660383 @@ -2971,7 +3001,7 @@ packages: license: Apache-2.0 license_family: APACHE purls: - - pkg:pypi/argcomplete?source=compressed-mapping + - pkg:pypi/argcomplete?source=hash-mapping run_exports: {} size: 46445 timestamp: 1782888485558 @@ -2985,7 +3015,7 @@ packages: license: Apache-2.0 license_family: Apache purls: - - pkg:pypi/asttokens?source=compressed-mapping + - pkg:pypi/asttokens?source=hash-mapping run_exports: {} size: 34639 timestamp: 1783975742052 @@ -3058,8 +3088,7 @@ packages: - cached_property >=1.5.2,<1.5.3.0a0 license: BSD-3-Clause license_family: BSD - purls: - - pkg:pypi/cached-property?source=compressed-mapping + purls: [] run_exports: {} size: 6836 timestamp: 1783242914545 @@ -3071,7 +3100,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/cached-property?source=compressed-mapping + - pkg:pypi/cached-property?source=hash-mapping run_exports: {} size: 14752 timestamp: 1783242913845 @@ -3094,7 +3123,7 @@ packages: - python >=3.10 license: ISC purls: - - pkg:pypi/certifi?source=compressed-mapping + - pkg:pypi/certifi?source=hash-mapping run_exports: {} size: 133877 timestamp: 1781719949728 @@ -3106,7 +3135,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/charset-normalizer?source=compressed-mapping + - pkg:pypi/charset-normalizer?source=hash-mapping run_exports: {} size: 61418 timestamp: 1783505332569 @@ -3135,7 +3164,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/click?source=compressed-mapping + - pkg:pypi/click?source=hash-mapping run_exports: {} size: 107155 timestamp: 1783085363526 @@ -3222,7 +3251,7 @@ packages: license: BSD-2-Clause license_family: BSD purls: - - pkg:pypi/decorator?source=compressed-mapping + - pkg:pypi/decorator?source=hash-mapping run_exports: {} size: 16102 timestamp: 1779115228886 @@ -3250,7 +3279,7 @@ packages: license: Apache-2.0 license_family: APACHE purls: - - pkg:pypi/distlib?source=compressed-mapping + - pkg:pypi/distlib?source=hash-mapping run_exports: {} size: 303705 timestamp: 1781320269259 @@ -3361,8 +3390,7 @@ packages: - python-multipart - uvicorn-standard license: MIT - purls: - - pkg:pypi/fastapi?source=compressed-mapping + purls: [] run_exports: {} size: 4845 timestamp: 1784297339154 @@ -3379,7 +3407,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/fastapi-cli?source=compressed-mapping + - pkg:pypi/fastapi-cli?source=hash-mapping run_exports: {} size: 20791 timestamp: 1784211753962 @@ -3406,7 +3434,7 @@ packages: - uvicorn-standard >=0.12.0 license: MIT purls: - - pkg:pypi/fastapi?source=compressed-mapping + - pkg:pypi/fastapi?source=hash-mapping run_exports: {} size: 104185 timestamp: 1784297339154 @@ -3417,7 +3445,7 @@ packages: - python >=3.10 license: Unlicense purls: - - pkg:pypi/filelock?source=compressed-mapping + - pkg:pypi/filelock?source=hash-mapping run_exports: {} size: 74480 timestamp: 1784269447478 @@ -3559,7 +3587,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/hpack?source=compressed-mapping + - pkg:pypi/hpack?source=hash-mapping run_exports: {} size: 32884 timestamp: 1782283986153 @@ -3630,7 +3658,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/idna?source=compressed-mapping + - pkg:pypi/idna?source=hash-mapping run_exports: {} size: 163869 timestamp: 1781620148226 @@ -3677,7 +3705,7 @@ packages: license: Apache-2.0 license_family: APACHE purls: - - pkg:pypi/importlib-metadata?source=compressed-mapping + - pkg:pypi/importlib-metadata?source=hash-mapping run_exports: {} size: 34766 timestamp: 1779714582554 @@ -3732,7 +3760,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/ipykernel?source=compressed-mapping + - pkg:pypi/ipykernel?source=hash-mapping run_exports: {} size: 138635 timestamp: 1781101665847 @@ -3757,7 +3785,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/ipython?source=compressed-mapping + - pkg:pypi/ipython?source=hash-mapping run_exports: {} size: 658801 timestamp: 1782482716877 @@ -3783,7 +3811,7 @@ packages: - python license: Apache-2.0 AND MIT purls: - - pkg:pypi/jedi?source=compressed-mapping + - pkg:pypi/jedi?source=hash-mapping run_exports: {} size: 2715215 timestamp: 1782251948616 @@ -4066,7 +4094,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/platformdirs?source=compressed-mapping + - pkg:pypi/platformdirs?source=hash-mapping run_exports: {} size: 26308 timestamp: 1779972894916 @@ -4143,7 +4171,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/pycparser?source=compressed-mapping + - pkg:pypi/pycparser?source=hash-mapping run_exports: {} size: 55886 timestamp: 1779293633166 @@ -4198,7 +4226,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/pydantic-settings?source=compressed-mapping + - pkg:pypi/pydantic-settings?source=hash-mapping run_exports: {} size: 52920 timestamp: 1781884990751 @@ -4294,7 +4322,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/pytest?source=compressed-mapping + - pkg:pypi/pytest?source=hash-mapping run_exports: {} size: 306724 timestamp: 1782127176429 @@ -4375,7 +4403,7 @@ packages: license: Apache-2.0 license_family: APACHE purls: - - pkg:pypi/python-multipart?source=compressed-mapping + - pkg:pypi/python-multipart?source=hash-mapping run_exports: {} size: 38132 timestamp: 1780610429919 @@ -4434,7 +4462,7 @@ packages: license: Apache-2.0 license_family: APACHE purls: - - pkg:pypi/requests?source=compressed-mapping + - pkg:pypi/requests?source=hash-mapping run_exports: {} size: 68709 timestamp: 1778851103479 @@ -4466,7 +4494,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/rich-toolkit?source=compressed-mapping + - pkg:pypi/rich-toolkit?source=hash-mapping run_exports: {} size: 34824 timestamp: 1784024058693 @@ -4516,7 +4544,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/setuptools?source=compressed-mapping + - pkg:pypi/setuptools?source=hash-mapping run_exports: {} size: 642081 timestamp: 1783619174976 @@ -4577,7 +4605,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/soupsieve?source=compressed-mapping + - pkg:pypi/soupsieve?source=hash-mapping run_exports: {} size: 38802 timestamp: 1779635534390 @@ -4810,7 +4838,7 @@ packages: - python license: MIT purls: - - pkg:pypi/tomlkit?source=compressed-mapping + - pkg:pypi/tomlkit?source=hash-mapping run_exports: {} size: 49054 timestamp: 1784287707171 @@ -4850,7 +4878,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/traitlets?source=compressed-mapping + - pkg:pypi/traitlets?source=hash-mapping run_exports: {} size: 115158 timestamp: 1780507822178 @@ -4904,7 +4932,7 @@ packages: license: PSF-2.0 license_family: PSF purls: - - pkg:pypi/typing-extensions?source=compressed-mapping + - pkg:pypi/typing-extensions?source=hash-mapping run_exports: {} size: 52631 timestamp: 1783002732887 @@ -4958,7 +4986,7 @@ packages: license: BSD-3-Clause license_family: BSD purls: - - pkg:pypi/uvicorn?source=compressed-mapping + - pkg:pypi/uvicorn?source=hash-mapping run_exports: {} size: 57699 timestamp: 1783518883165 @@ -5007,7 +5035,7 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/wcwidth?source=compressed-mapping + - pkg:pypi/wcwidth?source=hash-mapping run_exports: {} size: 132415 timestamp: 1782771807703 @@ -5020,19 +5048,17 @@ packages: license: MIT license_family: MIT purls: - - pkg:pypi/zipp?source=compressed-mapping + - pkg:pypi/zipp?source=hash-mapping run_exports: {} size: 24190 timestamp: 1779159948016 -- conda_source: autostorage[dcc1f98c] @ . +- conda_source: autostorage[e06c168f] @ . variants: target_platform: noarch depends: - python >=3.12 - python * - - click >=8.0 - - matplotlib-base >=3.8 - - pint >=0.25.2 + - pynauty >=2.8.8.1 - sqlmodel >=0.0.31 host_packages: - conda: https://conda.anaconda.org/conda-forge/linux-64/_openmp_mutex-4.5-20_gnu.conda diff --git a/pixi.toml b/pixi.toml index b0a2d72..c550230 100644 --- a/pixi.toml +++ b/pixi.toml @@ -6,7 +6,7 @@ preview = ["pixi-build"] [dependencies] python = ">=3.12,<3.14" autostorage = { path = "." } -automol = "==0.0.19" # local:false +automol = "==0.0.22" # local:false # automol = { path = "../automol" } # local:true matplotlib-base = ">=3.8" @@ -70,8 +70,6 @@ lint = "ruff check . --fix" types = "ty check" imports = "lint-imports" test = "pytest" -# Apply database migrations (set AUTOSTORAGE_DATABASE_URL to target a specific DB) -migrate = "alembic upgrade head" # Run pre-commit hooks (see lefthook.yaml) pre-commit = "lefthook run pre-commit --all-files" local-pre-commit = "lefthook run local-pre-commit --all-files" diff --git a/pyproject.toml b/pyproject.toml index 92a4340..eeca2d9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,18 +7,12 @@ authors = [ ] requires-python = ">= 3.12" dependencies = [ - "automol==0.0.19", - "click>=8.0", - "matplotlib>=3.8", - "pint>=0.25.2", + "automol==0.0.22", + "pynauty>=2.8.8.1", "sqlmodel>=0.0.31", - "irmsd>=0.1.1", "stereomolgraph>=0.0.22b0", ] -[project.scripts] -autostorage-merge = "autostorage.database:merge_databases" - [build-system] build-backend = "uv_build" requires = ["uv_build<0.12"] @@ -32,7 +26,6 @@ module-name = "autostorage" exclude = [ "docs", "**/*.ipynb", - "migrations", ] [tool.ruff.lint] @@ -71,12 +64,10 @@ exclude_type_checking_imports = true name = "Autostorage Layering" type = "layers" layers = [ - "autostorage.utils", "autostorage.database", - "autostorage.merge", "autostorage.events", "autostorage.models", - "autostorage.types | autostorage.exc", + "autostorage.types", ] # Pytest configurations diff --git a/schema/full_schema.pintora b/schema/full_schema.pintora new file mode 100644 index 0000000..a6d6f4d --- /dev/null +++ b/schema/full_schema.pintora @@ -0,0 +1,149 @@ +erDiagram + @param fill #FFFFFF + @param stroke #000000 + @param labelBackground #FFFFFF + @param fontSize 24 + @param fontFamily "Georgia" + @param layoutDirection "LR" + @param entityPaddingX 30 + @param entityPaddingY 30 + @param ranksep 300 + @param edgesep 50 + + MODEL { + int id PK + string program UK + string program_version UK + string method UK + string basis UK + } + CALCULATION { + int id PK + int model_id FK + enum calc_type + json input_provenance + json output_provenance + } + GEOMETRY { + int id PK + json symbols + json coordinates + int charge + int spin + } + CALCULATION_GEOMETRY_LINK { + int geometry_id FK + int calculation_id FK + enum role + } + TRAJECTORY { + int id PK + } + CALCULATION_TRAJECTORY_LINK { + int trajectory_id FK + int calculation_id FK + enum role + } + TRAJECTORY_GEOMETRY_LINK { + int geometry_id FK + int calculation_id FK + json index + } + ENERGY { + int id PK + int geometry_id FK + int calculation_id FK + int value + } + GRADIENT { + int id PK + int geometry_id FK + int calculation_id FK + json value + } + HESSIAN { + int id PK + int geometry_id FK + int calculation_id FK + json value + } + STATIONARY_POINT { + int id PK + int geometry_id FK + int calculation_id FK + int order + bool is_pseudo + bool is_valid + } + IDENTITY { + int id PK + string kind + string algorithm + string value + } + IDENTITY_EXTRA { + int id PK + int identity_id FK + string attribute + string value + } + STATIONARY_IDENTITY_LINK { + int identity_id FK + int stationary_id FK + } + STAGE { + int id PK + bool is_ts + } + STATIONARY_STAGE_LINK { + int stationary_id FK + int stage_id FK + } + STEP { + int id PK + int stage_id1 FK + int stage_id2 FK + int stage_id_ts FK + bool is_barrierless + } + VALIDATION { + int id PK + int calculation_id FK + str method + json extras + } + STEP_VALIDATION_LINK { + int step_id FK + int validation_id FK + } + %%Calculation + MODEL }o--|| CALCULATION : specifies + CALCULATION_GEOMETRY_LINK }o--|| CALCULATION : links + CALCULATION_TRAJECTORY_LINK }o--|| CALCULATION : links + CALCULATION ||..o{ ENERGY : produces + CALCULATION ||..o{ GRADIENT : produces + CALCULATION ||..o{ HESSIAN : produces + %%Geometry + CALCULATION_GEOMETRY_LINK }o--|| GEOMETRY : links + TRAJECTORY_GEOMETRY_LINK }o--|| GEOMETRY : links + ENERGY }o..|| GEOMETRY : "evaluated from" + GRADIENT }o..|| GEOMETRY : "evaluated from" + HESSIAN }o..|| GEOMETRY : "evaluated from" + %%Trajectory + CALCULATION_TRAJECTORY_LINK }o--|| TRAJECTORY : "links" + TRAJECTORY_GEOMETRY_LINK }o--|| TRAJECTORY : "links" + %%Stationary + STATIONARY_POINT }o..|| GEOMETRY : "has structure" + STATIONARY_IDENTITY_LINK }o--|| STATIONARY_POINT : "links" + STATIONARY_STAGE_LINK }o--|| STATIONARY_POINT : "links" + %%Identity + STATIONARY_IDENTITY_LINK }o--|| IDENTITY : "links" + IDENTITY ||--o{ IDENTITY_EXTRA : "has" + %%Stage + STATIONARY_STAGE_LINK }o--|| STAGE : "links" + STAGE ||..o{ STEP : "node of" + %%Step + STEP_VALIDATION_LINK }o--|| STEP : "links" + %%Validation + STEP_VALIDATION_LINK }o--|| VALIDATION : "links" + CALCULATION ||--o{ VALIDATION : "provides" diff --git a/schema/full_schema.svg b/schema/full_schema.svg new file mode 100644 index 0000000..b5dd768 --- /dev/null +++ b/schema/full_schema.svg @@ -0,0 +1 @@ +specifieslinkslinksproducesproducesproduceslinkslinksevaluated fromevaluated fromevaluated fromlinkslinkshas structurelinkslinkslinkshaslinksnode oflinkslinksprovidesMODELPKintidUKstringprogramUKstringprogram_versionUKstringmethodUKstringbasisCALCULATIONPKintidFKintmodel_idenumcalc_typejsoninput_provenancejsonoutput_provenanceGEOMETRYPKintidjsonsymbolsjsoncoordinatesintchargeintspinCALCULATION_GEOMETRY_LINKFKintgeometry_idFKintcalculation_idenumroleTRAJECTORYPKintidCALCULATION_TRAJECTORY_LINKFKinttrajectory_idFKintcalculation_idenumroleTRAJECTORY_GEOMETRY_LINKFKintgeometry_idFKintcalculation_idjsonindexENERGYPKintidFKintgeometry_idFKintcalculation_idintvalueGRADIENTPKintidFKintgeometry_idFKintcalculation_idjsonvalueHESSIANPKintidFKintgeometry_idFKintcalculation_idjsonvalueSTATIONARY_POINTPKintidFKintgeometry_idFKintcalculation_idintorderboolis_pseudoboolis_validIDENTITYPKintidstringkindstringalgorithmstringvalueIDENTITY_EXTRAPKintidFKintidentity_idstringattributestringvalueSTATIONARY_IDENTITY_LINKFKintidentity_idFKintstationary_idSTAGEPKintidboolis_tsSTATIONARY_STAGE_LINKFKintstationary_idFKintstage_idSTEPPKintidFKintstage_id1FKintstage_id2FKintstage_id_tsboolis_barrierlessVALIDATIONPKintidFKintcalculation_idstrmethodjsonextrasSTEP_VALIDATION_LINKFKintstep_idFKintvalidation_id \ No newline at end of file diff --git a/schema/simplified_schema.pintora b/schema/simplified_schema.pintora new file mode 100644 index 0000000..5fa34de --- /dev/null +++ b/schema/simplified_schema.pintora @@ -0,0 +1,121 @@ +@pre +@style { + .calculation { + backgroundColor: #44BB99; + } + .model { + backgroundColor: #DDDDDD; + } + .geometry { + backgroundColor: #99DDFF; + } + .trajectory { + backgroundColor: #BBCC33; + } + .result { + backgroundColor: #77AADD; + } + .stationary { + backgroundColor: #EE8866; + } + .stage { + backgroundColor: #FFAABB; + } + .step { + backgroundColor: #AA4499; + } +} +@endpre + +erDiagram + %% Universal styling + @param stroke #000000 + @param labelBackground #FFFFFF + @param fontSize 36 + @param fontFamily "Georgia" + @param entityPaddingX 30 + @param entityPaddingY 30 + @param ranksep 150 + @param edgesep 175 + + %% Style binds + @bindClass entity-model model + @bindClass entity-calculation calculation + @bindClass entity-geometry geometry + @bindClass entity-trajectory trajectory + @bindClass entity-energy result + @bindClass entity-gradient result + @bindClass entity-hessian result + @bindClass entity-stationary_point stationary + @bindClass entity-identity stationary + @bindClass entity-identity_extra stationary + @bindClass entity-stage stage + @bindClass entity-step step + @bindClass entity-validation step + + %% Attributes + model { + string program + string method + string basis + } + calculation { + enum calc_type + json input_provenance + json output_provenance + } + geometry { + json symbols + blob coordinates + int charge + int spin + } + energy { + } + gradient { + } + hessian { + } + stationary_point { + int order + bool is_pseudo + } + identity { + } + identity_extra { + } + stage { + bool is_ts + } + step { + bool is_barrierless + } + validation { + } + + + %% Entity Relationships + model }o--|| calculation : "specifies" + + calculation }o--o{ geometry : "I/O" + calculation }o--o{ trajectory : "I/O" + calculation ||..o{ energy : "produces" + calculation ||..o{ gradient : "produces" + calculation ||..o{ hessian : "produces" + calculation ||..o{ stationary_point : "produces" + calculation ||..o{ validation : "performs" + + geometry }o--o{ trajectory : "member of" + energy }o..|| geometry : "from" + gradient }o..|| geometry : "from" + hessian }o..|| geometry : "from" + + stationary_point }o..|| geometry : "has" + stationary_point }o--o{ identity : "has" + stationary_point }o--o{ stage : "species of" + + identity ||--o{ identity_extra : "has" + + stage ||..o{ step : "node of" + + validation }o--o{ step : "validates" \ No newline at end of file diff --git a/schema/simplified_schema.png b/schema/simplified_schema.png new file mode 100644 index 0000000..90f3407 Binary files /dev/null and b/schema/simplified_schema.png differ diff --git a/schema/simplified_schema.svg b/schema/simplified_schema.svg new file mode 100644 index 0000000..82b6c16 --- /dev/null +++ b/schema/simplified_schema.svg @@ -0,0 +1 @@ +specifiesI/OI/Oproducesproducesproducesproducesperformsmember offromfromfromhashasspecies ofhasnode ofvalidatesmodelstringprogramstringmethodstringbasiscalculationenumcalc_typejsoninput_provenancejsonoutput_provenancegeometryjsonsymbolsblobcoordinatesintchargeintspinenergygradienthessianstationary_pointintorderboolis_pseudoidentityidentity_extrastageboolis_tsstepboolis_barrierlessvalidationtrajectory \ No newline at end of file diff --git a/src/autostorage/__init__.py b/src/autostorage/__init__.py index 793579c..e54fe91 100644 --- a/src/autostorage/__init__.py +++ b/src/autostorage/__init__.py @@ -2,14 +2,15 @@ __version__ = "0.0.12" -from . import exc, merge, types, utils +from . import events, types from .database import Database -from .merge import MergeReport from .models import ( CalculationGeometryLink, CalculationRow, + CalculationTrajectoryLink, EnergyRow, GeometryRow, + GeometryTrajectoryLink, GradientRow, HessianRow, IdentityExtraRow, @@ -21,21 +22,20 @@ TrajectoryRow, ValidationRow, ) -from .types import CalcStatus, CalcType, Role +from .types import Role __all__ = [ - "CalcStatus", - "CalcType", "CalculationGeometryLink", "CalculationRow", + "CalculationTrajectoryLink", "Database", "EnergyRow", "GeometryRow", + "GeometryTrajectoryLink", "GradientRow", "HessianRow", "IdentityExtraRow", "IdentityRow", - "MergeReport", "ModelRow", "Role", "StageRow", @@ -43,8 +43,6 @@ "StepRow", "TrajectoryRow", "ValidationRow", - "exc", - "merge", + "events", "types", - "utils", ] diff --git a/src/autostorage/database.py b/src/autostorage/database.py index 2189907..3f3bfd5 100644 --- a/src/autostorage/database.py +++ b/src/autostorage/database.py @@ -1,29 +1,18 @@ -"""SQLite database connection and session management.""" +"""SQLite database connection management.""" import json -from collections.abc import Iterable, Iterator -from contextlib import contextmanager from functools import partial from pathlib import Path -from types import TracebackType -from typing import Self -import click -from sqlalchemy import event -from sqlalchemy import select as sa_select -from sqlalchemy.exc import MultipleResultsFound, NoResultFound -from sqlmodel import Session, SQLModel, create_engine -from sqlmodel.sql.expression import Select, SelectOfScalar +from sqlalchemy import create_engine, event +from sqlalchemy.orm import Session +from sqlmodel import SQLModel # Ensure all modules are loaded with the database -from .events import * # noqa: F403 -from .merge import MergeReport -from .merge import merge_databases as _merge_databases +from . import events # noqa: F401 from .models import * # noqa: F403 -type SelectStatement[T] = Select[T] | SelectOfScalar[T] - -__all__ = ["Database", "Select", "SelectOfScalar", "SelectStatement"] +__all__ = ["Database"] class Database: @@ -36,8 +25,6 @@ class Database: Path to SQLite database file. engine SQLAlchemy engine instance. - _session - Persistent database session. """ def __init__(self, path: str | Path, *, echo: bool = False) -> None: @@ -73,181 +60,18 @@ def _set_sqlite_pragma(dbapi_connection, _connection_record) -> None: # noqa: A cursor.close() SQLModel.metadata.create_all(self.engine) - self._session: Session = Session(self.engine) - - @contextmanager - def session(self) -> Iterator[Session]: - """Yield the persistent database session. - - Note - ---- - This yields the single, long-lived `Session` created in `__init__`, - not a fresh session per call — rows returned from queries stay - attached, so lazy-loaded relationships keep working after the `with` - block exits. As a result, a `Database` instance is not safe for - concurrent use by multiple threads: `check_same_thread=False` only - allows the underlying DBAPI connection to be used from a different - thread than it was created on (e.g. a single background worker), it - does not make the `Session` itself thread-safe. - """ - try: - yield self._session - except Exception: - self._session.rollback() - raise - - def add[RowT: SQLModel](self, row: RowT) -> None: - """Add row to session. - - Note - ---- - The row is not validated or written to the database until the next `flush()` or - `commit()`. Integrity/shape errors (unique constraints, the shape event - listeners) raise there, which may be far removed from this call. - """ - with self.session() as session: - session.add(row) - - def add_all[RowT: SQLModel](self, rows: Iterable[RowT]) -> None: - """Add multiple rows to session. - - Note - ---- - Bulk `add()`. Same staging-only caveat applies. - """ - with self.session() as session: - session.add_all(rows) - - def merge[RowT: SQLModel](self, row: RowT) -> RowT: - """Merge row into current session and commit, returning the merged row.""" - with self.session() as session: - merged = session.merge(row) - session.commit() - return merged - - def merge_from(self, source_db: "Database", *, commit: bool = True) -> MergeReport: - """Merge another database's contents into this one. - - Unlike `merge()` (a same-session upsert of a single row), this - copies every row from a separate `source_db` into this database, - remapping ids/foreign keys and deduplicating content-unique rows. - - See Also - -------- - autostorage.merge - """ - return _merge_databases(target=self, source=source_db, commit=commit) - def flush(self) -> None: - """Flush pending changes to the database without committing. + def session(self) -> Session: + """Return a fresh `Session` bound to this database's engine. Note ---- - Unlike `commit()`, this doesn't trigger SQLAlchemy's default `expire_on_commit` - behavior, so an already-loaded object whose row was removed by a DB-level - `ondelete="CASCADE"` during this flush would be read back stale. `expire_all()` - forces those objects to reload (or raise) on next access instead. + A new `Session` is created per call; use it as a context manager + (`with database.session() as session: ...`) to close it on exit. + Nothing is committed automatically — call `session.commit()` explicitly. """ - with self.session() as session: - session.flush() - session.expire_all() - - def commit(self) -> None: - """Commit database session.""" - with self.session() as session: - session.commit() - - def delete[RowT: SQLModel](self, row: RowT) -> None: - """Delete row from database.""" - with self.session() as session: - session.delete(row) - session.commit() - - def get_or_none[RowT: SQLModel]( - self, model: type[RowT], row_id: int - ) -> RowT | None: - """Get row from database, returning `None` instead of raising on a miss.""" - with self.session() as session: - return session.get(model, row_id) - - def get[RowT: SQLModel](self, model: type[RowT], row_id: int) -> RowT: - """Get row from database.""" - row = self.get_or_none(model, row_id) - if row is not None: - return row - - msg = f"{model} with {row_id = } not found." - raise LookupError(msg) - - def exec_first[RowT: SQLModel](self, stmt: SelectStatement[RowT]) -> RowT | None: - """Return the first match to a statement.""" - with self.session() as session: - return session.exec(stmt).first() - - def exec_one[RowT: SQLModel](self, stmt: SelectStatement[RowT]) -> RowT: - """Return the single match to a statement.""" - with self.session() as session: - try: - return session.exec(stmt).one() - except NoResultFound as exc: - msg = f"No row found matching {stmt}." - raise LookupError(msg) from exc - except MultipleResultsFound as exc: - msg = f"Multiple rows found matching {stmt}." - raise LookupError(msg) from exc - - def exec_all[RowT: SQLModel](self, stmt: SelectStatement[RowT]) -> list[RowT]: - """Return all matches to a statement.""" - with self.session() as session: - return list(session.exec(stmt).all()) - - def exists[RowT: SQLModel](self, stmt: SelectStatement[RowT]) -> bool: - """Return whether any row matches a statement. - - Executes as a single `EXISTS` subquery instead of `exec_first`, so a matching - row is never materialized. - """ - with self.session() as session: - return bool( - session.exec(sa_select(stmt.exists())).scalar() # ty:ignore[no-matching-overload] - ) + return Session(self.engine) def close(self) -> None: """Close the database connection.""" self.engine.dispose() - - def __enter__(self) -> Self: - """Enter a `with Database(...) as db:` block.""" - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc_value: object, - traceback: TracebackType | None, - ) -> None: - """Roll back on exception, then close the database connection.""" - del exc_value, traceback - if exc_type is not None: - self._session.rollback() - self.close() - - -@click.command(name="autostorage-merge") -@click.argument("target", type=click.Path(path_type=Path)) -@click.argument("source", type=click.Path(path_type=Path)) -def merge_databases(target: Path, source: Path) -> None: - """Merge one on-disk database's contents into another, as a CLI command. - - Installed as the ``autostorage-merge`` console script (see ``[project.scripts]`` in - ``pyproject.toml``); also runnable via ``python -m autostorage.database``. - """ - with Database(target) as target_db, Database(source) as source_db: - report = _merge_databases(target=target_db, source=source_db) - - click.echo(f"Copied: {report.copied}") - click.echo(f"Reused: {report.reused}") - - -if __name__ == "__main__": - merge_databases() diff --git a/src/autostorage/events.py b/src/autostorage/events.py index 1401ca3..a53bcf5 100644 --- a/src/autostorage/events.py +++ b/src/autostorage/events.py @@ -1,405 +1,389 @@ """SQLAlchemy ORM event listeners for validation and auto-managed identities.""" -from collections.abc import Iterable from typing import Any -import numpy as np -from automol import Algorithm, geom -from sqlalchemy import event, tuple_ +from automol import Identity +from automol.ident import Algorithm +from sqlalchemy import event from sqlalchemy.engine import Connection -from sqlalchemy.orm import Mapper, object_session -from sqlalchemy.orm.attributes import flag_modified, get_history -from sqlmodel import Integer, Session, cast, func, select +from sqlalchemy.orm import Mapper, Session -from .exc import DataIntegrityError, ResultShapeError from .models import ( GeometryRow, + GeometryTrajectoryLink, GradientRow, HessianRow, IdentityExtraRow, IdentityRow, - StageRow, StationaryPointRow, StepRow, - _geometry_hash, ) -def _resolve_geometry( - target: GradientRow | HessianRow | StationaryPointRow, -) -> GeometryRow | None: - """Return the target's geometry, resolving it via the session if unattached. - - Setting only `geometry_id` leaves `.geometry` unpopulated until the ORM syncs it. - """ - geometry = target.geometry - if geometry is None and target.geometry_id is not None: - session = object_session(target) - if session is not None: - geometry = session.get(GeometryRow, target.geometry_id) - return geometry - - -@event.listens_for(GradientRow, "before_insert") -@event.listens_for(GradientRow, "before_update") -def verify_gradient_shape( - mapper: Mapper, # noqa: ARG001 - connection: Connection, # noqa: ARG001 - target: GradientRow, -) -> None: - """Verify shape of the Gradient array before saving to the database.""" - geometry = _resolve_geometry(target) - if geometry is None: - return - - expected = (3 * geometry.atom_count,) - actual = np.shape(target.value) - - if actual != expected: - raise ResultShapeError(target, actual, expected) - - -@event.listens_for(HessianRow, "before_insert") -@event.listens_for(HessianRow, "before_update") -def verify_hessian_shape( - mapper: Mapper, # noqa: ARG001 +@event.listens_for(StepRow, "before_insert") +@event.listens_for(StepRow, "before_update") +def sort_step_stage_ids( + mapper: Mapper[StepRow], # noqa: ARG001 connection: Connection, # noqa: ARG001 - target: HessianRow, + target: StepRow, ) -> None: - """Verify shape of the Hessian matrix before saving to the database.""" - geometry = _resolve_geometry(target) - if geometry is None: - return - - expected_dim = 3 * geometry.atom_count - expected = (expected_dim, expected_dim) - actual = np.shape(target.value) - - if actual != expected: - raise ResultShapeError(target, actual, expected) + """Auto-sort stage_id1 and stage_id2 so stage_id1 < stage_id2.""" + if ( + target.stage_id1 is not None + and target.stage_id2 is not None + and target.stage_id1 > target.stage_id2 + ): + target.stage_id1, target.stage_id2 = target.stage_id2, target.stage_id1 -@event.listens_for(HessianRow, "before_update") -def invalidate_hessian_frequency_cache( - mapper: Mapper, # noqa: ARG001 +@event.listens_for(StepRow, "before_insert") +@event.listens_for(StepRow, "before_update") +def verify_step_barrierless_consistency( + mapper: Mapper[StepRow], # noqa: ARG001 connection: Connection, # noqa: ARG001 - target: HessianRow, + target: StepRow, ) -> None: - """Drop a cached `harmonic_frequencies` if `value` changed.""" - if get_history(target, "value").added: - target.__dict__.pop("harmonic_frequencies", None) + """Verify is_barrierless consistency with stage_id_ts.""" + if target.stage_id_ts is None: + if not target.is_barrierless: + msg = "Barrierless step (stage_id_ts=None) must have is_barrierless=True" + raise ValueError(msg) + elif target.is_barrierless: + msg = ( + "Step with transition state (stage_id_ts!=None) must have " + "is_barrierless=False" + ) + raise ValueError(msg) -def _recompute_geometry_stationary_validity( - geometry: GeometryRow, *, excluding: Iterable[HessianRow] = () +@event.listens_for(Session, "before_flush") +def verify_gradient_shapes_before_flush( + session: Session, + flush_context: Any, # noqa: ARG001, ANN401 + instances: Any, # noqa: ARG001, ANN401 ) -> None: - """Recompute `StationaryPointRow.is_valid` for a geometry from its Hessians. + """Verify gradient shapes match 3 * natoms for all gradients being flushed.""" + for obj in list(session.new) + list(session.dirty): + if not isinstance(obj, GradientRow): + continue - `excluding` skips Hessians pending deletion (still in `geometry.hessians` at - `before_flush` time since the DELETE hasn't been issued yet). - """ - excluded_ids = {id(h) for h in excluding} - hessians = [h for h in geometry.hessians if id(h) not in excluded_ids] - if not hessians: - return + if obj.geometry_id is None: + continue - # A shape-invalid `value` (still pending its own `verify_hessian_shape` - # before_insert check later in this same flush) can't have its order computed; - # skip it here rather than raising `vibrational_analysis`'s raw ValueError. - orders: set[int] = set() - for h in hessians: - try: - orders.add(h.order) - except ValueError: + # Load the geometry to get natoms + geometry_row = session.get(GeometryRow, obj.geometry_id) + if geometry_row is None: continue - if len(orders) > 1: - msg = f"Geometry Hessians do not agree on order. {orders = }." - raise DataIntegrityError(msg) - if orders and geometry.stationary_points: - expected_order = orders.pop() - for stationary in geometry.stationary_points: - stationary.is_valid = stationary.order == expected_order + natoms = len(geometry_row.symbols) + expected_shape = (3 * natoms,) + actual_shape = obj.value.shape + + if actual_shape != expected_shape: + msg = ( + f"Gradient shape {actual_shape} does not match expected " + f"shape {expected_shape} for geometry with {natoms} atoms" + ) + raise ValueError(msg) @event.listens_for(Session, "before_flush") -def revalidate_geometry_orders_on_insert_update( +def verify_hessian_shapes_before_flush( session: Session, - flush_context: Any, # noqa: ANN401, ARG001 - instances: Any, # noqa: ANN401, ARG001 + flush_context: Any, # noqa: ARG001, ANN401 + instances: Any, # noqa: ARG001, ANN401 ) -> None: - """Recompute order consensus for a geometry when a Stationary/Hessian changes. - - A session-level `before_flush` listener, not a per-instance `before_insert`/ - `before_update` mapper event: the recompute below mutates sibling - `StationaryPointRow`s that may already be clean going into this flush, and a - mapper event fires too late in the flush cycle for that mutation to be - included — SQLAlchemy silently drops it (the "Attribute history events... will - not result in database updates" warning) instead of writing it. - """ - candidates = ( - obj - for obj in session.new | session.dirty - if isinstance(obj, StationaryPointRow | HessianRow) - ) - geometries: dict[int, GeometryRow] = {} - for obj in candidates: - geometry = _resolve_geometry(obj) - if geometry is not None: - geometries[id(geometry)] = geometry + """Verify Hessian shapes match (3 * natoms, 3 * natoms) for all Hessians.""" + for obj in list(session.new) + list(session.dirty): + if not isinstance(obj, HessianRow): + continue + + if obj.geometry_id is None: + continue + + # Load the geometry to get natoms + geometry_row = session.get(GeometryRow, obj.geometry_id) + if geometry_row is None: + continue + + natoms = len(geometry_row.symbols) + expected_shape = (3 * natoms, 3 * natoms) + actual_shape = obj.value.shape - for geometry in geometries.values(): - _recompute_geometry_stationary_validity(geometry) + if actual_shape != expected_shape: + msg = ( + f"Hessian shape {actual_shape} does not match expected " + f"shape {expected_shape} for geometry with {natoms} atoms" + ) + raise ValueError(msg) @event.listens_for(Session, "before_flush") -def revalidate_geometry_orders_on_hessian_delete( +def verify_valid_stationary_has_hessian( session: Session, - flush_context: Any, # noqa: ANN401, ARG001 - instances: Any, # noqa: ANN401, ARG001 + flush_context: Any, # noqa: ARG001, ANN401 + instances: Any, # noqa: ARG001, ANN401 ) -> None: - """Recompute order consensus for a geometry when one of its Hessians is deleted.""" - deleted_hessians = [obj for obj in session.deleted if isinstance(obj, HessianRow)] - if not deleted_hessians: - return + """Verify that stationary points marked as valid have an associated Hessian.""" + for obj in list(session.new) + list(session.dirty): + if not isinstance(obj, StationaryPointRow): + continue - geometries = {h.geometry_id: h.geometry for h in deleted_hessians if h.geometry} - for geometry in geometries.values(): - excluding = [h for h in deleted_hessians if h.geometry_id == geometry.id] - _recompute_geometry_stationary_validity(geometry, excluding=excluding) + if not obj.is_valid: + continue + if obj.geometry_id is None: + continue -_IMMUTABLE_GEOMETRY_FIELDS = ("symbols", "coordinates") + # Load the geometry to check for hessians + geometry_row = session.get(GeometryRow, obj.geometry_id) + if geometry_row is None: + continue + + if not geometry_row.hessians: + msg = ( + f"StationaryPointRow cannot be marked as valid without a Hessian " + f"attached to its geometry (geometry_id={obj.geometry_id})" + ) + raise ValueError(msg) -@event.listens_for(GeometryRow, "before_update") -def verify_geometry_immutable_fields( - mapper: Mapper, # noqa: ARG001 +@event.listens_for(GeometryTrajectoryLink, "before_insert") +@event.listens_for(GeometryTrajectoryLink, "before_update") +def verify_trajectory_geometry_ndim_insert( + mapper: Mapper[GeometryTrajectoryLink], # noqa: ARG001 connection: Connection, # noqa: ARG001 - target: GeometryRow, + target: GeometryTrajectoryLink, ) -> None: - """Reject changes to `symbols`/`coordinates` on an already-persisted geometry. + """Ensure linked geometry's index length matches trajectory ndim.""" + if target.trajectory is None: + return - Otherwise an in-place edit could invalidate Gradient/Hessian shape checks - already run against it. - """ - for attr in _IMMUTABLE_GEOMETRY_FIELDS: - history = get_history(target, attr) - if history.added or history.deleted: - msg = f"GeometryRow.{attr} cannot be changed after insert." - raise DataIntegrityError(msg) + traj_ndim = target.trajectory.ndim + index_len = len(target.index) if target.index is not None else None + if index_len is not None and traj_ndim is not None and index_len != traj_ndim: + msg = ( + f"Geometry index length {index_len} does not match " + f"trajectory ndim {traj_ndim}" + ) + raise ValueError(msg) -@event.listens_for(GeometryRow, "before_insert") -@event.listens_for(GeometryRow, "before_update") -def compute_geometry_hash( - mapper: Mapper, # noqa: ARG001 - connection: Connection, # noqa: ARG001 - target: GeometryRow, -) -> None: - """Populate `geometry_hash` from bit-identical geometry fields before saving. + if traj_ndim is None and index_len is not None: + target.trajectory.ndim = index_len + elif index_len is None and traj_ndim is not None: + msg = f"Geometry index is missing but trajectory ndim is {traj_ndim}" + raise ValueError(msg) + + +def _find_or_create_identity(session: Session, identity: Identity) -> IdentityRow: + """Find or create an IdentityRow for the given Identity. - Written via `target.__dict__[...]` + `flag_modified`, not `target.geometry_hash = - ...`: `Geometry`'s `validate_assignment=True` pydantic config corrupts SQLAlchemy's - flush-time identity-key bookkeeping when a plain attribute assignment happens - inside a mapper event, breaking every GeometryRow insert. + Checks both the database and pending session inserts. """ - target.__dict__["geometry_hash"] = _geometry_hash( - target.symbols, target.coordinates, target.charge, target.spin + # First check the database + existing = ( + session.query(IdentityRow) + .filter_by( + kind=identity.kind, + algorithm=str(identity.algorithm), + value=identity.value, + ) + .first() ) - flag_modified(target, "geometry_hash") + # If not in database, check session.new for pending inserts + if existing is None: + for new_obj in session.new: + if ( + isinstance(new_obj, IdentityRow) + and new_obj.kind == identity.kind + and new_obj.algorithm == str(identity.algorithm) + and new_obj.value == identity.value + ): + existing = new_obj + break + + if existing is None: + # Create new identity row + new_identity = IdentityRow( + kind=identity.kind, + algorithm=str(identity.algorithm), + value=identity.value, + ) + session.add(new_identity) + return new_identity -# Identity algorithms managed here, so other code (e.g. `merge.py`) knows not to -# copy/attach them explicitly. -AUTO_MANAGED_IDENTITY_ALGORITHMS: frozenset[Algorithm] = frozenset( - {Algorithm.RDKIT_INCHI, Algorithm.IRMSD} -) + return existing @event.listens_for(Session, "before_flush") -def add_inchi_identities(session: Session, flush_context: Any, instances: Any) -> None: # noqa: ANN401, ARG001 - """Attach InChI and SMILES identities to new stationary point rows before flush.""" - pending_items = [] - inchi_lookups = [] - +def add_inchi_identities_before_flush( + session: Session, + flush_context: Any, # noqa: ARG001, ANN401 + instances: Any, # noqa: ARG001, ANN401 +) -> None: + """Automatically attach InChI identities to newly inserted stationary points.""" for obj in session.new: if not isinstance(obj, StationaryPointRow): continue - geometry = _resolve_geometry(obj) - if geometry is None: - continue - try: - inchi = IdentityRow.from_geometry( - geo=geometry, algorithm=Algorithm.RDKIT_INCHI - ) - pending_items.append((obj, inchi, geometry)) - inchi_lookups.append((inchi.algorithm, inchi.value)) - except ValueError: - continue - - if not pending_items: - return - - stmt = select(IdentityRow).where( - tuple_(IdentityRow.algorithm, IdentityRow.value).in_(inchi_lookups) # ty:ignore[invalid-argument-type] - ) - existing_rows = session.exec(stmt).all() - - identity_map = {(r.algorithm, r.value): r for r in existing_rows} - for obj, inchi, geometry in pending_items: - lookup_key = (inchi.algorithm, inchi.value) - existing = identity_map.get(lookup_key) - - if existing: - obj.identities.append(existing) + if obj.geometry_id is None or obj.identities: continue - obj.identities.append(inchi) - identity_map[lookup_key] = inchi - - try: - smiles = IdentityRow.from_geometry( - geometry, algorithm=Algorithm.RDKIT_SMILES - ) - smiles_extra = IdentityExtraRow( - identity=inchi, attribute="smiles", value=smiles.value - ) - - session.add(smiles_extra) - - except ValueError: + # Load the geometry + geometry_row = session.get(GeometryRow, obj.geometry_id) + if geometry_row is None: continue + # Generate InChI identity from geometry + identity = Identity.from_geometry(geometry_row, algorithm=Algorithm.RDKIT_INCHI) -def _matching_conformer_identity( - obj: StationaryPointRow, inchi: IdentityRow -) -> IdentityRow | None: - """Find the conformer identity of a geometric duplicate among InChI matches.""" - peers = [c for c in inchi.stationary_points if c is not obj] - if not peers: - return None + # Find or create the identity row + identity_row = _find_or_create_identity(session, identity) - geometry = _resolve_geometry(obj) - if geometry is None: - return None + # Link the identity to the stationary point + if identity_row not in obj.identities: + obj.identities.append(identity_row) - # Peers may have only `geometry_id` set (e.g. a bulk loader); resolve via session. - resolved_peers = [(c, _resolve_geometry(c)) for c in peers] - resolved_peers = [(c, g) for c, g in resolved_peers if g is not None] - if not resolved_peers: - return None - matches = geom.is_duplicate_conformer(geometry, [g for _, g in resolved_peers]) - match_idx = next((i for i, m in enumerate(matches) if m), None) - if match_idx is None: - return None +def _find_or_create_identity_extra( + session: Session, identity: IdentityRow, attribute: str, value: str +) -> IdentityExtraRow | None: + """Find or create an IdentityExtraRow. - matched_peer, _ = resolved_peers[match_idx] - return matched_peer.identity(kind=Algorithm.IRMSD.kind) + Checks both the database and pending session inserts. Returns None if the + extra already exists. + """ + # Check if the identity has an ID (already in database) + if identity.id is not None: + existing = ( + session.query(IdentityExtraRow) + .filter_by( + identity_id=identity.id, + attribute=attribute, + value=value, + ) + .first() + ) + if existing is not None: + return None # Already exists in database + + # Check session.new for pending inserts (both for this identity and in general) + for new_obj in session.new: + if ( + isinstance(new_obj, IdentityExtraRow) + and (new_obj.identity is identity or new_obj.identity_id == identity.id) + and new_obj.attribute == attribute + and new_obj.value == value + ): + return None # Already pending insertion + + # Create new extra, using the identity relationship + return IdentityExtraRow( + identity=identity, + attribute=attribute, + value=value, + ) @event.listens_for(Session, "before_flush") -def assign_conformer_ids(session: Session, flush_context: Any, instances: Any) -> None: # noqa: ANN401, ARG001 - """Assign a shared conformer-group identity to duplicate stationary points.""" - pending_items: list[tuple[StationaryPointRow, IdentityRow]] = [] - +def add_smiles_extras_before_flush( + session: Session, + flush_context: Any, # noqa: ARG001, ANN401 + instances: Any, # noqa: ARG001, ANN401 +) -> None: + """Automatically attach SMILES as IdentityExtraRow to stationary points.""" for obj in session.new: if not isinstance(obj, StationaryPointRow): continue - if obj.identity(kind=Algorithm.IRMSD.kind) is not None: - continue - inchi = obj.identity(algorithm=Algorithm.RDKIT_INCHI) - if inchi is not None: - pending_items.append((obj, inchi)) + if obj.geometry_id is None or not obj.identities: + continue - if not pending_items: - return + # Load the geometry + geometry_row = session.get(GeometryRow, obj.geometry_id) + if geometry_row is None: + continue - next_group_id: int | None = None + # Generate SMILES from geometry + try: + smiles_identity = Identity.from_geometry( + geometry_row, algorithm=Algorithm.RDKIT_SMILES + ) + smiles_value = smiles_identity.value + except Exception: # noqa: BLE001, S112 + # Skip if SMILES generation fails (e.g., invalid structure) + continue - for obj, inchi in pending_items: - match_ident = _matching_conformer_identity(obj, inchi) + # Get the InChI identity (should exist from add_inchi_identities_before_flush) + inchi_identity = next( + ( + ident + for ident in obj.identities + if ident.algorithm == str(Algorithm.RDKIT_INCHI) + ), + None, + ) - if match_ident is not None: - obj.identities.append(match_ident) + if inchi_identity is None: continue - if next_group_id is None: - # Assumes single-writer; concurrent writers rely on the DB's uniqueness - # constraint to fail one session's commit instead. - current_max = session.exec( - select(func.max(cast(IdentityRow.value, Integer))).where( - IdentityRow.kind == Algorithm.IRMSD.kind - ) - ).first() - next_group_id = (current_max or 0) + 1 - else: - next_group_id += 1 - - conformer = IdentityRow.from_value( - str(next_group_id), algorithm=Algorithm.IRMSD + # Find or create the SMILES extra + smiles_extra = _find_or_create_identity_extra( + session, inchi_identity, "rdkit_smiles", smiles_value ) - obj.identities.append(conformer) - -@event.listens_for(StepRow, "before_insert") -@event.listens_for(StepRow, "before_update") -def verify_stage_order_and_barrierless( - mapper: Mapper, # noqa: ARG001 - connection: Connection, # noqa: ARG001 - target: StepRow, -) -> None: - """Verify order of stage ids in a `StepRow` and whether it's barrierless.""" - stg_id1 = target.stage_id1 or target.stage1.id - stg_id2 = target.stage_id2 or target.stage2.id + if smiles_extra is not None: + session.add(smiles_extra) - if not stg_id1 or not stg_id2: - msg = "Cannot sort stage IDs; IDs aren't assigned to stages." - raise DataIntegrityError(msg) - if stg_id1 > stg_id2: - target.stage_id1, target.stage_id2 = stg_id2, stg_id1 +@event.listens_for(Session, "before_flush") +def add_hill_extras_before_flush( + session: Session, + flush_context: Any, # noqa: ARG001, ANN401 + instances: Any, # noqa: ARG001, ANN401 +) -> None: + """Automatically attach Hill formula as IdentityExtraRow to stationary points.""" + for obj in session.new: + if not isinstance(obj, StationaryPointRow): + continue - target.is_barrierless = not target.stage_id_ts + if obj.geometry_id is None or not obj.identities: + continue + # Load the geometry + geometry_row = session.get(GeometryRow, obj.geometry_id) + if geometry_row is None: + continue -def _resolve_stage(target: StepRow, id_attr: str, rel_attr: str) -> StageRow | None: - """Return one of a `StepRow`'s stages, resolving via session if unattached. + # Generate Hill formula from geometry + try: + hill_identity = Identity.from_geometry( + geometry_row, algorithm=Algorithm.HILL_FORMULA + ) + hill_value = hill_identity.value + except Exception: # noqa: BLE001, S112 + # Skip if Hill formula generation fails (e.g., invalid structure) + continue - Mirrors `_resolve_geometry`. - """ - stage = getattr(target, rel_attr) - if stage is None: - stage_id = getattr(target, id_attr) - if stage_id is not None: - session = object_session(target) - if session is not None: - stage = session.get(StageRow, stage_id) - return stage + # Get the InChI identity (should exist from add_inchi_identities_before_flush) + inchi_identity = next( + ( + ident + for ident in obj.identities + if ident.algorithm == str(Algorithm.RDKIT_INCHI) + ), + None, + ) + if inchi_identity is None: + continue -@event.listens_for(StepRow, "before_insert") -@event.listens_for(StepRow, "before_update") -def verify_stage_ts_consistency( - mapper: Mapper, # noqa: ARG001 - connection: Connection, # noqa: ARG001 - target: StepRow, -) -> None: - """Verify `is_ts` agreement between a `StepRow` and its linked stages. + # Find or create the Hill formula extra + hill_extra = _find_or_create_identity_extra( + session, inchi_identity, "hill_formula", hill_value + ) - `stage1`/`stage2` must not be transition-state stages. When `stage_id_ts` is set, - the referenced stage must be one. - """ - stage1 = _resolve_stage(target, "stage_id1", "stage1") - stage2 = _resolve_stage(target, "stage_id2", "stage2") - stage_ts = _resolve_stage(target, "stage_id_ts", "stage_ts") - - if (stage1 is not None and stage1.is_ts) or (stage2 is not None and stage2.is_ts): - msg = "Step's stage1/stage2 cannot be a transition-state stage." - raise DataIntegrityError(msg) - if stage_ts is not None and not stage_ts.is_ts: - msg = "Step's stage_ts must reference a transition-state stage." - raise DataIntegrityError(msg) + if hill_extra is not None: + session.add(hill_extra) diff --git a/src/autostorage/exc.py b/src/autostorage/exc.py deleted file mode 100644 index abe56c5..0000000 --- a/src/autostorage/exc.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Autostorage exceptions.""" - -from typing import Self - -from sqlmodel import SQLModel - -__all__ = ["DataIntegrityError", "MissingPrimaryKeyError", "ResultShapeError"] - - -class DataIntegrityError(Exception): - """Raise when an ORM event detects a data integrity violation.""" - - -class ResultShapeError(Exception): - """Raise when a result violates expected shape.""" - - def __init__( - self: Self, model: SQLModel, actual: tuple[int, ...], expected: tuple[int, ...] - ) -> None: - """Initialize exception.""" - class_name = model.__class__.__name__ - msg = f"{class_name} shape ({actual}) does not match expected ({expected})." - super().__init__(msg) - - -class MissingPrimaryKeyError(Exception): - """Raise when primary keys weren't provided to a query method.""" - - def __init__(self: Self, rows: list[SQLModel]) -> None: - row_ids = [ - f"{row.__class__.__name__}: {getattr(row, 'id', None)}" for row in rows - ] - msg = ( - f"Cannot perform operation using unpersisted database instance(s).\n" - f"Try Database.add(row) or Database.merge(row) before querying.\n" - f"({','.join(row_ids)})." - ) - super().__init__(msg) diff --git a/src/autostorage/merge.py b/src/autostorage/merge.py deleted file mode 100644 index 6f22ed6..0000000 --- a/src/autostorage/merge.py +++ /dev/null @@ -1,374 +0,0 @@ -"""Merge one database's contents into another, with validation at merge time.""" - -from dataclasses import dataclass -from typing import TYPE_CHECKING - -from sqlalchemy import inspect as sa_inspect -from sqlmodel import SQLModel, func, select - -from .events import AUTO_MANAGED_IDENTITY_ALGORITHMS -from .models import ( - GeometryRow, - IdentityExtraRow, - IdentityRow, - ModelRow, - StationaryIdentityLink, -) - -if TYPE_CHECKING: - from .database import Database - -__all__ = ["MergeReport", "merge_databases"] - - -@dataclass(frozen=True, slots=True) -class MergeReport: - """Summary of a `merge_databases` call. - - Attributes - ---------- - copied - New rows created in the target database, by table name. - reused - Source rows deduplicated onto an existing target row, by table name - (only ``model``, ``geometry``, and ``identity`` support dedup). - """ - - copied: dict[str, int] - reused: dict[str, int] - - -def merge_databases( - target: "Database", source: "Database", *, commit: bool = True -) -> MergeReport: - """Copy `source`'s contents into `target`, validating and deduplicating. - - Rows are copied with freshly-assigned primary keys and remapped foreign - keys. `ModelRow`, `GeometryRow`, and non-auto-managed `IdentityRow`s are - deduplicated against `target`'s existing content; everything else is - copied fresh. - - InChI/conformer identities are deliberately skipped here: inserting each - source `StationaryPointRow` fresh lets `autostorage.events`'s flush - listeners regenerate and dedup them against `target`'s live state, so a - conformer shared by both databases collapses onto one identity. Every - copied row passes through normal ORM validation. - - Only `flush()` is used until the end, so an error partway through - leaves `target` unchanged. - - Parameters - ---------- - target - Database to copy rows into. - source - Database to copy rows from. Only ever read, never modified. - commit, optional - If True (default), commit once every table has copied successfully. - If False, leave flushed-but-uncommitted for the caller. - - Returns - ------- - MergeReport - Per-table counts of rows copied and reused. - """ - _check_mergeable(target, source) - - id_map: dict[type[SQLModel], dict[int, int]] = {} - copied: dict[str, int] = {} - reused: dict[str, int] = {} - skipped_identity_ids: set[int] = set() - - for cls in _ordered_models(): - if cls is ModelRow: - _copy_models( - target=target, - source=source, - id_map=id_map, - copied=copied, - reused=reused, - ) - elif cls is GeometryRow: - _copy_geometries( - target=target, - source=source, - id_map=id_map, - copied=copied, - reused=reused, - ) - elif cls is IdentityRow: - skipped_identity_ids = _copy_identities( - target=target, - source=source, - id_map=id_map, - copied=copied, - reused=reused, - ) - elif cls is IdentityExtraRow: - rows = [ - row - for row in source.exec_all(select(IdentityExtraRow)) - if row.identity_id not in skipped_identity_ids - ] - _copy_table( - IdentityExtraRow, rows, target=target, id_map=id_map, copied=copied - ) - elif cls is StationaryIdentityLink: - rows = [ - row - for row in source.exec_all(select(StationaryIdentityLink)) - if row.identity_id not in skipped_identity_ids - ] - _copy_table( - StationaryIdentityLink, - rows, - target=target, - id_map=id_map, - copied=copied, - ) - else: - _copy_table( - cls, - source.exec_all(select(cls)), - target=target, - id_map=id_map, - copied=copied, - ) - - if commit: - target.commit() - else: - target.flush() - - return MergeReport(copied=copied, reused=reused) - - -def _is_same_database(a: "Database", b: "Database") -> bool: - """Return whether `a` and `b` refer to the same underlying database. - - In-memory databases (path ``":memory:"``) never count as a match; - on-disk `Database`s match if they resolve to the same file. - """ - if a is b: - return True - if str(a.path) == ":memory:" or str(b.path) == ":memory:": - return False - return a.path.resolve() == b.path.resolve() - - -def _reflected_schema(db: "Database") -> dict[str, set[str]]: - """Return the actual on-disk table/column names for `db`.""" - inspector = sa_inspect(db.engine) - return { - table_name: {column["name"] for column in inspector.get_columns(table_name)} - for table_name in inspector.get_table_names() - } - - -def _check_mergeable(target: "Database", source: "Database") -> None: - """Reject a merge that can't safely proceed. - - Raises - ------ - ValueError - If `source` and `target` are the same database, or either is - missing an expected table/column (e.g. an outdated schema). - """ - if _is_same_database(target, source): - msg = "Cannot merge a database into itself." - raise ValueError(msg) - - expected = { - table.name: {column.name for column in table.columns} - for table in SQLModel.metadata.tables.values() - } - for db, label in ((source, "source"), (target, "target")): - actual = _reflected_schema(db) - for table_name, expected_columns in expected.items(): - if table_name not in actual: - msg = f"{label} database is missing table {table_name!r}." - raise ValueError(msg) - missing_columns = expected_columns - actual[table_name] - if missing_columns: - msg = ( - f"{label} database's {table_name!r} table is missing " - f"column(s) {sorted(missing_columns)!r}." - ) - raise ValueError(msg) - - -def _mapped_classes() -> dict[str, type[SQLModel]]: - """Map each table name to its mapped `SQLModel` class.""" - mapping: dict[str, type[SQLModel]] = {} - for mapper in SQLModel._sa_registry.mappers: # noqa: SLF001 - table = getattr(mapper.class_, "__table__", None) - if table is not None: - mapping[table.name] = mapper.class_ - return mapping - - -def _ordered_models() -> list[type[SQLModel]]: - """Return every mapped table's `SQLModel` class, in FK-safe insertion order.""" - mapping = _mapped_classes() - return [mapping[table.name] for table in SQLModel.metadata.sorted_tables] - - -def _fk_targets(cls: type[SQLModel]) -> list[tuple[str, type[SQLModel]]]: - """Return `(column name, target class)` for each foreign key column on `cls`.""" - mapping = _mapped_classes() - return [ - (column.name, mapping[fk.column.table.name]) - for column in cls.__table__.columns # ty:ignore[unresolved-attribute] - for fk in column.foreign_keys - ] - - -def _copy_row( - row: SQLModel, *, id_map: dict[type[SQLModel], dict[int, int]] -) -> SQLModel: - """Build a new, unsaved row copying `row`'s content with FKs remapped.""" - cls = type(row) - content = row.model_dump(exclude={"id", "created_at", "updated_at"}) - for column_name, target_cls in _fk_targets(cls): - old_id = content.get(column_name) - if old_id is not None: - content[column_name] = id_map[target_cls][old_id] - return cls(**content) - - -def _copy_table( - cls: type[SQLModel], - rows: list[SQLModel], - *, - target: "Database", - id_map: dict[type[SQLModel], dict[int, int]], - copied: dict[str, int], -) -> None: - """Copy `rows` (all of type `cls`) into `target`, remapping FKs via `id_map`.""" - if not rows: - return - - new_rows = [_copy_row(row, id_map=id_map) for row in rows] - target.add_all(new_rows) - target.flush() - - if "id" in cls.model_fields: - id_map[cls] = { - row.id: new_row.id # ty:ignore[unresolved-attribute] - for row, new_row in zip(rows, new_rows, strict=True) - } - - copied[cls.__tablename__] = len(new_rows) # ty:ignore[invalid-assignment] - - -def _table_count(db: "Database", cls: type[SQLModel]) -> int: - """Return the number of rows currently in `cls`'s table.""" - return db.exec_first(select(func.count()).select_from(cls)) or 0 - - -def _copy_models( - *, - target: "Database", - source: "Database", - id_map: dict[type[SQLModel], dict[int, int]], - copied: dict[str, int], - reused: dict[str, int], -) -> None: - """Find-or-create every source `ModelRow` against `target`.""" - rows = source.exec_all(select(ModelRow)) - if not rows: - return - - before = _table_count(target, ModelRow) - mapping: dict[int, int] = {} - for row in rows: - new_row = ModelRow.find_or_create( - target, - program=row.program, - method=row.method, - program_version=row.program_version, - basis=row.basis, - commit=False, - ) - mapping[row.id] = new_row.id # ty:ignore[invalid-assignment] - - id_map[ModelRow] = mapping - created = _table_count(target, ModelRow) - before - copied["model"] = created - reused["model"] = len(rows) - created - - -def _copy_geometries( - *, - target: "Database", - source: "Database", - id_map: dict[type[SQLModel], dict[int, int]], - copied: dict[str, int], - reused: dict[str, int], -) -> None: - """Find-or-create every source `GeometryRow` against `target`.""" - rows = source.exec_all(select(GeometryRow)) - if not rows: - return - - before = _table_count(target, GeometryRow) - mapping: dict[int, int] = {} - for row in rows: - new_row = GeometryRow.find_or_create( - target, - symbols=row.symbols, - coordinates=row.coordinates, - charge=row.charge, - spin=row.spin, - commit=False, - ) - mapping[row.id] = new_row.id # ty:ignore[invalid-assignment] - - id_map[GeometryRow] = mapping - created = _table_count(target, GeometryRow) - before - copied["geometry"] = created - reused["geometry"] = len(rows) - created - - -def _copy_identities( - *, - target: "Database", - source: "Database", - id_map: dict[type[SQLModel], dict[int, int]], - copied: dict[str, int], - reused: dict[str, int], -) -> set[int]: - """Find-or-create every non-auto-managed source `IdentityRow` against `target`. - - Auto-managed identities (`AUTO_MANAGED_IDENTITY_ALGORITHMS`) are left for - `autostorage.events`'s flush listeners to regenerate instead. - - Returns - ------- - set[int] - Source-side ids of skipped identities, so callers can filter out - rows referencing them (identity extras, stationary-identity links). - """ - rows = source.exec_all(select(IdentityRow)) - skipped_ids: set[int] = set() - if not rows: - return skipped_ids - - before = _table_count(target, IdentityRow) - mapping: dict[int, int] = {} - handled = 0 - for row in rows: - if row.algorithm in AUTO_MANAGED_IDENTITY_ALGORITHMS: - skipped_ids.add(row.id) - continue - handled += 1 - new_row = IdentityRow.find_or_create( - target, algorithm=row.algorithm, value=row.value, commit=False - ) - mapping[row.id] = new_row.id # ty:ignore[invalid-assignment] - - id_map[IdentityRow] = mapping - created = _table_count(target, IdentityRow) - before - copied["identity"] = created - reused["identity"] = handled - created - return skipped_ids diff --git a/src/autostorage/models.py b/src/autostorage/models.py index 3663e24..c9696c0 100644 --- a/src/autostorage/models.py +++ b/src/autostorage/models.py @@ -1,16 +1,11 @@ -"""SQLModel row definitions for autostorage's persistence schema.""" +"""SQLModel row definitions for autostorage's schema.""" -import hashlib -import json -from datetime import datetime -from functools import cached_property -from typing import TYPE_CHECKING, Any, Self, dataclass_transform +from typing import Any import numpy as np -from automol import Algorithm, Geometry, Identity, geom +from automol import Geometry, Identity from automol.utils.types import FloatArray -from sqlalchemy import inspect as sa_inspect -from sqlalchemy import text +from pydantic import field_validator from sqlmodel import ( JSON, CheckConstraint, @@ -21,157 +16,251 @@ Relationship, SQLModel, UniqueConstraint, - func, - select, + text, ) from sqlmodel.main import SQLModelConfig -from stereomolgraph.algorithms.symmetry import ( - symmetry_number as _stereo_symmetry_number, -) -from autostorage.exc import MissingPrimaryKeyError +from .types import CompressedArrayTypeDecorator, Role, _fk_field -from .types import CalcStatus, CalcType, CompressedArrayTypeDecorator, Role -if TYPE_CHECKING: - from .database import Database +# 0. Link rows +# NOTE: Link tables are named by the two entities they connect, in alphabetical order. +class CalculationGeometryLink(SQLModel, table=True): + """Association table linking geometries to a calculation. + Attributes + ---------- + geometry_id + Foreign key to the linked geometry. + calculation_id + Foreign key to the linked calculation. + role + Role the geometry plays for this calculation (input/output). + geometry + The linked geometry (back-populated from `GeometryRow.calculation_links`). + calculation + The linked calculation (back-populated from `CalculationRow.geometry_links`). + """ -def _fk_field(target: str, *, nullable: bool = False, index: bool = True) -> Any: # noqa: ANN401 - """Build a standard foreign-key Field with ON DELETE CASCADE.""" - return Field( + __tablename__ = "calculation_geometry_link" + __table_args__ = ( + # The composite primary key only serves lookups keyed by `geometry_id` + # (its leading column); this adds a matching index for `calculation_id`. + Index("ix_calculation_geometry_link_calculation_id", "calculation_id"), + ) + + geometry_id: int | None = Field( + default=None, + foreign_key="geometry.id", + ondelete="CASCADE", + nullable=False, + primary_key=True, + ) + calculation_id: int | None = Field( default=None, - foreign_key=target, + foreign_key="calculation.id", ondelete="CASCADE", - nullable=nullable, - index=index, + nullable=False, + primary_key=True, + ) + role: Role = Field( + sa_column=Column(Enum(Role, values_callable=lambda x: [e.value for e in x])) ) + geometry: "GeometryRow" = Relationship(back_populates="calculation_links") + calculation: "CalculationRow" = Relationship(back_populates="geometry_links") -@dataclass_transform(kw_only_default=True, field_specifiers=(Field,)) -class TimestampMixin(SQLModel): - """Mixin adding server-managed creation/update timestamps. - Annotated as `datetime | None` since the value is unset in Python until the - database fills it in via `server_default`/`onupdate`; `nullable=False` - overrides the `NULL`-by-default column that an Optional annotation would - otherwise produce, since the DB always has a value once the row is flushed. +class GeometryTrajectoryLink(SQLModel, table=True): + """Association table linking geometries to a trajectory. + + Attributes + ---------- + geometry_id + Foreign key to the linked geometry. + trajectory_id + Foreign key to the linked trajectory. + index + Position of the geometry within the trajectory. + geometry + The linked geometry (back-populated from `GeometryRow.trajectory_links`). + trajectory + The linked trajectory (back-populated from `TrajectoryRow.geometry_links`). """ - created_at: datetime | None = Field( + __tablename__ = "geometry_trajectory_link" + __table_args__ = ( + Index("ix_geometry_trajectory_link_trajectory_id", "trajectory_id"), + ) + + geometry_id: int | None = Field( default=None, + foreign_key="geometry.id", + primary_key=True, + ondelete="CASCADE", nullable=False, - sa_column_kwargs={"server_default": func.now()}, ) - updated_at: datetime | None = Field( + trajectory_id: int | None = Field( default=None, + foreign_key="trajectory.id", + primary_key=True, + ondelete="CASCADE", nullable=False, - sa_column_kwargs={"server_default": func.now(), "onupdate": func.now()}, ) + index: list[int] | None = Field(default=None, sa_column=Column(JSON)) + + geometry: "GeometryRow" = Relationship(back_populates="trajectory_links") + trajectory: "TrajectoryRow" = Relationship(back_populates="geometry_links") -@dataclass_transform(kw_only_default=True, field_specifiers=(Field,)) -class BaseRow(TimestampMixin, SQLModel): - """Base for models with a primary ID.""" +class CalculationTrajectoryLink(SQLModel, table=True): + """Association table linking trajectories to a calculation. - id: int | None = Field(default=None, primary_key=True) + Attributes + ---------- + trajectory_id + Foreign key to the linked trajectory. + calculation_id + Foreign key to the linked calculation. + role + Role the trajectory plays for this calculation (input/output). + trajectory + The linked trajectory (back-populated from `TrajectoryRow.calculation_links`). + calculation + The linked calculation (back-populated from `CalculationRow.trajectory_links`). + """ + + __tablename__ = "calculation_trajectory_link" + __table_args__ = ( + Index("ix_calculation_trajectory_link_calculation_id", "calculation_id"), + ) + + trajectory_id: int | None = Field( + default=None, + foreign_key="trajectory.id", + ondelete="CASCADE", + nullable=False, + primary_key=True, + ) + calculation_id: int | None = Field( + default=None, + foreign_key="calculation.id", + ondelete="CASCADE", + nullable=False, + primary_key=True, + ) + role: Role = Field( + sa_column=Column(Enum(Role, values_callable=lambda x: [e.value for e in x])) + ) + trajectory: "TrajectoryRow" = Relationship(back_populates="calculation_links") + calculation: "CalculationRow" = Relationship(back_populates="trajectory_links") -class BaseResultRow(BaseRow): - """Base for result models.""" - geometry_id: int | None +class StageStationaryLink(SQLModel, table=True): + """Association table linking stationary points to reaction stages. - @classmethod - def query( - cls, - db: "Database", - *, - geo: "GeometryRow", - model: "ModelRow", - prov: dict[str, Any] | None = None, - ) -> Self | None: - """Query for result matching geometry, model, and provenance.""" - if not geo.id or not model.id: - raise MissingPrimaryKeyError([geo, model]) - - prov = prov or {} - stmt = ( - select(cls) - .join(CalculationRow) - .where( - cls.geometry_id == geo.id, - CalculationRow.model_id == model.id, - CalculationRow.input_provenance == prov, - ) - ) - return db.exec_first(stmt) - - -@dataclass_transform(kw_only_default=True, field_specifiers=(Field,)) -class BaseLink(SQLModel): - """Base for models without a primary ID.""" + Attributes + ---------- + stationary_id + Foreign key to the linked stationary point. + stage_id + Foreign key to the linked reaction stage. + """ - @classmethod - def create(cls, *rows: BaseRow, **attrs: object) -> Self: - """Construct a link, matching each row to its relationship by type. - - Parameters - ---------- - *rows - The rows to link (e.g. a ``GeometryRow`` and a ``CalculationRow``), - in any order. - **attrs - Extra attributes to set on the link (e.g. ``role``). - - Returns - ------- - Self - The constructed (unsaved) link row. - """ - relationships = sa_inspect(cls, raiseerr=True).relationships - fields: dict[str, BaseRow] = {} - for row in rows: - matches = [ - rel.key - for rel in relationships - if rel.key not in fields and isinstance(row, rel.mapper.class_) - ] - if not matches: - msg = f"{cls.__name__} has no unmatched relationship for {row!r}." - raise ValueError(msg) - if len(matches) > 1: - # Ambiguous: two+ unfilled relationships share this row's type, - # so matching by type alone can't tell them apart (e.g. a link - # table with two relationships to the same row model). Raise - # rather than silently picking one by declaration order. - msg = ( - f"{cls.__name__} has multiple unmatched relationships " - f"{matches} for {row!r}; construct this link directly instead." - ) - raise ValueError(msg) - fields[matches[0]] = row - return cls(**fields, **attrs) - - -def _geometry_hash( - symbols: list[str], coordinates: FloatArray, charge: int, spin: int -) -> str: - """Compute a hash identifying bit-identical geometry content.""" - hasher = hashlib.sha256() - hasher.update(json.dumps(symbols).encode()) - hasher.update(np.asarray(coordinates, dtype=np.float64).tobytes()) - hasher.update(charge.to_bytes(8, "big", signed=True)) - hasher.update(spin.to_bytes(8, "big", signed=True)) - return hasher.hexdigest() - - -# Geometry table -class GeometryRow(BaseRow, Geometry, table=True): + __tablename__ = "stage_stationary_link" + __table_args__ = (Index("ix_stage_stationary_link_stage_id", "stage_id"),) + + stationary_id: int | None = Field( + default=None, + foreign_key="stationary_point.id", + primary_key=True, + ondelete="CASCADE", + nullable=False, + ) + stage_id: int | None = Field( + default=None, + foreign_key="stage.id", + primary_key=True, + ondelete="CASCADE", + nullable=False, + ) + + +class StepValidationLink(SQLModel, table=True): + """Association table linking validations to a step. + + Attributes + ---------- + step_id + Foreign key to the linked step. + validation_id + Foreign key to the linked validation. + + Notes + ----- + Relationships are managed bidirectionally via `ValidationRow.step` and + `StepRow.validations` using this table's `link_model`. + """ + + __tablename__ = "step_validation_link" + __table_args__ = (Index("ix_step_validation_link_validation_id", "validation_id"),) + + step_id: int = Field( + foreign_key="step.id", + primary_key=True, + ondelete="CASCADE", + nullable=False, + ) + validation_id: int = Field( + foreign_key="validation.id", + primary_key=True, + ondelete="CASCADE", + nullable=False, + ) + + +class IdentityStationaryLink(SQLModel, table=True): + """Association table linking chemical identities to stationary points. + + Attributes + ---------- + stationary_id + Foreign key to the linked stationary point. + identity_id + Foreign key to the linked chemical identity. + + Notes + ----- + Relationships are managed bidirectionally via `StationaryPointRow.identities` + and `IdentityRow.stationary_points` using this table's `link_model`. + """ + + __tablename__ = "identity_stationary_link" + __table_args__ = (Index("ix_identity_stationary_link_identity_id", "identity_id"),) + + stationary_id: int = Field( + foreign_key="stationary_point.id", + primary_key=True, + ondelete="CASCADE", + nullable=False, + ) + identity_id: int = Field( + foreign_key="identity.id", + primary_key=True, + ondelete="CASCADE", + nullable=False, + ) + + +# 1. Existential data rows +class GeometryRow(SQLModel, Geometry, table=True): """Molecular geometry definition and metadata. Attributes ---------- + id + Primary key. symbols Atomic symbols in order. coordinates @@ -180,9 +269,6 @@ class GeometryRow(BaseRow, Geometry, table=True): Total molecular charge. spin Number of unpaired electrons (2S). - geometry_hash - Content hash of `symbols`/`coordinates`/`charge`/`spin`, used to reject - exactly-duplicate geometries (see `find_or_create`). energies Energy results computed at this geometry. gradients @@ -198,13 +284,13 @@ class GeometryRow(BaseRow, Geometry, table=True): """ __tablename__ = "geometry" - __table_args__ = (UniqueConstraint("geometry_hash", name="unique_geometry_hash"),) + model_config = SQLModelConfig(arbitrary_types_allowed=True) + id: int | None = Field(default=None, primary_key=True) symbols: list[str] = Field(sa_column=Column(JSON)) coordinates: FloatArray = Field(sa_column=Column(CompressedArrayTypeDecorator())) charge: int spin: int - geometry_hash: str | None = Field(default=None, nullable=False) energies: list["EnergyRow"] = Relationship(back_populates="geometry") gradients: list["GradientRow"] = Relationship(back_populates="geometry") @@ -212,67 +298,157 @@ class GeometryRow(BaseRow, Geometry, table=True): stationary_points: list["StationaryPointRow"] = Relationship( back_populates="geometry" ) - trajectory_links: list["TrajectoryGeometryLink"] = Relationship( + trajectory_links: list["GeometryTrajectoryLink"] = Relationship( back_populates="geometry" ) calculation_links: list["CalculationGeometryLink"] = Relationship( back_populates="geometry" ) - @cached_property - def symmetry_number(self) -> int: - """Symmetry number from stereo-preserving graph automorphisms. + @field_validator("coordinates", mode="before") + @classmethod + def _validate_coordinates(cls, value: list[list[float]] | FloatArray) -> FloatArray: + """Convert list to numpy array if needed.""" + if isinstance(value, list): + return np.asarray(value, dtype=float) + return value - Cached per instance since counting graph isomorphisms is expensive. - """ - graph = geom.stereo_mol_graph(self) - return _stereo_symmetry_number(graph) - @classmethod - def find_or_create( # noqa: PLR0913 - cls, - db: "Database", - *, - symbols: list[str], - coordinates: FloatArray, - charge: int, - spin: int, - commit: bool = True, - ) -> Self: - """Return the matching geometry row, creating and saving one if absent. - - Matches on exact content via `geometry_hash`, so this only reuses - bit-identical geometries. - - Parameters - ---------- - commit, optional - If True (default), commit a newly-created row immediately. If - False, only flush it (still assigns `.id`), leaving the caller's - transaction open — for a caller staging several dedup lookups - that must succeed or fail together. - """ - geometry_hash = _geometry_hash(symbols, coordinates, charge, spin) - stmt = select(cls).where(cls.geometry_hash == geometry_hash) - existing = db.exec_first(stmt) - if existing is not None: - return existing - - row = cls(symbols=symbols, coordinates=coordinates, charge=charge, spin=spin) - db.add(row) - if commit: - db.commit() - else: - db.flush() - return row - - -# Result tables -class EnergyRow(BaseResultRow, table=True): +class TrajectoryRow(SQLModel, table=True): + """Ordered sequence of geometries from a calculation trajectory. + + Attributes + ---------- + id + Primary key. + geometry_links + Raw link rows connecting geometries to this trajectory. + calculation_links + Raw link rows connecting calculations to this trajectory. + """ + + __tablename__ = "trajectory" + + id: int | None = Field(default=None, primary_key=True) + ndim: int | None = Field(default=None, nullable=True) + + geometry_links: list["GeometryTrajectoryLink"] = Relationship( + back_populates="trajectory" + ) + calculation_links: list["CalculationTrajectoryLink"] = Relationship( + back_populates="trajectory" + ) + + +class ModelRow(SQLModel, table=True): + """Calculation model specification. + + Attributes + ---------- + id + Primary key. + calc_type + Type of calculation (energy, gradient, hessian, etc.). + program + Quantum chemistry program used (psi4, ORCA, ...) + program_version + Quantum chemistry program version. + method + Computational method (B3LYP, MP2, ...) + basis + Orbital basis set. + keywords + Additional keywords and options for the calculation. + calculations + Calculations performed using this model. + """ + + __tablename__ = "model" + + id: int | None = Field(default=None, primary_key=True) + calc_type: str + program: str + program_version: str | None = None + method: str + basis: str | None = None + keywords: dict[str, Any] | None = Field( + default_factory=dict, sa_column=Column(JSON) + ) + + calculations: list["CalculationRow"] = Relationship(back_populates="model") + + +class CalculationRow(SQLModel, table=True): + """Quantum chemistry calculation and its associated data. + + Attributes + ---------- + id + Primary key. + model_id + Foreign key to the model used for this calculation. + input_provenance + Metadata describing how the input was generated. + output_provenance + Metadata describing how the output was produced. + model + Model used for this calculation. + geometry_links + Raw link rows connecting geometries to this calculation. + trajectory_links + Raw link rows connecting trajectories to this calculation. + energies + Energy results produced by this calculation. + gradients + Gradient results produced by this calculation. + hessians + Hessian results produced by this calculation. + validations + Validation results performed by this calculation. + stationary_points + Stationary points identified by this calculation. + """ + + __tablename__ = "calculation" + + id: int | None = Field(default=None, primary_key=True) + model_id: int | None = Field( + default=None, + foreign_key="model.id", + ondelete="CASCADE", + nullable=False, + index=True, + ) + input_provenance: dict[str, Any] | None = Field( + default_factory=dict, sa_column=Column(JSON) + ) + output_provenance: dict[str, Any] | None = Field( + default_factory=dict, sa_column=Column(JSON) + ) + + model: "ModelRow" = Relationship(back_populates="calculations") + energies: list["EnergyRow"] = Relationship(back_populates="calculation") + gradients: list["GradientRow"] = Relationship(back_populates="calculation") + hessians: list["HessianRow"] = Relationship(back_populates="calculation") + validations: list["ValidationRow"] = Relationship(back_populates="calculation") + stationary_points: list["StationaryPointRow"] = Relationship( + back_populates="calculation" + ) + geometry_links: list["CalculationGeometryLink"] = Relationship( + back_populates="calculation" + ) + trajectory_links: list["CalculationTrajectoryLink"] = Relationship( + back_populates="calculation" + ) + + +class EnergyRow(SQLModel, table=True): """Energy result for a specific geometry and calculation. Attributes ---------- + id + Primary key. geometry_id Foreign key to the geometry this energy was evaluated at. calculation_id @@ -287,19 +463,22 @@ class EnergyRow(BaseResultRow, table=True): __tablename__ = "energy" + id: int | None = Field(default=None, primary_key=True) geometry_id: int | None = _fk_field("geometry.id") calculation_id: int | None = _fk_field("calculation.id") value: float - calculation: "CalculationRow" = Relationship() + calculation: "CalculationRow" = Relationship(back_populates="energies") geometry: "GeometryRow" = Relationship(back_populates="energies") -class GradientRow(BaseResultRow, table=True): +class GradientRow(SQLModel, table=True): """Energy gradient result for a specific geometry and calculation. Attributes ---------- + id + Primary key. geometry_id Foreign key to the geometry this gradient was evaluated at. calculation_id @@ -315,19 +494,22 @@ class GradientRow(BaseResultRow, table=True): __tablename__ = "gradient" model_config = SQLModelConfig(arbitrary_types_allowed=True) + id: int | None = Field(default=None, primary_key=True) geometry_id: int | None = _fk_field("geometry.id") calculation_id: int | None = _fk_field("calculation.id") value: FloatArray = Field(sa_column=Column(CompressedArrayTypeDecorator())) - calculation: "CalculationRow" = Relationship() + calculation: "CalculationRow" = Relationship(back_populates="gradients") geometry: "GeometryRow" = Relationship(back_populates="gradients") -class HessianRow(BaseResultRow, table=True): +class HessianRow(SQLModel, table=True): """Hessian result for a specific geometry and calculation. Attributes ---------- + id + Primary key. geometry_id Foreign key to the geometry this Hessian was evaluated at. calculation_id @@ -343,6 +525,7 @@ class HessianRow(BaseResultRow, table=True): __tablename__ = "hessian" model_config = SQLModelConfig(arbitrary_types_allowed=True) + id: int | None = Field(default=None, primary_key=True) geometry_id: int | None = _fk_field("geometry.id") calculation_id: int | None = _fk_field("calculation.id") @@ -350,165 +533,51 @@ class HessianRow(BaseResultRow, table=True): sa_column=Column(CompressedArrayTypeDecorator(dtype=np.float32)) ) - calculation: "CalculationRow" = Relationship() + calculation: "CalculationRow" = Relationship(back_populates="hessians") geometry: "GeometryRow" = Relationship(back_populates="hessians") - @cached_property - def harmonic_frequencies(self) -> tuple[float, ...]: - """Harmonic frequencies derived from the Hessian. - - Cached per instance, since vibrational analysis re-diagonalizes the - Hessian on every call and `.order` (used by `_recompute_geometry_ - stationary_validity` for every sibling Hessian of a geometry, on - every relevant flush) depends on it. Invalidated on `value` update - by `invalidate_hessian_frequency_cache` in `events.py`. - """ - freqs, _ = geom.vibrational_analysis(geo=self.geometry, hess=self.value) - return freqs - - @property - def order(self) -> int: - """Hessian order.""" - return sum(1 for f in self.harmonic_frequencies if f < 0.0) - - -# Trajectory table -class TrajectoryRow(BaseRow, table=True): - """Ordered sequence of geometries from a calculation trajectory. - - Attributes - ---------- - geometry_links - Raw link rows connecting geometries to this trajectory. - calculation_links - Raw link rows connecting calculations to this trajectory. - """ - - __tablename__ = "trajectory" - - geometry_links: list["TrajectoryGeometryLink"] = Relationship( - back_populates="trajectory" - ) - calculation_links: list["CalculationTrajectoryLink"] = Relationship( - back_populates="trajectory" - ) - - -class TrajectoryGeometryLink(BaseLink, table=True): - """Association table linking geometries to a trajectory. - - Attributes - ---------- - geometry_id - Foreign key to the linked geometry. - trajectory_id - Foreign key to the linked trajectory. - index - Position of the geometry within the trajectory. - geometry - The linked geometry. - trajectory - The linked trajectory. - """ - - __tablename__ = "trajectory_geometry_link" - __table_args__ = ( - Index("ix_trajectory_geometry_link_trajectory_id", "trajectory_id"), - ) - geometry_id: int | None = Field( - default=None, - foreign_key="geometry.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - trajectory_id: int | None = Field( - default=None, - foreign_key="trajectory.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - index: list[int] | None = Field(default=None, sa_column=Column(JSON)) - - geometry: "GeometryRow" = Relationship(back_populates="trajectory_links") - trajectory: "TrajectoryRow" = Relationship(back_populates="geometry_links") - - -# Link tables declared here, ahead of the StationaryPointRow/IdentityRow and -# StationaryPointRow/StageRow entities they connect, because SQLModel's -# `link_model=` kwarg needs the actual class object at class-body-evaluation -# time — unlike every other cross-model reference in this file, it can't be -# satisfied by a lazily-resolved string forward ref. -class StationaryIdentityLink(BaseLink, table=True): - """Association table linking stationary points to chemical identities. +class ValidationRow(SQLModel, table=True): + """Validation result for a specific step and calculation. Attributes ---------- - stationary_id - Foreign key to the linked stationary point. - identity_id - Foreign key to the linked identity. + id + Primary key. + calculation_id + Foreign key to the calculation that performed this validation. + method + Type of validation step (e.g., ``irc``) + extras + Additional metadata attached to this validation. + calculation + Calculation that performed this validation. + step + Reaction step this validation belongs to. """ - __tablename__ = "stationary_identity_link" - __table_args__ = (Index("ix_stationary_identity_link_identity_id", "identity_id"),) - - stationary_id: int = Field( - foreign_key="stationary_point.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - identity_id: int = Field( - foreign_key="identity.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - - -class StationaryStageLink(BaseLink, table=True): - """Association table linking stationary points to reaction stages. + __tablename__ = "validation" - Attributes - ---------- - stationary_id - Foreign key to the linked stationary point. - stage_id - Foreign key to the linked reaction stage. - stationary - The linked stationary point. - stage - The linked reaction stage. - """ + id: int | None = Field(default=None, primary_key=True) + calculation_id: int | None = _fk_field("calculation.id") - __tablename__ = "stationary_stage_link" - __table_args__ = (Index("ix_stationary_stage_link_stage_id", "stage_id"),) + method: str + extras: dict[str, Any] = Field(default_factory=dict, sa_column=Column(JSON)) - stationary_id: int | None = Field( - default=None, - foreign_key="stationary_point.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - stage_id: int | None = Field( - default=None, - foreign_key="stage.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, + calculation: "CalculationRow" = Relationship(back_populates="validations") + step: "StepRow" = Relationship( + back_populates="validations", link_model=StepValidationLink ) -# Stationary point rows -class StationaryPointRow(BaseRow, table=True): +# 2. Stationary point rows +class StationaryPointRow(SQLModel, table=True): """A stationary point on a potential energy surface. Attributes ---------- + id + Primary key. geometry_id Foreign key to the underlying molecular geometry. calculation_id @@ -518,221 +587,63 @@ class StationaryPointRow(BaseRow, table=True): is_pseudo Whether this point is not a true stationary point (e.g. constrained). is_valid - Whether `order` agrees with the consensus order of its geometry's - Hessians (see `autostorage.events.revalidate_geometry_orders_on_insert_update`). + Whether `order` agrees with the consensus Hessian order of its geometry. geometry Geometry defining the coordinates of this point. - calculation - Calculation that identified this point. - identities - Chemical identifiers (e.g. InChI, SMILES) for this point. - stages - Reaction stages this stationary point belongs to. - """ - - __tablename__ = "stationary_point" - - geometry_id: int | None = _fk_field("geometry.id") - calculation_id: int | None = _fk_field("calculation.id") - order: int = 0 - is_pseudo: bool = False - is_valid: bool = False - - geometry: "GeometryRow" = Relationship(back_populates="stationary_points") - calculation: "CalculationRow" = Relationship() - identities: list["IdentityRow"] = Relationship( - back_populates="stationary_points", link_model=StationaryIdentityLink - ) - stages: list["StageRow"] = Relationship( - back_populates="stationaries", link_model=StationaryStageLink - ) - - @classmethod - def query( - cls, - db: "Database", - *, - ident: Identity, - model: "ModelRow | None" = None, - prov: dict[Any, Any] | None = None, - calc_type: CalcType | None = None, - ) -> Self | None: - """Query for stationary point matching geometry, model, and provenance.""" - stmt = ( - select(cls) - .join( - StationaryIdentityLink, - cls.id == StationaryIdentityLink.stationary_id, # ty:ignore[invalid-argument-type] - ) - .join( - IdentityRow, - IdentityRow.id == StationaryIdentityLink.identity_id, # ty:ignore[invalid-argument-type] - ) - .where( - IdentityRow.kind == ident.kind, - IdentityRow.algorithm == ident.algorithm, - IdentityRow.value == ident.value, - ) - ) - - if model or prov or calc_type: - stmt = stmt.join( - CalculationRow, - cls.calculation_id == CalculationRow.id, # ty:ignore[invalid-argument-type] - ) - - if model: - if not model.id: - raise MissingPrimaryKeyError([model]) - stmt = stmt.where(CalculationRow.model_id == model.id) - - if prov: - stmt = stmt.where(CalculationRow.input_provenance == prov) - - if calc_type: - stmt = stmt.where(CalculationRow.calc_type == calc_type) - - return db.exec_first(stmt) - - def identity( - self, - *, - kind: str | None = None, - algorithm: Any | None = None, # noqa: ANN401 - ) -> "IdentityRow | None": - """Return the first loaded identity matching kind and/or algorithm. - - Searches `self.identities` (the already-loaded relationship list), - not the database — use `StationaryPointRow.query` for a DB lookup. - """ - return next( - ( - i - for i in self.identities - if (kind is None or i.kind == kind) - and (algorithm is None or i.algorithm == algorithm) - ), - None, - ) - - -class IdentityRow(BaseRow, Identity, table=True): - """A chemical identifier associated with one or more stationary points. - - Attributes - ---------- - kind - Category of identifier (e.g. ``stereoisomer``, ``formula``). - algorithm - Method used to generate the identifier (e.g. ``rdkit inchi``, ``rdkit smiles``). - value - The resulting identifier string. - stationary_points - Stationary points sharing this identity. - identity_extras - Additional key-value metadata attached to this identity. - """ - - __tablename__ = "identity" - __table_args__ = ( - UniqueConstraint("kind", "algorithm", "value", name="unique_identity"), - ) - - stationary_points: list["StationaryPointRow"] = Relationship( - back_populates="identities", link_model=StationaryIdentityLink - ) - identity_extras: list["IdentityExtraRow"] = Relationship(back_populates="identity") - - @classmethod - def find_or_create( - cls, - db: "Database", - *, - algorithm: Algorithm, - value: str, - commit: bool = True, - ) -> Self: - """Return the matching identity row, creating and saving one if absent. - - `kind` isn't a parameter here since it's fully determined by - `algorithm` (see `Identity.from_value`), so matching on - `(algorithm, value)` is equivalent to `unique_identity`'s full - `(kind, algorithm, value)` constraint. - - Parameters - ---------- - commit, optional - If True (default), commit a newly-created row immediately. If - False, only flush it (still assigns `.id`), leaving the caller's - transaction open — for a caller staging several dedup lookups - that must succeed or fail together. - """ - stmt = select(cls).where(cls.algorithm == algorithm, cls.value == value) - existing = db.exec_first(stmt) - if existing is not None: - return existing - - row = cls.from_value(value, algorithm=algorithm) - db.add(row) - if commit: - db.commit() - else: - db.flush() - return row - - -class IdentityExtraRow(BaseRow, table=True): - """Additional key-value metadata attached to a chemical identity. - - Attributes - ---------- - identity_id - Foreign key to the parent identity. - attribute - Name of the extra attribute. - value - Value of the extra attribute. - identity - The parent identity this extra belongs to. + calculation + Calculation that identified this point. + identities + Chemical identifiers (e.g. InChI, SMILES) for this point. + stages + Reaction stages this stationary point belongs to. """ - __tablename__ = "identity_extras" - - identity_id: int | None = Field( - default=None, - foreign_key="identity.id", - ondelete="CASCADE", - nullable=False, - index=True, - ) + __tablename__ = "stationary_point" - attribute: str - value: str + id: int | None = Field(default=None, primary_key=True) + geometry_id: int | None = _fk_field("geometry.id") + calculation_id: int | None = _fk_field("calculation.id") + order: int = 0 + is_pseudo: bool = False + is_valid: bool = False - identity: "IdentityRow" = Relationship(back_populates="identity_extras") + geometry: "GeometryRow" = Relationship(back_populates="stationary_points") + calculation: "CalculationRow" = Relationship(back_populates="stationary_points") + identities: list["IdentityRow"] = Relationship( + back_populates="stationary_points", link_model=IdentityStationaryLink + ) + stages: list["StageRow"] = Relationship( + back_populates="stationaries", link_model=StageStationaryLink + ) -# Reaction rows -class StageRow(BaseRow, table=True): +# 3. Reaction network rows +class StageRow(SQLModel, table=True): """A chemical state (reactant, product, or transition state) in a reaction. Attributes ---------- + id + Primary key. is_ts Whether this stage represents a transition state. stationaries - Stationary points that make up this stage. + Stationary points that make up this stage (bidirectional via + `link_model=StageStationaryLink`). steps Reaction steps referencing this stage as `stage1`, `stage2`, or - `stage_ts` (read-only; derived from `StepRow`'s foreign keys). + `stage_ts`. Read-only view derived from `StepRow`'s foreign keys; + use `stage1`, `stage2`, `stage_ts` relationships on `StepRow` for + writing. """ __tablename__ = "stage" + id: int | None = Field(default=None, primary_key=True) is_ts: bool = False stationaries: list["StationaryPointRow"] = Relationship( - back_populates="stages", link_model=StationaryStageLink + back_populates="stages", link_model=StageStationaryLink ) steps: list["StepRow"] = Relationship( sa_relationship_kwargs={ @@ -745,98 +656,14 @@ class StageRow(BaseRow, table=True): } ) - @classmethod - def query( - cls, - db: "Database", - stationaries: list["StationaryPointRow"], - *, - is_ts: bool = False, - ) -> Self | None: - """Query for existing stage with stationaries.""" - target_ids = [s.id for s in stationaries] - if len(target_ids) != len(stationaries): - raise MissingPrimaryKeyError(list(stationaries)) - - stmt = ( - select(cls) - .join(StationaryStageLink) - .where(cls.is_ts == is_ts) - .group_by(cls.id) # ty:ignore[invalid-argument-type] - .having( - func.count(StationaryStageLink.stationary_id) == len(target_ids), # ty:ignore[invalid-argument-type] - func.count( - func.nullif( - StationaryStageLink.stationary_id.in_(target_ids), # ty:ignore[unresolved-attribute] - False, # noqa: FBT003 - ) - ) - == len(target_ids), - ) - ) - return db.exec_first(stmt) - - @classmethod - def find_or_create( - cls, - db: "Database", - stationaries: list["StationaryPointRow"], - *, - is_ts: bool = False, - ) -> Self: - """Return the matching stage row, creating and saving one if absent. - - Note - ---- - Unlike `ModelRow`/`StepRow`, there is no DB-level uniqueness - constraint backing this dedup, so it relies entirely on - `StageRow.query`'s app-level lookup. - """ - existing = cls.query(db, stationaries, is_ts=is_ts) - if existing is not None: - return existing - - row = cls(stationaries=stationaries, is_ts=is_ts) - db.add(row) - db.commit() - return row - - -# Declared here, ahead of StepRow, for the same `link_model=` reason as -# StationaryIdentityLink/StationaryStageLink above. -class StepValidationLink(BaseLink, table=True): - """Association table linking validations to a step. - - Attributes - ---------- - step_id - Foreign key to the linked step. - validation_id - Foreign key to the linked validation. - """ - - __tablename__ = "step_validation_link" - __table_args__ = (Index("ix_step_validation_link_validation_id", "validation_id"),) - - step_id: int = Field( - foreign_key="step.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - validation_id: int = Field( - foreign_key="validation.id", - primary_key=True, - ondelete="CASCADE", - nullable=False, - ) - -class StepRow(BaseRow, table=True): +class StepRow(SQLModel, table=True): """An elementary reaction step connecting a reactant, transition state, and product. Attributes ---------- + id + Primary key. stage_id1, stage_id2 Foreign keys to the step's two non-TS stages (stored with `stage_id1 < stage_id2`). @@ -845,10 +672,13 @@ class StepRow(BaseRow, table=True): barrierless step. is_barrierless Whether this step proceeds without a formal transition state. - stage1, stage2 - The step's two non-TS stages. + stage1 + The step's first non-TS stage (reactant or product). + stage2 + The step's second non-TS stage (reactant or product). stage_ts The step's transition-state stage, or `None` if barrierless. + Note: not back-populated from StageRow.steps (which is read-only). validations Validation calculations performed on this step. """ @@ -861,8 +691,7 @@ class StepRow(BaseRow, table=True): CheckConstraint("stage_id1 < stage_id2", name="chk_stage_order"), # `unq_step_stages` doesn't catch duplicate barrierless steps (stage_id_ts # NULL), since SQL never treats NULL as equal to itself in a unique - # constraint. This expression index closes that gap at the DB level, - # defense-in-depth alongside `StepRow.query`'s app-level lookup. + # constraint. This expression index closes that gap at the DB level. Index( "unq_step_stages_null_safe", "stage_id1", @@ -877,6 +706,7 @@ class StepRow(BaseRow, table=True): Index("ix_step_stage_id_ts", "stage_id_ts"), ) + id: int | None = Field(default=None, primary_key=True) stage_id1: int | None = Field( default=None, foreign_key="stage.id", @@ -911,353 +741,69 @@ class StepRow(BaseRow, table=True): sa_relationship_kwargs={"foreign_keys": "[StepRow.stage_id_ts]"} ) - @classmethod - def query( - cls, - db: "Database", - stage1: "StageRow", - stage2: "StageRow", - stage_ts: "StageRow | None" = None, - ) -> Self | None: - """Query for an existing step connecting specific stages.""" - if not stage1.id or not stage2.id or (stage_ts and not stage_ts.id): - raise MissingPrimaryKeyError( - [s for s in [stage1, stage2, stage_ts] if s is not None] - ) - - # Enforce the database CheckConstraint: stage_id1 < stage_id2 - id1, id2 = sorted([stage1.id, stage2.id]) - ts_id = stage_ts.id if stage_ts else None - - stmt = select(cls).where( - cls.stage_id1 == id1, - cls.stage_id2 == id2, - cls.stage_id_ts == ts_id, - ) - return db.exec_first(stmt) - @classmethod - def find_or_create( - cls, - db: "Database", - stage1: "StageRow", - stage2: "StageRow", - stage_ts: "StageRow | None" = None, - ) -> Self: - """Return the matching step row, creating and saving one if absent.""" - existing = cls.query(db, stage1, stage2, stage_ts) - if existing is not None: - return existing - - row = cls(stage1=stage1, stage2=stage2, stage_ts=stage_ts) - db.add(row) - db.commit() - return row - - -# Calculation rows -class ModelRow(BaseRow, table=True): - """Calculation model specification. +# 4. Identity rows +class IdentityRow(SQLModel, Identity, table=True): + """A chemical identifier associated with one or more stationary points. Attributes ---------- - program - Quantum chemistry program used (psi4, ORCA, ...) - program_version - Quantum chemistry program version. - method - Computational method (B3LYP, MP2, ...) - basis - Orbital basis set. + id + Primary key. + kind + Category of identifier (e.g. ``stereoisomer``, ``formula``). + algorithm + Method used to generate the identifier (e.g. ``rdkit inchi``, ``rdkit smiles``). + value + The resulting identifier string. + stationary_points + Stationary points sharing this identity. + identity_extras + Additional key-value metadata attached to this identity. """ - __tablename__ = "model" + __tablename__ = "identity" __table_args__ = ( - UniqueConstraint( - "program", - "program_version", - "method", - "basis", - name="unique_model", - ), - # `unique_model` doesn't catch duplicates when `program_version` or `basis` - # is NULL (see `find_or_create` below). This expression index closes that - # gap at the DB level, defense-in-depth alongside the app-level lookup. - Index( - "unique_model_null_safe", - "program", - text("coalesce(program_version, '')"), - "method", - text("coalesce(basis, '')"), - unique=True, - ), - ) - - program: str - program_version: str | None = None - method: str - basis: str | None = None - - @classmethod - def find_or_create( # noqa: PLR0913 - cls, - db: "Database", - *, - program: str, - method: str, - program_version: str | None = None, - basis: str | None = None, - commit: bool = True, - ) -> Self: - """Return the matching model row, creating and saving one if absent. - - ``unique_model`` doesn't catch duplicates when ``program_version`` - or ``basis`` is NULL, since SQL treats NULL as distinct from itself - in unique constraints. Callers that don't always supply both should - use this instead of constructing and adding a ``ModelRow`` directly, - to avoid silently accumulating duplicate rows for the same model. - - Parameters - ---------- - commit, optional - If True (default), commit a newly-created row immediately. If - False, only flush it (still assigns `.id`), leaving the caller's - transaction open — for a caller staging several dedup lookups - that must succeed or fail together. - """ - stmt = select(cls).where( - cls.program == program, - cls.program_version == program_version, - cls.method == method, - cls.basis == basis, - ) - existing = db.exec_first(stmt) - if existing is not None: - return existing - - row = cls( - program=program, - program_version=program_version, - method=method, - basis=basis, - ) - db.add(row) - if commit: - db.commit() - else: - db.flush() - return row - - -class CalculationRow(BaseRow, table=True): - """Quantum chemistry calculation and its associated data. - - Attributes - ---------- - model_id - Foreign key to the model used for this calculation. - calc_type - Type of calculation performed. - status - Lifecycle status of this calculation. - error_message - Error message recorded for a failed calculation, if any. - input_provenance - Metadata describing how the input was generated. - output_provenance - Metadata describing how the output was produced. - model - Model used for this calculation. - geometry_links - Raw link rows connecting geometries to this calculation. - trajectory_links - Raw link rows connecting trajectories to this calculation. - """ - - __tablename__ = "calculation" - - model_id: int | None = Field( - default=None, - foreign_key="model.id", - ondelete="CASCADE", - nullable=False, - index=True, - ) - calc_type: CalcType = Field( - sa_column=Column(Enum(CalcType, values_callable=lambda x: [e.value for e in x])) - ) - status: CalcStatus = Field( - default=CalcStatus.PENDING, - sa_column=Column( - Enum(CalcStatus, values_callable=lambda x: [e.value for e in x]) - ), - ) - error_message: str | None = Field(default=None) - # Intentionally unbounded free-form JSON; add a size/schema guardrail if - # these are ever populated from a less-trusted input path. - input_provenance: dict[str, Any] | None = Field( - default_factory=dict, sa_column=Column(JSON) - ) - output_provenance: dict[str, Any] | None = Field( - default_factory=dict, sa_column=Column(JSON) - ) - - model: "ModelRow" = Relationship() - geometry_links: list["CalculationGeometryLink"] = Relationship( - back_populates="calculation" - ) - trajectory_links: list["CalculationTrajectoryLink"] = Relationship( - back_populates="calculation" + UniqueConstraint("kind", "algorithm", "value", name="unique_identity"), ) - @property - def input_geometries(self) -> list["GeometryRow"]: - """Geometries linked to this calculation with an INPUT role.""" - return [ - link.geometry for link in self.geometry_links if link.role == Role.INPUT - ] - - @property - def output_geometries(self) -> list["GeometryRow"]: - """Geometries linked to this calculation with an OUTPUT role.""" - return [ - link.geometry for link in self.geometry_links if link.role == Role.OUTPUT - ] - - @property - def input_trajectories(self) -> list["TrajectoryRow"]: - """Trajectories linked to this calculation with an INPUT role.""" - return [ - link.trajectory for link in self.trajectory_links if link.role == Role.INPUT - ] - - @property - def output_trajectories(self) -> list["TrajectoryRow"]: - """Trajectories linked to this calculation with an OUTPUT role.""" - return [ - link.trajectory - for link in self.trajectory_links - if link.role == Role.OUTPUT - ] - - -class CalculationGeometryLink(BaseLink, table=True): - """Association table linking geometries to a calculation. - - Attributes - ---------- - geometry_id - Foreign key to the linked geometry. - calculation_id - Foreign key to the linked calculation. - role - Role the geometry plays for this calculation (input/output). - geometry - The linked geometry. - calculation - The linked calculation. - """ - - __tablename__ = "calculation_geometry_link" - __table_args__ = ( - # The composite primary key only serves lookups keyed by `geometry_id` - # (its leading column); this adds a matching index for `calculation_id`. - Index("ix_calculation_geometry_link_calculation_id", "calculation_id"), - ) + id: int | None = Field(default=None, primary_key=True) - geometry_id: int | None = Field( - default=None, - foreign_key="geometry.id", - ondelete="CASCADE", - nullable=False, - primary_key=True, - ) - calculation_id: int | None = Field( - default=None, - foreign_key="calculation.id", - ondelete="CASCADE", - nullable=False, - primary_key=True, - ) - role: Role = Field( - sa_column=Column(Enum(Role, values_callable=lambda x: [e.value for e in x])) + stationary_points: list["StationaryPointRow"] = Relationship( + back_populates="identities", link_model=IdentityStationaryLink ) - - geometry: "GeometryRow" = Relationship(back_populates="calculation_links") - calculation: "CalculationRow" = Relationship(back_populates="geometry_links") + identity_extras: list["IdentityExtraRow"] = Relationship(back_populates="identity") -class CalculationTrajectoryLink(BaseLink, table=True): - """Association table linking trajectories to a calculation. +class IdentityExtraRow(SQLModel, table=True): + """Additional key-value metadata attached to a chemical identity. Attributes ---------- - trajectory_id - Foreign key to the linked trajectory. - calculation_id - Foreign key to the linked calculation. - role - Role the trajectory plays for this calculation (input/output). - trajectory - The linked trajectory. - calculation - The linked calculation. + id + Primary key. + identity_id + Foreign key to the parent identity. + attribute + Name of the extra attribute. + value + Value of the extra attribute. + identity + The parent identity this extra belongs to. """ - __tablename__ = "calculation_trajectory_link" - __table_args__ = ( - Index("ix_calculation_trajectory_link_calculation_id", "calculation_id"), - ) + __tablename__ = "identity_extras" - trajectory_id: int | None = Field( - default=None, - foreign_key="trajectory.id", - ondelete="CASCADE", - nullable=False, - primary_key=True, - ) - calculation_id: int | None = Field( + id: int | None = Field(default=None, primary_key=True) + identity_id: int | None = Field( default=None, - foreign_key="calculation.id", + foreign_key="identity.id", ondelete="CASCADE", nullable=False, - primary_key=True, - ) - role: Role = Field( - sa_column=Column(Enum(Role, values_callable=lambda x: [e.value for e in x])) + index=True, ) - trajectory: "TrajectoryRow" = Relationship(back_populates="calculation_links") - calculation: "CalculationRow" = Relationship(back_populates="trajectory_links") - - -class ValidationRow(BaseRow, table=True): - """Validation result for a specific step and calculation. - - Attributes - ---------- - calculation_id - Foreign key to the calculation that performed this validation. - method - Type of validation step (e.g., ``irc``) - extras - Additional metadata attached to this validation. - calculation - Calculation that performed this validation. - step - Reaction step this validation belongs to. - """ - - __tablename__ = "validation" - - calculation_id: int | None = _fk_field("calculation.id") - - method: str - # Intentionally unbounded free-form JSON; add a size/schema guardrail if - # this is ever populated from a less-trusted input path. - extras: dict[str, Any] = Field(default_factory=dict, sa_column=Column(JSON)) + attribute: str + value: str - calculation: "CalculationRow" = Relationship() - step: "StepRow" = Relationship( - back_populates="validations", link_model=StepValidationLink - ) + identity: "IdentityRow" = Relationship(back_populates="identity_extras") diff --git a/src/autostorage/types.py b/src/autostorage/types.py index 9222d1a..e37fcda 100644 --- a/src/autostorage/types.py +++ b/src/autostorage/types.py @@ -8,13 +8,20 @@ import numpy as np from sqlalchemy import LargeBinary from sqlalchemy.types import TypeDecorator +from sqlmodel import Field -__all__ = [ - "CalcStatus", - "CalcType", - "CompressedArrayTypeDecorator", - "Role", -] +__all__ = ["CompressedArrayTypeDecorator", "Role"] + + +def _fk_field(target: str, *, nullable: bool = False, index: bool = True) -> Any: # noqa: ANN401 + """Build a standard foreign-key Field with ON DELETE CASCADE.""" + return Field( + default=None, + foreign_key=target, + ondelete="CASCADE", + nullable=nullable, + index=index, + ) class CompressedArrayTypeDecorator(TypeDecorator): @@ -55,74 +62,3 @@ class Role(StrEnum): INPUT = "input" OUTPUT = "output" - - -class CalcType(StrEnum): - """Primary calculation types. - - Attributes - ---------- - OPT - Geometry optimization to find a local minimum on the PES. - OPT_TS - Saddle-point geometry optimization to locate a transition state structure. - CONFORMER - Conformational search/sampling to identify low-energy spatial arrangements. - SCAN - PES scan across user-defined geometric coordinates. - IRC - Intrinsic Reaction Coordinate mapping minimum energy pathway from TS to - its connected reactants and products. - MEP - Minimum Energy Path multi-image chain searches (e.g., Nudged Elastic Band, - String Methods) to discover reaction trajectories and TS guesses. - ENERGY - Single-point electronic energy evaluation at a fixed molecular geometry. - GRADIENT - Nuclear gradient evaluation to compute forces acting on the atoms. - FREQUENCY - Vibrational frequency analysis to verify stationary point order and compute - a Hessian/zero-point energy. - THERMO - Statistical mechanics/thermochemical parsing to determine enthalpy (`H`), - entropy (`S`), and Gibbs free energy (`G`). - UNDEFINED - Placeholder for generic or unclassified calculation types. - """ - - # Structural exploration - OPT = "optimization" - OPT_TS = "transition_optimization" - CONFORMER = "conformer_search" - # Path generation - SCAN = "scan" - IRC = "intrinsic_reaction_coordinate" - MEP = "minimum_energy_path_search" - # Properties - ENERGY = "energy" - GRADIENT = "gradient" - FREQUENCY = "frequency" - THERMO = "thermochemistry" - # Fallback - UNDEFINED = "undefined" - - -class CalcStatus(StrEnum): - """Lifecycle status of a calculation. - - Attributes - ---------- - PENDING - Queued but not yet started. - RUNNING - Currently executing. - SUCCEEDED - Completed successfully. - FAILED - Terminated with an error. - """ - - PENDING = "pending" - RUNNING = "running" - SUCCEEDED = "succeeded" - FAILED = "failed" diff --git a/src/autostorage/utils.py b/src/autostorage/utils.py deleted file mode 100644 index a4c31e8..0000000 --- a/src/autostorage/utils.py +++ /dev/null @@ -1,945 +0,0 @@ -"""MESS input export and potential energy surface plotting.""" - -import io -from collections.abc import Sequence -from dataclasses import dataclass -from pathlib import Path -from typing import TYPE_CHECKING, Any, cast - -from automol import geom -from matplotlib.axes import Axes -from matplotlib.backends.backend_agg import FigureCanvasAgg -from matplotlib.figure import Figure - -from .exc import MissingPrimaryKeyError -from .models import ( - EnergyRow, - GeometryRow, - HessianRow, - ModelRow, - StageRow, - StationaryPointRow, - StepRow, -) - -if TYPE_CHECKING: - from .database import Database - -# CODATA Hartree -> kcal/mol (automol.utils.constants has no molar-energy conversion). -HARTREE_TO_KCAL_PER_MOL = 627.5094740631 - -# Non-default ground-state electronic levels (energy 1/cm, degeneracy), keyed by Hill -# formula. Species not listed here default to a single level at 0 1/cm, degeneracy -# spin + 1. -_ELECTRONIC_LEVELS_BY_HILL_FORMULA: dict[str, tuple[tuple[float, int], ...]] = { - "HO": ((0.0, 2), (140.0, 2)), -} - -# Fallback for a barrierless step's TS block, which has no geometry to compute a -# symmetry number from; must be checked by hand. -_SYMMETRY_NUMBER_PLACEHOLDER = ( - "1 ! TODO(autostorage): placeholder -- " - "no TS geometry to compute a symmetry number from, verify manually" -) - -# PES plot rendering constants. Grayscale only; dashed + muted gray flags missing data. -_LEVEL_HALF_WIDTH = 0.3 -_LEVEL_COLOR = "black" -_LEVEL_LINEWIDTH = 2.5 -_CONNECTOR_LINEWIDTH = 1.25 -_FLAGGED_COLOR = "0.6" -_FLAGGED_LINESTYLE = "--" -_MISSING_ENERGY_SENTINEL_KCAL = 0.0 -_MISSING_ENERGY_SUFFIX = " (no energy data)" -_LABEL_FONTSIZE = 9 -_Y_AXIS_LABEL = "Relative Energy (kcal/mol)" -_X_AXIS_LABEL = "Reaction Coordinate" -_DEFAULT_FIGSIZE = (6.0, 4.5) - - -@dataclass(frozen=True, slots=True) -class _FragmentData: - """Resolved geometry/frequency data for one stationary point of a species.""" - - stationary: StationaryPointRow - geometry: GeometryRow - frequencies: tuple[float, ...] | None - - -@dataclass(frozen=True, slots=True) -class _SpeciesData: - """Resolved rendering data for one non-TS stage (a well or bimolecular state).""" - - stage: StageRow - label: str - name: str - zero_energy_kcal: float | None - fragments: tuple[_FragmentData, ...] - - -def _require_stage_id(stage: StageRow) -> int: - """Return `stage.id`, raising if the stage hasn't been persisted.""" - if stage.id is None: - raise MissingPrimaryKeyError([stage]) - return stage.id - - -def _collect_stages(steps: Sequence[StepRow]) -> list[StageRow]: - """Return unique stages referenced by `steps`, in first-encounter order.""" - stages: list[StageRow] = [] - seen_ids: set[int] = set() - for step in steps: - for stage in (step.stage1, step.stage2, step.stage_ts): - if stage is None: - continue - stage_id = _require_stage_id(stage) - if stage_id not in seen_ids: - seen_ids.add(stage_id) - stages.append(stage) - return stages - - -def _auto_labels(stages: Sequence[StageRow]) -> dict[int, str]: - """Assign auto-generated `W#`/`P#` labels to stages, in first-encounter order.""" - labels: dict[int, str] = {} - well_count = 0 - bimolecular_count = 0 - for stage in stages: - if stage.is_ts: - continue - stage_id = _require_stage_id(stage) - if len(stage.stationaries) == 1: - well_count += 1 - labels[stage_id] = f"W{well_count}" - else: - bimolecular_count += 1 - labels[stage_id] = f"P{bimolecular_count}" - return labels - - -def _auto_barrier_labels(steps: Sequence[StepRow]) -> list[str]: - """Assign sequential `B#` labels, one per step, in `steps` order.""" - return [f"B{i}" for i in range(1, len(steps) + 1)] - - -def _auto_name(stage: StageRow) -> str: - """Return a Hill-formula-based comment name for a stage's fragment(s).""" - fragments = sorted(stage.stationaries, key=lambda s: s.id or 0) - return " + ".join(geom.hill_formula(f.geometry) for f in fragments) - - -def _resolve_label( - stage: StageRow, auto: dict[int, str], override: dict[int, str] -) -> str: - """Return the override label for `stage`, falling back to its auto label.""" - stage_id = _require_stage_id(stage) - return override.get(stage_id, auto[stage_id]) - - -def _resolve_name(stage: StageRow, override: dict[int, str]) -> str: - """Return the override comment name for `stage`, falling back to `_auto_name`.""" - stage_id = _require_stage_id(stage) - if stage_id in override: - return override[stage_id] - return _auto_name(stage) - - -def _energy_hartree(db: "Database", geo: GeometryRow, model: ModelRow) -> float | None: - """Return the Hartree energy of `geo` at `model`, or `None` if not found.""" - if geo.id is None or model.id is None: - raise MissingPrimaryKeyError([geo, model]) - energy = EnergyRow.query(db, geo=geo, model=model) - return energy.value if energy is not None else None - - -def _relative_energy_kcal( - value_hartree: float | None, ref_hartree: float -) -> float | None: - """Convert a Hartree energy to kcal/mol relative to `ref_hartree`.""" - if value_hartree is None: - return None - return (value_hartree - ref_hartree) * HARTREE_TO_KCAL_PER_MOL - - -def _zpe_hartree(geo: GeometryRow, frequencies: tuple[float, ...] | None) -> float: - """Return the harmonic ZPE correction for `frequencies`, or 0.0 if `None`.""" - if not frequencies: - return 0.0 - return geom.harmonic_zpv(geo, hess=[], freqs=frequencies) - - -def _zpe_corrected_energy_hartree( - db: "Database", - geo: GeometryRow, - model: ModelRow, - *, - frequencies: tuple[float, ...] | None = None, -) -> float | None: - """Return the ZPE-corrected Hartree energy of `geo` at `model`. - - Falls back to the bare electronic energy if `frequencies` is `None` and - no `HessianRow` is found at `model`. Returns `None` if no `EnergyRow` is - found. - """ - value_hartree = _energy_hartree(db, geo, model) - if value_hartree is None: - return None - if frequencies is None: - hessian = HessianRow.query(db, geo=geo, model=model) - frequencies = hessian.harmonic_frequencies if hessian is not None else None - return value_hartree + _zpe_hartree(geo, frequencies) - - -def _resolve_ref_hartree( - db: "Database", ref: StationaryPointRow, model: ModelRow -) -> float: - """Return the ZPE-corrected Hartree energy of `ref`, raising if not found.""" - ref_hartree = _zpe_corrected_energy_hartree(db, ref.geometry, model) - if ref_hartree is None: - msg = ( - f"No EnergyRow found for reference geometry {ref.geometry.id} " - f"at model {model.id}." - ) - raise ValueError(msg) - return ref_hartree - - -def _electronic_levels(geo: GeometryRow) -> tuple[tuple[float, int], ...]: - """Return `(energy_cm1, degeneracy)` ground-state electronic levels for `geo`.""" - formula = geom.hill_formula(geo) - if formula in _ELECTRONIC_LEVELS_BY_HILL_FORMULA: - return _ELECTRONIC_LEVELS_BY_HILL_FORMULA[formula] - return ((0.0, geo.spin + 1),) - - -def _indent(text: str, spaces: int) -> str: - """Indent every line of `text` by `spaces` spaces.""" - prefix = " " * spaces - return "\n".join(prefix + line for line in text.splitlines()) - - -def _format_number_columns( - values: Sequence[float], *, per_line: int = 3, width: int = 10, precision: int = 2 -) -> str: - """Render `values` right-aligned in fixed-width columns, `per_line` per row.""" - lines = [] - for i in range(0, len(values), per_line): - chunk = values[i : i + per_line] - lines.append("".join(f"{v:>{width}.{precision}f}" for v in chunk)) - return "\n".join(lines) - - -def _render_geometry_block(geo: GeometryRow) -> str: - """Render a MESS `Geometry[angstrom]` block.""" - lines = [f"Geometry[angstrom] {geo.atom_count}"] - lines.extend( - f"{symbol} {x:.6f} {y:.6f} {z:.6f}" - for symbol, (x, y, z) in zip(geo.symbols, geo.coordinates, strict=True) - ) - return "\n".join(lines) - - -def _render_frequencies_block(frequencies: tuple[float, ...]) -> str: - """Render a MESS `Frequencies[1/cm]` block.""" - header = f"Frequencies[1/cm] {len(frequencies)}" - if not frequencies: - return header - return f"{header}\n{_format_number_columns(frequencies)}" - - -def _render_electronic_levels_block(geo: GeometryRow) -> str: - """Render a MESS `ElectronicLevels[1/cm]` block.""" - levels = _electronic_levels(geo) - lines = [f"ElectronicLevels[1/cm] {len(levels)}"] - lines.extend(f"{energy:.1f} {degeneracy}" for energy, degeneracy in levels) - return "\n".join(lines) - - -def _render_zero_energy_block( - energy_kcal: float | None, *, keyword: str = "ZeroEnergy" -) -> str: - """Render a MESS `ZeroEnergy`/`GroundEnergy[kcal/mol]` line.""" - if energy_kcal is None: - return ( - f"{keyword}[kcal/mol] 0.00 ! TODO(autostorage): no EnergyRow found " - "at requested model -- fill in manually" - ) - return f"{keyword}[kcal/mol] {energy_kcal:.2f}" - - -def _render_fragment_zero_energy_block() -> str: - """Render a fragment's `ZeroEnergy[1/cm]` line. - - Always 0; the enclosing `Bimolecular` block's `GroundEnergy` line carries - the pair's relative energy. - """ - return "ZeroEnergy[1/cm] 0" - - -def _render_core_rigidrotor_block(symmetry_number: int | None) -> str: - """Render a MESS `Core RigidRotor` block. - - Falls back to a flagged placeholder when `symmetry_number` is `None`. - """ - factor = ( - _SYMMETRY_NUMBER_PLACEHOLDER if symmetry_number is None else symmetry_number - ) - return f"Core RigidRotor\n SymmetryFactor {factor}\nEnd" - - -def _render_fragment_block(fragment: _FragmentData, label: str) -> str: - """Render a `Fragment` sub-block within a `Bimolecular` species.""" - parts = [ - f"Fragment {label}", - _indent("RRHO", 2), - _indent(_render_geometry_block(fragment.geometry), 4), - _indent(_render_core_rigidrotor_block(fragment.geometry.symmetry_number), 4), - ] - if fragment.frequencies: - parts.append(_indent(_render_frequencies_block(fragment.frequencies), 4)) - parts.append(_indent(_render_fragment_zero_energy_block(), 4)) - parts.append(_indent(_render_electronic_levels_block(fragment.geometry), 4)) - parts.append(_indent("End", 2)) - return "\n".join(parts) - - -def _render_species_block( - fragment: _FragmentData, zero_energy_kcal: float | None -) -> str: - """Render a well's `Species` sub-block.""" - parts = [ - "Species", - _indent("RRHO", 2), - _indent(_render_geometry_block(fragment.geometry), 4), - _indent(_render_core_rigidrotor_block(fragment.geometry.symmetry_number), 4), - ] - if fragment.frequencies: - parts.append(_indent(_render_frequencies_block(fragment.frequencies), 4)) - parts.append(_indent(_render_zero_energy_block(zero_energy_kcal), 4)) - parts.append(_indent(_render_electronic_levels_block(fragment.geometry), 4)) - parts.append(_indent("End", 2)) - return "\n".join(parts) - - -def _render_well_block(species: _SpeciesData) -> str: - """Render a `Well` block for a single-fragment stage.""" - (fragment,) = species.fragments - header = f"Well {species.label} # {species.name}" - body = _render_species_block(fragment, species.zero_energy_kcal) - return f"{header}\n{_indent(body, 2)}\nEnd" - - -def _render_bimolecular_block(species: _SpeciesData) -> str: - """Render a `Bimolecular` block for a multi-fragment stage.""" - header = f"Bimolecular {species.label} # {species.name}" - fragment_blocks = "\n".join( - _indent(_render_fragment_block(f, geom.hill_formula(f.geometry)), 2) - for f in species.fragments - ) - ground_energy = _indent( - _render_zero_energy_block(species.zero_energy_kcal, keyword="GroundEnergy"), 2 - ) - return f"{header}\n{fragment_blocks}\n{ground_energy}\nEnd" - - -def _build_fragment_data( - db: "Database", stationary: StationaryPointRow, *, model: ModelRow -) -> _FragmentData: - """Resolve geometry and frequency data for one fragment stationary point.""" - geometry = stationary.geometry - if geometry.id is None or model.id is None: - raise MissingPrimaryKeyError([geometry, model]) - hessian = HessianRow.query(db, geo=geometry, model=model) - frequencies = hessian.harmonic_frequencies if hessian is not None else None - return _FragmentData( - stationary=stationary, geometry=geometry, frequencies=frequencies - ) - - -def _build_species_data( # noqa: PLR0913 - db: "Database", - stage: StageRow, - *, - model: ModelRow, - ref_hartree: float, - label: str, - name: str, -) -> _SpeciesData: - """Resolve all rendering data for a well/bimolecular stage.""" - fragments = tuple( - _build_fragment_data(db, s, model=model) - for s in sorted(stage.stationaries, key=lambda s: s.id or 0) - ) - values_hartree = [ - _zpe_corrected_energy_hartree( - db, fragment.geometry, model, frequencies=fragment.frequencies - ) - for fragment in fragments - ] - zero_energy_kcal = ( - None - if any(value is None for value in values_hartree) - else _relative_energy_kcal(sum(values_hartree), ref_hartree) - ) - return _SpeciesData( - stage=stage, - label=label, - name=name, - zero_energy_kcal=zero_energy_kcal, - fragments=fragments, - ) - - -def _render_barrier_block( # noqa: PLR0913 - db: "Database", - step: StepRow, - barrier_label: str, - species1: _SpeciesData, - species2: _SpeciesData, - *, - model: ModelRow, - ref_hartree: float, -) -> str: - """Render a `Barrier` block for a step with a transition state.""" - ts_stage = step.stage_ts - (ts_stationary,) = ts_stage.stationaries - ts_geometry = ts_stationary.geometry - hessian = HessianRow.query(db, geo=ts_geometry, model=model) - real_frequencies = ( - tuple(f for f in hessian.harmonic_frequencies if f > 0.0) - if hessian is not None - else None - ) - ts_frequencies = hessian.harmonic_frequencies if hessian is not None else None - value_hartree = _zpe_corrected_energy_hartree( - db, ts_geometry, model, frequencies=ts_frequencies - ) - zero_energy_kcal = _relative_energy_kcal(value_hartree, ref_hartree) - - header = ( - f"Barrier {barrier_label} {species1.label} {species2.label}" - f" # {species1.name} = {species2.name}" - ) - parts = [ - _indent("RRHO", 2), - _indent(_render_geometry_block(ts_geometry), 4), - _indent(_render_core_rigidrotor_block(ts_geometry.symmetry_number), 4), - ] - if real_frequencies: - parts.append(_indent(_render_frequencies_block(real_frequencies), 4)) - parts.append(_indent(_render_zero_energy_block(zero_energy_kcal), 4)) - parts.append(_indent(_render_electronic_levels_block(ts_geometry), 4)) - body = "\n".join(parts) - return f"{header}\n{body}\nEnd" - - -def _render_barrierless_placeholder_block( - step: StepRow, - barrier_label: str, - species1: _SpeciesData, - species2: _SpeciesData, -) -> str: - """Render a placeholder `Barrier` block for a step with no transition state. - - Not directly MESS-runnable; must be completed by hand. - """ - del step # kept in the signature for symmetry with `_render_barrier_block` - energies = [ - e - for e in (species1.zero_energy_kcal, species2.zero_energy_kcal) - if e is not None - ] - zero_energy_kcal = max(energies) if energies else None - - header = ( - f"Barrier {barrier_label} {species1.label} {species2.label}" - f" # {species1.name} = {species2.name}" - ) - lines = [ - "! TODO(autostorage): barrierless step -- no transition state exists " - "in the database.", - "! Fill in a PhaseSpaceTheory/Variational flux-parameter model by hand,", - "! or replace this Barrier block with the appropriate MESS " - "barrierless-channel construct.", - _indent("RRHO", 2), - _indent("Geometry[angstrom]", 4), - _indent( - "! TODO(autostorage): no TS geometry -- supply variational " - "geometries manually", - 6, - ), - _indent(_render_core_rigidrotor_block(None), 4), - _indent("Frequencies[1/cm] 0", 4), - _indent( - "! TODO(autostorage): no TS frequencies -- supply manually or " - "replace with a Variational/PST model", - 6, - ), - _indent(_render_zero_energy_block(zero_energy_kcal), 4), - _indent("ElectronicLevels[1/cm] 1", 4), - _indent("0.0 1", 6), - ] - return f"{header}\n" + "\n".join(lines) + "\nEnd" - - -def export_mess_input( # noqa: PLR0913 - db: "Database", - steps: Sequence[StepRow], - *, - ref: StationaryPointRow, - model: ModelRow, - labels: dict[int, str] | None = None, - names: dict[int, str] | None = None, -) -> str: - """Render `steps` as MESS `Well`/`Bimolecular`/`Barrier` input blocks. - - Non-TS stages become `Well`/`Bimolecular` blocks, auto-labeled - `W1, W2, ...`/`P1, P2, ...` and named by Hill formula, unless overridden - via `labels`/`names` (keyed by `StageRow.id`). Each step becomes one - `Barrier` block, auto-labeled `B1, B2, ...`; a barrierless step instead - gets a `TODO(autostorage)`-flagged placeholder. - - Energies are the electronic energy at `model` plus the harmonic - zero-point energy from each geometry's `HessianRow` at `model` (falling - back to the bare electronic energy where no such `HessianRow` exists), - relative to `ref`. Symmetry numbers come from `GeometryRow.symmetry_ - number`, except a barrierless step's TS block, which has no geometry. - Does not emit `Model`/`EnergyRelaxation`/`CollisionFrequency` blocks; the - caller must prepend those. - - Parameters - ---------- - db - Database to query energies and Hessians from. - steps - Elementary reaction steps to include, in output order. - ref - Stationary point defining the zero of energy (0.0 kcal/mol). - model - Level of theory used for every energy and frequency lookup. - labels, optional - Override MESS labels, keyed by `StageRow.id`. - names, optional - Override comment names, keyed by `StageRow.id`. - - Returns - ------- - The full MESS input text for the given steps. - - Raises - ------ - ValueError - No `EnergyRow` found for `ref` at `model`. - - Examples - -------- - >>> import numpy as np - >>> from autostorage import ( - ... CalcType, - ... CalculationRow, - ... Database, - ... EnergyRow, - ... GeometryRow, - ... ModelRow, - ... StageRow, - ... StationaryPointRow, - ... StepRow, - ... ) - >>> db = Database(":memory:") - >>> model = ModelRow(program="ORCA", method="b3lyp") - >>> calc = CalculationRow(model=model, calc_type=CalcType.OPT) - >>> geo1 = GeometryRow( - ... symbols=["O", "H"], - ... coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.97]]), - ... charge=0, - ... spin=1, - ... ) - >>> geo2 = GeometryRow( - ... symbols=["O", "H"], - ... coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.50]]), - ... charge=0, - ... spin=1, - ... ) - >>> s1 = StationaryPointRow(geometry=geo1, calculation=calc) - >>> s2 = StationaryPointRow(geometry=geo2, calculation=calc) - >>> db.add_all([s1, s2]) - >>> db.commit() - >>> db.add_all( - ... [ - ... EnergyRow(geometry=geo1, calculation=calc, value=-75.0), - ... EnergyRow(geometry=geo2, calculation=calc, value=-74.9), - ... ] - ... ) - >>> db.commit() - >>> stage1 = StageRow(stationaries=[s1]) - >>> stage2 = StageRow(stationaries=[s2]) - >>> step = StepRow(stage1=stage1, stage2=stage2) - >>> db.add(step) - >>> db.commit() - >>> text = export_mess_input(db, [step], ref=s1, model=model) - >>> "Well W1" in text and "Well W2" in text - True - >>> "TODO(autostorage): barrierless step" in text - True - >>> db.close() - """ - labels = labels or {} - names = names or {} - - ref_hartree = _resolve_ref_hartree(db, ref, model) - - stages = _collect_stages(steps) - auto_labels = _auto_labels(stages) - - species_by_stage_id: dict[int, _SpeciesData] = {} - for stage in stages: - if stage.is_ts: - continue - species_by_stage_id[_require_stage_id(stage)] = _build_species_data( - db, - stage, - model=model, - ref_hartree=ref_hartree, - label=_resolve_label(stage, auto_labels, labels), - name=_resolve_name(stage, names), - ) - - blocks = [] - for stage in stages: - if stage.is_ts: - continue - species = species_by_stage_id[_require_stage_id(stage)] - if len(species.fragments) == 1: - blocks.append(_render_well_block(species)) - else: - blocks.append(_render_bimolecular_block(species)) - - barrier_labels = _auto_barrier_labels(steps) - for step, barrier_label in zip(steps, barrier_labels, strict=True): - species1 = species_by_stage_id[_require_stage_id(step.stage1)] - species2 = species_by_stage_id[_require_stage_id(step.stage2)] - if step.is_barrierless: - blocks.append( - _render_barrierless_placeholder_block( - step, barrier_label, species1, species2 - ) - ) - else: - blocks.append( - _render_barrier_block( - db, - step, - barrier_label, - species1, - species2, - model=model, - ref_hartree=ref_hartree, - ) - ) - - return "\n".join(blocks) + "\n" - - -@dataclass(frozen=True, slots=True) -class _LevelPlacement: - """Resolved y-position and flagged-status of one drawn level.""" - - y_kcal: float - flagged: bool - - -@dataclass(frozen=True, slots=True) -class PESPlot: - """A rendered potential energy surface diagram. - - Attributes - ---------- - figure - The rendered figure. - axes - The axes the diagram was drawn into. - """ - - figure: Figure - axes: Axes - - def _repr_png_(self) -> bytes: - """Return a PNG-encoded snapshot of `figure`, for Jupyter's rich display.""" - buffer = io.BytesIO() - self.figure.savefig(buffer, format="png", dpi=150, bbox_inches="tight") - return buffer.getvalue() - - def save(self, path: str | Path, **savefig_kwargs: Any) -> None: # noqa: ANN401 - """Save `figure` to `path`; format is inferred from its suffix.""" - savefig_kwargs.setdefault("bbox_inches", "tight") - self.figure.savefig(path, **savefig_kwargs) - - -def _draw_level( - axes: Axes, x: float, energy_kcal: float | None, *, label: str, name: str -) -> _LevelPlacement: - """Draw one flat energy-level segment (species well/bimolecular or TS peak).""" - flagged = energy_kcal is None - y = _MISSING_ENERGY_SENTINEL_KCAL if flagged else energy_kcal - color = _FLAGGED_COLOR if flagged else _LEVEL_COLOR - axes.plot( - [x - _LEVEL_HALF_WIDTH, x + _LEVEL_HALF_WIDTH], - [y, y], - color=color, - linewidth=_LEVEL_LINEWIDTH, - linestyle=_FLAGGED_LINESTYLE if flagged else "-", - solid_capstyle="butt", - zorder=3, - ) - text = f"{label}\n{name}{_MISSING_ENERGY_SUFFIX if flagged else ''}" - axes.annotate( - text, - xy=(x, y), - xytext=(0, 6), - textcoords="offset points", - ha="center", - va="bottom", - fontsize=_LABEL_FONTSIZE, - color=color, - annotation_clip=False, - ) - return _LevelPlacement(y_kcal=y, flagged=flagged) - - -def _draw_connector( # noqa: PLR0913 - axes: Axes, - x1: float, - placement1: _LevelPlacement, - x2: float, - placement2: _LevelPlacement, - *, - linestyle: str = "-", -) -> None: - """Draw a connector line between two levels' inner edges. - - Sorts by x first, since `x1 < x2` cannot be assumed. - """ - (left_x, left), (right_x, right) = sorted( - [(x1, placement1), (x2, placement2)], key=lambda pair: pair[0] - ) - flagged = left.flagged or right.flagged - axes.plot( - [left_x + _LEVEL_HALF_WIDTH, right_x - _LEVEL_HALF_WIDTH], - [left.y_kcal, right.y_kcal], - color=_FLAGGED_COLOR if flagged else _LEVEL_COLOR, - linewidth=_CONNECTOR_LINEWIDTH, - linestyle=_FLAGGED_LINESTYLE if flagged else linestyle, - zorder=2, - ) - - -def _annotate_barrier_label(axes: Axes, x: float, y: float, label: str) -> None: - """Annotate a barrierless connector's midpoint with its `B#` label.""" - axes.annotate( - label, - xy=(x, y), - xytext=(0, 6), - textcoords="offset points", - ha="center", - va="bottom", - fontsize=_LABEL_FONTSIZE, - color=_FLAGGED_COLOR, - style="italic", - annotation_clip=False, - ) - - -def plot_pes( # noqa: PLR0913 - db: "Database", - steps: Sequence[StepRow], - *, - ref: StationaryPointRow, - model: ModelRow, - labels: dict[int, str] | None = None, - names: dict[int, str] | None = None, - ax: Axes | None = None, -) -> PESPlot: - r"""Render `steps` as a potential energy surface diagram. - - Wells/bimolecular species are derived the same way as - `export_mess_input` and drawn as flat "level" segments along an - unlabeled, ordinal x-axis (first-encounter order, not necessarily - reactant-to-product). Each step becomes a peak at its transition state - (labeled `B1, B2, ...`), connected to both stages by diagonal lines; a - barrierless step instead draws one direct dashed connector with no peak. - - Energies are the electronic energy at `model` plus the harmonic - zero-point energy from each geometry's `HessianRow` at `model` (falling - back to the bare electronic energy where no such `HessianRow` exists), - relative to `ref`, in kcal/mol. A species/TS with no `EnergyRow` at - `model` is drawn at 0.0 kcal/mol in a flagged (dashed, muted gray) style - with " (no energy data)" appended to its label. - - Parameters - ---------- - db - Database to query energies from. - steps - Elementary reaction steps to include, in output order. - ref - Stationary point defining the zero of energy (0.0 kcal/mol). - model - Level of theory used for every energy lookup. - labels, optional - Override labels, keyed by `StageRow.id`. - names, optional - Override names, keyed by `StageRow.id`. - ax, optional - Axes to draw into. If `None`, a new figure/axes is created. - - Returns - ------- - Wrapper holding the rendered figure/axes. - - Raises - ------ - ValueError - No `EnergyRow` found for `ref` at `model`. - - Examples - -------- - >>> import numpy as np - >>> from autostorage import ( - ... CalcType, - ... CalculationRow, - ... Database, - ... EnergyRow, - ... GeometryRow, - ... ModelRow, - ... StageRow, - ... StationaryPointRow, - ... StepRow, - ... ) - >>> db = Database(":memory:") - >>> model = ModelRow(program="ORCA", method="b3lyp") - >>> calc = CalculationRow(model=model, calc_type=CalcType.OPT) - >>> geo1 = GeometryRow( - ... symbols=["O", "H"], - ... coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.97]]), - ... charge=0, - ... spin=1, - ... ) - >>> geo2 = GeometryRow( - ... symbols=["O", "H"], - ... coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.50]]), - ... charge=0, - ... spin=1, - ... ) - >>> s1 = StationaryPointRow(geometry=geo1, calculation=calc) - >>> s2 = StationaryPointRow(geometry=geo2, calculation=calc) - >>> db.add_all([s1, s2]) - >>> db.commit() - >>> db.add_all( - ... [ - ... EnergyRow(geometry=geo1, calculation=calc, value=-75.0), - ... EnergyRow(geometry=geo2, calculation=calc, value=-74.9), - ... ] - ... ) - >>> db.commit() - >>> stage1 = StageRow(stationaries=[s1]) - >>> stage2 = StageRow(stationaries=[s2]) - >>> step = StepRow(stage1=stage1, stage2=stage2) - >>> db.add(step) - >>> db.commit() - >>> plot = plot_pes(db, [step], ref=s1, model=model) - >>> [line.get_linestyle() for line in plot.axes.lines].count("--") - 1 - >>> sorted(text.get_text().splitlines()[0] for text in plot.axes.texts) - ['B1', 'W1', 'W2'] - >>> plot._repr_png_().startswith(b"\x89PNG") - True - >>> db.close() - """ - labels = labels or {} - names = names or {} - - ref_hartree = _resolve_ref_hartree(db, ref, model) - - stages = _collect_stages(steps) - auto_labels = _auto_labels(stages) - non_ts_stages = [stage for stage in stages if not stage.is_ts] - x_by_stage_id = { - _require_stage_id(stage): float(i) for i, stage in enumerate(non_ts_stages) - } - - species_by_stage_id: dict[int, _SpeciesData] = {} - for stage in non_ts_stages: - species_by_stage_id[_require_stage_id(stage)] = _build_species_data( - db, - stage, - model=model, - ref_hartree=ref_hartree, - label=_resolve_label(stage, auto_labels, labels), - name=_resolve_name(stage, names), - ) - - if ax is None: - figure = Figure(figsize=_DEFAULT_FIGSIZE) - FigureCanvasAgg(figure) - axes = figure.add_subplot() - else: - axes = ax - figure = cast("Figure", ax.figure) - - placements_by_stage_id: dict[int, _LevelPlacement] = {} - for stage in non_ts_stages: - stage_id = _require_stage_id(stage) - species = species_by_stage_id[stage_id] - placements_by_stage_id[stage_id] = _draw_level( - axes, - x_by_stage_id[stage_id], - species.zero_energy_kcal, - label=species.label, - name=species.name, - ) - - barrier_labels = _auto_barrier_labels(steps) - for step, barrier_label in zip(steps, barrier_labels, strict=True): - id1 = _require_stage_id(step.stage1) - id2 = _require_stage_id(step.stage2) - x1, x2 = x_by_stage_id[id1], x_by_stage_id[id2] - placement1 = placements_by_stage_id[id1] - placement2 = placements_by_stage_id[id2] - - if step.is_barrierless: - _draw_connector( - axes, x1, placement1, x2, placement2, linestyle=_FLAGGED_LINESTYLE - ) - _annotate_barrier_label( - axes, - (x1 + x2) / 2, - (placement1.y_kcal + placement2.y_kcal) / 2, - barrier_label, - ) - else: - ts_stage = step.stage_ts - (ts_stationary,) = ts_stage.stationaries - ts_value_hartree = _zpe_corrected_energy_hartree( - db, ts_stationary.geometry, model - ) - ts_energy_kcal = _relative_energy_kcal(ts_value_hartree, ref_hartree) - x_ts = (x1 + x2) / 2 - ts_placement = _draw_level( - axes, - x_ts, - ts_energy_kcal, - label=barrier_label, - name=_resolve_name(ts_stage, names), - ) - _draw_connector(axes, x1, placement1, x_ts, ts_placement) - _draw_connector(axes, x_ts, ts_placement, x2, placement2) - - axes.set_ylabel(_Y_AXIS_LABEL) - axes.set_xlabel(_X_AXIS_LABEL) - axes.set_xticks([]) - axes.spines[["top", "right", "bottom"]].set_visible(False) - axes.grid(axis="y", color="0.85", linewidth=0.8, zorder=0) - axes.set_axisbelow(True) - axes.margins(x=0.15, y=0.15) - - return PESPlot(figure=figure, axes=axes) diff --git a/tests/conftest.py b/tests/conftest.py deleted file mode 100644 index f7dd919..0000000 --- a/tests/conftest.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Autostorage test fixtures.""" - -from collections.abc import Iterator - -import numpy as np -import pytest -from numpy.random import Generator - -from autostorage import ( - CalcType, - CalculationGeometryLink, - CalculationRow, - Database, - GeometryRow, - ModelRow, -) -from autostorage.types import Role - - -@pytest.fixture -def rng() -> Generator: - """Fixture for numpy rng.""" - return np.random.default_rng(seed=679) - - -@pytest.fixture -def database() -> Iterator[Database]: - """In-memory database fixture.""" - db = Database(":memory:") - - try: - yield db - - finally: - db.close() - - -@pytest.fixture -def model_row() -> ModelRow: - """Fixture for ModelRow.""" - return ModelRow( - program="ORCA", - program_version="6.1.1", - method="b3lyp", - basis="def2-SVP", - ) - - -@pytest.fixture -def geometry_row() -> GeometryRow: - """Fixture for GeometryRow.""" - return GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.8], [0, 0, 0], [0.8, 0, 0]]), - charge=0, - spin=0, - ) - - -@pytest.fixture -def calculation_row(model_row: ModelRow) -> CalculationRow: - """Fixture for CalculationRow.""" - return CalculationRow(model=model_row, calc_type=CalcType.UNDEFINED) - - -@pytest.fixture -def calc_geo_link( - calculation_row: CalculationRow, geometry_row: GeometryRow -) -> CalculationGeometryLink: - """Fixture for CalculationGeometryLink.""" - return CalculationGeometryLink.create( - calculation_row, geometry_row, role=Role.INPUT - ) diff --git a/tests/data/propyl_oxirane.xyz b/tests/data/propyl_oxirane.xyz deleted file mode 100644 index eaf8d47..0000000 --- a/tests/data/propyl_oxirane.xyz +++ /dev/null @@ -1,17 +0,0 @@ -15 -Coordinates from ORCA-job freq E -271.046152099005 - C -1.96724663066026 0.80558754753944 -0.26984185877179 - C -1.29117726250631 -0.36918071893157 0.34150331476840 - C -0.04004321603558 -0.80422018365461 -0.42635994313307 - C 1.03214563151731 0.24863206570885 -0.38677606439848 - O 2.34959758787060 -0.17269436687856 -0.68722367928961 - C 2.07345817914486 0.23122270310282 0.64338371469441 - H -2.06876374685845 0.86895595068965 -1.34453061196824 - H -2.52840627498046 1.50521946095658 0.33021365027168 - H -1.98290858765521 -1.22174853490473 0.38141017294270 - H -1.02615105441698 -0.15208251033875 1.37942492009934 - H 0.36599658366260 -1.72653521419848 -0.00740760059068 - H -0.30318264941652 -1.01073961341585 -1.46734173779766 - H 0.76769398303030 1.21217226524433 -0.81453115079735 - H 2.54681052278233 1.15523188532636 0.95299265615060 - H 2.07217694452179 -0.56982074624546 1.37508422781971 diff --git a/tests/data/propyl_oxirane_frequencies.gz b/tests/data/propyl_oxirane_frequencies.gz deleted file mode 100644 index d525ed5..0000000 Binary files a/tests/data/propyl_oxirane_frequencies.gz and /dev/null differ diff --git a/tests/data/propyl_oxirane_hessian.gz b/tests/data/propyl_oxirane_hessian.gz deleted file mode 100644 index d3703c8..0000000 Binary files a/tests/data/propyl_oxirane_hessian.gz and /dev/null differ diff --git a/tests/data/test.xyz b/tests/data/test.xyz deleted file mode 100644 index d07e12f..0000000 --- a/tests/data/test.xyz +++ /dev/null @@ -1,12 +0,0 @@ -10 - -C 1.02331523 -0.46639020 0.21542267 -C 0.06901154 0.72077423 0.13887406 -C -1.08374249 -0.15756372 -0.33604711 -H 1.80979527 -0.50460426 -0.52211178 -H 1.31739380 -0.78886787 1.20208171 -H -0.11131780 1.20106024 1.09818765 -H 0.36827831 1.47934510 -0.58115241 -H -1.88835969 -0.32096402 0.36369830 -H -1.39599296 -0.03447919 -1.36127092 -H -0.10838122 -1.12831031 -0.21768215 \ No newline at end of file diff --git a/tests/test_database.py b/tests/test_database.py index 1e03894..ad83cc5 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -1,236 +1,211 @@ -"""Test for database module.""" +"""Database module tests.""" + +import tempfile +from collections.abc import Generator +from pathlib import Path import pytest -from numpy.random import Generator -from sqlalchemy import inspect from sqlalchemy.exc import IntegrityError -from sqlmodel import select +from sqlalchemy.orm import Session -from autostorage import ( +from autostorage.database import Database +from autostorage.models import ( CalculationGeometryLink, CalculationRow, - Database, - GeometryRow, - GradientRow, + IdentityRow, + ModelRow, ) -from autostorage.database import ModelRow, Select, SelectStatement -from autostorage.exc import ResultShapeError - - -def test__add(database: Database, model_row: ModelRow) -> None: - """Test add to database.""" - database.add(model_row) - database.commit() - - assert model_row.id - - -def test__invalid_add(database: Database, model_row: ModelRow) -> None: - """Test invalid add to database.""" - model_row2 = model_row.model_copy(deep=True) - - database.add(model_row) - database.commit() - - # Violate ModelRow's (program, program_version, method, basis) unique constraint - database.add(model_row2) - with pytest.raises(IntegrityError): - database.commit() - - -def test__get(database: Database, model_row: ModelRow) -> None: - """Test get from database.""" - database.add(model_row) - database.commit() - assert model_row.id - - match = database.get(ModelRow, model_row.id) - assert match == model_row -def test__invalid_get(database: Database) -> None: - """Test invalid get from database.""" - with pytest.raises(LookupError): - database.get(ModelRow, 679) - - -def test__get_or_none_returns_row_or_none( - database: Database, model_row: ModelRow -) -> None: - """Test get_or_none returns the row on a hit and None on a miss.""" - database.add(model_row) - database.commit() - assert model_row.id - - assert database.get_or_none(ModelRow, model_row.id) == model_row - assert database.get_or_none(ModelRow, 679) is None - - -def test__add_all(database: Database) -> None: - """Test add_all stages multiple rows for the next flush/commit.""" - rows = [ModelRow(program="orca", method="xtb", basis=f"basis{i}") for i in range(3)] - database.add_all(rows) - database.commit() - - assert all(row.id is not None for row in rows) - assert len({row.id for row in rows}) == len(rows) - - -def test__delete(database: Database, model_row: ModelRow) -> None: - """Test delete from database.""" - database.add(model_row) - database.commit() - assert model_row.id - - database.delete(model_row) - database.commit() - with pytest.raises(LookupError, match=r"with row_id = 1 not found."): - database.get(ModelRow, model_row.id) +@pytest.fixture +def db_path() -> Generator[Path, None, None]: + """Create a temporary database path for testing.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) / "test.db" @pytest.fixture -def orca_model_statement() -> SelectStatement: - """Fixture for Statement.""" - return Select(ModelRow).where(ModelRow.program == "ORCA") - - -def test__exec_first( - database: Database, - model_row: ModelRow, - orca_model_statement: SelectStatement, -) -> None: - """Test exec first from database.""" - database.add(model_row) - match = database.exec_first(orca_model_statement) - assert match - - -def test__exec_one( - database: Database, model_row: ModelRow, orca_model_statement: SelectStatement -) -> None: - """Test exec one from database.""" - database.add(model_row) - match = database.exec_one(orca_model_statement) - assert match - - -def test__invalid_exec_one( - database: Database, orca_model_statement: SelectStatement -) -> None: - """Test delete and invalid exec one from database.""" - with pytest.raises(LookupError): - database.exec_one(orca_model_statement) - - -def test__exec_all( - database: Database, model_row: ModelRow, orca_model_statement: SelectStatement -) -> None: - """Test exec all from database.""" - database.add(model_row) - for match in database.exec_all(orca_model_statement): - assert match - - -def test__exists_true_and_false( - database: Database, model_row: ModelRow, orca_model_statement: SelectStatement -) -> None: - """Test exists() returns True for a match and False otherwise.""" - database.add(model_row) - database.commit() - - assert database.exists(orca_model_statement) is True - missing_stmt = select(ModelRow).where(ModelRow.program == "nonexistent") - assert database.exists(missing_stmt) is False - - -def test__select_statement_chaining(database: Database, model_row: ModelRow) -> None: - """Test that native SQLModel statement chaining works through exec_*.""" - database.add(model_row) - database.commit() - - stmt = select(ModelRow).where(ModelRow.program == "ORCA") - assert database.exec_first(stmt) == model_row - assert database.exec_one(stmt) == model_row - assert list(database.exec_all(stmt)) - - missing_stmt = select(ModelRow).where(ModelRow.program == "nonexistent") - assert database.exec_first(missing_stmt) is None - - -def test__select_statement_offset_and_distinct(database: Database) -> None: - """Test offset and distinct on a plain select() statement.""" - rows = [ - ModelRow(program="ORCA", method="b3lyp", basis=f"basis{i}") for i in range(3) - ] - for row in rows: - database.add(row) - database.commit() - - ordered_stmt = select(ModelRow).order_by(ModelRow.basis).offset(1) - ordered = list(database.exec_all(ordered_stmt)) - assert [r.basis for r in ordered] == ["basis1", "basis2"] - - distinct_stmt = select(ModelRow).distinct() - programs = list(database.exec_all(distinct_stmt)) - assert {p.program for p in programs} == {"ORCA"} - - -def test__merge_commits(database: Database, model_row: ModelRow) -> None: - """Test that merge() commits immediately.""" - merged = database.merge(model_row) - assert merged.id - - # A rollback after merge() must not undo it, since merge() already committed. - database._session.rollback() # noqa: SLF001 - assert database.get(ModelRow, merged.id) == merged - - -def test__session_rolls_back_on_generic_error( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, - rng: Generator, -) -> None: - """A non-IntegrityError failure rolls back, leaving the session usable.""" - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - database.add( - GradientRow( - calculation=calculation_row, - geometry=geometry_row, - value=rng.uniform(size=2), - ) - ) - - with pytest.raises(ResultShapeError): - database.commit() - - # The session must still be usable for subsequent, unrelated operations. - unrelated = ModelRow(program="ORCA", method="b3lyp") - database.add(unrelated) - database.commit() - assert unrelated.id - - -def test__link_table_reverse_lookup_indexes_exist(database: Database) -> None: - """Test that the trailing column of each composite-PK link table is indexed. - - The composite primary key on each of these tables only serves lookups keyed - by its leading column; each also needs its own index for the other direction. - """ - expected = { - "calculation_geometry_link": "calculation_id", - "calculation_trajectory_link": "calculation_id", - "trajectory_geometry_link": "trajectory_id", - "stationary_identity_link": "identity_id", - "stationary_stage_link": "stage_id", - "step_validation_link": "validation_id", - } - inspector = inspect(database.engine) - for table, column in expected.items(): - indexed_columns = { - name for idx in inspector.get_indexes(table) for name in idx["column_names"] - } - assert column in indexed_columns +def database(db_path: Path) -> Generator[Database, None, None]: + """Create a Database instance for testing.""" + db = Database(db_path) + yield db + db.close() + + +class TestDatabaseInit: + """Tests for Database initialization.""" + + def test_init_with_string_path(self, db_path: Path) -> None: + """Database can be initialized with a string path.""" + db = Database(str(db_path)) + assert db.path == db_path + assert db.engine is not None + db.close() + + def test_init_with_path_object(self, db_path: Path) -> None: + """Database can be initialized with a Path object.""" + db = Database(db_path) + assert db.path == db_path + assert db.engine is not None + db.close() + + def test_init_creates_schema(self, db_path: Path) -> None: + """Database initialization creates all tables.""" + db = Database(db_path) + # Check that tables exist by attempting to create a session and query + with db.session() as session: + # This should not raise an error if schema is created + assert session.is_active + db.close() + + def test_init_with_echo_false(self, db_path: Path) -> None: + """Database can be initialized with echo=False.""" + db = Database(db_path, echo=False) + assert db.engine.echo is False + db.close() + + def test_init_with_echo_true(self, db_path: Path) -> None: + """Database can be initialized with echo=True.""" + db = Database(db_path, echo=True) + assert db.engine.echo is True + db.close() + + def test_path_attribute(self, database: Database, db_path: Path) -> None: + """Database stores path as Path object.""" + assert isinstance(database.path, Path) + assert database.path == db_path + + def test_engine_attribute(self, database: Database) -> None: + """Database creates a SQLAlchemy engine.""" + assert database.engine is not None + assert "sqlite" in str(database.engine.url) + + +class TestDatabaseSession: + """Tests for Database.session() method.""" + + def test_session_returns_session_instance(self, database: Database) -> None: + """session() returns a SQLAlchemy Session.""" + sess = database.session() + assert isinstance(sess, Session) + sess.close() + + def test_session_is_bound_to_engine(self, database: Database) -> None: + """Session is bound to the database engine.""" + sess = database.session() + assert sess.get_bind() == database.engine + sess.close() + + def test_session_as_context_manager(self, database: Database) -> None: + """Session can be used as a context manager.""" + with database.session() as sess: + assert isinstance(sess, Session) + assert sess.is_active + + def test_multiple_sessions(self, database: Database) -> None: + """Multiple sessions can be created from same database.""" + sess1 = database.session() + sess2 = database.session() + assert sess1 is not sess2 + assert sess1.get_bind() == sess2.get_bind() + sess1.close() + sess2.close() + + def test_session_context_manager_rollback(self, database: Database) -> None: + """Session exits context manager cleanly.""" + sess = database.session() + with sess: + pass + # Session may still exist but transaction should be complete + assert sess is not None + + +class TestDatabaseClose: + """Tests for Database.close() method.""" + + def test_close_disposes_engine(self, db_path: Path) -> None: + """close() disposes the engine.""" + db = Database(db_path) + db.close() + # After dispose, new connections should be created fresh + # We can verify this indirectly by creating a new session + # (which would fail if engine was truly destroyed) + assert db.engine is not None + + def test_can_create_session_after_close(self, db_path: Path) -> None: + """A new session can be created after close() due to engine re-pooling.""" + db = Database(db_path) + db.close() + # Pool should be reset but engine still functional + sess = db.session() + assert isinstance(sess, Session) + sess.close() + + +class TestDatabaseIntegration: + """Integration tests for database operations.""" + + def test_insert_and_query_row(self, database: Database) -> None: + """Can insert and query a row from the database.""" + with database.session() as session: + # Create an identity row + identity = IdentityRow( + algorithm="RDKIT_INCHI", + kind="inchi", + value="InChI=1S/CH4/h1H4", + ) + session.add(identity) + session.commit() + + # Query it back + result = session.query(IdentityRow).filter_by(kind="inchi").first() + assert result is not None + assert result.value == "InChI=1S/CH4/h1H4" + + def test_json_serializer_sorts_keys(self, database: Database) -> None: + """JSON serializer sorts keys for consistent output.""" + with database.session() as session: + # Create a model first (required for CalculationRow) + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + # Create input_provenance with unordered keys + provenance = {"z_key": "z", "a_key": "a", "m_key": "m"} + + calc = CalculationRow( + model_id=model.id, + calc_type="energy", + input_provenance=provenance, + ) + session.add(calc) + session.commit() + + # Retrieve and verify keys are consistent + result = session.query(CalculationRow).first() + assert result is not None + # The serializer should have sorted keys during storage + assert result.input_provenance == provenance + + def test_foreign_keys_enabled(self, database: Database) -> None: + """Foreign key constraints are enforced.""" + with database.session() as session: + # Try to create a link with non-existent geometry_id + link = CalculationGeometryLink( + calculation_id=9999, # Non-existent + geometry_id=9999, # Non-existent + role="input", + ) + session.add(link) + # Foreign key constraint should prevent commit + with pytest.raises(IntegrityError): + session.commit() + + def test_concurrent_session_access(self, database: Database) -> None: + """Multiple concurrent sessions can access the database.""" + with database.session() as sess1, database.session() as sess2: + # Both sessions should be active simultaneously + assert sess1.is_active + assert sess2.is_active + # Both should access the same database + assert sess1.get_bind() == sess2.get_bind() diff --git a/tests/test_events.py b/tests/test_events.py new file mode 100644 index 0000000..3214799 --- /dev/null +++ b/tests/test_events.py @@ -0,0 +1,1765 @@ +"""Tests for SQLAlchemy ORM event listeners.""" + +import tempfile +from collections.abc import Callable, Generator +from pathlib import Path + +import numpy as np +import pytest +from sqlalchemy.exc import IntegrityError + +from autostorage.database import Database +from autostorage.models import ( + CalculationRow, + GeometryRow, + GeometryTrajectoryLink, + GradientRow, + HessianRow, + IdentityRow, + ModelRow, + StageRow, + StationaryPointRow, + StepRow, + TrajectoryRow, +) + +# Test data constants +NDIM_2 = 2 +NDIM_3 = 3 +EXPECTED_IDENTITY_COUNT_TWO = 2 +EXPECTED_EXTRAS_COUNT = 2 +NATOMS_THREE = 3 +NATOMS_TWO = 2 + + +@pytest.fixture +def db_path() -> Generator[Path, None, None]: + """Create a temporary database path for testing.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) / "test.db" + + +@pytest.fixture +def database(db_path: Path) -> Generator[Database, None, None]: + """Create a Database instance for testing.""" + db = Database(db_path) + yield db + db.close() + + +@pytest.fixture +def make_model_gradient() -> Callable[[], ModelRow]: + """Create factory for gradient calculation ModelRow.""" + + def _make() -> ModelRow: + return ModelRow(calc_type="gradient", program="psi4", method="B3LYP") + + return _make + + +@pytest.fixture +def make_model_frequency() -> Callable[[], ModelRow]: + """Create factory for frequency calculation ModelRow.""" + + def _make() -> ModelRow: + return ModelRow(calc_type="frequency", program="psi4", method="B3LYP") + + return _make + + +@pytest.fixture +def make_model_opt() -> Callable[[], ModelRow]: + """Create factory for optimization calculation ModelRow.""" + + def _make() -> ModelRow: + return ModelRow(calc_type="opt", program="psi4", method="B3LYP") + + return _make + + +@pytest.fixture +def make_calculation() -> Callable[[int], CalculationRow]: + """Create factory for CalculationRow with empty provenance.""" + + def _make(model_id: int) -> CalculationRow: + return CalculationRow( + model_id=model_id, input_provenance={}, output_provenance={} + ) + + return _make + + +@pytest.fixture +def make_geometry_2atom() -> Callable[[], GeometryRow]: + """Create factory for 2-atom geometry (C, H).""" + + def _make() -> GeometryRow: + return GeometryRow( + symbols=["C", "H"], + coordinates=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + + return _make + + +@pytest.fixture +def make_geometry_3atom() -> Callable[[], GeometryRow]: + """Create factory for 3-atom geometry (C, H, H).""" + + def _make() -> GeometryRow: + return GeometryRow( + symbols=["C", "H", "H"], + coordinates=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + charge=0, + spin=0, + ) + + return _make + + +@pytest.fixture +def make_geometry_5atom() -> Callable[[], GeometryRow]: + """Create factory for 5-atom geometry (C, H, H, H, H).""" + + def _make() -> GeometryRow: + return GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + + return _make + + +class TestSortStepStageIds: + """Tests for sort_step_stage_ids event listener.""" + + def test_stage_ids_sorted_on_insert(self, database: Database) -> None: + """stage_id1 and stage_id2 are auto-sorted on insert.""" + with database.session() as session: + # Create two stages (stage1 gets smaller id) + stage_small = StageRow(is_ts=False) + stage_large = StageRow(is_ts=False) + session.add_all([stage_small, stage_large]) + session.flush() + + # stage_large.id should be > stage_small.id + assert stage_large.id is not None + assert stage_small.id is not None + assert stage_large.id > stage_small.id + + # Create step with reversed stage IDs (large first, then small) + step = StepRow( + stage_id1=stage_large.id, + stage_id2=stage_small.id, + is_barrierless=True, + ) + session.add(step) + session.flush() + + # Verify IDs were auto-sorted (should be small < large) + assert step.stage_id1 is not None + assert step.stage_id2 is not None + assert step.stage_id1 < step.stage_id2 + assert step.stage_id1 == stage_small.id + assert step.stage_id2 == stage_large.id + + def test_stage_ids_sorted_on_update(self, database: Database) -> None: + """stage_id1 and stage_id2 are auto-sorted on update.""" + with database.session() as session: + # Create three stages + stage_a = StageRow(is_ts=False) + stage_b = StageRow(is_ts=False) + stage_c = StageRow(is_ts=False) + session.add_all([stage_a, stage_b, stage_c]) + session.flush() + + # Ensure A < B < C + assert stage_a.id is not None + assert stage_b.id is not None + assert stage_c.id is not None + assert stage_a.id < stage_b.id < stage_c.id + + # Create step with correct order (A, B) + step = StepRow( + stage_id1=stage_a.id, + stage_id2=stage_b.id, + is_barrierless=True, + ) + session.add(step) + session.flush() + + # Update with reversed order (C, A) -> should become (A, C) + step.stage_id1 = stage_c.id + step.stage_id2 = stage_a.id + session.flush() + + # Verify IDs were auto-sorted to (A, C) + assert step.stage_id1 is not None + assert step.stage_id2 is not None + assert step.stage_id1 < step.stage_id2 + assert step.stage_id1 == stage_a.id + assert step.stage_id2 == stage_c.id + + def test_equal_stage_ids_unchanged(self, database: Database) -> None: + """If stage_id1 == stage_id2, they remain unchanged.""" + with database.session() as session: + # Create a single stage + stage = StageRow(is_ts=False) + session.add(stage) + session.flush() + + # Create step with equal IDs (will fail constraint but that's OK) + step = StepRow( + stage_id1=stage.id, + stage_id2=stage.id, + is_barrierless=True, + ) + session.add(step) + + # The constraint check happens during commit + with pytest.raises(IntegrityError): + session.commit() + + def test_sorting_three_stage_ids(self, database: Database) -> None: + """Sorting works with three different stages.""" + with database.session() as session: + # Create three stages + stage_x = StageRow(is_ts=False) + stage_y = StageRow(is_ts=False) + stage_z = StageRow(is_ts=False) + session.add_all([stage_x, stage_y, stage_z]) + session.flush() + + # X < Y < Z + assert stage_x.id is not None + assert stage_y.id is not None + assert stage_z.id is not None + + # Create step with reversed order (Z, X) + step = StepRow( + stage_id1=stage_z.id, + stage_id2=stage_x.id, + is_barrierless=True, + ) + session.add(step) + session.flush() + + # Verify they were sorted to (X, Z) + assert step.stage_id1 is not None + assert step.stage_id2 is not None + assert step.stage_id1 < step.stage_id2 + assert step.stage_id1 == stage_x.id + assert step.stage_id2 == stage_z.id + + +class TestVerifyStepBarrierlessConsistency: + """Tests for verify_step_barrierless_consistency event listener.""" + + def test_barrierless_requires_no_ts(self, database: Database) -> None: + """Barrierless step (stage_id_ts=None) must have is_barrierless=True.""" + with database.session() as session: + # Create two stages + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + # Try to create barrierless step without is_barrierless=True + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=None, + is_barrierless=False, # Invalid + ) + session.add(step) + + with pytest.raises(ValueError, match="must have is_barrierless=True"): + session.flush() + + def test_non_barrierless_requires_ts(self, database: Database) -> None: + """Non-barrierless step must have stage_id_ts!=None.""" + with database.session() as session: + # Create three stages (two regular, one TS) + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + # Try to create non-barrierless step without transition state + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=None, + is_barrierless=False, + ) + session.add(step) + + with pytest.raises(ValueError, match="must have is_barrierless=True"): + session.flush() + + def test_barrierless_step_valid(self, database: Database) -> None: + """Barrierless step with is_barrierless=True and stage_id_ts=None is valid.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=None, + is_barrierless=True, + ) + session.add(step) + session.flush() + + # Should succeed + assert step.stage_id_ts is None + assert step.is_barrierless is True + + def test_step_with_ts_valid(self, database: Database) -> None: + """Step with transition state and is_barrierless=False is valid.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage1, stage2, stage_ts]) + session.flush() + + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step) + session.flush() + + # Should succeed + assert step.stage_id_ts == stage_ts.id + assert step.is_barrierless is False + + def test_step_with_ts_requires_is_barrierless_false( + self, database: Database + ) -> None: + """Step with stage_id_ts!=None must have is_barrierless=False.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage1, stage2, stage_ts]) + session.flush() + + # Try to create step with TS but is_barrierless=True (invalid) + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=stage_ts.id, + is_barrierless=True, # Invalid with TS present + ) + session.add(step) + + with pytest.raises(ValueError, match="must have is_barrierless=False"): + session.flush() + + def test_barrierless_consistency_on_update(self, database: Database) -> None: + """Barrierless consistency is checked on update.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + # Create a valid barrierless step + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=None, + is_barrierless=True, + ) + session.add(step) + session.flush() + + # Try to update to invalid state + step.is_barrierless = False + + with pytest.raises(ValueError, match="must have is_barrierless=True"): + session.flush() + + +class TestVerifyGradientShape: + """Tests for verify_gradient_shape event listener.""" + + def test_valid_gradient_shape_on_insert( + self, + database: Database, + make_model_gradient: Callable[[], ModelRow], + make_geometry_3atom: Callable[[], GeometryRow], + ) -> None: + """Gradient with correct shape (3 * natoms,) is accepted on insert.""" + with database.session() as session: + model = make_model_gradient() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = make_geometry_3atom() + session.add(geom) + session.flush() + + # Create gradient with correct shape (3 * 3 = 9 elements) + gradient = GradientRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]), + ) + gradient.geometry = geom + session.add(gradient) + session.flush() + + assert gradient.value.shape == (9,) + assert len(gradient.geometry.symbols) == NATOMS_THREE + + def test_invalid_gradient_shape_on_insert( + self, + database: Database, + make_model_gradient: Callable[[], ModelRow], + make_geometry_3atom: Callable[[], GeometryRow], + ) -> None: + """Gradient with incorrect shape raises ValueError on insert.""" + with database.session() as session: + model = make_model_gradient() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = make_geometry_3atom() + session.add(geom) + session.flush() + + # Create gradient with incorrect shape (only 6 elements instead of 9) + gradient = GradientRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), + ) + gradient.geometry = geom + session.add(gradient) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_valid_gradient_shape_on_update( + self, + database: Database, + make_model_gradient: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Gradient shape is validated on update.""" + with database.session() as session: + model = make_model_gradient() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = make_geometry_2atom() + session.add(geom) + session.flush() + + # Create gradient with correct shape (3 * 2 = 6 elements) + gradient = GradientRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), + ) + gradient.geometry = geom + session.add(gradient) + session.flush() + + # Update to new valid values (same shape) + gradient.value = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + session.flush() + + assert gradient.value.shape == (6,) + + def test_invalid_gradient_shape_on_update( + self, + database: Database, + make_model_gradient: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Updating gradient to incorrect shape raises ValueError.""" + with database.session() as session: + model = make_model_gradient() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = make_geometry_2atom() + session.add(geom) + session.flush() + + # Create gradient with correct shape (3 * 2 = 6 elements) + gradient = GradientRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), + ) + gradient.geometry = geom + session.add(gradient) + session.flush() + + # Update to invalid shape + gradient.value = np.array([1.0, 2.0, 3.0]) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_none_geometry_skipped( + self, database: Database, make_model_gradient: Callable[[], ModelRow] + ) -> None: + """Gradient with None geometry is skipped by validation.""" + with database.session() as session: + model = make_model_gradient() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create gradient with no geometry relationship loaded + gradient = GradientRow( + geometry_id=None, + calculation_id=calc.id, + value=np.array([0.1, 0.2, 0.3]), + ) + session.add(gradient) + + # Event should handle None geometry gracefully + # (the insert will fail on FK constraint, but event shouldn't crash) + + +class TestVerifyHessianShape: + """Tests for verify_hessian_shape event listener.""" + + def test_valid_hessian_shape_on_insert( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Hessian with correct shape (3*natoms, 3*natoms) is accepted on insert.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with correct shape (6x6 for 2 atoms) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + session.flush() + + assert hessian.value.shape == (6, 6) + assert len(hessian.geometry.symbols) == NATOMS_TWO + + def test_invalid_hessian_shape_on_insert( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_3atom: Callable[[], GeometryRow], + ) -> None: + """Hessian with incorrect shape raises ValueError on insert.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_3atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with incorrect shape (6x6 instead of 9x9) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_hessian_wrong_first_dimension( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Hessian with wrong first dimension raises ValueError.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with wrong first dimension (5x6 instead of 6x6) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((5, 6), dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_hessian_wrong_second_dimension( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Hessian with wrong second dimension raises ValueError.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with wrong second dimension (6x5 instead of 6x6) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 5), dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_hessian_1d_array_rejected( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Hessian with 1D array (wrong dimensionality) raises ValueError.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with 1D array (flattened, wrong dimensionality) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random(36, dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_valid_hessian_shape_on_update( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Hessian shape is validated on update.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with correct shape (6x6) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + session.flush() + + # Update to new valid values (same shape) + hessian.value = np.eye(6, dtype=np.float32) + session.flush() + + assert hessian.value.shape == (6, 6) + + def test_invalid_hessian_shape_on_update( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Updating Hessian to incorrect shape raises ValueError.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with correct shape (6x6) + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + hessian.geometry = geom + session.add(hessian) + session.flush() + + # Update to invalid shape (3x3) + hessian.value = rng.random((3, 3), dtype=np.float32) + + with pytest.raises(ValueError, match="does not match expected"): + session.flush() + + def test_none_geometry_skipped( + self, database: Database, make_model_frequency: Callable[[], ModelRow] + ) -> None: + """Hessian with None geometry is skipped by validation.""" + with database.session() as session: + model = make_model_frequency() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create Hessian with no geometry relationship loaded + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=None, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + session.add(hessian) + + # Event should handle None geometry gracefully + # (the insert will fail on FK constraint, but event shouldn't crash) + + +class TestVerifyTrajectoryGeometryNdim: + """Tests for verify_trajectory_geometry_ndim_insert event listener.""" + + def test_matching_index_and_ndim( + self, database: Database, make_geometry_5atom: Callable[[], GeometryRow] + ) -> None: + """Index length matching ndim is accepted.""" + with database.session() as session: + geom = make_geometry_5atom() + traj = TrajectoryRow(ndim=NDIM_2) + session.add_all([geom, traj]) + session.flush() + + link = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=traj.id, + index=[0, 1], # length matches ndim + ) + link.trajectory = traj + session.add(link) + session.flush() + + assert link.index == [0, 1] + assert link.trajectory.ndim == NDIM_2 + + def test_mismatched_index_and_ndim_raises( + self, database: Database, make_geometry_5atom: Callable[[], GeometryRow] + ) -> None: + """Index length not matching ndim raises ValueError.""" + with database.session() as session: + geom = make_geometry_5atom() + traj = TrajectoryRow(ndim=NDIM_3) + session.add_all([geom, traj]) + session.flush() + + link = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=traj.id, + index=[0, 1], # length doesn't match ndim + ) + link.trajectory = traj + session.add(link) + + with pytest.raises(ValueError, match="does not match"): + session.flush() + + def test_index_infers_ndim( + self, database: Database, make_geometry_5atom: Callable[[], GeometryRow] + ) -> None: + """Index length infers trajectory ndim if ndim is None.""" + with database.session() as session: + geom = make_geometry_5atom() + traj = TrajectoryRow(ndim=None) + session.add_all([geom, traj]) + session.flush() + + link = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=traj.id, + index=[0, 1, 2], # length 3 + ) + link.trajectory = traj + session.add(link) + session.flush() + + assert link.trajectory.ndim == NDIM_3 + + def test_missing_index_with_set_ndim_raises( + self, database: Database, make_geometry_5atom: Callable[[], GeometryRow] + ) -> None: + """Missing index when ndim is set raises ValueError.""" + with database.session() as session: + geom = make_geometry_5atom() + traj = TrajectoryRow(ndim=NDIM_2) + session.add_all([geom, traj]) + session.flush() + + link = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=traj.id, + index=None, # Missing index + ) + link.trajectory = traj + session.add(link) + + with pytest.raises(ValueError, match="index is missing"): + session.flush() + + def test_none_trajectory_skipped(self, database: Database) -> None: + """Link with None trajectory is skipped gracefully.""" + with database.session() as session: + link = GeometryTrajectoryLink( + geometry_id=1, # Will be invalid but event shouldn't crash + trajectory_id=None, + index=None, + ) + session.add(link) + + # Event should handle None trajectory gracefully + # (the insert will fail on FK constraint, but event shouldn't crash) + + def test_index_none_and_ndim_none( + self, database: Database, make_geometry_5atom: Callable[[], GeometryRow] + ) -> None: + """Both index and ndim None is allowed.""" + with database.session() as session: + geom = make_geometry_5atom() + traj = TrajectoryRow(ndim=None) + session.add_all([geom, traj]) + session.flush() + + link = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=traj.id, + index=None, + ) + link.trajectory = traj + session.add(link) + session.flush() + + assert link.index is None + assert link.trajectory.ndim is None + + def test_geometry_ndim_update(self, database: Database) -> None: + """Trajectory ndim is updated on link insert if previously None.""" + with database.session() as session: + geom = GeometryRow( + symbols=["C"], + coordinates=[[0.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + traj = TrajectoryRow(ndim=None) + session.add_all([geom, traj]) + session.flush() + + # First link infers ndim + link1 = GeometryTrajectoryLink( + geometry_id=geom.id, + trajectory_id=traj.id, + index=[0, 1], + ) + link1.trajectory = traj + session.add(link1) + session.flush() + + assert traj.ndim == NDIM_2 + + # Second link with same trajectory and matching index works + geom2 = GeometryRow( + symbols=["C"], + coordinates=[[1.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom2) + session.flush() + + link2 = GeometryTrajectoryLink( + geometry_id=geom2.id, + trajectory_id=traj.id, + index=[1, 2], # matches ndim + ) + link2.trajectory = traj + session.add(link2) + session.flush() + + assert link2.index == [1, 2] + + +class TestAddInchiIdentity: + """Tests for add_inchi_identity event listener.""" + + def test_inchi_identity_added_on_insert( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """InChI identity is automatically attached to a new stationary point.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C", "O"], + coordinates=[[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], + charge=0, + spin=1, + ) + session.add(geom) + session.flush() + + stat_point = StationaryPointRow( + geometry_id=geom.id, calculation_id=calc.id, order=0 + ) + session.add(stat_point) + session.flush() + + assert len(stat_point.identities) == 1 + identity = stat_point.identities[0] + assert identity.kind == "stereoisomer" + assert identity.algorithm == "rdkit inchi" + assert identity.value.startswith("InChI=") + + def test_existing_inchi_identity_reused( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Existing InChI identity is reused for duplicate geometries.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc1 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + calc2 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add_all([calc1, calc2]) + session.flush() + + # Create two identical geometries + geom1 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + geom2 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + session.add_all([geom1, geom2]) + session.flush() + + stat1 = StationaryPointRow( + geometry_id=geom1.id, calculation_id=calc1.id, order=0 + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, calculation_id=calc2.id, order=0 + ) + session.add_all([stat1, stat2]) + session.flush() + + assert len(stat1.identities) == 1 + assert len(stat2.identities) == 1 + assert stat1.identities[0].id == stat2.identities[0].id + assert stat1.identities[0].value == stat2.identities[0].value + + identity_count = session.query(IdentityRow).count() + assert identity_count == 1 + + def test_different_geometries_create_different_identities( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Different geometries create different InChI identities.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc1 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + calc2 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add_all([calc1, calc2]) + session.flush() + + geom1 = GeometryRow( + symbols=["C", "O"], + coordinates=[[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], + charge=0, + spin=1, + ) + geom2 = GeometryRow( + symbols=["C", "C"], + coordinates=[[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add_all([geom1, geom2]) + session.flush() + + stat1 = StationaryPointRow( + geometry_id=geom1.id, calculation_id=calc1.id, order=0 + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, calculation_id=calc2.id, order=0 + ) + session.add_all([stat1, stat2]) + session.flush() + + assert len(stat1.identities) == 1 + assert len(stat2.identities) == 1 + assert stat1.identities[0].id != stat2.identities[0].id + assert stat1.identities[0].value != stat2.identities[0].value + + identity_count = session.query(IdentityRow).count() + assert identity_count == EXPECTED_IDENTITY_COUNT_TWO + + +class TestAddSmilesExtras: + """Tests for add_smiles_extras_before_flush event listener.""" + + def test_smiles_extra_added_on_insert( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """SMILES is automatically attached as IdentityExtraRow to stationary point.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + session.add(geom) + session.flush() + + stat_point = StationaryPointRow( + geometry_id=geom.id, calculation_id=calc.id, order=0 + ) + session.add(stat_point) + session.flush() + + # Should have one InChI identity with one SMILES extra + assert len(stat_point.identities) == 1 + identity = stat_point.identities[0] + assert identity.algorithm == "rdkit inchi" + + # Reload to get the identity_extras relationship populated + session.expire_all() + identity = session.get(IdentityRow, identity.id) + assert identity is not None + assert len(identity.identity_extras) == EXPECTED_EXTRAS_COUNT + + # Check for SMILES extra + extras_by_attr = { + extra.attribute: extra.value for extra in identity.identity_extras + } + assert "rdkit_smiles" in extras_by_attr + assert extras_by_attr["rdkit_smiles"] == "C" # Methane SMILES + + def test_duplicate_smiles_not_created( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Duplicate SMILES are not created for the same geometry.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc1 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + calc2 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add_all([calc1, calc2]) + session.flush() + + # Create two identical geometries + geom1 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + geom2 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + session.add_all([geom1, geom2]) + session.flush() + + stat1 = StationaryPointRow( + geometry_id=geom1.id, calculation_id=calc1.id, order=0 + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, calculation_id=calc2.id, order=0 + ) + session.add_all([stat1, stat2]) + session.flush() + + # Both should share the same identity (InChI) + assert stat1.identities[0].id == stat2.identities[0].id + + # Should have two extras (SMILES + Hill) for the shared identity + session.expire_all() + identity = session.get(IdentityRow, stat1.identities[0].id) + assert identity is not None + assert len(identity.identity_extras) == EXPECTED_EXTRAS_COUNT + + # Check for SMILES extra + extras_by_attr = { + extra.attribute: extra.value for extra in identity.identity_extras + } + assert "rdkit_smiles" in extras_by_attr + assert extras_by_attr["rdkit_smiles"] == "C" + + def test_different_smiles_for_different_geometries( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Different geometries create different SMILES extras.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc1 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + calc2 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add_all([calc1, calc2]) + session.flush() + + # Methane + geom1 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + # Ethane + geom2 = GeometryRow( + symbols=["C", "C"], + coordinates=[[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add_all([geom1, geom2]) + session.flush() + + stat1 = StationaryPointRow( + geometry_id=geom1.id, calculation_id=calc1.id, order=0 + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, calculation_id=calc2.id, order=0 + ) + session.add_all([stat1, stat2]) + session.flush() + + # Should have different identities + assert stat1.identities[0].id != stat2.identities[0].id + + # Each identity should have two extras (SMILES + Hill) + session.expire_all() + identity1 = session.get(IdentityRow, stat1.identities[0].id) + identity2 = session.get(IdentityRow, stat2.identities[0].id) + assert identity1 is not None + assert identity2 is not None + + assert len(identity1.identity_extras) == EXPECTED_EXTRAS_COUNT + extras1_by_attr = { + extra.attribute: extra.value for extra in identity1.identity_extras + } + assert "rdkit_smiles" in extras1_by_attr + assert extras1_by_attr["rdkit_smiles"] == "C" + + assert len(identity2.identity_extras) == EXPECTED_EXTRAS_COUNT + extras2_by_attr = { + extra.attribute: extra.value for extra in identity2.identity_extras + } + assert "rdkit_smiles" in extras2_by_attr + # Ethane SMILES should be different from methane + assert extras2_by_attr["rdkit_smiles"] != "C" + + +class TestAddHillExtras: + """Tests for add_hill_extras_before_flush event listener.""" + + def test_hill_extra_added_on_insert( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Hill formula is attached as IdentityExtraRow to stationary point.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + session.add(geom) + session.flush() + + stat_point = StationaryPointRow( + geometry_id=geom.id, calculation_id=calc.id, order=0 + ) + session.add(stat_point) + session.flush() + + # Should have one InChI identity with two extras (SMILES + Hill formula) + assert len(stat_point.identities) == 1 + identity = stat_point.identities[0] + assert identity.algorithm == "rdkit inchi" + + # Reload to get the identity_extras relationship populated + session.expire_all() + identity = session.get(IdentityRow, identity.id) + assert identity is not None + assert len(identity.identity_extras) == EXPECTED_EXTRAS_COUNT + + # Check for both SMILES and Hill formula extras + extras_by_attr = { + extra.attribute: extra.value for extra in identity.identity_extras + } + assert "rdkit_smiles" in extras_by_attr + assert extras_by_attr["rdkit_smiles"] == "C" # Methane SMILES + assert "hill_formula" in extras_by_attr + assert extras_by_attr["hill_formula"] == "CH4" # Methane Hill formula + + def test_duplicate_hill_not_created( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Duplicate Hill formulas are not created for the same geometry.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc1 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + calc2 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add_all([calc1, calc2]) + session.flush() + + # Create two identical geometries + geom1 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + geom2 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + session.add_all([geom1, geom2]) + session.flush() + + stat1 = StationaryPointRow( + geometry_id=geom1.id, calculation_id=calc1.id, order=0 + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, calculation_id=calc2.id, order=0 + ) + session.add_all([stat1, stat2]) + session.flush() + + # Both should share the same identity (InChI) + assert stat1.identities[0].id == stat2.identities[0].id + + # Should have two extras (SMILES + Hill) for the shared identity + session.expire_all() + identity = session.get(IdentityRow, stat1.identities[0].id) + assert identity is not None + assert len(identity.identity_extras) == EXPECTED_EXTRAS_COUNT + + # Check for both SMILES and Hill formula extras + extras_by_attr = { + extra.attribute: extra.value for extra in identity.identity_extras + } + assert "rdkit_smiles" in extras_by_attr + assert extras_by_attr["rdkit_smiles"] == "C" + assert "hill_formula" in extras_by_attr + assert extras_by_attr["hill_formula"] == "CH4" + + def test_different_hill_for_different_geometries( + self, database: Database, make_model_opt: Callable[[], ModelRow] + ) -> None: + """Different geometries create different Hill formula extras.""" + with database.session() as session: + model = make_model_opt() + session.add(model) + session.flush() + + calc1 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + calc2 = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add_all([calc1, calc2]) + session.flush() + + # Methane + geom1 = GeometryRow( + symbols=["C", "H", "H", "H", "H"], + coordinates=[ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + [-1.0, 0.0, 0.0], + ], + charge=0, + spin=0, + ) + # Ethane + geom2 = GeometryRow( + symbols=["C", "C"], + coordinates=[[0.0, 0.0, 0.0], [1.5, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add_all([geom1, geom2]) + session.flush() + + stat1 = StationaryPointRow( + geometry_id=geom1.id, calculation_id=calc1.id, order=0 + ) + stat2 = StationaryPointRow( + geometry_id=geom2.id, calculation_id=calc2.id, order=0 + ) + session.add_all([stat1, stat2]) + session.flush() + + # Should have different identities + assert stat1.identities[0].id != stat2.identities[0].id + + # Each identity should have two extras (SMILES + Hill) + session.expire_all() + identity1 = session.get(IdentityRow, stat1.identities[0].id) + identity2 = session.get(IdentityRow, stat2.identities[0].id) + assert identity1 is not None + assert identity2 is not None + + assert len(identity1.identity_extras) == EXPECTED_EXTRAS_COUNT + extras1_by_attr = { + extra.attribute: extra.value for extra in identity1.identity_extras + } + assert "rdkit_smiles" in extras1_by_attr + assert extras1_by_attr["rdkit_smiles"] == "C" + assert "hill_formula" in extras1_by_attr + assert extras1_by_attr["hill_formula"] == "CH4" + + assert len(identity2.identity_extras) == EXPECTED_EXTRAS_COUNT + extras2_by_attr = { + extra.attribute: extra.value for extra in identity2.identity_extras + } + assert "rdkit_smiles" in extras2_by_attr + # Ethane SMILES should be different from methane + assert extras2_by_attr["rdkit_smiles"] != "C" + assert "hill_formula" in extras2_by_attr + # Ethane Hill formula should be different from methane + assert extras2_by_attr["hill_formula"] != "CH4" + + +class TestVerifyValidStationaryHasHessian: + """Tests for verify_valid_stationary_has_hessian event listener.""" + + def test_valid_stationary_with_hessian_accepted( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Stationary point marked valid with a Hessian is accepted.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Add a Hessian to the geometry + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + session.add(hessian) + session.flush() + + # Create a valid stationary point - should succeed + stat = StationaryPointRow( + geometry_id=geom.id, + calculation_id=calc.id, + order=0, + is_valid=True, + ) + session.add(stat) + session.flush() + + assert stat.is_valid is True + assert len(geom.hessians) == 1 + + def test_valid_stationary_without_hessian_rejected( + self, + database: Database, + make_model_opt: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Stationary point marked valid without a Hessian is rejected.""" + with database.session() as session: + model = make_model_opt() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create a valid stationary point without Hessian - should fail + stat = StationaryPointRow( + geometry_id=geom.id, + calculation_id=calc.id, + order=0, + is_valid=True, + ) + session.add(stat) + + with pytest.raises(ValueError, match="cannot be marked as valid"): + session.flush() + + def test_invalid_stationary_without_hessian_accepted( + self, + database: Database, + make_model_opt: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Stationary point not marked valid can exist without a Hessian.""" + with database.session() as session: + model = make_model_opt() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create stationary point with is_valid=False - should succeed + stat = StationaryPointRow( + geometry_id=geom.id, + calculation_id=calc.id, + order=0, + is_valid=False, + ) + session.add(stat) + session.flush() + + assert stat.is_valid is False + + def test_valid_stationary_update_to_valid_without_hessian_rejected( + self, + database: Database, + make_model_opt: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Updating stationary to valid without Hessian is rejected.""" + with database.session() as session: + model = make_model_opt() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create invalid stationary point + stat = StationaryPointRow( + geometry_id=geom.id, + calculation_id=calc.id, + order=0, + is_valid=False, + ) + session.add(stat) + session.flush() + + # Try to update to valid without Hessian - should fail + stat.is_valid = True + + with pytest.raises(ValueError, match="cannot be marked as valid"): + session.flush() + + def test_update_to_valid_with_hessian_accepted( + self, + database: Database, + make_model_frequency: Callable[[], ModelRow], + make_geometry_2atom: Callable[[], GeometryRow], + ) -> None: + """Updating stationary to valid with Hessian is accepted.""" + with database.session() as session: + model = make_model_frequency() + geom = make_geometry_2atom() + session.add_all([model, geom]) + session.flush() + + calc = CalculationRow( + model_id=model.id, input_provenance={}, output_provenance={} + ) + session.add(calc) + session.flush() + + # Create invalid stationary point + stat = StationaryPointRow( + geometry_id=geom.id, + calculation_id=calc.id, + order=0, + is_valid=False, + ) + session.add(stat) + session.flush() + + # Add a Hessian + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((6, 6), dtype=np.float32), + ) + session.add(hessian) + session.flush() + + # Update to valid with Hessian present - should succeed + stat.is_valid = True + session.flush() + + assert stat.is_valid is True diff --git a/tests/test_merge.py b/tests/test_merge.py deleted file mode 100644 index 3866009..0000000 --- a/tests/test_merge.py +++ /dev/null @@ -1,421 +0,0 @@ -"""Tests for autostorage.merge.""" - -import sqlite3 -from collections.abc import Iterator -from pathlib import Path - -import numpy as np -import pytest -from automol import Algorithm -from sqlalchemy import text -from sqlmodel import SQLModel, select - -from autostorage import ( - CalculationRow, - Database, - GeometryRow, - HessianRow, - IdentityExtraRow, - IdentityRow, - ModelRow, - StageRow, - StationaryPointRow, - StepRow, -) -from autostorage.exc import ResultShapeError -from autostorage.merge import _fk_targets, _ordered_models -from autostorage.models import StationaryIdentityLink -from autostorage.types import CalcType, CompressedArrayTypeDecorator - - -@pytest.fixture -def target() -> Iterator[Database]: - """In-memory database fixture, playing the role of a merge's target.""" - db = Database(":memory:") - try: - yield db - finally: - db.close() - - -@pytest.fixture -def source() -> Iterator[Database]: - """In-memory database fixture, playing the role of a merge's source.""" - db = Database(":memory:") - try: - yield db - finally: - db.close() - - -def _water_geometry() -> GeometryRow: - """Build a water GeometryRow.""" - return GeometryRow( - symbols=["O", "H", "H"], - coordinates=np.array([[0.0, 0.0, 0.0], [0.96, 0.0, 0.0], [-0.24, 0.93, 0.0]]), - charge=0, - spin=0, - ) - - -def _add_water_stationary(db: Database) -> StationaryPointRow: - """Add a model/calculation/geometry/stationary-point chain for water to `db`.""" - model = ModelRow.find_or_create(db, program="orca", method="xtb") - calculation = CalculationRow(model=model, calc_type=CalcType.OPT) - db.add(calculation) - geometry = _water_geometry() - db.add(geometry) - db.flush() - - stationary = StationaryPointRow( - calculation_id=calculation.id, geometry_id=geometry.id - ) - db.add(stationary) - db.commit() - return stationary - - -# --- Preconditions ----------------------------------------------------------- - - -def test__merge_from_self_raises(target: Database) -> None: - """Test that merging a database into itself is rejected.""" - with pytest.raises(ValueError, match="itself"): - target.merge_from(target) - - -def test__distinct_in_memory_databases_never_collide( - target: Database, source: Database -) -> None: - """Test that two distinct `:memory:` databases are never treated as the same.""" - report = target.merge_from(source) - assert report.copied == {} - - -def test__same_on_disk_file_collides(tmp_path: Path) -> None: - """Test that two `Database`s opened on the same on-disk file are rejected.""" - path = tmp_path / "shared.db" - a = Database(path) - b = Database(path) - try: - with pytest.raises(ValueError, match="itself"): - a.merge_from(b) - finally: - a.close() - b.close() - - -def test__missing_column_fails_precondition(tmp_path: Path) -> None: - """Test that a source DB missing an expected column fails before any row copies.""" - path = tmp_path / "stale.db" - Database(path).close() - - conn = sqlite3.connect(path) - conn.execute("ALTER TABLE geometry DROP COLUMN spin") - conn.commit() - conn.close() - - source_db = Database(path) - target_db = Database(":memory:") - try: - with pytest.raises(ValueError, match="spin"): - target_db.merge_from(source_db) - finally: - source_db.close() - target_db.close() - - -# --- FK remapping -------------------------------------------------------- - - -def test__multi_tier_fk_remapping(target: Database, source: Database) -> None: - """Test that FKs resolve to newly-copied rows across a full model->step chain. - - `target` is pre-populated with an unrelated model/calculation/geometry/ - stationary/stage chain of its own first, so every table's autoincrement - ids advance past 1 before the merge -- a broken remap that left a raw - source id in place could then plausibly resolve to one of these decoy - rows instead of failing outright, which the final content check catches - (the decoy geometry is a single helium atom, trivially distinguishable - from the water geometries actually being merged). - """ - decoy_model = ModelRow.find_or_create(target, program="decoy", method="decoy") - decoy_calculation = CalculationRow(model=decoy_model, calc_type=CalcType.OPT) - target.add(decoy_calculation) - decoy_geometry = GeometryRow( - symbols=["He"], coordinates=np.zeros((1, 3)), charge=0, spin=0 - ) - target.add(decoy_geometry) - target.flush() - decoy_stationary = StationaryPointRow( - calculation_id=decoy_calculation.id, geometry_id=decoy_geometry.id - ) - target.add(decoy_stationary) - target.commit() - target.add(StageRow(stationaries=[decoy_stationary])) - target.commit() - - model = ModelRow.find_or_create(source, program="orca", method="xtb") - calculation = CalculationRow(model=model, calc_type=CalcType.OPT) - source.add(calculation) - source.flush() - - geometry1 = _water_geometry() - geometry2 = GeometryRow( - symbols=["O", "H", "H"], - coordinates=np.array([[0.0, 0.0, 0.0], [1.1, 0.0, 0.0], [-0.3, 1.0, 0.0]]), - charge=0, - spin=0, - ) - source.add(geometry1) - source.add(geometry2) - source.flush() - - stationary1 = StationaryPointRow( - calculation_id=calculation.id, geometry_id=geometry1.id - ) - stationary2 = StationaryPointRow( - calculation_id=calculation.id, geometry_id=geometry2.id - ) - source.add(stationary1) - source.add(stationary2) - source.commit() - - stage1 = StageRow(stationaries=[stationary1]) - stage2 = StageRow(stationaries=[stationary2]) - source.add(stage1) - source.add(stage2) - source.commit() - - StepRow.find_or_create(source, stage1, stage2) - - target.merge_from(source) - - merged_steps = target.exec_all(select(StepRow)) - assert len(merged_steps) == 1 - (merged_step,) = merged_steps - assert merged_step.stage_id_ts is None - - merged_geometries = { - stationary.geometry.coordinates.tobytes() - for stage_id in (merged_step.stage_id1, merged_step.stage_id2) - for stationary in target.get(StageRow, stage_id).stationaries - } - assert merged_geometries == { - geometry1.coordinates.tobytes(), - geometry2.coordinates.tobytes(), - } - - -# --- ModelRow/IdentityRow dedup ---------------------------------------------- - - -def test__model_dedup_distinguishes_null_basis( - target: Database, source: Database -) -> None: - """Test that a differing basis is kept distinct, matching basis=None is reused.""" - ModelRow.find_or_create(target, program="orca", method="xtb") - before = len(target.exec_all(select(ModelRow))) - - ModelRow.find_or_create(source, program="orca", method="xtb") - ModelRow.find_or_create(source, program="orca", method="xtb", basis="def2-svp") - - report = target.merge_from(source) - - assert report.reused["model"] == 1 - assert report.copied["model"] == 1 - assert len(target.exec_all(select(ModelRow))) == before + report.copied["model"] - - -# --- GeometryRow dedup -------------------------------------------------------- - - -def test__geometry_dedup_reuses_identical_geometry( - target: Database, source: Database -) -> None: - """Test that an identical geometry is reused, and dependents resolve to it.""" - target_geometry = _water_geometry() - target.add(target_geometry) - target.commit() - before = len(target.exec_all(select(GeometryRow))) - - _add_water_stationary(source) - - report = target.merge_from(source) - - assert report.reused["geometry"] == 1 - assert report.copied["geometry"] == 0 - assert len(target.exec_all(select(GeometryRow))) == before - - (merged_stationary,) = target.exec_all(select(StationaryPointRow)) - assert merged_stationary.geometry_id == target_geometry.id - - -# --- Identity/events interaction --------------------------------------------- - - -def test__conformer_and_inchi_collapse_across_merged_databases( - target: Database, source: Database -) -> None: - """Test that a matching conformer from source and target share one identity.""" - stationary_t = _add_water_stationary(target) - _add_water_stationary(source) - - target.merge_from(source) - - stationaries = target.exec_all(select(StationaryPointRow)) - merged_stationary = next(s for s in stationaries if s.id != stationary_t.id) - - conformer_t = stationary_t.identity(kind=Algorithm.IRMSD.kind) - conformer_merged = merged_stationary.identity(kind=Algorithm.IRMSD.kind) - assert conformer_t is not None - assert conformer_merged is not None - assert conformer_t.id == conformer_merged.id - - inchi_t = stationary_t.identity(algorithm=Algorithm.RDKIT_INCHI) - inchi_merged = merged_stationary.identity(algorithm=Algorithm.RDKIT_INCHI) - assert inchi_t is not None - assert inchi_merged is not None - assert inchi_t.id == inchi_merged.id - - -def test__smiles_extra_and_identity_links_not_duplicated( - target: Database, source: Database -) -> None: - """Test that the auto-attached SMILES extra/links aren't duplicated by a merge.""" - _add_water_stationary(target) - _add_water_stationary(source) - - target.merge_from(source) - - (extra,) = target.exec_all(select(IdentityExtraRow)) - assert extra.attribute == "smiles" - - # Each stationary point shares the same two identities (InChI + conformer) - # rather than getting its own duplicate pair. - stationaries = target.exec_all(select(StationaryPointRow)) - shared_identity_kinds = {Algorithm.RDKIT_INCHI.kind, Algorithm.IRMSD.kind} - links = target.exec_all(select(StationaryIdentityLink)) - assert len(links) == len(stationaries) * len(shared_identity_kinds) - - -def test__non_auto_managed_identity_is_explicitly_deduped( - target: Database, source: Database -) -> None: - """Test that an identity kind other than InChI/conformer is find-or-created. - - `Algorithm.RDKIT_SMILES` is never persisted as a standalone `IdentityRow` - by any event (it's only ever folded into an `IdentityExtraRow`), so this - exercises the "explicit copy" branch of identity handling directly, - using a manually-attached identity the way a caller outside the - InChI/conformer auto-generation path might. - """ - stationary_t = _add_water_stationary(target) - smiles_t = IdentityRow.find_or_create( - target, algorithm=Algorithm.RDKIT_SMILES, value="O" - ) - assert stationary_t.id is not None - assert smiles_t.id is not None - target.add( - StationaryIdentityLink(stationary_id=stationary_t.id, identity_id=smiles_t.id) - ) - target.commit() - - stationary_s = _add_water_stationary(source) - smiles_s = IdentityRow.find_or_create( - source, algorithm=Algorithm.RDKIT_SMILES, value="O" - ) - assert stationary_s.id is not None - assert smiles_s.id is not None - source.add( - StationaryIdentityLink(stationary_id=stationary_s.id, identity_id=smiles_s.id) - ) - source.commit() - - report = target.merge_from(source) - - assert report.reused["identity"] == 1 - smiles_rows = [ - identity - for identity in target.exec_all(select(IdentityRow)) - if identity.algorithm == Algorithm.RDKIT_SMILES - ] - assert len(smiles_rows) == 1 - - -# --- Atomicity ----------------------------------------------------------- - - -def test__failed_merge_rolls_back_earlier_tiers( - target: Database, source: Database -) -> None: - """Test that a validation failure partway through leaves target unchanged. - - The invalid Hessian is planted via raw SQL rather than the ORM, since - `source`'s own shape-check event would otherwise reject it immediately - -- this simulates a stale/hand-edited source file, exactly the case - merge-time validation exists to catch. - """ - ModelRow.find_or_create(source, program="new-program", method="xtb") - - model = ModelRow.find_or_create(source, program="orca", method="xtb") - calculation = CalculationRow(model=model, calc_type=CalcType.FREQUENCY) - source.add(calculation) - geometry = _water_geometry() - source.add(geometry) - source.commit() - - bad_value = CompressedArrayTypeDecorator(dtype=np.float32).process_bind_param( - np.zeros((2, 2)), None - ) - with source.engine.begin() as conn: - conn.execute( - text( - "INSERT INTO hessian (geometry_id, calculation_id, value) " - "VALUES (:geo, :calc, :value)" - ), - {"geo": geometry.id, "calc": calculation.id, "value": bad_value}, - ) - - with pytest.raises(ResultShapeError): - target.merge_from(source) - - assert target.exec_all(select(ModelRow)) == [] - assert target.exec_all(select(HessianRow)) == [] - - -def test__commit_false_leaves_merge_uncommitted( - target: Database, source: Database -) -> None: - """Test that commit=False flushes (assigning ids) but doesn't commit.""" - ModelRow.find_or_create(source, program="orca", method="xtb") - - report = target.merge_from(source, commit=False) - - assert report.copied["model"] == 1 - (model,) = target.exec_all(select(ModelRow)) - assert model.id is not None - - target._session.rollback() # noqa: SLF001 - assert target.exec_all(select(ModelRow)) == [] - - -# --- Generic ordering (private-API safety net) -------------------------- - - -def test__ordered_models_covers_every_table() -> None: - """Test that `_ordered_models` maps every table to a class. - - Guards against `SQLModel._sa_registry` (a private attribute) silently - dropping tables after a SQLModel upgrade. - """ - assert len(_ordered_models()) == len(SQLModel.metadata.tables) - - -def test__ordered_models_is_topologically_valid() -> None: - """Test that every model's FK targets appear earlier in the returned order.""" - models = _ordered_models() - position = {cls: i for i, cls in enumerate(models)} - for cls in models: - for _, target_cls in _fk_targets(cls): - assert position[target_cls] < position[cls] diff --git a/tests/test_migrations.py b/tests/test_migrations.py deleted file mode 100644 index bb01672..0000000 --- a/tests/test_migrations.py +++ /dev/null @@ -1,57 +0,0 @@ -"""Smoke test for Alembic migrations.""" - -from pathlib import Path - -import pytest -from alembic.autogenerate import compare_metadata -from alembic.command import upgrade -from alembic.config import Config -from alembic.runtime.migration import MigrationContext -from sqlalchemy import create_engine, text -from sqlmodel import SQLModel - -import autostorage.database # noqa: F401 (registers models on SQLModel.metadata) - -REPO_ROOT = Path(__file__).resolve().parent.parent - -# Expression-based unique indexes SQLite can't reflect (see CLAUDE.md's -# migrations note); `compare_metadata` silently skips them too, so they need -# a direct `sqlite_master` check to guard against a future migration -# dropping one by accident. -_NULL_SAFE_EXPRESSION_INDEXES = ("unique_model_null_safe", "unq_step_stages_null_safe") - - -@pytest.mark.filterwarnings( - "ignore:Skipped unsupported reflection of expression-based index" - ":sqlalchemy.exc.SAWarning" -) -@pytest.mark.filterwarnings( - "ignore:autogenerate skipping metadata-specified expression-based index:UserWarning" -) -def test__migrations_upgrade_to_head_matches_current_models(tmp_path: Path) -> None: - """Test that running all migrations reproduces the schema `create_all()` builds.""" - db_path = tmp_path / "migrated.db" - config = Config(REPO_ROOT / "alembic.ini") - config.set_main_option("script_location", str(REPO_ROOT / "migrations")) - config.set_main_option("sqlalchemy.url", f"sqlite:///{db_path}") - - upgrade(config, "head") - - engine = create_engine(f"sqlite:///{db_path}") - try: - with engine.connect() as connection: - context = MigrationContext.configure(connection) - diff = compare_metadata(context, SQLModel.metadata) - - index_names = { - row[0] - for row in connection.execute( - text("SELECT name FROM sqlite_master WHERE type = 'index'") - ) - } - finally: - engine.dispose() - - assert diff == [] - for index_name in _NULL_SAFE_EXPRESSION_INDEXES: - assert index_name in index_names diff --git a/tests/test_models.py b/tests/test_models.py index 3b394e8..0e328f5 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,1117 +1,1058 @@ -"""Autostorage models tests.""" +"""Models module tests.""" -import time -from unittest import mock +import tempfile +from collections.abc import Generator +from pathlib import Path import numpy as np import pytest -from automol import Algorithm, Identity -from numpy.random import Generator -from scipy.spatial.transform import Rotation -from sqlalchemy import inspect as sa_inspect from sqlalchemy.exc import IntegrityError -from autostorage import ( - CalcStatus, - CalcType, +from autostorage.database import Database +from autostorage.models import ( CalculationGeometryLink, CalculationRow, - Database, + CalculationTrajectoryLink, EnergyRow, GeometryRow, + GeometryTrajectoryLink, GradientRow, HessianRow, + IdentityExtraRow, + IdentityRow, + IdentityStationaryLink, ModelRow, StageRow, + StageStationaryLink, StationaryPointRow, StepRow, + StepValidationLink, TrajectoryRow, ValidationRow, ) -from autostorage.exc import DataIntegrityError, MissingPrimaryKeyError, ResultShapeError -from autostorage.models import CalculationTrajectoryLink from autostorage.types import Role -def test__link_create_matches_rows_by_type( - calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that link.create() matches rows to relationships regardless of order.""" - link = CalculationGeometryLink.create( - geometry_row, calculation_row, role=Role.INPUT - ) - - assert link.calculation is calculation_row - assert link.geometry is geometry_row - assert link.role == Role.INPUT - - -def test__link_create_rejects_unmatched_row( - calculation_row: CalculationRow, model_row: ModelRow -) -> None: - """Test that link.create() raises when a row has no matching relationship.""" - with pytest.raises(ValueError, match="no unmatched relationship"): - CalculationGeometryLink.create(calculation_row, model_row, role=Role.INPUT) - - -def test__link_create_rejects_ambiguous_row_type( - calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that link.create() raises when 2+ unfilled relationships share a type. - - No current `BaseLink` subclass has two relationships to the same row - type, so this patches `sa_inspect` to simulate one, guarding the - ambiguity check against silently picking a relationship by declaration - order if such a link table is ever added. - """ - real_relationships = list(sa_inspect(CalculationGeometryLink).relationships) - duplicate_geometry_rel = next( - rel for rel in real_relationships if rel.key == "geometry" - ) - - with mock.patch("autostorage.models.sa_inspect") as mock_inspect: - mock_inspect.return_value.relationships = [ - *real_relationships, - duplicate_geometry_rel, - ] - with pytest.raises(ValueError, match="multiple unmatched relationships"): - CalculationGeometryLink.create( - geometry_row, calculation_row, role=Role.INPUT +@pytest.fixture +def db_path() -> Generator[Path, None, None]: + """Create a temporary database path for testing.""" + with tempfile.TemporaryDirectory() as tmpdir: + yield Path(tmpdir) / "test.db" + + +@pytest.fixture +def database(db_path: Path) -> Generator[Database, None, None]: + """Create a Database instance for testing.""" + db = Database(db_path) + yield db + db.close() + + +class TestGeometryRow: + """Tests for GeometryRow model.""" + + def test_create_geometry_with_list_coordinates(self, database: Database) -> None: + """GeometryRow can be created with list coordinates.""" + with database.session() as session: + geom = GeometryRow( + symbols=["C", "H"], + coordinates=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom) + session.commit() + + assert geom.id is not None + assert isinstance(geom.coordinates, np.ndarray) + assert geom.coordinates.shape == (2, 3) + + def test_create_geometry_with_numpy_coordinates(self, database: Database) -> None: + """GeometryRow can be created with numpy array coordinates.""" + with database.session() as session: + coords = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]) + geom = GeometryRow(symbols=["C", "H"], coordinates=coords, charge=0, spin=0) + session.add(geom) + session.commit() + + assert geom.id is not None + assert isinstance(geom.coordinates, np.ndarray) + np.testing.assert_array_equal(geom.coordinates, coords) + + def test_geometry_symbols_stored_as_json(self, database: Database) -> None: + """GeometryRow symbols are stored and retrieved correctly.""" + with database.session() as session: + geom = GeometryRow( + symbols=["C", "O", "H"], + coordinates=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom) + session.commit() + + result = session.query(GeometryRow).filter_by(id=geom.id).first() + assert result is not None + assert result.symbols == ["C", "O", "H"] + + def test_geometry_charge_and_spin(self, database: Database) -> None: + """GeometryRow stores charge and spin correctly.""" + with database.session() as session: + geom = GeometryRow( + symbols=["C"], + coordinates=[[0.0, 0.0, 0.0]], + charge=1, + spin=1, + ) + session.add(geom) + session.commit() + + assert geom.charge == 1 + assert geom.spin == 1 + + def test_geometry_relationships(self, database: Database) -> None: + """GeometryRow relationships are initially empty.""" + with database.session() as session: + geom = GeometryRow( + symbols=["C"], + coordinates=[[0.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom) + session.commit() + + assert geom.energies == [] + assert geom.gradients == [] + assert geom.hessians == [] + assert geom.stationary_points == [] + assert geom.trajectory_links == [] + assert geom.calculation_links == [] + + +class TestTrajectoryRow: + """Tests for TrajectoryRow model.""" + + def test_create_trajectory_with_ndim(self, database: Database) -> None: + """TrajectoryRow can be created with ndim specified.""" + with database.session() as session: + traj = TrajectoryRow(ndim=3) + session.add(traj) + session.commit() + + assert traj.id is not None + assert traj.ndim == 3 # noqa: PLR2004 + + def test_create_trajectory_without_ndim(self, database: Database) -> None: + """TrajectoryRow can be created without ndim.""" + with database.session() as session: + traj = TrajectoryRow(ndim=None) + session.add(traj) + session.commit() + + assert traj.id is not None + assert traj.ndim is None + + def test_trajectory_relationships(self, database: Database) -> None: + """TrajectoryRow relationships are initially empty.""" + with database.session() as session: + traj = TrajectoryRow(ndim=2) + session.add(traj) + session.commit() + + assert traj.geometry_links == [] + assert traj.calculation_links == [] + + +class TestModelRow: + """Tests for ModelRow model.""" + + def test_create_model_minimal(self, database: Database) -> None: + """ModelRow can be created with minimal required fields.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.commit() + + assert model.id is not None + assert model.calc_type == "energy" + assert model.program == "psi4" + assert model.method == "B3LYP" + assert model.basis is None + assert model.program_version is None + assert model.keywords == {} + + def test_create_model_complete(self, database: Database) -> None: + """ModelRow can be created with all fields specified.""" + with database.session() as session: + model = ModelRow( + calc_type="gradient", + program="orca", + program_version="5.0.3", + method="MP2", + basis="cc-pvdz", + keywords={"convergence": "tight", "scf_type": "df"}, + ) + session.add(model) + session.commit() + + assert model.id is not None + assert model.program_version == "5.0.3" + assert model.basis == "cc-pvdz" + assert model.keywords == {"convergence": "tight", "scf_type": "df"} + + def test_model_keywords_default_empty_dict(self, database: Database) -> None: + """ModelRow keywords default to empty dict.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="gaussian", method="HF") + session.add(model) + session.commit() + + assert model.keywords == {} + + +class TestCalculationRow: + """Tests for CalculationRow model.""" + + def test_create_calculation(self, database: Database) -> None: + """CalculationRow can be created with a model reference.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow( + model_id=model.id, + input_provenance={"source": "test"}, + output_provenance={"status": "success"}, + ) + session.add(calc) + session.commit() + + assert calc.id is not None + assert calc.model_id == model.id + assert calc.input_provenance == {"source": "test"} + assert calc.output_provenance == {"status": "success"} + + def test_calculation_relationships(self, database: Database) -> None: + """CalculationRow relationships are initially empty.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.commit() + + assert calc.energies == [] + assert calc.gradients == [] + assert calc.hessians == [] + assert calc.validations == [] + assert calc.stationary_points == [] + assert calc.geometry_links == [] + assert calc.trajectory_links == [] + + def test_calculation_model_relationship(self, database: Database) -> None: + """CalculationRow.model relationship works correctly.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.commit() + + assert calc.model is not None + assert calc.model.id == model.id + assert calc.model.method == "B3LYP" + + def test_calculation_requires_model(self, database: Database) -> None: + """CalculationRow requires a valid model_id.""" + with database.session() as session: + calc = CalculationRow(model_id=9999) + session.add(calc) + + with pytest.raises(IntegrityError): + session.commit() + + +class TestResultRows: + """Tests for result row models (Energy, Gradient, Hessian).""" + + def test_create_energy_row(self, database: Database) -> None: + """EnergyRow can be created with geometry and calculation.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + energy = EnergyRow( + geometry_id=geom.id, calculation_id=calc.id, value=-37.8422 + ) + session.add(energy) + session.commit() + + assert energy.id is not None + assert energy.value == -37.8422 # noqa: PLR2004 + assert energy.geometry_id == geom.id + assert energy.calculation_id == calc.id + + def test_create_gradient_row(self, database: Database) -> None: + """GradientRow can be created with numpy array.""" + with database.session() as session: + model = ModelRow(calc_type="gradient", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C", "H"], + coordinates=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom) + session.flush() + + grad_value = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]) + gradient = GradientRow( + geometry_id=geom.id, calculation_id=calc.id, value=grad_value ) + session.add(gradient) + session.commit() + + assert gradient.id is not None + np.testing.assert_array_equal(gradient.value, grad_value) + + def test_create_hessian_row(self, database: Database) -> None: + """HessianRow can be created with 2D numpy array.""" + with database.session() as session: + model = ModelRow(calc_type="frequency", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], + coordinates=[[0.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom) + session.flush() + + hess_value = np.eye(3, dtype=np.float32) + hessian = HessianRow( + geometry_id=geom.id, calculation_id=calc.id, value=hess_value + ) + session.add(hessian) + session.commit() + + assert hessian.id is not None + np.testing.assert_array_equal(hessian.value, hess_value) + + def test_result_relationships(self, database: Database) -> None: + """Result rows have correct relationships to geometry and calculation.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + energy = EnergyRow( + geometry_id=geom.id, calculation_id=calc.id, value=-37.8422 + ) + session.add(energy) + session.commit() + + assert energy.geometry is not None + assert energy.geometry.id == geom.id + assert energy.calculation is not None + assert energy.calculation.id == calc.id + + +class TestValidationRow: + """Tests for ValidationRow model.""" + + def test_create_validation(self, database: Database) -> None: + """ValidationRow can be created with calculation reference.""" + with database.session() as session: + model = ModelRow(calc_type="irc", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + validation = ValidationRow( + calculation_id=calc.id, + method="irc", + extras={"convergence": "tight"}, + ) + session.add(validation) + session.commit() + + assert validation.id is not None + assert validation.method == "irc" + assert validation.extras == {"convergence": "tight"} + + def test_validation_extras_default(self, database: Database) -> None: + """ValidationRow extras default to empty dict.""" + with database.session() as session: + model = ModelRow(calc_type="irc", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + validation = ValidationRow(calculation_id=calc.id, method="irc") + session.add(validation) + session.commit() + + assert validation.extras == {} + + +class TestStationaryPointRow: + """Tests for StationaryPointRow model.""" + + def test_create_stationary_point(self, database: Database) -> None: + """StationaryPointRow can be created with default values.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + stat_pt = StationaryPointRow(geometry_id=geom.id, calculation_id=calc.id) + session.add(stat_pt) + session.commit() + + assert stat_pt.id is not None + assert stat_pt.order == 0 + assert stat_pt.is_pseudo is False + assert stat_pt.is_valid is False + + def test_create_transition_state(self, database: Database) -> None: + """StationaryPointRow can represent a transition state (order=1).""" + with database.session() as session: + model = ModelRow(calc_type="opt_ts", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + # Add Hessian to make stationary point valid + rng = np.random.default_rng() + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=rng.random((3, 3), dtype=np.float32), + ) + session.add(hessian) + session.flush() + + stat_pt = StationaryPointRow( + geometry_id=geom.id, calculation_id=calc.id, order=1, is_valid=True + ) + session.add(stat_pt) + session.commit() + + assert stat_pt.order == 1 + assert stat_pt.is_valid is True + + def test_stationary_point_relationships(self, database: Database) -> None: + """StationaryPointRow relationships work correctly.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + stat_pt = StationaryPointRow(geometry_id=geom.id, calculation_id=calc.id) + session.add(stat_pt) + session.commit() + + assert stat_pt.geometry is not None + assert stat_pt.geometry.id == geom.id + assert stat_pt.calculation is not None + assert stat_pt.calculation.id == calc.id + # Note: auto-generated InChI identity from event listener + assert len(stat_pt.identities) >= 1 + assert stat_pt.stages == [] + + +class TestStageRow: + """Tests for StageRow model.""" + + def test_create_stage_non_ts(self, database: Database) -> None: + """StageRow can be created as a non-TS stage.""" + with database.session() as session: + stage = StageRow(is_ts=False) + session.add(stage) + session.commit() + + assert stage.id is not None + assert stage.is_ts is False + + def test_create_stage_ts(self, database: Database) -> None: + """StageRow can be created as a TS stage.""" + with database.session() as session: + stage = StageRow(is_ts=True) + session.add(stage) + session.commit() + + assert stage.id is not None + assert stage.is_ts is True + + def test_stage_relationships(self, database: Database) -> None: + """StageRow relationships are initially empty.""" + with database.session() as session: + stage = StageRow(is_ts=False) + session.add(stage) + session.commit() + + assert stage.stationaries == [] + assert stage.steps == [] + + +class TestStepRow: + """Tests for StepRow model.""" + + def test_create_barrierless_step(self, database: Database) -> None: + """StepRow can be created as a barrierless step.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=None, + is_barrierless=True, + ) + session.add(step) + session.commit() + + assert step.id is not None + assert step.stage_id_ts is None + assert step.is_barrierless is True + + def test_create_step_with_ts(self, database: Database) -> None: + """StepRow can be created with a transition state.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage1, stage2, stage_ts]) + session.flush() + + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step) + session.commit() + + assert step.id is not None + assert step.stage_id_ts == stage_ts.id + assert step.is_barrierless is False + + def test_step_unique_constraint(self, database: Database) -> None: + """StepRow enforces unique constraint on stage IDs.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage1, stage2, stage_ts]) + session.flush() + + step1 = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step1) + session.flush() + + # Try to create duplicate step + step2 = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step2) + + with pytest.raises(IntegrityError): + session.commit() + + def test_step_stage_relationships(self, database: Database) -> None: + """StepRow stage relationships work correctly.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + stage_ts = StageRow(is_ts=True) + session.add_all([stage1, stage2, stage_ts]) + session.flush() + + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + stage_id_ts=stage_ts.id, + is_barrierless=False, + ) + session.add(step) + session.commit() + + assert step.stage1 is not None + assert step.stage1.id == stage1.id + assert step.stage2 is not None + assert step.stage2.id == stage2.id + assert step.stage_ts is not None + assert step.stage_ts.id == stage_ts.id + + +class TestIdentityRow: + """Tests for IdentityRow model.""" + + def test_create_identity(self, database: Database) -> None: + """IdentityRow can be created with kind, algorithm, and value.""" + with database.session() as session: + identity = IdentityRow( + kind="stereoisomer", + algorithm="rdkit inchi", + value="InChI=1S/CH4/h1H4", + ) + session.add(identity) + session.commit() + + assert identity.id is not None + assert identity.kind == "stereoisomer" + assert identity.algorithm == "rdkit inchi" + assert identity.value == "InChI=1S/CH4/h1H4" + + def test_identity_unique_constraint(self, database: Database) -> None: + """IdentityRow enforces unique constraint on kind, algorithm, value.""" + with database.session() as session: + identity1 = IdentityRow( + kind="stereoisomer", + algorithm="rdkit inchi", + value="InChI=1S/CH4/h1H4", + ) + session.add(identity1) + session.flush() + + # Try to create duplicate identity + identity2 = IdentityRow( + kind="stereoisomer", + algorithm="rdkit inchi", + value="InChI=1S/CH4/h1H4", + ) + session.add(identity2) + + with pytest.raises(IntegrityError): + session.commit() + + def test_identity_relationships(self, database: Database) -> None: + """IdentityRow relationships are initially empty.""" + with database.session() as session: + identity = IdentityRow( + kind="stereoisomer", algorithm="rdkit inchi", value="InChI=1S/CH4/h1H4" + ) + session.add(identity) + session.commit() -def test__row_timestamps_set_on_create(database: Database) -> None: - """Test that created_at/updated_at are populated by the database on insert.""" - row = ModelRow(program="orca", method="xtb") - database.add(row) - database.commit() - - assert row.created_at is not None - assert row.updated_at is not None - - -def test__row_updated_at_advances_on_update(database: Database) -> None: - """Test that updated_at advances on a later commit while created_at doesn't.""" - row = ModelRow(program="orca", method="xtb") - database.add(row) - database.commit() - created_at, updated_at = row.created_at, row.updated_at - assert created_at is not None - assert updated_at is not None - - # SQLite's CURRENT_TIMESTAMP has one-second resolution. - time.sleep(1.1) - row.basis = "def2-svp" - database.add(row) - database.commit() - - assert row.updated_at is not None - assert row.created_at == created_at - assert row.updated_at > updated_at - - -def test__model_find_or_create_reuses_matching_row(database: Database) -> None: - """Test that find_or_create returns the same row for repeated calls.""" - first = ModelRow.find_or_create(database, program="orca", method="xtb") - second = ModelRow.find_or_create(database, program="orca", method="xtb") - - assert first.id is not None - assert first.id == second.id - - -def test__model_find_or_create_distinguishes_basis(database: Database) -> None: - """Test that find_or_create treats a differing basis as a distinct model.""" - no_basis = ModelRow.find_or_create(database, program="orca", method="xtb") - with_basis = ModelRow.find_or_create( - database, program="orca", method="xtb", basis="def2-svp" - ) - - assert no_basis.id != with_basis.id - - -def test__model_null_safe_index_catches_duplicate(database: Database) -> None: - """Test that a direct duplicate insert (bypassing find_or_create) is rejected. - - `unique_model` alone doesn't catch this, since SQL treats NULL as distinct - from itself; `unique_model_null_safe` is the defense-in-depth index that does. - """ - database.add(ModelRow(program="orca", method="xtb")) - database.add(ModelRow(program="orca", method="xtb")) - - with pytest.raises(IntegrityError): - database.commit() - - -def test__calculation_default_status_is_pending(model_row: ModelRow) -> None: - """Test that a bare CalculationRow defaults to PENDING with no error message.""" - calculation = CalculationRow(model=model_row, calc_type=CalcType.UNDEFINED) - - assert calculation.status == CalcStatus.PENDING - assert calculation.error_message is None - - -def test__calculation_status_transitions( - database: Database, model_row: ModelRow -) -> None: - """Test that status/error_message round-trip through the database.""" - calculation = CalculationRow( - model=model_row, - calc_type=CalcType.UNDEFINED, - status=CalcStatus.FAILED, - error_message="boom", - ) - database.add(calculation) - database.commit() - assert calculation.id is not None - - fetched = database.get(CalculationRow, calculation.id) - assert fetched.status == CalcStatus.FAILED - assert fetched.error_message == "boom" - - -def test__calculation_geometry_role_properties( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, -) -> None: - """Test that input_geometries/output_geometries filter links by role.""" - output_geometry = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.9], [0, 0, 0], [0.9, 0, 0]]), - charge=0, - spin=0, - ) - output_link = CalculationGeometryLink.create( - calculation_row, output_geometry, role=Role.OUTPUT - ) - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - database.add(output_geometry) - database.add(output_link) - database.commit() - - assert calculation_row.input_geometries == [geometry_row] - assert calculation_row.output_geometries == [output_geometry] - - -def test__calculation_trajectory_role_properties( - database: Database, calculation_row: CalculationRow -) -> None: - """Test that input_trajectories/output_trajectories filter links by role.""" - # Committed one at a time: TrajectoryRow has no non-base columns, and - # SQLite's batched multi-row insert can't apply the created_at - # server_default to two such rows in a single flush. - input_trajectory = TrajectoryRow() - database.add(input_trajectory) - database.commit() - output_trajectory = TrajectoryRow() - database.add(output_trajectory) - database.commit() - - input_link = CalculationTrajectoryLink.create( - calculation_row, input_trajectory, role=Role.INPUT - ) - output_link = CalculationTrajectoryLink.create( - calculation_row, output_trajectory, role=Role.OUTPUT - ) - database.add(calculation_row) - database.add(input_link) - database.add(output_link) - database.commit() - - assert calculation_row.input_trajectories == [input_trajectory] - assert calculation_row.output_trajectories == [output_trajectory] - - -def test__validation_requires_calculation(database: Database) -> None: - """Test that a ValidationRow without a calculation is rejected.""" - database.add(ValidationRow(method="irc")) - - with pytest.raises(IntegrityError): - database.commit() - - -def test__gradient_shape( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, - rng: Generator, -) -> None: - """Test gradient shape is validated before committing to database.""" - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - - gradient = GradientRow( - calculation=calculation_row, - geometry=geometry_row, - value=rng.uniform(size=2), - ) - database.add(gradient) - with pytest.raises(ResultShapeError): - database.commit() - - -def test__hessian_shape( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, - rng: Generator, -) -> None: - """Test hessian shape is validated before committing to database.""" - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - database.commit() - - hess = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=rng.uniform(size=(3, 2)), - ) - database.add(hess) - - with pytest.raises(ResultShapeError): - database.commit() - - -def test__geometry_symbols_immutable_after_insert( - database: Database, geometry_row: GeometryRow -) -> None: - """Test that mutating symbols after insert is rejected.""" - database.add(geometry_row) - database.commit() - - geometry_row.symbols = ["H", "O", "O"] - database.add(geometry_row) - with pytest.raises(DataIntegrityError, match="symbols"): - database.commit() - - -def test__geometry_coordinates_immutable_after_insert( - database: Database, geometry_row: GeometryRow -) -> None: - """Test that mutating coordinates after insert is rejected.""" - database.add(geometry_row) - database.commit() - - geometry_row.coordinates = np.array(geometry_row.coordinates) + 0.1 - database.add(geometry_row) - with pytest.raises(DataIntegrityError, match="coordinates"): - database.commit() - - -def test__geometry_charge_and_spin_remain_mutable( - database: Database, geometry_row: GeometryRow -) -> None: - """Test that charge/spin can still be updated after insert.""" - database.add(geometry_row) - database.commit() - assert geometry_row.id - - geometry_row.charge = 1 - geometry_row.spin = 1 - database.add(geometry_row) - database.commit() - - fetched = database.get(GeometryRow, geometry_row.id) - assert fetched.charge == 1 - assert fetched.spin == 1 - - -def test__geometry_find_or_create_reuses_matching_row( - database: Database, geometry_row: GeometryRow -) -> None: - """Test that find_or_create returns the same row for repeated calls.""" - first = GeometryRow.find_or_create( - database, - symbols=list(geometry_row.symbols), - coordinates=np.array(geometry_row.coordinates), - charge=geometry_row.charge, - spin=geometry_row.spin, - ) - second = GeometryRow.find_or_create( - database, - symbols=list(geometry_row.symbols), - coordinates=np.array(geometry_row.coordinates), - charge=geometry_row.charge, - spin=geometry_row.spin, - ) - - assert first.id is not None - assert first.id == second.id - - -def test__geometry_unique_hash_catches_direct_duplicate_insert( - database: Database, geometry_row: GeometryRow -) -> None: - """Test that a direct duplicate insert (bypassing find_or_create) is rejected.""" - duplicate = GeometryRow( - symbols=list(geometry_row.symbols), - coordinates=np.array(geometry_row.coordinates), - charge=geometry_row.charge, - spin=geometry_row.spin, - ) - database.add(geometry_row) - database.add(duplicate) - - with pytest.raises(IntegrityError): - database.commit() - - -def test__geometry_near_duplicate_is_not_deduped( - database: Database, geometry_row: GeometryRow, rng: Generator -) -> None: - """Test that a rotated/translated/jittered near-duplicate is a distinct row. - - `geometry_hash` only catches bit-identical content; chemically-equivalent - but numerically distinct conformers are handled separately (and more - coarsely) by `events.py`'s InChI/conformer identity matching. - """ - near_duplicate = _jittered_copy(geometry_row, rng) - - database.add(geometry_row) - database.add(near_duplicate) - database.commit() - - assert geometry_row.id is not None - assert near_duplicate.id is not None - assert geometry_row.id != near_duplicate.id - - -def test__hessian_properties( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, - rng: Generator, -) -> None: - """Test hessian harmonic frequencies and order properties.""" - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - - n = geometry_row.atom_count - hessian = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=rng.uniform(size=(3 * n, 3 * n)), - ) - assert hessian.harmonic_frequencies - assert hessian.order - - -def test__hessian_frequency_cache_invalidated_on_value_update( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, - rng: Generator, -) -> None: - """Test that updating `value` invalidates the cached harmonic frequencies. - - `harmonic_frequencies` is a `functools.cached_property`; `Session.commit()` - doesn't clear cached-property entries (only mapped attributes), so this - guards against `invalidate_hessian_frequency_cache` regressing and leaving - stale frequencies/order behind after an in-place `value` update. - """ - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - - n = geometry_row.atom_count - hessian = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=rng.uniform(size=(3 * n, 3 * n)), - ) - database.add(hessian) - database.commit() - - original_frequencies = hessian.harmonic_frequencies - assert "harmonic_frequencies" in hessian.__dict__ - - hessian.value = rng.uniform(size=(3 * n, 3 * n)) - database.add(hessian) - database.commit() - - assert hessian.harmonic_frequencies != original_frequencies - - -def test__geometry_symmetry_number() -> None: - """Test that symmetry_number is computed from a geometry's point group. - - Uses a properly symmetric (C2v) water geometry -- unlike `geometry_row`, - whose bond lengths/angle are arbitrary and so has no symmetry to detect -- - to guard against a broken computation silently returning 1 for everything. - """ - oh = 0.9584 - angle = np.radians(104.45) - o = np.array([0.0, 0.0, 0.0]) - h1 = np.array([0.0, oh * np.sin(angle / 2), oh * np.cos(angle / 2)]) - h2 = np.array([0.0, -oh * np.sin(angle / 2), oh * np.cos(angle / 2)]) - water = GeometryRow( - symbols=["O", "H", "H"], coordinates=np.array([o, h1, h2]), charge=0, spin=0 - ) - expected_symmetry_number = 2 - assert water.symmetry_number == expected_symmetry_number - - -def test__result_query( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - calc_geo_link: CalculationGeometryLink, - rng: Generator, -) -> None: - """Test querying of result tables.""" - database.add(calculation_row) - database.add(geometry_row) - database.add(calc_geo_link) - database.commit() - - n = geometry_row.atom_count - hess = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=rng.uniform(size=(3 * n, 3 * n)), - ) - database.add(hess) - - database.commit() - - hess2 = HessianRow.query(database, geo=geometry_row, model=calculation_row.model) - assert hess2 - assert hess2.id == hess.id - - -def test__provenance_query_matches_regardless_of_dict_key_order( - database: Database, geometry_row: GeometryRow, model_row: ModelRow -) -> None: - """Test that provenance-filtered queries ignore dict key insertion order. - - SQLite JSON columns compare by exact serialized text, so two dicts built with - different key insertion order would previously fail to match even though - they're equal in Python; `Database`'s `json_serializer` canonicalizes key - order on write to fix this. - """ - calculation = CalculationRow( - model=model_row, calc_type=CalcType.ENERGY, input_provenance={"b": 1, "a": 2} - ) - database.add(calculation) - database.add(geometry_row) - database.commit() - - energy = EnergyRow(calculation=calculation, geometry=geometry_row, value=-1.0) - database.add(energy) - database.commit() - - found = EnergyRow.query( - database, geo=geometry_row, model=model_row, prov={"a": 2, "b": 1} - ) - assert found - assert found.id == energy.id - - -def test__stationary_inchi( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test InChI is attached before committing to database.""" - database.add(calculation_row) - database.add(geometry_row) - - stationary = StationaryPointRow( - calculation=calculation_row, geometry=geometry_row, order=0 - ) - database.add(stationary) - database.commit() - - assert stationary.identities[0].value == "InChI=1S/H2O/h1H2" - - -def test__stationary_inchi_resolves_unattached_geometry_id( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test InChI/conformer identities attach when only `geometry_id` is set. - - Regression test: `add_inchi_identities`/`assign_conformer_ids` must - resolve the geometry via the session (like the shape/order validators - do), rather than reading the `.geometry` relationship directly, which - stays unpopulated until the ORM syncs it. - """ - database.add(calculation_row) - database.add(geometry_row) - database.flush() - - stationary = StationaryPointRow( - calculation_id=calculation_row.id, geometry_id=geometry_row.id, order=0 - ) - database.add(stationary) - database.commit() - - assert stationary.identities[0].value == "InChI=1S/H2O/h1H2" - assert stationary.identity(algorithm=Algorithm.IRMSD) is not None - - -def test__stationary_identity_matches_by_kind_and_algorithm( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test identity() lookup by kind, algorithm, both, and no match.""" - database.add(calculation_row) - database.add(geometry_row) - - stationary = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary) - database.commit() - - inchi = stationary.identity(kind="stereoisomer") - assert inchi - assert inchi.value == "InChI=1S/H2O/h1H2" - - conformer = stationary.identity(algorithm=Algorithm.IRMSD) - assert conformer - assert conformer.kind == "conformer" - - assert ( - stationary.identity(kind="stereoisomer", algorithm=Algorithm.RDKIT_INCHI) - is inchi - ) - assert stationary.identity(kind="nonexistent") is None - - -def test__stationary_order_hessian_first( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test stationary point order is validated when geometry Hessian is present. - - Corrects a valid StationaryPointRow marked as invalid. - """ - database.add(calculation_row) - database.add(geometry_row) - - n = geometry_row.atom_count - hessian_row = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=np.zeros((3 * n, 3 * n)), - ) - database.add(hessian_row) - - stationary = StationaryPointRow( - calculation=calculation_row, geometry=geometry_row, order=0, is_valid=False - ) - database.add(stationary) - assert not stationary.is_valid - - database.commit() - assert stationary.is_valid - - -def test__stationary_order_hessian_second( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test stationary point order is validated when geometry Hessian is present. - - Corrects an invalid StationaryPointRow marked as valid. - """ - database.add(calculation_row) - database.add(geometry_row) - - stationary = StationaryPointRow( - calculation=calculation_row, geometry=geometry_row, order=1, is_valid=True - ) - database.add(stationary) - assert stationary.is_valid - - n = geometry_row.atom_count - hessian_row = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=np.zeros((3 * n, 3 * n)), - ) - database.add(hessian_row) - - database.commit() - assert not stationary.is_valid - - -def test__hessian_delete_leaves_is_valid_correct_with_remaining_hessian( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that deleting one of two agreeing Hessians keeps is_valid correct.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - n = geometry_row.atom_count - hessian1 = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=np.zeros((3 * n, 3 * n)), - ) - hessian2 = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=np.zeros((3 * n, 3 * n)), - ) - database.add(hessian1) - database.add(hessian2) - - stationary = StationaryPointRow( - calculation=calculation_row, geometry=geometry_row, order=0 - ) - database.add(stationary) - database.commit() - assert stationary.is_valid - - database.delete(hessian1) - assert stationary.is_valid - - -def test__hessian_delete_leaves_is_valid_untouched_when_no_hessians_remain( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that deleting the last Hessian doesn't reset is_valid to False.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - n = geometry_row.atom_count - hessian = HessianRow( - calculation=calculation_row, - geometry=geometry_row, - value=np.zeros((3 * n, 3 * n)), - ) - database.add(hessian) - - stationary = StationaryPointRow( - calculation=calculation_row, geometry=geometry_row, order=0 - ) - database.add(stationary) - database.commit() - assert stationary.is_valid - - database.delete(hessian) - assert stationary.is_valid - - -def test__stationary_query( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test querying of stationary points.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary) - - ident = Identity.from_geometry(geo=geometry_row, algorithm=Algorithm.RDKIT_INCHI) - stationary2 = StationaryPointRow.query( - database, ident=ident, model=calculation_row.model - ) - - assert stationary2 - assert stationary2.id == stationary.id - - -def test__invalid_stationary_query( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test invalid querying of stationary points.""" - ident = Identity.from_geometry(geo=geometry_row, algorithm=Algorithm.RDKIT_INCHI) - with pytest.raises(MissingPrimaryKeyError): - StationaryPointRow.query(database, ident=ident, model=calculation_row.model) - - -def test__stage_and_step_query( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test querying of stages and steps built on the chainable Query API.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary1 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - stationary2 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary1) - database.add(stationary2) - database.commit() - - stage1 = StageRow(stationaries=[stationary1]) - stage2 = StageRow(stationaries=[stationary2]) - database.add(stage1) - database.add(stage2) - database.commit() - - stage_match = StageRow.query(database, [stationary1]) - assert stage_match - assert stage_match.id == stage1.id - - step = StepRow(stage1=stage1, stage2=stage2) - database.add(step) - database.commit() - - step_match = StepRow.query(database, stage1, stage2) - assert step_match - assert step_match.id == step.id - - -def test__stage_find_or_create_reuses_matching_row( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that find_or_create returns the same row for repeated calls.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary) - database.commit() - - first = StageRow.find_or_create(database, [stationary]) - second = StageRow.find_or_create(database, [stationary]) - - assert first.id is not None - assert first.id == second.id - - -def test__stage_find_or_create_distinguishes_is_ts( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that find_or_create treats differing is_ts as distinct stages.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary) - database.commit() - - non_ts = StageRow.find_or_create(database, [stationary], is_ts=False) - ts = StageRow.find_or_create(database, [stationary], is_ts=True) - - assert non_ts.id != ts.id - - -def test__step_find_or_create_reuses_matching_row( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that find_or_create returns the same row for repeated calls.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary1 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - stationary2 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary1) - database.add(stationary2) - database.commit() - - stage1 = StageRow(stationaries=[stationary1]) - stage2 = StageRow(stationaries=[stationary2]) - database.add(stage1) - database.add(stage2) - database.commit() - - first = StepRow.find_or_create(database, stage1, stage2) - second = StepRow.find_or_create(database, stage1, stage2) - - assert first.id is not None - assert first.id == second.id - - -def test__step_null_safe_index_catches_barrierless_duplicate( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a direct duplicate barrierless step (bypassing StepRow.query) fails. - - `unq_step_stages` alone doesn't catch this, since `stage_id_ts` is NULL for - both rows and SQL treats NULL as distinct from itself; `unq_step_stages_null_safe` - is the defense-in-depth index that does. - """ - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary1 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - stationary2 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary1) - database.add(stationary2) - database.commit() - - stage1 = StageRow(stationaries=[stationary1]) - stage2 = StageRow(stationaries=[stationary2]) - database.add(stage1) - database.add(stage2) - database.commit() - - database.add(StepRow(stage1=stage1, stage2=stage2)) - database.add(StepRow(stage1=stage1, stage2=stage2)) - - with pytest.raises(IntegrityError): - database.commit() - - -def test__step_rejects_ts_stage_as_stage1_or_stage2( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a TS stage cannot be used as stage1/stage2.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationary1 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - stationary2 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary1) - database.add(stationary2) - database.commit() - - stage_ts = StageRow(stationaries=[stationary1], is_ts=True) - stage2 = StageRow(stationaries=[stationary2]) - database.add(stage_ts) - database.add(stage2) - database.commit() - - database.add(StepRow(stage1=stage_ts, stage2=stage2)) - with pytest.raises(DataIntegrityError, match="transition-state"): - database.commit() - - -def test__step_rejects_non_ts_stage_as_stage_ts( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a non-TS stage cannot be used as stage_ts.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationaries = [ - StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - for _ in range(3) - ] - for stationary in stationaries: - database.add(stationary) - database.commit() - - stage1, stage2, stage3 = (StageRow(stationaries=[s]) for s in stationaries) - database.add(stage1) - database.add(stage2) - database.add(stage3) - database.commit() - - database.add(StepRow(stage1=stage1, stage2=stage2, stage_ts=stage3)) - with pytest.raises(DataIntegrityError, match="stage_ts"): - database.commit() - - -def test__step_accepts_consistent_ts_configuration( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a step with a genuine TS stage commits without error.""" - database.add(calculation_row) - database.add(geometry_row) - database.commit() - - stationaries = [ - StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - for _ in range(3) - ] - for stationary in stationaries: - database.add(stationary) - database.commit() - - stage1 = StageRow(stationaries=[stationaries[0]]) - stage2 = StageRow(stationaries=[stationaries[1]]) - stage_ts = StageRow(stationaries=[stationaries[2]], is_ts=True) - database.add(stage1) - database.add(stage2) - database.add(stage_ts) - database.commit() - - step = StepRow(stage1=stage1, stage2=stage2, stage_ts=stage_ts) - database.add(step) - database.commit() - - assert step.id is not None - assert not step.is_barrierless - - -def _hooh_geometry_row(dihedral_deg: float) -> GeometryRow: - """Build an HOOH GeometryRow at a given H-O-O-H dihedral angle.""" - roo, roh, hoo_ang = 1.45, 0.97, np.radians(100.0) - dih = np.radians(dihedral_deg) - o1 = np.array([0, 0, 0]) - o2 = np.array([roo, 0, 0]) - h1 = o1 + roh * np.array([-np.cos(hoo_ang), np.sin(hoo_ang), 0]) - base = roh * np.array([np.cos(hoo_ang), np.sin(hoo_ang), 0]) - rot = np.array( - [ - [1, 0, 0], - [0, np.cos(dih), -np.sin(dih)], - [0, np.sin(dih), np.cos(dih)], - ] - ) - h2 = o2 + rot @ base - coordinates = np.array([h1, o1, o2, h2]) - return GeometryRow( - symbols=["H", "O", "O", "H"], coordinates=coordinates, charge=0, spin=0 - ) - - -def _jittered_copy(geometry_row: GeometryRow, rng: Generator) -> GeometryRow: - """Build a small-noise, rotated, translated copy of a geometry.""" - coordinates = np.array(geometry_row.coordinates) - coordinates = coordinates + rng.normal(scale=0.01, size=coordinates.shape) - rot = Rotation.from_euler("xyz", [30, 20, 10], degrees=True) - coordinates = rot.apply(coordinates) + np.array([3.0, 3.0, 3.0]) - return GeometryRow( - symbols=list(geometry_row.symbols), - coordinates=coordinates, - charge=geometry_row.charge, - spin=geometry_row.spin, - ) - - -def test__conformer_identity_merge_on_duplicate_geometry( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - rng: Generator, -) -> None: - """Test that near-identical geometries share one conformer identity.""" - duplicate_geometry = _jittered_copy(geometry_row, rng) - - database.add(calculation_row) - database.add(geometry_row) - database.add(duplicate_geometry) - - stationary1 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - stationary2 = StationaryPointRow( - calculation=calculation_row, geometry=duplicate_geometry - ) - database.add(stationary1) - database.add(stationary2) - database.commit() - - conformer1 = next(i for i in stationary1.identities if i.kind == "conformer") - conformer2 = next(i for i in stationary2.identities if i.kind == "conformer") - assert conformer1.id == conformer2.id - - -def test__conformer_identity_merge_with_unattached_geometry_ids( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - rng: Generator, -) -> None: - """Test conformer dedup when peers are built with only `geometry_id` set. - - Regression test: `_matching_conformer_identity` must resolve each InChI - peer's geometry the same way it resolves the target row's own geometry - (via the session, for a row with only `geometry_id` set), rather than - reading a peer's `.geometry` relationship directly, which stays - unpopulated until the ORM syncs it — as is the case for rows built by a - bulk loader like a database merge rather than via `.geometry=`. - """ - duplicate_geometry = _jittered_copy(geometry_row, rng) - - database.add(calculation_row) - database.add(geometry_row) - database.add(duplicate_geometry) - database.flush() - - stationary1 = StationaryPointRow( - calculation_id=calculation_row.id, geometry_id=geometry_row.id - ) - stationary2 = StationaryPointRow( - calculation_id=calculation_row.id, geometry_id=duplicate_geometry.id - ) - database.add(stationary1) - database.add(stationary2) - database.commit() - - conformer1 = next(i for i in stationary1.identities if i.kind == "conformer") - conformer2 = next(i for i in stationary2.identities if i.kind == "conformer") - assert conformer1.id == conformer2.id - - -def test__conformer_identity_split_on_distinct_conformer( - database: Database, calculation_row: CalculationRow -) -> None: - """Test that geometrically distinct conformers of the same species split.""" - anti = _hooh_geometry_row(180) - gauche = _hooh_geometry_row(60) - - database.add(calculation_row) - database.add(anti) - database.add(gauche) - - stationary1 = StationaryPointRow(calculation=calculation_row, geometry=anti) - stationary2 = StationaryPointRow(calculation=calculation_row, geometry=gauche) - database.add(stationary1) - database.add(stationary2) - database.commit() - - conformer1 = next(i for i in stationary1.identities if i.kind == "conformer") - conformer2 = next(i for i in stationary2.identities if i.kind == "conformer") - assert conformer1.value != conformer2.value - - -def test__conformer_group_id_increments_across_distinct_species_in_one_flush( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - rng: Generator, -) -> None: - """Test that one flush assigns distinct group ids to distinct species.""" - duplicate_geometry = _jittered_copy(geometry_row, rng) - hooh_geometry = _hooh_geometry_row(180) - - database.add(calculation_row) - database.add(geometry_row) - database.add(duplicate_geometry) - database.add(hooh_geometry) - - water1 = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - water2 = StationaryPointRow( - calculation=calculation_row, geometry=duplicate_geometry - ) - hooh = StationaryPointRow(calculation=calculation_row, geometry=hooh_geometry) - database.add(water1) - database.add(water2) - database.add(hooh) - database.commit() - - water1_conf = next(i for i in water1.identities if i.kind == "conformer") - water2_conf = next(i for i in water2.identities if i.kind == "conformer") - hooh_conf = next(i for i in hooh.identities if i.kind == "conformer") - - assert water1_conf.id == water2_conf.id - assert hooh_conf.id != water1_conf.id - - -def test__conformer_identity_idempotent_on_second_flush( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that re-flushing an already-committed stationary point is idempotent.""" - database.add(calculation_row) - database.add(geometry_row) - - stationary = StationaryPointRow(calculation=calculation_row, geometry=geometry_row) - database.add(stationary) - database.commit() - - conformer_id = next(i for i in stationary.identities if i.kind == "conformer").id - - stationary.order = 1 - database.add(stationary) - database.commit() - - conformer_identities = [i for i in stationary.identities if i.kind == "conformer"] - assert len(conformer_identities) == 1 - assert conformer_identities[0].id == conformer_id + assert identity.stationary_points == [] + assert identity.identity_extras == [] + + +class TestIdentityExtraRow: + """Tests for IdentityExtraRow model.""" + + def test_create_identity_extra(self, database: Database) -> None: + """IdentityExtraRow can be created with attribute and value.""" + with database.session() as session: + identity = IdentityRow( + kind="stereoisomer", + algorithm="rdkit inchi", + value="InChI=1S/CH4/h1H4", + ) + session.add(identity) + session.flush() + + extra = IdentityExtraRow( + identity_id=identity.id, + attribute="molecular_weight", + value="16.04", + ) + session.add(extra) + session.commit() + + assert extra.id is not None + assert extra.attribute == "molecular_weight" + assert extra.value == "16.04" + + def test_identity_extra_relationship(self, database: Database) -> None: + """IdentityExtraRow.identity relationship works correctly.""" + with database.session() as session: + identity = IdentityRow( + kind="stereoisomer", algorithm="rdkit inchi", value="InChI=1S/CH4/h1H4" + ) + session.add(identity) + session.flush() + + extra = IdentityExtraRow( + identity_id=identity.id, attribute="molecular_weight", value="16.04" + ) + session.add(extra) + session.commit() + + assert extra.identity is not None + assert extra.identity.id == identity.id + assert extra.identity.value == "InChI=1S/CH4/h1H4" + + +class TestLinkModels: + """Tests for link/association table models.""" + + def test_calculation_geometry_link(self, database: Database) -> None: + """CalculationGeometryLink can be created with role.""" + with database.session() as session: + model = ModelRow(calc_type="energy", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + link = CalculationGeometryLink( + calculation_id=calc.id, geometry_id=geom.id, role=Role.INPUT + ) + session.add(link) + session.commit() + + assert link.role == Role.INPUT + assert link.geometry_id == geom.id + assert link.calculation_id == calc.id + + def test_geometry_trajectory_link(self, database: Database) -> None: + """GeometryTrajectoryLink can be created with index.""" + with database.session() as session: + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + traj = TrajectoryRow(ndim=2) + session.add_all([geom, traj]) + session.flush() + + link = GeometryTrajectoryLink( + geometry_id=geom.id, trajectory_id=traj.id, index=[0, 0] + ) + session.add(link) + session.commit() + + assert link.index == [0, 0] + + def test_calculation_trajectory_link(self, database: Database) -> None: + """CalculationTrajectoryLink can be created.""" + with database.session() as session: + model = ModelRow(calc_type="irc", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + traj = TrajectoryRow(ndim=1) + session.add_all([calc, traj]) + session.flush() + + link = CalculationTrajectoryLink( + calculation_id=calc.id, trajectory_id=traj.id, role="output" + ) + session.add(link) + session.commit() + + assert link.role == "output" + + def test_stage_stationary_link(self, database: Database) -> None: + """StageStationaryLink can be created.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + stat_pt = StationaryPointRow(geometry_id=geom.id, calculation_id=calc.id) + stage = StageRow(is_ts=False) + session.add_all([stat_pt, stage]) + session.flush() + + link = StageStationaryLink(stationary_id=stat_pt.id, stage_id=stage.id) + session.add(link) + session.commit() + + assert link.stationary_id == stat_pt.id + assert link.stage_id == stage.id + + def test_step_validation_link(self, database: Database) -> None: + """StepValidationLink can be created.""" + with database.session() as session: + stage1 = StageRow(is_ts=False) + stage2 = StageRow(is_ts=False) + session.add_all([stage1, stage2]) + session.flush() + + step = StepRow( + stage_id1=stage1.id, + stage_id2=stage2.id, + is_barrierless=True, + ) + session.add(step) + session.flush() + + model = ModelRow(calc_type="irc", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + validation = ValidationRow(calculation_id=calc.id, method="irc") + session.add(validation) + session.flush() + + assert step.id is not None + assert validation.id is not None + link = StepValidationLink(step_id=step.id, validation_id=validation.id) + session.add(link) + session.commit() + + assert link.step_id == step.id + assert link.validation_id == validation.id + + def test_identity_stationary_link(self, database: Database) -> None: + """IdentityStationaryLink can be created.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + stat_pt = StationaryPointRow(geometry_id=geom.id, calculation_id=calc.id) + identity = IdentityRow( + kind="stereoisomer", algorithm="rdkit smiles", value="C" + ) + session.add_all([stat_pt, identity]) + session.flush() + + assert stat_pt.id is not None + assert identity.id is not None + link = IdentityStationaryLink( + stationary_id=stat_pt.id, identity_id=identity.id + ) + session.add(link) + session.commit() + + assert link.stationary_id == stat_pt.id + assert link.identity_id == identity.id + + +class TestModelIntegration: + """Integration tests for model interactions.""" + + def test_geometry_with_multiple_results(self, database: Database) -> None: + """Geometry can have multiple result types attached.""" + with database.session() as session: + model = ModelRow(calc_type="frequency", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C", "H"], + coordinates=[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], + charge=0, + spin=0, + ) + session.add(geom) + session.flush() + + energy = EnergyRow( + geometry_id=geom.id, calculation_id=calc.id, value=-37.8422 + ) + gradient = GradientRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]), + ) + hessian = HessianRow( + geometry_id=geom.id, + calculation_id=calc.id, + value=np.eye(6, dtype=np.float32), + ) + session.add_all([energy, gradient, hessian]) + session.commit() + + assert len(geom.energies) == 1 + assert len(geom.gradients) == 1 + assert len(geom.hessians) == 1 + + def test_calculation_with_multiple_geometries(self, database: Database) -> None: + """Calculation can be linked to multiple geometries.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom1 = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + geom2 = GeometryRow( + symbols=["C"], coordinates=[[0.1, 0.0, 0.0]], charge=0, spin=0 + ) + session.add_all([geom1, geom2]) + session.flush() + + link1 = CalculationGeometryLink( + calculation_id=calc.id, geometry_id=geom1.id, role=Role.INPUT + ) + link2 = CalculationGeometryLink( + calculation_id=calc.id, geometry_id=geom2.id, role=Role.OUTPUT + ) + session.add_all([link1, link2]) + session.commit() + + assert len(calc.geometry_links) == 2 # noqa: PLR2004 + + def test_stationary_point_with_identity(self, database: Database) -> None: + """Stationary point can be linked to an identity.""" + with database.session() as session: + model = ModelRow(calc_type="opt", program="psi4", method="B3LYP") + session.add(model) + session.flush() + + calc = CalculationRow(model_id=model.id) + session.add(calc) + session.flush() + + geom = GeometryRow( + symbols=["C"], coordinates=[[0.0, 0.0, 0.0]], charge=0, spin=0 + ) + session.add(geom) + session.flush() + + stat_pt = StationaryPointRow(geometry_id=geom.id, calculation_id=calc.id) + identity = IdentityRow( + kind="stereoisomer", algorithm="rdkit smiles", value="C" + ) + session.add_all([stat_pt, identity]) + session.flush() + + assert stat_pt.id is not None + assert identity.id is not None + link = IdentityStationaryLink( + stationary_id=stat_pt.id, identity_id=identity.id + ) + session.add(link) + session.commit() + + # Test bidirectional relationship + # Note: stat_pt will have auto-generated identities from event listener + # plus the manually linked one + assert len(stat_pt.identities) >= 2 # noqa: PLR2004 + assert identity.id in [i.id for i in stat_pt.identities] + assert len(identity.stationary_points) == 1 + assert identity.stationary_points[0].id == stat_pt.id diff --git a/tests/test_plot.py b/tests/test_plot.py deleted file mode 100644 index 892a5a6..0000000 --- a/tests/test_plot.py +++ /dev/null @@ -1,270 +0,0 @@ -"""Autostorage PES plot tests.""" - -from pathlib import Path - -import numpy as np -import pytest -from matplotlib.axes import Axes -from matplotlib.backends.backend_agg import FigureCanvasAgg -from matplotlib.figure import Figure - -from autostorage import CalculationRow, Database, GeometryRow, StageRow, StepRow -from autostorage.utils import PESPlot, plot_pes -from tests.test_utils import _diatomic, _stationary, _with_energy - -PNG_MAGIC_BYTES = b"\x89PNG\r\n\x1a\n" - -EXPECTED_TS_PATH_LINE_COUNT = 5 # 3 levels + 2 connectors -EXPECTED_TS_PATH_TEXT_COUNT = 3 # 2 species + 1 TS peak -EXPECTED_BARRIERLESS_LINE_COUNT = 3 # 2 levels + 1 connector - - -def _new_figure_and_axes() -> tuple[Figure, Axes]: - """Build a bare Figure/Axes without touching pyplot's global state.""" - figure = Figure() - FigureCanvasAgg(figure) - return figure, figure.add_subplot() - - -def test__plot_normal_path_with_ts( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a well + bimolecular + barrier network draws all segments.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - _with_energy(database, calculation_row, geometry_row, -76.0) - - frag1_geo = GeometryRow( - symbols=["N", "H", "H"], - coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.8, -0.5], [0.0, -0.8, -0.5]]), - charge=0, - spin=1, - ) - frag2_geo = GeometryRow( - symbols=["H", "C", "O"], - coordinates=np.array([[0.0, 1.2, 0.0], [0.0, 0.0, 0.0], [1.1, -0.6, 0.0]]), - charge=0, - spin=1, - ) - frag1 = _stationary(database, calculation_row, frag1_geo) - frag2 = _stationary(database, calculation_row, frag2_geo) - _with_energy(database, calculation_row, frag1_geo, -55.6) - _with_energy(database, calculation_row, frag2_geo, -20.3) - - ts_geo = GeometryRow( - symbols=["N", "H", "H"], - coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.9, -0.5], [0.0, -0.9, -0.5]]), - charge=0, - spin=1, - ) - ts = _stationary(database, calculation_row, ts_geo, order=1) - _with_energy(database, calculation_row, ts_geo, -75.5) - - # All three stages are built together, right before `step`, so no intervening - # `database.commit()` (from the helpers above) flushes the session while one is - # linked via backref to a persistent stationary but not yet in the session itself. - well_stage = StageRow(stationaries=[ref]) - bimolecular_stage = StageRow(stationaries=[frag1, frag2]) - ts_stage = StageRow(stationaries=[ts], is_ts=True) - - step = StepRow(stage1=well_stage, stage2=bimolecular_stage, stage_ts=ts_stage) - database.add(step) - database.commit() - - result = plot_pes(database, [step], ref=ref, model=calculation_row.model) - - assert isinstance(result, PESPlot) - assert isinstance(result.figure, Figure) - assert isinstance(result.axes, Axes) - assert len(result.axes.lines) == EXPECTED_TS_PATH_LINE_COUNT - assert len(result.axes.texts) == EXPECTED_TS_PATH_TEXT_COUNT - assert any(t.get_text().startswith("B1") for t in result.axes.texts) - - -def test__plot_barrierless_step( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a barrierless step draws a dashed connector, no TS peak.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - assert step.is_barrierless - - result = plot_pes(database, [step], ref=ref, model=calculation_row.model) - - assert len(result.axes.lines) == EXPECTED_BARRIERLESS_LINE_COUNT - connector = result.axes.lines[-1] - assert connector.get_linestyle() == "--" - assert any(t.get_text() == "B1" for t in result.axes.texts) - - -def test__plot_missing_energy_flagged( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a species missing an EnergyRow is flagged, not raised.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - result = plot_pes(database, [step], ref=ref, model=calculation_row.model) - - flagged_lines = [line for line in result.axes.lines if line.get_linestyle() == "--"] - assert any(np.asarray(line.get_ydata())[0] == 0.0 for line in flagged_lines) - assert any(t.get_text().endswith("(no energy data)") for t in result.axes.texts) - - -def test__plot_repr_png_returns_valid_png( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that _repr_png_ returns bytes recognizable as a PNG (Jupyter hook).""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - result = plot_pes(database, [step], ref=ref, model=calculation_row.model) - png = result._repr_png_() - - assert isinstance(png, bytes) - assert png.startswith(PNG_MAGIC_BYTES) - - -def test__plot_save_writes_png( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - tmp_path: Path, -) -> None: - """Test that save() writes a real PNG file, inferring format from suffix.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - result = plot_pes(database, [step], ref=ref, model=calculation_row.model) - path = tmp_path / "pes.png" - result.save(path) - - assert path.exists() - assert path.read_bytes()[:8] == PNG_MAGIC_BYTES - - -def test__plot_save_writes_svg( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - tmp_path: Path, -) -> None: - """Test that save() writes a real SVG file, inferring format from suffix.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - result = plot_pes(database, [step], ref=ref, model=calculation_row.model) - path = tmp_path / "pes.svg" - result.save(path) - - assert path.exists() - assert " None: - """Test that a caller-supplied Axes is drawn into and returned as-is.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - figure, axes = _new_figure_and_axes() - result = plot_pes(database, [step], ref=ref, model=calculation_row.model, ax=axes) - - assert result.axes is axes - assert result.figure is figure - - -def test__plot_raises_on_missing_reference_energy( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a missing reference energy raises rather than silently defaulting.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - with pytest.raises(ValueError, match="No EnergyRow found for reference"): - plot_pes(database, [step], ref=ref, model=calculation_row.model) diff --git a/tests/test_utils.py b/tests/test_utils.py deleted file mode 100644 index b4882df..0000000 --- a/tests/test_utils.py +++ /dev/null @@ -1,435 +0,0 @@ -"""Autostorage MESS export tests.""" - -import re - -import numpy as np -import pytest -from numpy.random import Generator - -from autostorage import ( - CalculationRow, - Database, - EnergyRow, - GeometryRow, - HessianRow, - StageRow, - StationaryPointRow, - StepRow, -) -from autostorage.utils import HARTREE_TO_KCAL_PER_MOL, export_mess_input - - -def _stationary( - database: Database, - calculation: CalculationRow, - geometry: GeometryRow, - *, - order: int = 0, -) -> StationaryPointRow: - """Create, persist, and return a stationary point for `geometry`.""" - stationary = StationaryPointRow( - calculation=calculation, geometry=geometry, order=order - ) - database.add(stationary) - database.commit() - return stationary - - -def _with_energy( - database: Database, calculation: CalculationRow, geometry: GeometryRow, value: float -) -> None: - """Persist an `EnergyRow` for `geometry` at `calculation`'s model.""" - database.add(EnergyRow(geometry=geometry, calculation=calculation, value=value)) - database.commit() - - -def _diatomic(symbols: list[str], distance: float, *, spin: int = 0) -> GeometryRow: - """Build a simple two-atom `GeometryRow` along the z-axis.""" - return GeometryRow( - symbols=symbols, - coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, distance]]), - charge=0, - spin=spin, - ) - - -def _species_section(text: str, label: str) -> str: - """Return the `Well`/`Bimolecular` block in `text` for the given label.""" - for keyword in ("Well", "Bimolecular"): - marker = f"{keyword} {label}" - if marker in text: - start = text.index(marker) - break - else: - msg = f"No species block found for label {label!r}." - raise AssertionError(msg) - - rest = text[start + len(marker) :] - boundaries = [ - start + len(marker) + rest.index(kw) - for kw in ("\nWell ", "\nBimolecular ", "\nBarrier ") - if kw in rest - ] - end = min(boundaries) if boundaries else len(text) - return text[start:end] - - -def _barrier_section(text: str, label: str) -> str: - """Return the `Barrier` block in `text` for the given label (to end of text).""" - marker = f"Barrier {label}" - return text[text.index(marker) :] - - -def test__export_well_bimolecular_barrier_round_trip( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a well + bimolecular + barrier network renders all block types.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - _with_energy(database, calculation_row, geometry_row, -76.0) - - frag1_geo = GeometryRow( - symbols=["N", "H", "H"], - coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.8, -0.5], [0.0, -0.8, -0.5]]), - charge=0, - spin=1, - ) - frag2_geo = GeometryRow( - symbols=["H", "C", "O"], - coordinates=np.array([[0.0, 1.2, 0.0], [0.0, 0.0, 0.0], [1.1, -0.6, 0.0]]), - charge=0, - spin=1, - ) - frag1 = _stationary(database, calculation_row, frag1_geo) - frag2 = _stationary(database, calculation_row, frag2_geo) - _with_energy(database, calculation_row, frag1_geo, -55.6) - _with_energy(database, calculation_row, frag2_geo, -20.3) - - ts_geo = GeometryRow( - symbols=["N", "H", "H"], - coordinates=np.array([[0.0, 0.0, 0.0], [0.0, 0.9, -0.5], [0.0, -0.9, -0.5]]), - charge=0, - spin=1, - ) - ts = _stationary(database, calculation_row, ts_geo, order=1) - _with_energy(database, calculation_row, ts_geo, -75.5) - - # All three stages are built together, right before `step`, so no intervening - # `database.commit()` (from the helpers above) flushes the session while one is - # linked via backref to a persistent stationary but not yet in the session itself. - well_stage = StageRow(stationaries=[ref]) - bimolecular_stage = StageRow(stationaries=[frag1, frag2]) - ts_stage = StageRow(stationaries=[ts], is_ts=True) - - step = StepRow(stage1=well_stage, stage2=bimolecular_stage, stage_ts=ts_stage) - database.add(step) - database.commit() - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - expected_fragment_count = 2 - - assert "Well W1" in text - assert "Bimolecular P1" in text - assert text.count("Fragment ") == expected_fragment_count - barrier_section = _barrier_section(text, "B1") - assert "W1" in barrier_section.splitlines()[0] - assert "P1" in barrier_section.splitlines()[0] - - -def test__export_zero_energy_relative_to_reference( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that ZeroEnergy values are computed relative to the reference energy.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - - ref_value = -76.000 - other_value = -75.950 - _with_energy(database, calculation_row, geometry_row, ref_value) - _with_energy(database, calculation_row, other_geo, other_value) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - expected = (other_value - ref_value) * HARTREE_TO_KCAL_PER_MOL - ref_section = _species_section(text, "W1") - other_section = _species_section(text, "W2") - ref_match = re.search(r"ZeroEnergy\[kcal/mol\]\s+(-?\d+\.\d+)", ref_section) - other_match = re.search(r"ZeroEnergy\[kcal/mol\]\s+(-?\d+\.\d+)", other_section) - assert ref_match is not None - assert other_match is not None - assert float(ref_match.group(1)) == pytest.approx(0.0, abs=1e-2) - assert float(other_match.group(1)) == pytest.approx(expected, abs=1e-2) - - -def test__export_oh_electronic_levels_special_case( - database: Database, calculation_row: CalculationRow -) -> None: - """Test the OH spin-orbit electronic-levels special case and the fallback.""" - database.add(calculation_row) - database.commit() - - oh_geo = _diatomic(["O", "H"], 0.97, spin=1) - other_geo = _diatomic(["H", "H"], 0.74) - oh = _stationary(database, calculation_row, oh_geo) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, oh_geo, -75.7) - _with_energy(database, calculation_row, other_geo, -1.2) - - oh_stage = StageRow(stationaries=[oh]) - other_stage = StageRow(stationaries=[other]) - step = StepRow(stage1=oh_stage, stage2=other_stage) - database.add(step) - database.commit() - - assert oh_stage.id is not None - assert other_stage.id is not None - labels = {oh_stage.id: "WOH", other_stage.id: "WOTHER"} - text = export_mess_input( - database, [step], ref=oh, model=calculation_row.model, labels=labels - ) - - oh_section = _species_section(text, "WOH") - other_section = _species_section(text, "WOTHER") - - assert "ElectronicLevels[1/cm] 2" in oh_section - assert "0.0 2" in oh_section - assert "140.0 2" in oh_section - - assert "ElectronicLevels[1/cm] 1" in other_section - assert "0.0 1" in other_section - - -def test__export_barrierless_placeholder( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a barrierless step renders a flagged placeholder Barrier block.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - assert step.is_barrierless - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - barrier_section = _barrier_section(text, "B1") - assert "TODO(autostorage): barrierless step" in barrier_section - assert "Frequencies[1/cm] 0" in barrier_section - - -def test__export_symmetry_factor_computed_and_barrierless_placeholder( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that SymmetryFactor is a real computed value, except for barrierless TSs.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - other_stage = StageRow(stationaries=[other]) - step = StepRow(stage1=StageRow(stationaries=[ref]), stage2=other_stage) - database.add(step) - database.commit() - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - other_section = _species_section(text, "W2") - assert f"SymmetryFactor {other_geo.symmetry_number}" in other_section - assert "TODO(autostorage)" not in other_section - - barrier_section = _barrier_section(text, "B1") - assert "SymmetryFactor 1 ! TODO(autostorage)" in barrier_section - - -def test__export_custom_labels_and_names_override( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that label/name overrides apply only to the given stage ids.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -1.0) - - ref_stage = StageRow(stationaries=[ref]) - step = StepRow(stage1=ref_stage, stage2=StageRow(stationaries=[other])) - database.add(step) - database.commit() - - assert ref_stage.id is not None - text = export_mess_input( - database, - [step], - ref=ref, - model=calculation_row.model, - labels={ref_stage.id: "X1"}, - names={ref_stage.id: "custom name"}, - ) - - assert "Well X1 # custom name" in text - - well_labels = re.findall(r"^Well {2}(\S+)", text, re.MULTILINE) - assert "X1" in well_labels - other_labels = [label for label in well_labels if label != "X1"] - assert len(other_labels) == 1 - assert re.fullmatch(r"W\d+", other_labels[0]) - - -def test__export_ts_excludes_imaginary_frequency( - database: Database, - calculation_row: CalculationRow, - geometry_row: GeometryRow, - rng: Generator, -) -> None: - """Test that a Barrier block's Frequencies exclude the imaginary TS mode.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - - ts_geo = GeometryRow( - symbols=["H", "O", "H"], - coordinates=np.array([[0, 0, 0.9], [0, 0, 0], [0.9, 0, 0]]), - charge=0, - spin=0, - ) - ts = _stationary(database, calculation_row, ts_geo, order=1) - - n = ts_geo.atom_count - ts_hessian = HessianRow( - calculation=calculation_row, - geometry=ts_geo, - value=rng.uniform(size=(3 * n, 3 * n)), - ) - database.add(ts_hessian) - database.commit() - - _with_energy(database, calculation_row, geometry_row, -76.0) - _with_energy(database, calculation_row, other_geo, -75.9) - _with_energy(database, calculation_row, ts_geo, -75.8) - - expected_positive_count = sum(1 for f in ts_hessian.harmonic_frequencies if f > 0.0) - assert ts_hessian.order >= 1 - - step = StepRow( - stage1=StageRow(stationaries=[ref]), - stage2=StageRow(stationaries=[other]), - stage_ts=StageRow(stationaries=[ts], is_ts=True), - ) - database.add(step) - database.commit() - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - barrier_section = _barrier_section(text, "B1") - match = re.search(r"Frequencies\[1/cm\]\s+(\d+)", barrier_section) - assert match is not None - assert int(match.group(1)) == expected_positive_count - - -def test__export_missing_energy_renders_todo( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a species missing an EnergyRow renders a TODO instead of raising.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, geometry_row, -76.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - assert "TODO(autostorage): no EnergyRow found" in text - - -def test__export_raises_on_missing_reference_energy( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that a missing reference energy raises rather than silently defaulting.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - other_geo = _diatomic(["H", "H"], 0.74) - other = _stationary(database, calculation_row, other_geo) - _with_energy(database, calculation_row, other_geo, -1.0) - - step = StepRow( - stage1=StageRow(stationaries=[ref]), stage2=StageRow(stationaries=[other]) - ) - database.add(step) - database.commit() - - with pytest.raises(ValueError, match="No EnergyRow found for reference"): - export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - -def test__export_fragment_order_deterministic( - database: Database, calculation_row: CalculationRow, geometry_row: GeometryRow -) -> None: - """Test that bimolecular fragments render in ascending-id order, not list order.""" - database.add(calculation_row) - database.commit() - - ref = _stationary(database, calculation_row, geometry_row) - _with_energy(database, calculation_row, geometry_row, -76.0) - - frag_a_geo = _diatomic(["C", "O"], 1.13) - frag_b_geo = _diatomic(["H", "H"], 0.74) - frag_a = _stationary(database, calculation_row, frag_a_geo) - frag_b = _stationary(database, calculation_row, frag_b_geo) - _with_energy(database, calculation_row, frag_a_geo, -10.0) - _with_energy(database, calculation_row, frag_b_geo, -1.0) - - assert frag_a.id is not None - assert frag_b.id is not None - assert frag_a.id < frag_b.id - - bimolecular_stage = StageRow(stationaries=[frag_b, frag_a]) - step = StepRow(stage1=StageRow(stationaries=[ref]), stage2=bimolecular_stage) - database.add(step) - database.commit() - - text = export_mess_input(database, [step], ref=ref, model=calculation_row.model) - - assert text.index("Fragment CO") < text.index("Fragment H2")