rc-foundry 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

foundry/inference_engines/checkpoint_registry.py

@@ -3,20 +3,62 @@
 import os
 from dataclasses import dataclass
 from pathlib import Path
+from typing import Iterable, List
 
+import dotenv
 
-def get_default_checkpoint_dir() -> Path:
-    """Get the default checkpoint directory.
+DEFAULT_CHECKPOINT_DIR = Path.home() / ".foundry" / "checkpoints"
+
+
+def _normalize_paths(paths: Iterable[Path]) -> list[Path]:
+    """Return absolute, deduplicated paths in order."""
+    seen = set()
+    normalized: List[Path] = []
+    for path in paths:
+        resolved = path.expanduser().absolute()
+        if resolved not in seen:
+            normalized.append(resolved)
+            seen.add(resolved)
+    return normalized
 
-    Priority:
-    1. FOUNDRY_CHECKPOINTS_DIR environment variable
-    2. ~/.foundry/checkpoints
+
+def get_default_checkpoint_dirs() -> list[Path]:
+    """Return checkpoint search paths.
+
+    Always starts with the default ~/.foundry/checkpoints directory and then
+    appends any additional directories from the colon-separated
+    FOUNDRY_CHECKPOINT_DIRS environment variable.
     """
-    if "FOUNDRY_CHECKPOINTS_DIR" in os.environ and os.environ.get(
-        "FOUNDRY_CHECKPOINTS_DIR"
-    ):
-        return Path(os.environ["FOUNDRY_CHECKPOINTS_DIR"]).absolute()
-    return Path.home() / ".foundry" / "checkpoints"
+    env_dirs = os.environ.get("FOUNDRY_CHECKPOINT_DIRS", "")
+
+    # For backward compatibility, also check FOUNDRY_CHECKPOINTS_DIR
+    if not env_dirs:
+        env_dirs = os.environ.get("FOUNDRY_CHECKPOINTS_DIR", "")
+
+    extra_dirs: list[Path] = []
+    if env_dirs:
+        extra_dirs = [Path(p.strip()) for p in env_dirs.split(":") if p.strip()]
+    return _normalize_paths([*extra_dirs, DEFAULT_CHECKPOINT_DIR])
+
+
+def get_default_checkpoint_dir() -> Path:
+    """Backward-compatible helper returning the primary checkpoint directory."""
+    return get_default_checkpoint_dirs()[0]
+
+
+def append_checkpoint_to_env(checkpoint_dirs: list[Path]) -> bool:
+    dotenv_path = dotenv.find_dotenv()
+    if dotenv_path:
+        checkpoint_dirs = _normalize_paths(checkpoint_dirs)
+        dotenv.set_key(
+            dotenv_path=dotenv_path,
+            key_to_set="FOUNDRY_CHECKPOINT_DIRS",
+            value_to_set=":".join(str(path) for path in checkpoint_dirs),
+            export=False,
+        )
+        return True
+    else:
+        return False
 
 
 @dataclass
@@ -27,7 +69,12 @@ class RegisteredCheckpoint:
     sha256: None = None  # Optional: add checksum for verification
 
     def get_default_path(self):
-        return get_default_checkpoint_dir() / self.filename
+        checkpoint_dirs = get_default_checkpoint_dirs()
+        for checkpoint_dir in checkpoint_dirs:
+            candidate = checkpoint_dir / self.filename
+            if candidate.exists():
+                return candidate
+        return checkpoint_dirs[0] / self.filename
 
 
 REGISTERED_CHECKPOINTS = {

foundry/utils/alignment.py

@@ -18,14 +18,19 @@ def weighted_rigid_align(
     Returns:
       X_align_L: [B, L, 3]
     """
-    assert X_L.shape == X_gt_L.shape
-    assert X_L.shape[:-1] == w_L.shape
 
+    # Canonicalize dimensions
+    if X_L.ndim == 2:
+        X_L = X_L[None]
+    if X_gt_L.ndim == 2:
+        X_gt_L = X_gt_L[None]
     if X_exists_L is None:
         X_exists_L = torch.ones((X_L.shape[-2]), dtype=torch.bool)
     if w_L is None:
         w_L = torch.ones_like(X_L[..., 0])
     else:
+        if w_L.ndim == 1:
+            w_L = w_L[None]
         w_L = w_L.to(torch.float32)
 
     # Assert `X_exists_L` is a boolean mask
@@ -33,6 +38,9 @@ def weighted_rigid_align(
         X_exists_L.dtype == torch.bool
     ), "X_exists_L should be a boolean mask! Otherwise, the alignment will be incorrect (silent failure)!"
 
