Add MPS support, upgrade PyTorch, refactor packaging, and fix CI + correctness/security bugs

.github/workflows/ci.yml

@@ -0,0 +1,79 @@
+name: CI
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+
+jobs:
+  quick-tests:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - name: data-utility
+            path: "tests/data tests/utility"
+          - name: models
+            path: "tests/models"
+          - name: layers
+            path: "tests/layers"
+    name: ${{ matrix.name }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-activate-base: false
+          miniforge-version: latest
+      - name: Create Locked Environment
+        shell: bash -el {0}
+        run: conda create -y -n chroma --file environment.linux-64.lock
+      - name: Install Package
+        shell: bash -el {0}
+        run: conda run -n chroma python -m pip install --no-deps -e .
+      - name: Run Quick Tests
+        shell: bash -el {0}
+        run: conda run -n chroma python -m pytest -q -m "not integration and not slow" ${{ matrix.path }}
+
+  integration:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-activate-base: false
+          miniforge-version: latest
+      - name: Create Locked Environment
+        shell: bash -el {0}
+        run: conda create -y -n chroma --file environment.linux-64.lock
+      - name: Install Package
+        shell: bash -el {0}
+        run: conda run -n chroma python -m pip install --no-deps -e .
+      - name: Run Integration Tests
+        shell: bash -el {0}
+        run: conda run -n chroma python -m pytest -q tests/models/test_chroma.py tests/utility/test_api.py tests/data/test_system.py
+
+  release:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          auto-activate-base: false
+          miniforge-version: latest
+      - name: Create Locked Environment
+        shell: bash -el {0}
+        run: conda create -y -n chroma --file environment.linux-64.lock
+      - name: Install Package
+        shell: bash -el {0}
+        run: conda run -n chroma python -m pip install --no-deps -e .
+      - name: Build Artifacts
+        shell: bash -el {0}
+        run: conda run -n chroma python -m build
+      - name: Check Artifacts
+        shell: bash -el {0}
+        run: conda run -n chroma python -m twine check dist/*
+      - name: Packaging Smoke Test
+        shell: bash -el {0}
+        run: conda run -n chroma python -m pytest -q tests/utility/test_packaging.py

.gitignore

@@ -8,6 +8,7 @@
 *.tsv
 *.pdb
 *.cif
+!tests/resources/*.cif
 *.pk
 *.pt
 *.fasta
@@ -24,6 +25,8 @@ __mmtf__
 __pycache__
 public
 htmlcov
+build
+dist
 make.bat
 examples
 chroma/layers/structure/params/centering_2g3n.params
@@ -32,3 +35,4 @@ config.json
 
 # ides
 .vscode
+/personal/

CONTRIBUTING.md

@@ -9,3 +9,36 @@ Note that your contributions will be governed by the Apache 2.0 license, meaning
 Please use GitHub pull requests to contribute code. See
 [GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
 information on using pull requests. We will try to monitor incoming requests with some regularity, but cannot promise a specific timeframe within which we will review your request. 
+
+## Tests and Fixtures
+
+Use the `chroma` Conda environment for local testing. The repository CI workflow in
+[`.github/workflows/ci.yml`](.github/workflows/ci.yml) runs the data and utility,
+model, layer, integration, and release smoke suites separately. The shared
+cross-platform CI environment is [`environment.yml`](environment.yml), while
+[`environment_Mac.yaml`](environment_Mac.yaml) keeps the Apple Silicon local setup
+used during development. CI installs from
+[`environment.linux-64.lock`](environment.linux-64.lock), and Apple Silicon
+local development should use
+[`environment.osx-arm64.lock`](environment.osx-arm64.lock).
+
+To regenerate the explicit lockfiles from [`environment.yml`](environment.yml),
+run:
+
+```bash
+python scripts/generate_lockfiles.py
+```
+
+The Linux lock targets `__glibc=2.17`, which matches the minimum ABI required
+by the pinned `linux-64` packages and is sufficient for the GitHub Actions
+`ubuntu-latest` runners used by CI.
+
+Fixture provenance and regeneration notes live in
+[`tests/resources/README.md`](tests/resources/README.md), including the generated
+`steps200_seed42_len100.cif` sample used by `tests/models/test_chroma.py`.
+
+For quick local iteration, you can skip the heavier tiers with:
+
+```bash
+pytest -m "not integration and not slow"
+```

README.md

@@ -28,15 +28,30 @@ Internally, Chroma uses diffusion modeling, equivariant graph neural networks, a
 * [Chroma API Tutorial](https://colab.research.google.com/github/generatebio/chroma/blob/main/notebooks/ChromaAPI.ipynb): Code notebook demonstrating protein I/O, sampling, and design configuration directly in `python`.
 * [Chroma Conditioner API Tutorial](https://colab.research.google.com/github/generatebio/chroma/blob/main/notebooks/ChromaConditioners.ipynb): A deeper dive under the hood for implementing new Chroma [Conditioners](#conditioner-api).
 
-**PyPi package**.You can install the latest release of Chroma with:
+
+**Install Chroma from GitHub**
+
+Create a locked Conda environment for your platform
+
 ```
-pip install generate-chroma
+conda create -n chroma --file environment.osx-arm64.lock
 ```
 
-**Install latest Chroma from github**
+On Linux, use:
+
+```
+conda create -n chroma --file environment.linux-64.lock
+```
+
+The shared source spec is `environment.yml`, and `environment_Mac.yaml`
+remains available as the Apple Silicon YAML variant used during development.
+
+Clone from github and install as a package
+
 ```
 git clone https://github.com/generatebio/chroma.git
-pip install -e chroma # use `-e` for it to be editable locally. 
+conda activate chroma
+python -m pip install --no-deps -e . # use `-e` for it to be editable locally.
 ```
 
 ## Sampling
@@ -101,7 +116,7 @@ The value of diffusion time conditioning $t$ can be set via the `design_t` param
 from chroma import Protein, Chroma
 chroma = Chroma()
 
-protein = Protein('1GFP')
+protein = Protein.from_PDBID('1GFP')
 protein = chroma.design(protein)
 
 protein.to("1GFP-redesign.cif")

chroma/__init__.py

@@ -12,8 +12,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from importlib import import_module
+
 __version__ = "1.0.0"
-from chroma.data.protein import Protein
-from chroma.layers.structure import conditioners
-from chroma.models.chroma import Chroma
-from chroma.utility import api
+__all__ = ["__version__", "Protein", "conditioners", "Chroma", "api"]
+
+
+def __getattr__(name):
+    if name == "Protein":
+        return import_module("chroma.data.protein").Protein
+    if name == "conditioners":
+        return import_module("chroma.layers.structure.conditioners")
+    if name == "Chroma":
+        return import_module("chroma.models.chroma").Chroma
+    if name == "api":
+        return import_module("chroma.utility.api")
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")