diff --git a/.github/workflows/build-images.yml b/.github/workflows/build-images.yml index 544285b..df630be 100644 --- a/.github/workflows/build-images.yml +++ b/.github/workflows/build-images.yml @@ -1,12 +1,27 @@ name: Build & push md-workflows images -# Builds the three-stage image chain on the self-hosted diffuse builder and pushes to Harbor. +# Builds the three-stage image chain on the self-hosted diffuse builder and publishes it: +# - base + gromacs -> ghcr.io/prism-science/md-workflows (public; tags prefixed base-/gromacs-), +# pushed with the workflow's own GITHUB_TOKEN; +# - actl (the Astera ACTL overlay) -> Astera's internal registry, library/md-workflows. The +# registry host is the ASTERA_REGISTRY secret, never written in this public repo; GitHub masks +# it wherever it is echoed, including inside the longer image refs built from it. # # Security: this workflow has NO pull_request trigger, so fork PRs can never reach the -# self-hosted runner or the Harbor credentials. Heavy builds run only on push to `astera` -# (path-filtered) and manual dispatch. Credentials live in the `harbor` GitHub Environment, -# which must be restricted to the `astera` deployment branch; the `diffuse-sh-builder` runner -# should sit in a runner group scoped to this repo. See the runbook in the PR description. +# self-hosted runner or any credentials. Heavy builds run only on push to `astera` +# (path-filtered) and manual dispatch. All three publishing jobs run in the `harbor` GitHub +# Environment, which must be restricted to the `astera` deployment branch: that restriction is +# what limits publishing -- including the public GHCR images -- to `astera`, and the environment +# holds the Astera registry secrets (ASTERA_REGISTRY, HARBOR_USERNAME, HARBOR_PASSWORD). The +# `diffuse-sh-builder` runner should sit in a runner group scoped to this repo. See the runbook +# in the PR description. +# +# Manual dispatch runs only the `test-build` job unless `publish` is ticked: all three stages +# built on the runner's local Docker daemon, smoke-tested, then untagged. It pushes nothing and +# uses no registry credentials, so it needs no `harbor` environment and can be dispatched from +# any branch to try Dockerfile changes before they reach `astera`: +# gh workflow run build-images.yml --ref # build + smoke-test +# gh workflow run build-images.yml --ref astera -f publish=true # build + publish # # NOTE: runs-on matches only the unique custom label `diffuse-sh-builder` (not `self-hosted`), # because that runner was registered with --no-default-labels and so lacks the implicit @@ -22,6 +37,11 @@ on: - md_workflows/** - .github/workflows/build-images.yml workflow_dispatch: + inputs: + publish: + description: "Publish the images (GHCR + Astera registry). Leave unticked to build + smoke-test only." + type: boolean + default: false permissions: contents: read @@ -31,15 +51,15 @@ concurrency: cancel-in-progress: true env: - REGISTRY: harbor.astera.sh - # base + gromacs (bake the non-redistributable ChimeraX) live in the proprietary project, - # distinguished by a stage tag prefix. The final consumable actl image lands in library. - PROPRIETARY_IMAGE: harbor.astera.sh/diffuse-proprietary/md-workflows - LIBRARY_IMAGE: harbor.astera.sh/library/md-workflows + # Public stages, distinguished by a stage tag prefix. + PUBLIC_IMAGE: ghcr.io/prism-science/md-workflows + # Path of the actl overlay on the Astera registry; the host comes from secrets.ASTERA_REGISTRY. + ASTERA_IMAGE_PATH: library/md-workflows jobs: version: name: Compute version + if: github.event_name == 'push' || inputs.publish runs-on: ubuntu-latest outputs: semver: ${{ steps.v.outputs.semver }} @@ -61,20 +81,25 @@ jobs: base: name: Build base + if: github.event_name == 'push' || inputs.publish needs: version runs-on: [diffuse-sh-builder] + # Uses no environment secrets; the environment's branch restriction gates publishing. environment: harbor + permissions: + contents: read + packages: write outputs: digest: ${{ steps.build.outputs.digest }} steps: - uses: actions/checkout@v4 - uses: docker/setup-buildx-action@v3 - - name: Login to Harbor + - name: Login to GHCR uses: docker/login-action@v3 with: - registry: ${{ env.REGISTRY }} - username: ${{ secrets.HARBOR_USERNAME }} - password: ${{ secrets.HARBOR_PASSWORD }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - id: build uses: docker/build-push-action@v6 with: @@ -84,29 +109,34 @@ jobs: push: true provenance: false tags: | - ${{ env.PROPRIETARY_IMAGE }}:base-${{ needs.version.outputs.semver }} - ${{ env.PROPRIETARY_IMAGE }}:base-${{ needs.version.outputs.build }} - ${{ env.PROPRIETARY_IMAGE }}:base-${{ needs.version.outputs.sha }} - ${{ env.PROPRIETARY_IMAGE }}:base - cache-from: type=registry,ref=${{ env.PROPRIETARY_IMAGE }}:base-buildcache - cache-to: type=registry,ref=${{ env.PROPRIETARY_IMAGE }}:base-buildcache,mode=max + ${{ env.PUBLIC_IMAGE }}:base-${{ needs.version.outputs.semver }} + ${{ env.PUBLIC_IMAGE }}:base-${{ needs.version.outputs.build }} + ${{ env.PUBLIC_IMAGE }}:base-${{ needs.version.outputs.sha }} + ${{ env.PUBLIC_IMAGE }}:base + cache-from: type=registry,ref=${{ env.PUBLIC_IMAGE }}:base-buildcache + cache-to: type=registry,ref=${{ env.PUBLIC_IMAGE }}:base-buildcache,mode=max gromacs: name: Build gromacs + if: github.event_name == 'push' || inputs.publish needs: [version, base] runs-on: [diffuse-sh-builder] + # Uses no environment secrets; the environment's branch restriction gates publishing. environment: harbor + permissions: + contents: read + packages: write outputs: digest: ${{ steps.build.outputs.digest }} steps: - uses: actions/checkout@v4 - uses: docker/setup-buildx-action@v3 - - name: Login to Harbor + - name: Login to GHCR uses: docker/login-action@v3 with: - registry: ${{ env.REGISTRY }} - username: ${{ secrets.HARBOR_USERNAME }} - password: ${{ secrets.HARBOR_PASSWORD }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - id: build uses: docker/build-push-action@v6 with: @@ -116,29 +146,48 @@ jobs: push: true provenance: false build-args: | - BASE_IMAGE=${{ env.PROPRIETARY_IMAGE }}@${{ needs.base.outputs.digest }} + BASE_IMAGE=${{ env.PUBLIC_IMAGE }}@${{ needs.base.outputs.digest }} tags: | - ${{ env.PROPRIETARY_IMAGE }}:gromacs-${{ needs.version.outputs.semver }} - ${{ env.PROPRIETARY_IMAGE }}:gromacs-${{ needs.version.outputs.build }} - ${{ env.PROPRIETARY_IMAGE }}:gromacs-${{ needs.version.outputs.sha }} - ${{ env.PROPRIETARY_IMAGE }}:gromacs - cache-from: type=registry,ref=${{ env.PROPRIETARY_IMAGE }}:gromacs-buildcache - cache-to: type=registry,ref=${{ env.PROPRIETARY_IMAGE }}:gromacs-buildcache,mode=max + ${{ env.PUBLIC_IMAGE }}:gromacs-${{ needs.version.outputs.semver }} + ${{ env.PUBLIC_IMAGE }}:gromacs-${{ needs.version.outputs.build }} + ${{ env.PUBLIC_IMAGE }}:gromacs-${{ needs.version.outputs.sha }} + ${{ env.PUBLIC_IMAGE }}:gromacs + cache-from: type=registry,ref=${{ env.PUBLIC_IMAGE }}:gromacs-buildcache + cache-to: type=registry,ref=${{ env.PUBLIC_IMAGE }}:gromacs-buildcache,mode=max actl: - name: Build actl (consumable) - needs: [version, gromacs] + name: Build actl (Astera overlay) + if: github.event_name == 'push' || inputs.publish + needs: [version, base, gromacs] runs-on: [diffuse-sh-builder] environment: harbor + permissions: + contents: read + packages: read # pull the gromacs parent from GHCR, even while the package is private + env: + # Bare registry host only; the image path stays in the workflow env above. + ASTERA_REGISTRY: ${{ secrets.ASTERA_REGISTRY }} outputs: digest: ${{ steps.build.outputs.digest }} steps: + - name: Check the registry secret is set + run: | + if [ -z "${ASTERA_REGISTRY}" ]; then + echo "::error::Set ASTERA_REGISTRY (the bare Astera registry host) in the harbor environment." + exit 1 + fi - uses: actions/checkout@v4 - uses: docker/setup-buildx-action@v3 - - name: Login to Harbor + - name: Login to GHCR uses: docker/login-action@v3 with: - registry: ${{ env.REGISTRY }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + - name: Login to the Astera registry + uses: docker/login-action@v3 + with: + registry: ${{ env.ASTERA_REGISTRY }} username: ${{ secrets.HARBOR_USERNAME }} password: ${{ secrets.HARBOR_PASSWORD }} - id: build @@ -150,18 +199,63 @@ jobs: push: true provenance: false build-args: | - GROMACS_IMAGE=${{ env.PROPRIETARY_IMAGE }}@${{ needs.gromacs.outputs.digest }} + GROMACS_IMAGE=${{ env.PUBLIC_IMAGE }}@${{ needs.gromacs.outputs.digest }} tags: | - ${{ env.LIBRARY_IMAGE }}:${{ needs.version.outputs.semver }} - ${{ env.LIBRARY_IMAGE }}:${{ needs.version.outputs.build }} - ${{ env.LIBRARY_IMAGE }}:${{ needs.version.outputs.sha }} - ${{ env.LIBRARY_IMAGE }}:actl - ${{ env.LIBRARY_IMAGE }}:latest - cache-from: type=registry,ref=${{ env.LIBRARY_IMAGE }}:actl-buildcache - cache-to: type=registry,ref=${{ env.LIBRARY_IMAGE }}:actl-buildcache,mode=max + ${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:${{ needs.version.outputs.semver }} + ${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:${{ needs.version.outputs.build }} + ${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:${{ needs.version.outputs.sha }} + ${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:actl + ${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:latest + cache-from: type=registry,ref=${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:actl-buildcache + cache-to: type=registry,ref=${{ env.ASTERA_REGISTRY }}/${{ env.ASTERA_IMAGE_PATH }}:actl-buildcache,mode=max - name: Report digests + env: + BUILD: ${{ needs.version.outputs.build }} run: | - echo "base digest: ${{ needs.base.outputs.digest }}" || true + echo "base digest: ${{ needs.base.outputs.digest }}" echo "gromacs digest: ${{ needs.gromacs.outputs.digest }}" echo "actl digest: ${{ steps.build.outputs.digest }}" - echo "Published ${{ env.LIBRARY_IMAGE }}:${{ needs.version.outputs.build }}" + echo "Published ${PUBLIC_IMAGE}:gromacs-${BUILD}" + echo "Published ${ASTERA_REGISTRY}/${ASTERA_IMAGE_PATH}:${BUILD}" + + # publish unticked: the chain above pulls each parent stage from a registry by digest, so + # without a push the stages must instead be built in one job and chained through local tags, + # as in the README. `--builder default` pins the daemon's own BuildKit (docker driver), which + # can see locally built images in FROM, regardless of any buildx builder left selected on the + # runner. + test-build: + name: Build all stages (no push) + if: github.event_name == 'workflow_dispatch' && !inputs.publish + runs-on: [diffuse-sh-builder] + env: + TAG: test-${{ github.run_id }} + steps: + - uses: actions/checkout@v4 + - name: Build base + run: >- + docker buildx build --builder default --platform linux/amd64 + -f Dockerfile.base -t "md-base:${TAG}" . + - name: Build gromacs + run: >- + docker buildx build --builder default --platform linux/amd64 + -f Dockerfile.gromacs --build-arg BASE_IMAGE="md-base:${TAG}" -t "md-gromacs:${TAG}" . + - name: Build actl + run: >- + docker buildx build --builder default --platform linux/amd64 + -f Dockerfile.actl --build-arg GROMACS_IMAGE="md-gromacs:${TAG}" -t "md-actl:${TAG}" . + # GPU-free: tools the pipeline shells out to, the gemmi crystal code, and no ChimeraX. + - name: Smoke test + run: | + docker run --rm "md-actl:${TAG}" bash -c ' + set -euo pipefail + md-workflows --help > /dev/null + python -c "import gemmi, md_workflows.core.crystal; print(gemmi.__version__)" + for cmd in gmx pdb4amber; do command -v "${cmd}"; done + if command -v chimerax ChimeraX > /dev/null 2>&1; then + echo "ChimeraX is still installed" >&2 + exit 1 + fi + ' + - name: Remove test tags + if: always() + run: docker image rm -f "md-actl:${TAG}" "md-gromacs:${TAG}" "md-base:${TAG}" || true diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 42030c6..22f3765 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -49,6 +49,33 @@ jobs: # Contract + CLI tests are gmx-free, so they run on the hosted runner. run: pytest -q + # The image's lunus env pins Python 3.10 (Dockerfile.base) while `lint` runs on 3.12, so + # 3.11+-only code -- an unconditional `import tomllib` -- once passed CI and broke the image's + # CLI. A separate job rather than a matrix on `lint` leaves that job's check name unchanged for + # any branch-protection rule that requires it. + test-py310: + name: Tests (Python 3.10, as in the image) + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.10' + + - name: Install package + run: | + python -m pip install --upgrade pip + pip install '.[dev]' + + - name: Run tests + run: pytest -q + + - name: Check the CLI starts + run: md-workflows --help > /dev/null + build-check: name: Dockerfile build checks (no push) runs-on: ubuntu-latest diff --git a/.gitleaksignore b/.gitleaksignore new file mode 100644 index 0000000..b38d580 --- /dev/null +++ b/.gitleaksignore @@ -0,0 +1,21 @@ +# Accepted findings for the org secret scan (prism-science/github-workflows, secret-scan.yml). +# +# Rule astera-internal-hostname: historical references to Astera's internal container registry +# host, in the build workflows, Dockerfiles and README of these commits (on main, astera and +# dev). The repo is public, so they are already disclosed and git history keeps them; the host +# was removed from the working tree in PR #15, where build-images.yml now reads it from the +# ASTERA_REGISTRY secret. The host is a registry endpoint, not a credential; registry access +# stays behind the Harbor credentials in the `harbor` environment. +119f1f5ad21e6a60a6d599e2ddee5ebd2ea73d30:.github/workflows/build-images.yml:astera-internal-hostname:36 +f91aa04f88471ec7b71fc3e94c95821e05ed7057:Dockerfile.actl:astera-internal-hostname:11 +f91aa04f88471ec7b71fc3e94c95821e05ed7057:.github/workflows/build-images.yml:astera-internal-hostname:30 +f91aa04f88471ec7b71fc3e94c95821e05ed7057:.github/workflows/build-images.yml:astera-internal-hostname:33 +f91aa04f88471ec7b71fc3e94c95821e05ed7057:.github/workflows/build-images.yml:astera-internal-hostname:34 +7578642a282b3ce46bd55d0293459c008c9a2884:Dockerfile.astera:astera-internal-hostname:14 +425a67755927313d8a065654fd7935c2c6e6ab27:Dockerfile.astera:astera-internal-hostname:13 +425a67755927313d8a065654fd7935c2c6e6ab27:README.md:astera-internal-hostname:10 +425a67755927313d8a065654fd7935c2c6e6ab27:README.md:astera-internal-hostname:40 +9f2ddc1b545a90e828dd97c9ca359d60d522f7ac:Dockerfile.astera:astera-internal-hostname:13 +9f2ddc1b545a90e828dd97c9ca359d60d522f7ac:.github/workflows/astera-docker.yml:astera-internal-hostname:12 +9f2ddc1b545a90e828dd97c9ca359d60d522f7ac:README.md:astera-internal-hostname:10 +9f2ddc1b545a90e828dd97c9ca359d60d522f7ac:README.md:astera-internal-hostname:40 diff --git a/Dockerfile.actl b/Dockerfile.actl index 17f4722..7805b42 100644 --- a/Dockerfile.actl +++ b/Dockerfile.actl @@ -8,7 +8,7 @@ # Build locally (after building md-gromacs:local): # docker buildx build --platform linux/amd64 \ # -f Dockerfile.actl --build-arg GROMACS_IMAGE=md-gromacs:local \ -# -t harbor.astera.sh/library/md-workflows:local-actl . +# -t md-actl:local . ARG GROMACS_IMAGE=md-gromacs:local FROM ${GROMACS_IMAGE} AS actl diff --git a/Dockerfile.base b/Dockerfile.base index 2ae4fdd..7cdc1ba 100644 --- a/Dockerfile.base +++ b/Dockerfile.base @@ -1,6 +1,6 @@ # ==================== md-workflows base (stage 1 of 3) ==================== # Architecture-neutral foundation: CUDA 12.6 *devel* toolchain + the full scientific -# conda stack (AmberTools / DIALS / cctbx / mdtraj / …), lunus, and ChimeraX. This layer +# conda stack (AmberTools / DIALS / cctbx / mdtraj / …) and lunus. This layer # carries everything that is portable across compute; the non-portable GROMACS compile and # the md-workflows app are added downstream in Dockerfile.gromacs, and Astera workspace # conventions in Dockerfile.actl. @@ -106,33 +106,9 @@ RUN $MAMBA_EXE remove -n lunus -y scons \ rm -rf /opt/packages/lunus/.git; \ true -# ---------- ChimeraX ---------- -ARG CHIMERAX_URL="https://www.cgl.ucsf.edu/chimerax/cgi-bin/secure/chimerax-get.py?file=current/ubuntu-22.04/chimerax-daily.deb" -# The download is a two-step token dance behind www.cgl.ucsf.edu (POST choice=Accept -> -# meta-refresh URL carrying a session ident -> follow with the cookie). Run it under bash with -# `set -euo pipefail` and `curl -fsS` so a failed request aborts loudly instead of feeding an -# empty body into grep -- a silent `curl -s | grep` masks both curl's exit code (pipe takes the -# last command's status) and the real error (grep just exits 1 on no match), which is exactly how -# this step failed in CI with a bare "exit code 1" and zero diagnostics. -RUN apt-get update \ - && bash -euo pipefail -c '\ -url="$1"; \ -echo "ChimeraX: requesting download token..."; \ -resp="$(curl -fsS --retry 5 --retry-all-errors --retry-delay 5 --connect-timeout 30 -c /tmp/cx_cookies -d "choice=Accept" "$url")"; \ -redirect="$(printf "%s" "$resp" | grep -oP "url=\K[^\"]*" | head -n1 || true)"; \ -if [ -z "$redirect" ]; then echo "ChimeraX ERROR: no download URL in response (network block or page change). First 2KB:"; printf "%s" "$resp" | head -c 2048; echo; exit 1; fi; \ -echo "ChimeraX: downloading $redirect"; \ -curl -fsS --retry 5 --retry-all-errors --retry-delay 5 --connect-timeout 30 -b /tmp/cx_cookies -o /tmp/chimerax.deb "https://www.cgl.ucsf.edu$redirect"; \ -dpkg-deb --info /tmp/chimerax.deb >/dev/null 2>&1 || { echo "ChimeraX ERROR: downloaded file is not a valid .deb. First 512B:"; head -c 512 /tmp/chimerax.deb; echo; exit 1; }; \ -apt-get install -y /tmp/chimerax.deb \ -' _ "${CHIMERAX_URL}" \ - && rm -f /tmp/chimerax.deb /tmp/cx_cookies \ - && rm -rf /var/lib/apt/lists/* \ - && mkdir -p /home/dev/.config/ChimeraX - # ---------- interactive shells & pager ---------- # Kept out of the build-time list at the top so runtime additions don't invalidate the unpinned -# conda solve and ChimeraX fetch above. tcsh also provides /bin/csh (via update-alternatives). +# conda solve above. tcsh also provides /bin/csh (via update-alternatives). RUN apt-get update && apt-get install -y --no-install-recommends \ less \ tcsh \ diff --git a/README.md b/README.md index b7de6e5..223f988 100644 --- a/README.md +++ b/README.md @@ -11,12 +11,12 @@ without rebuilding the base: | Stage | Dockerfile | Contents | |-------|-----------|----------| -| base | `Dockerfile.base` | CUDA 12.6 devel toolchain, micromamba/conda `lunus` env, lunus, ChimeraX. Architecture-neutral. | +| base | `Dockerfile.base` | CUDA 12.6 devel toolchain, micromamba/conda `lunus` env, lunus. Architecture-neutral. | | gromacs | `Dockerfile.gromacs` | GROMACS (CUDA, tuned for H100 / AVX-512 by default) + the `md_workflows` package. The consumable image. | (An Astera-specific `Dockerfile.actl` overlay adds workspace conventions on top of `gromacs`. It is tracked here so it stays under CI lint coverage, but is built and published only from the -`astera` deployment branch.) +`astera` deployment branch, to Astera's internal registry.) ## 1) Build the images @@ -32,9 +32,16 @@ docker build -f Dockerfile.base -t md-base:local . docker build -f Dockerfile.gromacs --build-arg BASE_IMAGE=md-base:local -t md-gromacs:local . ``` -CI builds these stages and pushes versioned tags (derived from the `version` in -`pyproject.toml`) to the Astera Harbor registry; see `.github/workflows/build-images.yml` on -the `astera` branch. +Or pull the published images instead of building them. CI publishes `base` and `gromacs` from +the `astera` branch to `ghcr.io/prism-science/md-workflows`, as stage-prefixed tags derived from +the `version` in `pyproject.toml` (`gromacs-`, `gromacs--b`, +`gromacs-sha-`, and `gromacs` for the latest; likewise `base-…`): + +```bash +docker pull ghcr.io/prism-science/md-workflows:gromacs +``` + +See `.github/workflows/build-images.yml`. ## 2) Start a container @@ -90,6 +97,15 @@ run_profiles: min: { ntomp: 16, nb: gpu, pme: cpu, bonded: cpu, tunepme: false } ``` +`make_crystal` expands the asymmetric unit to the unit cell and propagates the lattice +with [gemmi](https://gemmi.readthedocs.io) (`md_workflows/core/crystal.py`) instead of +ChimeraX and AmberTools `PropPDB`, so neither is needed at run time. `crystal.spacegroup` +overrides the CRYST1 space group, `crystal.op_order` reorders the symmetry operations +(e.g. `[1, 2, 4, 3]` reproduces the chain labelling of the old ChimeraX route), and +`crystal.numbering` picks the supercell's atom/residue numbering — `continuous` +(PropPDB-compatible), `per-cell` (counters restart in every cell, for supercells past +the 99999-atom / 9999-residue PDB fields) or `auto`. + ```bash md-workflows -w run_dir -c config.yaml run-pipeline md-workflows -w run_dir -c config.yaml --resume run-pipeline # skip completed steps diff --git a/md_workflows/cli.py b/md_workflows/cli.py index ab9d6e9..83bace4 100644 --- a/md_workflows/cli.py +++ b/md_workflows/cli.py @@ -88,6 +88,19 @@ def _ntomp_override(ntomp: int | None) -> dict[str, Any]: return {"run_profiles": {k: {"ntomp": ntomp} for k in RUN_PROFILE_KEYS}} +def _parse_op_order(value: str | None) -> list[int] | None: + """Parse a comma/space separated 1-based symmetry-operation order.""" + if value is None: + return None + try: + return [int(token) for token in value.replace(",", " ").split()] + except ValueError as exc: + typer.secho( + f"invalid --op-order {value!r}: expected integers", fg=typer.colors.RED, err=True + ) + raise typer.Exit(1) from exc + + def _emit(result: StepResult) -> None: color = { StepStatus.COMPLETED: typer.colors.GREEN, @@ -141,13 +154,34 @@ def make_crystal_cmd( ix: int | None = typer.Option(None, "--ix", help="Supercell replication in x (and y/z)."), iy: int | None = typer.Option(None, "--iy"), iz: int | None = typer.Option(None, "--iz"), - chimerax_exec: str | None = typer.Option(None, "--chimerax-exec"), + spacegroup: str | None = typer.Option( + None, "--spacegroup", help="Override the space group in CRYST1, e.g. 'P 21 21 21'." + ), + op_order: str | None = typer.Option( + None, + "--op-order", + help="1-based symmetry-operation order, e.g. '1,2,4,3'; reorders the cell's chains.", + ), + numbering: str | None = typer.Option( + None, "--numbering", help="Supercell numbering: auto (default), continuous or per-cell." + ), + chimerax_exec: str | None = typer.Option( + None, "--chimerax-exec", help="Unused: the expansion is gemmi-based. Kept for callers." + ), ) -> None: """Build the crystal supercell.""" common: Common = ctx.obj crystal = { k: v - for k, v in {"ix": ix, "iy": iy, "iz": iz, "chimerax_exec": chimerax_exec}.items() + for k, v in { + "ix": ix, + "iy": iy, + "iz": iz, + "spacegroup": spacegroup, + "op_order": _parse_op_order(op_order), + "numbering": numbering, + "chimerax_exec": chimerax_exec, + }.items() if v is not None } cfg = _load_cfg(common, {"crystal": crystal} if crystal else None) diff --git a/md_workflows/core/config.py b/md_workflows/core/config.py index acb9b96..8bf6602 100644 --- a/md_workflows/core/config.py +++ b/md_workflows/core/config.py @@ -12,15 +12,23 @@ from __future__ import annotations +import sys from pathlib import Path from typing import Any, Literal -import tomllib import yaml from pydantic import BaseModel, ConfigDict, Field +if sys.version_info >= (3, 11): + import tomllib +else: # tomllib is 3.11+; the image's lunus env pins Python 3.10 + import tomli as tomllib + Offload = Literal["auto", "cpu", "gpu"] +# Atom/residue numbering of a propagated supercell; see ``core.crystal``. +CrystalNumbering = Literal["auto", "continuous", "per-cell"] + # Keys for the per-invocation mdrun profiles. Each GROMACS mdrun call in the pipeline # looks up its profile by one of these names. RUN_PROFILE_KEYS = ( @@ -106,6 +114,13 @@ class CrystalParams(BaseModel): ix: int = 1 iy: int | None = None # falls back to ix iz: int | None = None # falls back to ix + spacegroup: str | None = None # None => take it from the CRYST1 record + # 1-based reordering of the space group's symmetry operations, e.g. [1, 2, 4, 3] to + # reproduce the chain labelling of the old ChimeraX route; see ``core.crystal``. + op_order: list[int] | None = None + numbering: CrystalNumbering = "auto" + # Retained so existing configs and --chimerax-exec keep validating; the expansion is + # gemmi-based now and never launches ChimeraX. chimerax_exec: str = "chimerax" @@ -113,6 +128,10 @@ class WaterboxParams(BaseModel): model_config = ConfigDict(extra="forbid") nc_scale: int = 5 # cell subdivision / reservoir tiling factor (taylor) conc: float = 60.0 # gmx insert-molecules water concentration (mol/L) + # Numbering for the nc_scale^3 tiling; see ``core.crystal.propagate_cell``. "auto" + # keeps PropPDB-compatible continuous numbering until the water count overruns the + # PDB resSeq field, then restarts the counters in each tile. + numbering: CrystalNumbering = "auto" min_mdp: str = "min_water.mdp" equil_mdp: str = "equil_water.mdp" diff --git a/md_workflows/core/crystal.py b/md_workflows/core/crystal.py new file mode 100644 index 0000000..2f0e5a1 --- /dev/null +++ b/md_workflows/core/crystal.py @@ -0,0 +1,399 @@ +"""Crystallographic primitives: symmetry expansion and lattice propagation. + +These replace the two external tools ``make_crystal`` used to shell out to — ChimeraX +(``changechains`` / ``unitcell`` / ``combine``) and AmberTools ``PropPDB`` — with gemmi, +so building the crystal needs neither a graphical nor an Amber installation. + +Both primitives rewrite the fixed-column ``ATOM``/``HETATM`` text of the input rather +than round-tripping through a :class:`gemmi.Structure`: atom order, residue names and +numbers, and the ATOM/HETATM record types then survive byte-for-byte, which the +atom-order-dependent downstream topologies (tleap, GROMACS) require. For the same +reason the serial/resSeq fields stay plain decimal — gemmi's own ``write_pdb`` switches +to hybrid-36 past 99999 atoms and ``gmx solvate``, the consumer of ``xtal.pdb``, cannot +read that. Supercells too large for those fields are handled by +:func:`propagate_cell`'s ``per-cell`` numbering instead. + +All failures raise structured exceptions; nothing here calls ``sys.exit``. +""" + +from __future__ import annotations + +import math +import string +from collections.abc import Sequence +from pathlib import Path + +import gemmi +from pydantic import BaseModel + +from .config import CrystalNumbering +from .exceptions import CrystalSymmetryError + +_COORD_RECORDS = ("ATOM ", "HETATM") + +# Single-character chain ids handed out to symmetry copies and to propagated cells. +_CHAIN_IDS = string.ascii_uppercase + string.ascii_lowercase + string.digits + +# Capacity of the fixed-column PDB serial (5 cols) and resSeq (4 cols) fields. +_MAX_SERIAL = 99999 +_MAX_RESSEQ = 9999 + + +class ExpandResult(BaseModel): + """Unit cell written by :func:`expand_to_unit_cell`.""" + + cell_pdb: Path + spacegroup: str # space group the asymmetric unit was expanded with + chain_ids: list[str] # chain id per symmetry copy, in output order + natoms: int + + +class PropagateResult(BaseModel): + """Supercell written by :func:`propagate_cell`.""" + + out_pdb: Path + numbering: CrystalNumbering # scheme actually used ("auto" already resolved) + ncells: int + natoms: int + nresidues: int + + +def gemmi_version() -> str: + """gemmi version string, for a step's ``tool_versions`` metadata.""" + return gemmi.__version__ + + +# --------------------------------------------------------------------------- # +# Fixed-column PDB text helpers +# --------------------------------------------------------------------------- # +def _pad(line: str, width: int = 80) -> str: + """Right-pad a record so every fixed-column field can be sliced and rewritten.""" + return line.rstrip("\n").ljust(width) + + +def _is_coord(line: str) -> bool: + return line.startswith(_COORD_RECORDS) + + +def _get_xyz(line: str) -> tuple[float, float, float]: + return float(line[30:38]), float(line[38:46]), float(line[46:54]) + + +def _set_xyz(line: str, x: float, y: float, z: float) -> str: + line = _pad(line) + return f"{line[:30]}{x:8.3f}{y:8.3f}{z:8.3f}{line[54:]}" + + +def _set_chain(line: str, chain: str) -> str: + line = _pad(line) + return line[:21] + chain[0] + line[22:] + + +def _set_serial(line: str, serial: int) -> str: + line = _pad(line) + return f"{line[:6]}{serial % (_MAX_SERIAL + 1):5d}{line[11:]}" + + +def _set_resseq(line: str, resseq: int) -> str: + line = _pad(line) + return f"{line[:22]}{resseq % (_MAX_RESSEQ + 1):4d}{line[26:]}" + + +def _set_segid(line: str, segid: str) -> str: + """Write the segID field (columns 73-76).""" + line = _pad(line) + return f"{line[:72]}{segid[:4]:<4}{line[76:]}" + + +def _residue_key(line: str) -> tuple[str, str, str, str]: + """Chain id, residue name, residue sequence number and insertion code.""" + line = _pad(line) + return line[21], line[17:20], line[22:26], line[26] + + +def _count_residues(lines: Sequence[str]) -> int: + """Number of residues a PDB reader would delimit in ``lines``.""" + count = 0 + previous = None + for line in lines: + if _is_coord(line): + key = _residue_key(line) + if key != previous: + count += 1 + previous = key + return count + + +def _format_cryst1(cell: gemmi.UnitCell, spacegroup: str = "P 1") -> str: + return ( + f"CRYST1{cell.a:9.3f}{cell.b:9.3f}{cell.c:9.3f}" + f"{cell.alpha:7.2f}{cell.beta:7.2f}{cell.gamma:7.2f} {spacegroup:<11}" + ) + + +def _read_pdb(pdb: Path) -> tuple[str, list[str]]: + """Return the CRYST1 record and the ATOM/HETATM/TER body of a PDB file. + + Headers, CONECT and secondary-structure records are dropped: ChimeraX carried some + of them through, but nothing downstream in this pipeline reads them and they would + be stale after expansion. + """ + cryst1: str | None = None + body: list[str] = [] + with open(pdb) as fh: + for line in fh: + if line.startswith("CRYST1"): + if cryst1 is None: + cryst1 = _pad(line) + elif _is_coord(line) or line.startswith("TER"): + body.append(_pad(line)) + if cryst1 is None: + raise CrystalSymmetryError(f"{pdb}: no CRYST1 record, so the unit cell is unknown") + return cryst1, body + + +def _write_pdb(pdb: Path, lines: Sequence[str]) -> None: + pdb.write_text("\n".join(line.rstrip() for line in lines) + "\n") + + +def _cell_from_cryst1(cryst1: str) -> gemmi.UnitCell: + return gemmi.UnitCell( + float(cryst1[6:15]), + float(cryst1[15:24]), + float(cryst1[24:33]), + float(cryst1[33:40]), + float(cryst1[40:47]), + float(cryst1[47:54]), + ) + + +def _spacegroup_from_cryst1(cryst1: str, override: str | None = None) -> gemmi.SpaceGroup: + name = override if override else cryst1[55:66].strip() + if not name: + raise CrystalSymmetryError("CRYST1 carries no space group; set crystal.spacegroup") + spacegroup = gemmi.find_spacegroup_by_name(name) + if spacegroup is None: + raise CrystalSymmetryError(f"unknown space group {name!r}") + return spacegroup + + +def _ordered_operations(spacegroup: gemmi.SpaceGroup, op_order: Sequence[int] | None) -> list: + """Symmetry operations of ``spacegroup``, optionally reordered. + + gemmi lists the operations in Hall-symbol order. For P 21 21 21 that swaps the third + and fourth relative to the International Tables order ChimeraX used, which swaps the + contents of chains C and D; ``op_order=[1, 2, 4, 3]`` reproduces the ChimeraX + A/B/C/D labelling exactly. + """ + ops = list(spacegroup.operations()) + if op_order is None: + return ops + if sorted(op_order) != list(range(1, len(ops) + 1)): + raise CrystalSymmetryError( + f"op_order must be a permutation of 1..{len(ops)} for {spacegroup.xhm()}, " + f"got {list(op_order)}" + ) + return [ops[i - 1] for i in op_order] + + +def _lattice_vectors(cell: gemmi.UnitCell) -> list[gemmi.Position]: + """Cartesian a, b, c vectors of the cell. + + Orthogonalization is linear, so the fractional basis vectors map straight onto them. + """ + bases = ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0)) + return [cell.orthogonalize(gemmi.Fractional(*basis)) for basis in bases] + + +# --------------------------------------------------------------------------- # +# Asymmetric unit -> unit cell (replaces ChimeraX) +# --------------------------------------------------------------------------- # +def expand_to_unit_cell( + asu_pdb: str | Path, + cell_pdb: str | Path, + *, + spacegroup: str | None = None, + chain_ids: str | None = None, + op_order: Sequence[int] | None = None, +) -> ExpandResult: + """Expand one asymmetric unit into the full unit cell, written in space group P 1. + + Stands in for ChimeraX's ``changechains #1 A; unitcell #1; combine #2``: every + symmetry operation of the CRYST1 space group is applied to the input coordinates and + each copy is then displaced by the integer lattice vector that brings its + (unweighted) atom centroid into ``[0, 1)^3`` — the packing rule ``unitcell`` uses. + Copies are labelled ``A``, ``B``, ``C``, ... and the output CRYST1 already carries + ``P 1``, so no separate space-group rewrite is needed. + """ + asu_pdb, cell_pdb = Path(asu_pdb), Path(cell_pdb) + cryst1, body = _read_pdb(asu_pdb) + cell = _cell_from_cryst1(cryst1) + group = _spacegroup_from_cryst1(cryst1, spacegroup) + ops = _ordered_operations(group, op_order) + + ids = chain_ids if chain_ids else _CHAIN_IDS + if len(ops) > len(ids): + raise CrystalSymmetryError( + f"{group.xhm()} needs {len(ops)} chain ids for its symmetry copies but only " + f"{len(ids)} are available" + ) + + fractional = [ + cell.fractionalize(gemmi.Position(*_get_xyz(line))).tolist() + for line in body + if _is_coord(line) + ] + + out = [_format_cryst1(cell)] + used_ids: list[str] = [] + serial = 0 + for icopy, op in enumerate(ops): + chain = ids[icopy] + used_ids.append(chain) + moved = [op.apply_to_xyz(frac) for frac in fractional] + centroid = [sum(values) / len(moved) for values in zip(*moved, strict=True)] + shift = [-math.floor(value) for value in centroid] + packed = [ + cell.orthogonalize( + gemmi.Fractional(frac[0] + shift[0], frac[1] + shift[1], frac[2] + shift[2]) + ) + for frac in moved + ] + icoord = 0 + for line in body: + if _is_coord(line): + pos = packed[icoord] + icoord += 1 + serial += 1 + line = _set_xyz(line, pos.x, pos.y, pos.z) + out.append(_set_serial(_set_chain(line, chain), serial)) + out.append("END") + + _write_pdb(cell_pdb, out) + return ExpandResult( + cell_pdb=cell_pdb, + spacegroup=group.xhm(), + chain_ids=used_ids, + natoms=serial, + ) + + +# --------------------------------------------------------------------------- # +# Unit cell -> supercell (replaces AmberTools PropPDB) +# --------------------------------------------------------------------------- # +def propagate_cell( + cell_pdb: str | Path, + out_pdb: str | Path, + *, + ix: int = 1, + iy: int = 1, + iz: int = 1, + numbering: CrystalNumbering = "auto", +) -> PropagateResult: + """Replicate a P 1 unit cell ``ix`` x ``iy`` x ``iz`` times and scale CRYST1. + + Stands in for AmberTools ``PropPDB``, including its replication order (z fastest, + then y, then x), so the output is layout-compatible with it. + + ``numbering`` selects how atoms and residues are counted: + + - ``continuous`` — sequential across the whole supercell, which is what PropPDB + writes. The serial field holds 99999 atoms and resSeq 9999 residues; past that + the counters wrap. + - ``per-cell`` — restart both counters in every unit cell, give each (cell, input + chain) pair its own chain id and record the cell index in the segID field, so + every field stays inside its column width however large the supercell grows. + - ``auto`` (default) — ``continuous`` while it fits those fields, ``per-cell`` + beyond that; identical to ``continuous`` for every supercell it can represent. + """ + cell_pdb, out_pdb = Path(cell_pdb), Path(out_pdb) + if min(ix, iy, iz) < 1: + raise CrystalSymmetryError(f"propagation counts must be >= 1, got {ix}, {iy}, {iz}") + + cryst1, body = _read_pdb(cell_pdb) + cell = _cell_from_cryst1(cryst1) + avec, bvec, cvec = _lattice_vectors(cell) + supercell = gemmi.UnitCell( + cell.a * ix, cell.b * iy, cell.c * iz, cell.alpha, cell.beta, cell.gamma + ) + + ncells = ix * iy * iz + natoms_in = sum(1 for line in body if _is_coord(line)) + nres_in = _count_residues(body) + natoms = natoms_in * ncells + nresidues = nres_in * ncells + + mode = numbering + if mode == "auto": + fits = natoms <= _MAX_SERIAL and nresidues <= _MAX_RESSEQ + mode = "continuous" if fits else "per-cell" + + in_chains: list[str] = [] + for line in body: + if _is_coord(line) and line[21] not in in_chains: + in_chains.append(line[21]) + if mode == "per-cell" and (natoms_in > _MAX_SERIAL or nres_in > _MAX_RESSEQ): + raise CrystalSymmetryError( + f"{cell_pdb}: one unit cell ({natoms_in} atoms, {nres_in} residues) already " + "overflows the fixed-column PDB serial/resSeq fields" + ) + + out = [_format_cryst1(supercell)] + serial = 0 + resseq = 0 + previous_key = None + chain_map: dict[str, str] = {chain: chain for chain in in_chains} + segid = "" + icell = -1 + # z fastest, then y, then x — the order PropPDB writes its copies in. + for cx in range(ix): + for cy in range(iy): + for cz in range(iz): + icell += 1 + dx = cx * avec.x + cy * bvec.x + cz * cvec.x + dy = cx * avec.y + cy * bvec.y + cz * cvec.y + dz = cx * avec.z + cy * bvec.z + cz * cvec.z + + if mode == "per-cell": + serial = 0 + resseq = 0 + previous_key = None + base = icell * len(in_chains) + chain_map = { + chain: _CHAIN_IDS[(base + offset) % len(_CHAIN_IDS)] + for offset, chain in enumerate(in_chains) + } + segid = f"{icell:04d}" + + for line in body: + if mode == "per-cell": + line = _set_segid(_set_chain(line, chain_map[line[21]]), segid) + if _is_coord(line): + key = (cx, cy, cz, *_residue_key(line)) + if key != previous_key: + resseq += 1 + previous_key = key + serial += 1 + x, y, z = _get_xyz(line) + line = _set_xyz(line, x + dx, y + dy, z + dz) + out.append(_set_resseq(_set_serial(line, serial), resseq)) + out.append("END") + + # A wrapped resSeq can merge two neighbouring residues into one as far as any PDB + # reader is concerned, which silently corrupts the topology; refuse to write that. + delimited = _count_residues(out) + if delimited != nresidues: + hint = "; use crystal.numbering='per-cell'" if mode == "continuous" else "" + raise CrystalSymmetryError( + f"{out_pdb}: {nresidues} residues were propagated but a PDB reader would " + f"delimit {delimited} — adjacent residues share chain, name and number{hint}" + ) + + _write_pdb(out_pdb, out) + return PropagateResult( + out_pdb=out_pdb, + numbering=mode, + ncells=ncells, + natoms=natoms, + nresidues=nresidues, + ) diff --git a/md_workflows/core/exceptions.py b/md_workflows/core/exceptions.py index 4a21af4..8d0df97 100644 --- a/md_workflows/core/exceptions.py +++ b/md_workflows/core/exceptions.py @@ -28,7 +28,7 @@ def __init__(self, step: str, missing: list[Path]): class StepToolError(MDWorkflowError): - """An external tool (gmx, tleap, ChimeraX, AmberTools, ...) exited non-zero.""" + """An external tool (gmx, tleap, AmberTools, ...) exited non-zero.""" def __init__( self, @@ -55,6 +55,16 @@ def __init__( super().__init__("\n".join(parts)) +class CrystalSymmetryError(MDWorkflowError): + """A PDB lacks the symmetry metadata needed to build the crystal, or the requested + crystal cannot be written as a PDB file. + + Raised by the gemmi primitives in ``core.crystal`` — missing/unknown CRYST1 space + group, a bad symmetry-operation order, or a supercell whose numbering would not fit + the fixed-column serial/resSeq fields. + """ + + class GmxOutputParseError(MDWorkflowError): """A value expected in GROMACS stdout/log (e.g. a molecule count) was not found. diff --git a/md_workflows/core/steps/make_crystal.py b/md_workflows/core/steps/make_crystal.py index d50c885..1aecbee 100644 --- a/md_workflows/core/steps/make_crystal.py +++ b/md_workflows/core/steps/make_crystal.py @@ -1,8 +1,10 @@ """Build a crystal supercell from the parameterized protein. Corresponds to ``make_crystal.sh``: dry the protein, restore the CRYST1 cell, expand -the asymmetric unit to the full unit cell with ChimeraX, set the spacegroup to P1, then -replicate into a supercell with PropPDB. +the asymmetric unit to the full unit cell in space group P1, then replicate it into a +supercell. The expansion and the replication are done with gemmi (``core.crystal``), +which replaces the ChimeraX ``unitcell`` and AmberTools ``PropPDB`` calls the shell +recipe shelled out to; ``pdb4amber`` is the only external tool this step still needs. """ from __future__ import annotations @@ -50,7 +52,9 @@ def make_crystal( *, resume: bool = False, ) -> MakeCrystalResult: - from ..gmx import run_tool # local import keeps module import light for tests + # local imports keep module import (and the contract tests) light + from ..crystal import expand_to_unit_cell, gemmi_version, propagate_cell + from ..gmx import run_tool check_inputs(inputs) wd = inputs.workdir @@ -86,32 +90,33 @@ def make_crystal( # 2. Restore the CRYST1 cell from pdb_clean; drop stray ions / any wrong CRYST1. _prepend_cryst1(inputs.pdb_clean, prot_dry) - # 3. Expand the asymmetric unit to the full unit cell with ChimeraX. - _expand_unit_cell(wd, prot_dry, prot_dry_cell, params.chimerax_exec) - - # 4. Rewrite the CRYST1 spacegroup to P1 and prepend to the cell PDB. - _set_p1_spacegroup(prot_dry, prot_dry_cell) + # 3. Expand the asymmetric unit to the full unit cell, already in spacegroup P1. + expanded = expand_to_unit_cell( + prot_dry, + prot_dry_cell, + spacegroup=params.spacegroup, + op_order=params.op_order, + ) - # 5. Replicate the P1 cell into the requested supercell. + # 4. Replicate the P1 cell into the requested supercell. + metrics: dict[str, float | int] = { + "ix": ix, + "iy": iy, + "iz": iz, + "cell_atoms": expanded.natoms, + "cell_copies": len(expanded.chain_ids), + } if ix > 0 or iy > 0 or iz > 0: - run_tool( - [ - "PropPDB", - "-p", - str(prot_dry_cell), - "-o", - str(xtal), - "-ix", - str(ix), - "-iy", - str(iy), - "-iz", - str(iz), - ], - tool="PropPDB", - cwd=wd, - log_path=wd / "proppdb.log", + propagated = propagate_cell( + prot_dry_cell, + xtal, + ix=ix, + iy=iy, + iz=iz, + numbering=params.numbering, ) + metrics["xtal_atoms"] = propagated.natoms + metrics["xtal_residues"] = propagated.nresidues else: shutil.copy(prot_dry_cell, xtal) @@ -122,7 +127,8 @@ def make_crystal( outputs=outputs, params=params.model_dump(), input_checksums=checksums(consumed), - metrics={"ix": ix, "iy": iy, "iz": iz}, + tool_versions={"gemmi": gemmi_version()}, + metrics=metrics, ) @@ -145,34 +151,3 @@ def _prepend_cryst1(source_pdb: Path, target_pdb: Path) -> None: with open(target_pdb, "w") as fh: fh.write(cryst1) fh.writelines(filtered) - - -def _expand_unit_cell(workdir: Path, dry_pdb: Path, cell_pdb: Path, chimerax_exec: str) -> None: - from ..gmx import run_tool - - cxc = workdir / "expand.cxc" - cxc.write_text( - f"open {dry_pdb}\nchangechains #1 A\nunitcell #1\ncombine #2\nsave {cell_pdb} #3\nquit\n" - ) - run_tool( - [chimerax_exec, "--offscreen", "--nogui", str(cxc)], - tool="ChimeraX", - cwd=workdir, - log_path=workdir / "chimerax_expand.log", - ) - - -def _set_p1_spacegroup(dry_pdb: Path, cell_pdb: Path) -> None: - """Rewrite the CRYST1 line with a P 1 spacegroup and prepend to the cell PDB.""" - cryst1_p1 = "" - with open(dry_pdb) as fh: - for line in fh: - if line.startswith("CRYST1"): - cryst1_p1 = line[:55] + "P 1\n" - break - - with open(cell_pdb) as fh: - cell_content = fh.read() - with open(cell_pdb, "w") as fh: - fh.write(cryst1_p1) - fh.write(cell_content) diff --git a/md_workflows/core/steps/make_waterbox.py b/md_workflows/core/steps/make_waterbox.py index 26e6f01..9554706 100644 --- a/md_workflows/core/steps/make_waterbox.py +++ b/md_workflows/core/steps/make_waterbox.py @@ -4,6 +4,12 @@ ``nc_scale``, fill the sub-cell with water at ``conc`` mol/L, tile it back up ``nc_scale``x, restore the full crystal CRYST1, then minimize + equilibrate the box. +The tiling is done with gemmi (``core.crystal.propagate_cell``), which replaces the +AmberTools ``PropPDB`` call the shell recipe shelled out to; ``gmx insert-molecules`` is +the only external tool this step still needs. ``nc_scale``^3 cells of water overrun the +fixed-column PDB resSeq field easily (PropPDB wrapped it silently), so the default +``auto`` numbering restarts the counters per cell once the totals no longer fit. + The water count written to the topology is the exact molecule count of the expanded box (``count_waters``), superseding the scripts' log-derived estimates. """ @@ -62,6 +68,9 @@ def make_waterbox( *, resume: bool = False, ) -> MakeWaterboxResult: + # local import keeps module import (and the contract tests) light + from ..crystal import propagate_cell + check_inputs(inputs) wd = inputs.workdir wb = wd / "waterbox" @@ -118,23 +127,13 @@ def make_waterbox( cwd=wb, log_path=wb / "insert-molecules.log", ) - run_tool( - [ - "PropPDB", - "-p", - str(wb / "box_solv.pdb"), - "-o", - str(box_solv_expand), - "-ix", - str(nc), - "-iy", - str(nc), - "-iz", - str(nc), - ], - tool="PropPDB", - cwd=wb, - log_path=wb / "proppdb.log", + propagated = propagate_cell( + wb / "box_solv.pdb", + box_solv_expand, + ix=nc, + iy=nc, + iz=nc, + numbering=params.numbering, ) _restore_cryst1(box_solv_expand, cryst1_xtal) @@ -155,7 +154,7 @@ def make_waterbox( params=params.model_dump(), input_checksums=checksums(consumed), run_profiles_used={"waterbox_min": min_profile, "waterbox_equil": equil_profile}, - metrics={"nwat": nwat, "nc_scale": nc}, + metrics={"nwat": nwat, "nc_scale": nc, "box_atoms": propagated.natoms}, log_paths=[min_run.log, equil_run.log], ) diff --git a/pyproject.toml b/pyproject.toml index 689a5f7..2a0104e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,8 +8,10 @@ version = "0.2.0" description = "Molecular dynamics workflow orchestration scripts" requires-python = ">=3.10" dependencies = [ + "gemmi>=0.6", "pydantic>=2", "pyyaml>=6", + "tomli>=1.1; python_version < '3.11'", "typer>=0.12", ] diff --git a/tests/test_crystal.py b/tests/test_crystal.py new file mode 100644 index 0000000..d489057 --- /dev/null +++ b/tests/test_crystal.py @@ -0,0 +1,189 @@ +from pathlib import Path + +import gemmi +import pytest + +from md_workflows.core.crystal import expand_to_unit_cell, propagate_cell +from md_workflows.core.exceptions import CrystalSymmetryError + +CELL = (20.0, 30.0, 40.0, 90.0, 90.0, 90.0) + + +def _cryst1(spacegroup: str) -> str: + a, b, c, alpha, beta, gamma = CELL + return f"CRYST1{a:9.3f}{b:9.3f}{c:9.3f}{alpha:7.2f}{beta:7.2f}{gamma:7.2f} {spacegroup:<11}\n" + + +def _atom(serial: int, name: str, resname: str, resseq: int, xyz, chain: str = "A") -> str: + x, y, z = xyz + return ( + f"ATOM {serial:5d} {name:<4} {resname:>3} {chain}{resseq:4d} " + f"{x:8.3f}{y:8.3f}{z:8.3f} 1.00 0.00\n" + ) + + +def _write_asu(tmp_path: Path, spacegroup: str, name: str = "asu.pdb") -> Path: + """Three one-atom residues inside the cell, plus a TER record.""" + pdb = tmp_path / name + pdb.write_text( + _cryst1(spacegroup) + + _atom(1, "N", "ALA", 1, (1.0, 2.0, 3.0)) + + _atom(2, "CA", "GLY", 2, (4.0, 5.0, 6.0)) + + _atom(3, "CB", "SER", 3, (7.0, 8.0, 9.0)) + + "TER 3 SER A 3\n" + + "END\n" + ) + return pdb + + +def _coords(pdb: Path) -> list[tuple[float, float, float]]: + return [ + (float(line[30:38]), float(line[38:46]), float(line[46:54])) + for line in pdb.read_text().splitlines() + if line.startswith(("ATOM", "HETATM")) + ] + + +def _fields(pdb: Path, start: int, stop: int) -> list[str]: + return [ + line[start:stop] + for line in pdb.read_text().splitlines() + if line.startswith(("ATOM", "HETATM")) + ] + + +def test_expand_p1_writes_p1_cryst1_and_keeps_coordinates(tmp_path: Path): + asu = _write_asu(tmp_path, "P 1") + out = tmp_path / "cell.pdb" + result = expand_to_unit_cell(asu, out) + + assert result.spacegroup == "P 1" + assert result.chain_ids == ["A"] + assert result.natoms == 3 + assert out.read_text().splitlines()[0].endswith("P 1") + assert _coords(out) == [(1.0, 2.0, 3.0), (4.0, 5.0, 6.0), (7.0, 8.0, 9.0)] + + +def test_expand_labels_one_chain_per_symmetry_copy(tmp_path: Path): + asu = _write_asu(tmp_path, "P 21 21 21") + out = tmp_path / "cell.pdb" + result = expand_to_unit_cell(asu, out) + + assert result.chain_ids == ["A", "B", "C", "D"] + assert result.natoms == 12 + assert set(_fields(out, 21, 22)) == {"A", "B", "C", "D"} + # atom order, residue names and serials run through the copies without gaps + assert _fields(out, 17, 20) == ["ALA", "GLY", "SER"] * 4 + assert [int(s) for s in _fields(out, 6, 11)] == list(range(1, 13)) + + +def test_expand_packs_every_copy_centroid_into_the_cell(tmp_path: Path): + asu = _write_asu(tmp_path, "P 21 21 21") + out = tmp_path / "cell.pdb" + expand_to_unit_cell(asu, out) + + cell = gemmi.UnitCell(*CELL) + copies = _coords(out) + for start in range(0, len(copies), 3): + fractional = [cell.fractionalize(gemmi.Position(*xyz)) for xyz in copies[start : start + 3]] + for axis in ("x", "y", "z"): + centroid = sum(getattr(f, axis) for f in fractional) / 3 + assert 0.0 <= centroid < 1.0 + + +def test_expand_op_order_permutes_the_copies(tmp_path: Path): + asu = _write_asu(tmp_path, "P 21 21 21") + default = tmp_path / "default.pdb" + swapped = tmp_path / "swapped.pdb" + expand_to_unit_cell(asu, default) + expand_to_unit_cell(asu, swapped, op_order=[1, 2, 4, 3]) + + got, expected = _coords(swapped), _coords(default) + assert got[:6] == expected[:6] + assert got[6:9] == expected[9:12] # chain C now holds the fourth operation + assert got[9:12] == expected[6:9] + + +def test_expand_rejects_an_op_order_that_is_not_a_permutation(tmp_path: Path): + asu = _write_asu(tmp_path, "P 21 21 21") + with pytest.raises(CrystalSymmetryError, match="permutation"): + expand_to_unit_cell(asu, tmp_path / "cell.pdb", op_order=[1, 2, 3]) + + +def test_expand_honors_a_spacegroup_override(tmp_path: Path): + asu = _write_asu(tmp_path, "P 1") + result = expand_to_unit_cell(asu, tmp_path / "cell.pdb", spacegroup="P 2") + assert result.spacegroup == "P 1 2 1" # gemmi reports the full Hermann-Mauguin name + assert result.natoms == 6 + + +def test_expand_without_cryst1_raises(tmp_path: Path): + asu = tmp_path / "nocell.pdb" + asu.write_text(_atom(1, "N", "ALA", 1, (1.0, 2.0, 3.0))) + with pytest.raises(CrystalSymmetryError, match="no CRYST1"): + expand_to_unit_cell(asu, tmp_path / "cell.pdb") + + +def test_expand_with_an_unknown_spacegroup_raises(tmp_path: Path): + asu = _write_asu(tmp_path, "P 1") + with pytest.raises(CrystalSymmetryError, match="unknown space group"): + expand_to_unit_cell(asu, tmp_path / "cell.pdb", spacegroup="Q 7") + + +def test_propagate_shifts_by_lattice_vectors_and_scales_cryst1(tmp_path: Path): + cell = _write_asu(tmp_path, "P 1") + out = tmp_path / "xtal.pdb" + result = propagate_cell(cell, out, ix=2, iy=1, iz=1) + + assert (result.ncells, result.natoms, result.nresidues) == (2, 6, 6) + assert result.numbering == "continuous" + coords = _coords(out) + assert coords[:3] == [(1.0, 2.0, 3.0), (4.0, 5.0, 6.0), (7.0, 8.0, 9.0)] + assert coords[3:] == [(x + CELL[0], y, z) for x, y, z in coords[:3]] + cryst1 = out.read_text().splitlines()[0] + assert float(cryst1[6:15]) == 2 * CELL[0] + assert float(cryst1[15:24]) == CELL[1] + assert [int(s) for s in _fields(out, 6, 11)] == list(range(1, 7)) + assert [int(s) for s in _fields(out, 22, 26)] == list(range(1, 7)) + + +def test_propagate_replicates_with_z_fastest_like_proppdb(tmp_path: Path): + cell = _write_asu(tmp_path, "P 1") + out = tmp_path / "xtal.pdb" + propagate_cell(cell, out, ix=2, iy=1, iz=2) + + first_atoms = _coords(out)[::3] # one atom per replicated cell + assert first_atoms == [ + (1.0, 2.0, 3.0), + (1.0, 2.0, 3.0 + CELL[2]), + (1.0 + CELL[0], 2.0, 3.0), + (1.0 + CELL[0], 2.0, 3.0 + CELL[2]), + ] + + +def test_propagate_per_cell_restarts_counters_and_tags_the_segid(tmp_path: Path): + cell = _write_asu(tmp_path, "P 1") + out = tmp_path / "xtal.pdb" + result = propagate_cell(cell, out, ix=2, numbering="per-cell") + + assert result.numbering == "per-cell" + assert [int(s) for s in _fields(out, 6, 11)] == [1, 2, 3, 1, 2, 3] + assert [int(s) for s in _fields(out, 22, 26)] == [1, 2, 3, 1, 2, 3] + assert _fields(out, 21, 22) == ["A"] * 3 + ["B"] * 3 + assert _fields(out, 72, 76) == ["0000"] * 3 + ["0001"] * 3 + + +def test_propagate_auto_switches_to_per_cell_when_the_pdb_fields_overflow(tmp_path: Path): + cell = _write_asu(tmp_path, "P 1") + out = tmp_path / "xtal.pdb" + result = propagate_cell(cell, out, ix=15, iy=15, iz=15) + + assert result.nresidues == 3 * 15**3 > 9999 + assert result.numbering == "per-cell" + assert max(int(s) for s in _fields(out, 22, 26)) == 3 + + +def test_propagate_rejects_counts_below_one(tmp_path: Path): + cell = _write_asu(tmp_path, "P 1") + with pytest.raises(CrystalSymmetryError, match=">= 1"): + propagate_cell(cell, tmp_path / "xtal.pdb", ix=0)