+    assert X_L.shape == X_gt_L.shape
+    assert X_L.shape[:-1] == w_L.shape
+
     X_resolved = X_L[:, X_exists_L]
     X_gt_resolved = X_gt_L[:, X_exists_L]
     w_resolved = w_L[:, X_exists_L]

foundry/version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.5'
-__version_tuple__ = version_tuple = (0, 1, 5)
+__version__ = version = '0.1.7'
+__version_tuple__ = version_tuple = (0, 1, 7)
 
 __commit_id__ = commit_id = None

foundry_cli/download_checkpoints.py

@@ -6,7 +6,7 @@ from typing import Optional
 from urllib.request import urlopen
 
 import typer
-from dotenv import find_dotenv, load_dotenv, set_key
+from dotenv import load_dotenv
 from rich.console import Console
 from rich.progress import (
     BarColumn,
@@ -20,7 +20,8 @@ from rich.progress import (
 
 from foundry.inference_engines.checkpoint_registry import (
     REGISTERED_CHECKPOINTS,
-    get_default_checkpoint_dir,
+    append_checkpoint_to_env,
+    get_default_checkpoint_dirs,
 )
 
 load_dotenv(override=True)
@@ -29,6 +30,27 @@ app = typer.Typer(help="Foundry model checkpoint installation utilities")
 console = Console()
 
 
+def _resolve_checkpoint_dirs(checkpoint_dir: Optional[Path]) -> list[Path]:
+    """Return checkpoint search path with defaults first."""
+    checkpoint_dirs = get_default_checkpoint_dirs()
+    if checkpoint_dir is not None:
+        resolved = checkpoint_dir.expanduser().absolute()
+        if resolved not in checkpoint_dirs:
+            checkpoint_dirs.insert(0, resolved)
+        else:
+            # Move to front
+            checkpoint_dirs.remove(resolved)
+            checkpoint_dirs.insert(0, resolved)
+
+        # Try to persist checkpoint dir to .env (optional, may not exist in Colab etc.)
+        if append_checkpoint_to_env(checkpoint_dirs):
+            console.print(
+                f"Tracked checkpoint directories: {':'.join(str(path) for path in checkpoint_dirs)}"
+            )
+
+    return checkpoint_dirs
+
+
 def download_file(url: str, dest: Path, verify_hash: Optional[str] = None) -> None:
     """Download a file with progress bar and optional hash verification.
 
@@ -123,134 +145,112 @@ def install_model(model_name: str, checkpoint_dir: Path, force: bool = False) ->
 def install(
     models: list[str] = typer.Argument(
         ...,
-        help="Models to install: 'all', 'rfd3', 'rf3', 'mpnn', or combination",
+        help="Models to install: 'all', 'rfd3', 'rf3', 'mpnn', or a combination thereof",
     ),
     checkpoint_dir: Optional[Path] = typer.Option(
         None,
         "--checkpoint-dir",
         "-d",
-        help="Directory to save checkpoints (default: $FOUNDRY_CHECKPOINTS_DIR or ~/.foundry/checkpoints)",
+        help="Directory to save checkpoints (default search path: ~/.foundry/checkpoints plus any $FOUNDRY_CHECKPOINT_DIRS entries)",
     ),
     force: bool = typer.Option(
         False, "--force", "-f", help="Overwrite existing checkpoints"
     ),
 ):
     """Install model checkpoints for foundry.
-
     Examples:
-
         foundry install all
-
         foundry install rfd3 rf3
-
         foundry install proteinmpnn --checkpoint-dir ./checkpoints
     """
     # Determine checkpoint directory
-    if checkpoint_dir is None:
-        checkpoint_dir = get_default_checkpoint_dir()
+    checkpoint_dirs = _resolve_checkpoint_dirs(checkpoint_dir)
+    primary_checkpoint_dir = checkpoint_dirs[0]
 
-    console.print(f"[bold]Checkpoint directory:[/bold] {checkpoint_dir}")
-    console.print()
+    console.print(f"[bold]Install target:[/bold] {primary_checkpoint_dir}\n")
 
     # Expand 'all' to all available models
     if "all" in models:
+        models_to_install = list(REGISTERED_CHECKPOINTS.keys())
+    elif "base-models" in models:
         models_to_install = ["rfd3", "proteinmpnn", "ligandmpnn", "rf3"]
     else:
         models_to_install = models
 
     # Install each model
     for model_name in models_to_install:
-        install_model(model_name, checkpoint_dir, force)
+        install_model(model_name, primary_checkpoint_dir, force)
         console.print()
 
-    # Try to persist checkpoint dir to .env (optional, may not exist in Colab etc.)
-    dotenv_path = find_dotenv()
-    if dotenv_path:
-        set_key(
-            dotenv_path=dotenv_path,
-            key_to_set="FOUNDRY_CHECKPOINTS_DIR",
-            value_to_set=str(checkpoint_dir),
-            export=False,
-        )
-        console.print(f"Saved FOUNDRY_CHECKPOINTS_DIR to {dotenv_path}")
-
     console.print("[bold green]Installation complete![/bold green]")
 
 
-@app.command(name="list")
-def list_models():
+@app.command(name="list-available")
+def list_available():
     """List available model checkpoints."""
     console.print("[bold]Available models:[/bold]\n")
     for name, info in REGISTERED_CHECKPOINTS.items():
         console.print(f"  [cyan]{name:8}[/cyan] - {info.description}")
 
 
-@app.command()
-def show(
-    checkpoint_dir: Optional[Path] = typer.Option(
-        None,
-        "--checkpoint-dir",
-        "-d",
-        help="Checkpoint directory to show",
-    ),
-):
-    """Show installed checkpoints."""
-    if checkpoint_dir is None:
-        checkpoint_dir = get_default_checkpoint_dir()
+@app.command(name="list-installed")
+def list_installed():
+    """List installed checkpoints and their sizes."""
+    checkpoint_dirs = _resolve_checkpoint_dirs(None)
 
-    if not checkpoint_dir.exists():
-        console.print(
-            f"[yellow]No checkpoints directory found at {checkpoint_dir}[/yellow]"
-        )
-        raise typer.Exit(0)
+    checkpoint_files: list[tuple[Path, float]] = []
+    for checkpoint_dir in checkpoint_dirs:
+        if not checkpoint_dir.exists():
+            continue
+        ckpts = list(checkpoint_dir.glob("*.ckpt")) + list(checkpoint_dir.glob("*.pt"))
+        for ckpt in ckpts:
+            size = ckpt.stat().st_size / (1024**3)  # GB
+            checkpoint_files.append((ckpt, size))
 
-    checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
     if not checkpoint_files:
-        console.print(f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]")
+        console.print(
+            "[yellow]No checkpoint files found in any checkpoint directory[/yellow]"
+        )
         raise typer.Exit(0)
 
-    console.print(f"[bold]Installed checkpoints in {checkpoint_dir}:[/bold]\n")
+    console.print("[bold]Installed checkpoints:[/bold]\n")
     total_size = 0
-    for ckpt in sorted(checkpoint_files):
-        size = ckpt.stat().st_size / (1024**3)  # GB
+    for ckpt, size in sorted(checkpoint_files, key=lambda item: str(item[0])):
         total_size += size
-        console.print(f"  {ckpt.name:30} {size:8.2f} GB")
+        console.print(f"  {ckpt} {size:8.2f} GB")
 
     console.print(f"\n[bold]Total:[/bold] {total_size:.2f} GB")
 
 
-@app.command()
+@app.command(name="clean")
 def clean(
-    checkpoint_dir: Optional[Path] = typer.Option(
-        None,
-        "--checkpoint-dir",
-        "-d",
-        help="Checkpoint directory to clean",
-    ),
     confirm: bool = typer.Option(
         True, "--confirm/--no-confirm", help="Ask for confirmation before deleting"
     ),
 ):
     """Remove all downloaded checkpoints."""
-    if checkpoint_dir is None:
-        checkpoint_dir = get_default_checkpoint_dir()
-
-    if not checkpoint_dir.exists():
-        console.print(f"[yellow]No checkpoints found at {checkpoint_dir}[/yellow]")
-        raise typer.Exit(0)
+    checkpoint_dirs = _resolve_checkpoint_dirs(None)
 
     # List files to delete
-    checkpoint_files = list(checkpoint_dir.glob("*.ckpt"))
+    checkpoint_files: list[Path] = []
+    for checkpoint_dir in checkpoint_dirs:
+        if not checkpoint_dir.exists():
+            continue
+        checkpoint_files.extend(checkpoint_dir.glob("*.ckpt"))
+        checkpoint_files.extend(checkpoint_dir.glob("*.pt"))
+
     if not checkpoint_files:
-        console.print(f"[yellow]No checkpoint files found in {checkpoint_dir}[/yellow]")
+        console.print(
+            "[yellow]No checkpoint files found in any checkpoint directory[/yellow]"
+        )
         raise typer.Exit(0)
 
     console.print("[bold]Files to delete:[/bold]")
     total_size = 0
-    for ckpt in checkpoint_files:
+    for ckpt in sorted(checkpoint_files, key=str):
         size = ckpt.stat().st_size / (1024**3)  # GB
         total_size += size
-        console.print(f"  {ckpt.name} ({size:.2f} GB)")
+        console.print(f"  {ckpt} ({size:.2f} GB)")
 
     console.print(f"\n[bold]Total:[/bold] {total_size:.2f} GB")
 

{rc_foundry-0.1.5.dist-info → rc_foundry-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rc-foundry
-Version: 0.1.5
+Version: 0.1.7
 Summary: Shared utilities and training infrastructure for biomolecular structure prediction models.
 Author-email: Institute for Protein Design <contact@ipd.uw.edu>
 License: BSD 3-Clause License
@@ -104,33 +104,34 @@ All models within Foundry rely on [AtomWorks](https://github.com/RosettaCommons/
 pip install rc-foundry[all]
 ```
 
-**Downloading weights** All models can be downloaded to a target folder with:
-
+**Downloading weights** Models can be downloaded to a target folder with:
+```
+foundry install base-models --checkpoint-dir <path/to/ckpt/dir>
+```
+where `checkpoint-dir` will be `~/.foundry/checkpoints` by default. Foundry always searches `~/.foundry/checkpoints` plus any colon-separated entries in `$FOUNDRY_CHECKPOINT_DIRS` during inference or subsequent commands to find checkpoints. `base-models` installs the latest RFD3, RF3 and MPNN variants - you can also download all of the models supported (including multiple checkpoints of RF3) with `all`, or by listing the models sequentially (e.g. `foundry install rfd3 rf3 ...`).
+To list the registry of available checkpoints:
 ```
-foundry install all --checkpoint_dir <path/to/ckpt/dir>
+foundry list-available
 ```
-This will download all the models supported (including multiple checkpoints of RF3) but as a beginner you can start with:
+To check what you already have downloaded (searches `~/.foundry/checkpoints` plus `$FOUNDRY_CHECKPOINT_DIRS` if set):
 ```
-foundry install rfd3 ligandmpnn rf3 --checkpoint_dir  <path/to/ckpt/dir>
+foundry list-installed
 ```
 
->*See `examples/all.ipynb` for how to run each model in a notebook.*
+>*See `examples/all.ipynb` for how to run each model and design proteins end-to-end in a notebook.*
+
+### Google Colab
+For an interactive Google Colab notebook walking through a basic design pipeline with RFD3, MPNN, and RF3, please see the [IPD Design Pipeline Tutorial](https://colab.research.google.com/drive/1ZwIMV3n9h0ZOnIXX0GyKUuoiahgifBxh?usp=sharing).
 
 ### RFdiffusion3 (RFD3)
 
 [RFdiffusion3](https://www.biorxiv.org/content/10.1101/2025.09.18.676967v2) is an all-atom generative model capable of designing protein structures under complex constraints. 
 
-> *See [models/rfd3/README.md](models/rfd3/README.md) for complete documentation.*
-
 <div align="center">
-  <img src="models/rfd3/docs/.assets/trajectory.png" alt="RFdiffusion3 generation trajectory." width="700">
+  <img src="docs/_static/cover.png" alt="RFdiffusion3 generation trajectory." width="700">
 </div>
 
-### ProteinMPNN
-[ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187) and [LigandMPNN](https://www.nature.com/articles/s41592-025-02626-1) are lightweight inverse-folding models which can be use to design diverse sequences for backbones under constrained conditions.
-
-> *See [models/mpnn/README.md](models/mpnn/README.md) for complete documentation.*
-
+> *See [models/rfd3/README.md](models/rfd3/README.md) for complete documentation.*
 
 ### RosettaFold3 (RF3)
 
@@ -142,6 +143,11 @@ foundry install rfd3 ligandmpnn rf3 --checkpoint_dir  <path/to/ckpt/dir>
 
 > *See [models/rf3/README.md](models/rf3/README.md) for complete documentation.*
 
+### ProteinMPNN
+[ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187) and [LigandMPNN](https://www.nature.com/articles/s41592-025-02626-1) are lightweight inverse-folding models which can be use to design diverse sequences for backbones under constrained conditions.
+
+> *See [models/mpnn/README.md](models/mpnn/README.md) for complete documentation.*
+
 ---
 
 ## Development
@@ -159,11 +165,7 @@ foundry install rfd3 ligandmpnn rf3 --checkpoint_dir  <path/to/ckpt/dir>
 Install both `foundry` and models in editable mode for development:
 
 ```bash
-# Install foundry and RF3 in editable mode
-uv pip install -e . -e ./models/rf3 -e ./models/rfd3 -e ./models/mpnn
-
-# Or install only foundry (no models)
-uv pip install -e .
+uv pip install -e '.[all,dev]'
 ```
 
 This approach allows you to:
@@ -171,6 +173,9 @@ This approach allows you to:
 - Work on specific models without installing all models
 - Add new models as independent packages in `models/`
 
+> [!NOTE]
+> Running tests is not currently supported, test files may be missing.
+
 ### Adding New Models
 
 To add a new model: