rc-foundry 0.1.6__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.6'
-__version_tuple__ = version_tuple = (0, 1, 6)
+__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.6.dist-info → rc_foundry-0.1.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rc-foundry
-Version: 0.1.6
+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:

{rc_foundry-0.1.6.dist-info → rc_foundry-0.1.7.dist-info}/RECORD

@@ -1,7 +1,7 @@
 foundry/__init__.py,sha256=H8S1nl5v6YeW8ggn1jKy4GdtH7c-FGS-j7CqUCAEnAU,1926
 foundry/common.py,sha256=Aur8mH-CNmcUqSsw7VgaCQSW5sH1Bqf8Da91jzxPV1Y,3035
 foundry/constants.py,sha256=0n1wBKCvNuw3QaQehSbmsHYkIdaGn3tLeRFItBrdeHY,913
-foundry/version.py,sha256=riGXiVTWXmtdoju9hVCWvTxpszEMAAIK0sZZWoLKlnU,704
+foundry/version.py,sha256=szvPIs2C82UunpzuvVg3MbF4QhzbBYTsVJ8DmPfq6_E,704
 foundry/callbacks/__init__.py,sha256=VsRT1e4sqlJHPcTCsfupMEx82Iz-LoOAGPpwvf_OJeE,126
 foundry/callbacks/callback.py,sha256=xZBo_suP4bLrP6gl5uJPbaXm00DXigePa6dMeDxucgg,3890
 foundry/callbacks/health_logging.py,sha256=tEtkByOlaAA7nnelxb7PbM9_dcIgOsdbxCdQY3K5pMc,16664
@@ -10,7 +10,7 @@ foundry/callbacks/timing_logging.py,sha256=u-r0hKp7fWOY3mLk7CcuIwHgZbhte13m5M09x
 foundry/callbacks/train_logging.py,sha256=Xs3tmZA88qLxmdSOwt-x8YKN4NKb1kVm59uptNXl4Qo,10399
 foundry/hydra/resolvers.py,sha256=xyJzo6OeWAc_LOu8RiHhX7_CRNoLZ22626AvYHXYl4U,2186
 foundry/inference_engines/base.py,sha256=ZHdlmGUqH4-p3v4RdrLH-Ps8_zalr7j5mQ4x-S53N4M,8375
-foundry/inference_engines/checkpoint_registry.py,sha256=kt2Z1JhrAjoOiEpkIIQ0sLttie1ceL8OgXUBmmyA6iw,2544
+foundry/inference_engines/checkpoint_registry.py,sha256=c_me8Uz2NWXAaELhQ4bT1HMPfY8XrH67kvCKdDPrD8g,4149
 foundry/metrics/__init__.py,sha256=qL4wwaiQ7EtR30pmZ9MCknqx909BJcNvHVmNJUaz_WM,236
 foundry/metrics/losses.py,sha256=2CLUmf7oCdFUCvgJukdNkff0FVG3BlATI-NI60TtpVY,903
 foundry/metrics/metric.py,sha256=23pKh_Ra0EcHGo5cSzYQQrUGr5zWRxeufKSJ58tfXXo,12687
@@ -22,7 +22,7 @@ foundry/trainers/fabric.py,sha256=cjaTHbGuJEQwaGBvIAXD_il4bHtY-crsTY14Xn77uXA,40
 foundry/training/EMA.py,sha256=3OWA9Pz7XuDr-SRxbz24tZf55DmhSa2fKy9r5v2IXqA,2651
 foundry/training/checkpoint.py,sha256=mUiObg-qcF3tvMfVu77sD9m3yVRp71czv07ccliU7qQ,1791
 foundry/training/schedulers.py,sha256=StmXegPfIdLAv31FreCTrDh9dsOvNUfzG4YGa61Y4oE,3647
-foundry/utils/alignment.py,sha256=OAN7H2TqraGxP1uMXUpwLO7g0qS0cxUVjuV33pY16z0,2316
+foundry/utils/alignment.py,sha256=2anqy0mn9zeFEiVWS_EG7zHiyPk1C_gbUu-SRvQ5mAM,2502
 foundry/utils/components.py,sha256=Piw2TfQF26uuxC3hXG3iv_4rgud1lVO-cv6N-p05EDY,15200
 foundry/utils/datasets.py,sha256=pLBxVezm-TSrYuC5gFnJZdGnNWV7aPH2QiWIVE2hkdQ,16629
 foundry/utils/ddp.py,sha256=ydHrO6peGbRnWAwgH5rmpHuQd55g2gFzzoZJYypn7GU,3970
@@ -34,7 +34,7 @@ foundry/utils/squashfs.py,sha256=QlcwuJyVe-QVfIOS7o1QfLhaCQPNzzox7ln4n8dcYEg,523
 foundry/utils/torch.py,sha256=OLsqoxw4CTXbGzWUHernLUT7uQjLu0tVPtD8h8747DI,11211
 foundry/utils/weights.py,sha256=btz4S02xff2vgiq4xMfiXuhK1ERafqQPtmimo1DmoWY,10381
 foundry_cli/__init__.py,sha256=0BxY2RUKJLaMXUGgypPCwlTskTEFdVnkhTR4C4ft2Kw,52
-foundry_cli/download_checkpoints.py,sha256=UCNdy4VZyJe1PH_lnVLqy-VSMuTu875mGGd99ma7fTQ,8426
+foundry_cli/download_checkpoints.py,sha256=CxU9dKBa1vAkVd450tfH5aZAlQIUTrHsDGTbmxzd_JQ,8922
 mpnn/__init__.py,sha256=hgQcXFaCbAxFrhydVAy0xj8yC7UJF-GCCFhqD0sZ7I4,57
 mpnn/inference.py,sha256=wPtGR325eVRVeesXoWtBK6b_-VcU8BZae5IfQN3-mvA,1669
 mpnn/train.py,sha256=9eQGBd3rdNF5Zr2w8oUgETbqxBavNBajtA6Vbc5zESE,10239
@@ -119,18 +119,18 @@ rfd3/__init__.py,sha256=2Wto2IsUIj2lGag9m_gqgdCwBNl5p21-Xnr7W_RpU3c,348
 rfd3/callbacks.py,sha256=Zjt8RiaYWquoKOwRmC_wCUbRbov-V4zd2_73zjhgDHE,2783
 rfd3/cli.py,sha256=ka3K5H117fzDYIDXFpOpJV21w_XBrHYJZdFE0thsGBI,1644
 rfd3/constants.py,sha256=wLvDzrThpOrK8T3wGFNQeGrhAXOJQze8l3v_7pjIdMM,13141
-rfd3/engine.py,sha256=La_dB48Ewz0IdY1ocxvSWg-PXVAsySm0OGvwyz42lI8,20824
-rfd3/run_inference.py,sha256=ljzsCKEtrlfAvP0SDFPeQwTM3rV_X3ewHOhcRFVI37c,1258
+rfd3/engine.py,sha256=NwATrhYFyqT7C9Bie8mWtUiqqzXgs9x6nOCkmZYPiT4,21224
+rfd3/run_inference.py,sha256=HfRMQ30_SAHfc-VFzBV52F-aLaNdG6PW8VkdMyB__wE,1264
 rfd3/train.py,sha256=rHswffIUhOae3_iYyvAiQ3jALoFuzrcRUgMlbJLinlI,7947
 rfd3/inference/datasets.py,sha256=u-2U7deHXu-iOs7doiKKynewP-NEyJfdORSTDzUSaQI,6538
-rfd3/inference/input_parsing.py,sha256=mk3HBvo7MPTFEET7NagCo5TSjb47w-hxUDoeQxUW_h4,45449
-rfd3/inference/legacy_input_parsing.py,sha256=1wf_KF7qWnGLaVM8IXDl8fIsWCmxtOi2YlAiHEVELqw,28046
-rfd3/inference/parsing.py,sha256=Nq8CYmimnql4RM-5ZfPAvOFvCae4_CC2pYDzE6iCpWU,5290
-rfd3/inference/symmetry/atom_array.py,sha256=HH50Z07bTUnNUgCwAGslADbvMYHgsXn9s-fqwx6BvKw,11034
-rfd3/inference/symmetry/checks.py,sha256=wb7K327GnMwGG9bgOvvDAbaPsFj4nGZpEAolICUapNc,8908
+rfd3/inference/input_parsing.py,sha256=TyEzCzeCaNhuNi0RjMcq9fF2j3Sp36KbuZ1FUjlBTZ8,45442
+rfd3/inference/legacy_input_parsing.py,sha256=G2XxkrjdIpL6i1YY7xEmkFitVv__Pc45ow6IKKPHw64,28855
+rfd3/inference/parsing.py,sha256=ktAMUuZE3Pe4bKAjjV3zjqcEDmGlMZ-cotIUhJsEQQA,5402
+rfd3/inference/symmetry/atom_array.py,sha256=HfFagFUB5yB-Y4IfUM5nuVGWHC5AEkyHqt0JcIqTQ_E,10922
+rfd3/inference/symmetry/checks.py,sha256=y-Kq0l5OhEmmxsPBBsMMB0qaAt18FeEicD3-jSMQFa0,9900
 rfd3/inference/symmetry/contigs.py,sha256=6OvbZ2dJg-a0mvvKAC0VkzUH5HpUDxOJvkByIst_roU,2127
-rfd3/inference/symmetry/frames.py,sha256=G55p-aOXqEYG4kCyKxrgWAsS-gW9-gOTlBME6nhbKyU,10716
-rfd3/inference/symmetry/symmetry_utils.py,sha256=KwgxrdfO766RCEwF3VElAE85oEKiopPGRQDhJbKZaUA,15810
+rfd3/inference/symmetry/frames.py,sha256=aEwkmlUsYexERX9hu09JMhisC8QTpHPVhfITbL80-EE,10819
+rfd3/inference/symmetry/symmetry_utils.py,sha256=p_PkxU3sw6gYGO2EmZTrbNQdLjz1mdTWEIl5MjQdIuY,14664
 rfd3/metrics/design_metrics.py,sha256=O1RqZdjQPNlAWYRg6UJTERYg_gUI1_hVleKsm9xbWBY,16836
 rfd3/metrics/hbonds_hbplus_metrics.py,sha256=Sewy9KzmrA1OnfkasN-fmWrQ9IRx9G7Yyhe2ua0mk28,11518
 rfd3/metrics/hbonds_metrics.py,sha256=SIR4BnDhYdpVSqwXXRYpQ_tB-M0_fVyugGl08WivCmE,15257
@@ -170,7 +170,7 @@ rfd3/transforms/symmetry.py,sha256=GSnMF7oAnUxPozfafsRuHEv0yKXW0BpLTI6wsKGZrbc,2
 rfd3/transforms/training_conditions.py,sha256=UXiUPjDwrNKM95tRe0eXrMeRN8XlTPc_MXUvo6UpePo,19510
 rfd3/transforms/util_transforms.py,sha256=2AcLkzx-73ZFgcWD1cIHv7NyniRPI4_zThHK8azyQaY,18119
 rfd3/transforms/virtual_atoms.py,sha256=UpmxzPPd5FaJigcRoxgLSHHrLLOqsCvZ5PPZfQSGqII,12547
-rfd3/utils/inference.py,sha256=RQp5CCy6Z6uHVZ2Mx0zmmGluYEOrASke4bABtfRjpy0,26448
+rfd3/utils/inference.py,sha256=-8IKzkB9ulhLEJgapvnZSdIaIPQDPMpyPpHTQlFS7r0,27317
 rfd3/utils/io.py,sha256=wbdjUTQkDc3RCSM7gdogA-XOKR68HeQ-cfvyN4pP90w,9849
 rfd3/utils/vizualize.py,sha256=HPlczrA3zkOuxV5X05eOvy_Oga9e3cPnFUXOEP4RR_g,11046
 rf3/configs/inference.yaml,sha256=JmEZdkAnbnOrX79lGS5xrYYho9aBFfVxfUp-8KjJV5I,309
@@ -248,7 +248,7 @@ rfd3/configs/datasets/conditions/sequence_design.yaml,sha256=D1K6WOysmSAQ4LogltU
 rfd3/configs/datasets/conditions/tipatom.yaml,sha256=0010o7UUL-l75qI8HCjC_tdBXFWysm2dgVXzE7bQyZ0,650
 rfd3/configs/datasets/conditions/unconditional.yaml,sha256=z1eVHylswLyludXWFs1AMt3mTMu3EbAUHrP8J3XBsRU,446
 rfd3/configs/datasets/train/rfd3_monomer_distillation.yaml,sha256=1f61uFeRB8OD6sifFuIKFov8D7PcHpqRT4Z-M5EzO4w,1207
-rfd3/configs/datasets/train/pdb/af3_train_interface.yaml,sha256=mwbdGJQ9SXc8WvO3qqSWzS--K4rvbFsM0MR371FUrr0,1552
+rfd3/configs/datasets/train/pdb/af3_train_interface.yaml,sha256=DSIpXW2SQ3drDp12490y0tFbjbugecyA7TI_x3WrKng,1546
 rfd3/configs/datasets/train/pdb/af3_train_pn_unit.yaml,sha256=DPoEhLlyBu0RdBkkJeWB8pkOV4z0DBc6XmclLgww9II,1324
 rfd3/configs/datasets/train/pdb/base.yaml,sha256=2VUEAKADyvjJmWP4FeOJwRat9r6F3_GXuyGYjvMvArw,291
 rfd3/configs/datasets/train/pdb/base_no_weights.yaml,sha256=8HchN7DqYESBK520vShdg7xidWBSogGRAxfaxa5pKdE,554
@@ -285,7 +285,7 @@ rfd3/configs/hydra/default.yaml,sha256=SYDTSU8bAw20QssrtTi7lptiBD5H3XNyzApsyy0br
 rfd3/configs/hydra/no_logging.yaml,sha256=MUXDFcw-QwaRPz9HcE-1tdZwbNha1mexTe31G-Zt9_w,120
 rfd3/configs/inference_engine/base.yaml,sha256=ekP5U7bAALpeJGpwyj1v0N5LiEtptl5loRCtM8FRzRM,246
 rfd3/configs/inference_engine/dev.yaml,sha256=-2snClOTwj5TQt7jnwSrI4pzAiI4nFulXKJflmgIyUw,304
-rfd3/configs/inference_engine/rfdiffusion3.yaml,sha256=h2e9U9RFCcvXjKAJ6U8puj-8O-U57ZxeZLA0HLB2txA,2161
+rfd3/configs/inference_engine/rfdiffusion3.yaml,sha256=3bHIAhzFhFDIag0xQWYxHBUMSc71fjClHXKbZ-tpHzA,2112
 rfd3/configs/logger/csv.yaml,sha256=DtcywAIS4OxLXP2QxSEvqdrjhMpT6xHiGspoYw5qkus,245
 rfd3/configs/logger/default.yaml,sha256=pSyHyxT-J_T-g4_6TtD2yzN3rzxgY6rOG_Vh4RjZeFY,17
 rfd3/configs/logger/wandb.yaml,sha256=RhCnFtO0hNc3R75ts417l5ICZeGm74lOj9Bfe7ZvRNA,652
@@ -294,7 +294,7 @@ rfd3/configs/model/components/ema.yaml,sha256=AIzf4RZLKP8AcfaxdvZBS1rFw3AlSo431r
 rfd3/configs/model/components/rfd3_net.yaml,sha256=95FF4U7aWmLCoHvyxsRoE74n-bxTPD6KlAhPKNemVH4,3275
 rfd3/configs/model/optimizers/adam.yaml,sha256=cTRNo4_4lNgLv0b329v-KiC_MCQtTVVTxeer5Au_FIM,145
 rfd3/configs/model/samplers/edm.yaml,sha256=QycHAIrfhRgx0mJygTOs56FT93tGCWTGxrQSKBOA7Mc,483
-rfd3/configs/model/samplers/symmetry.yaml,sha256=BZZOIhk2ndAvIntf-16nnqCuOW43iWTB7iDU-RsxOcc,214
+rfd3/configs/model/samplers/symmetry.yaml,sha256=pI0Ens6jmbpAIl8E4eYsJR1SqIppe5OsWh91KfpjNjs,214
 rfd3/configs/model/schedulers/af3.yaml,sha256=xEtRb--KPjg_5pW_IJvN9AHWVqCtOM4QOnXlMH2KrEg,149
 rfd3/configs/paths/default.yaml,sha256=bjB04SNu_5E6W_v4mRBjwce0xmdKwO5wsVf4gfaRl0Y,1045
 rfd3/configs/paths/data/default.yaml,sha256=jfs1dbbcOqHja4_6lXheyRg4t0YExqVn2w0rZEWL6XE,788
@@ -304,8 +304,8 @@ rfd3/configs/trainer/rfd3_base.yaml,sha256=R3lZxdyjUirjlLU31qWlnZgHaz4GcWTGGIz4f
 rfd3/configs/trainer/loss/losses/diffusion_loss.yaml,sha256=FE4FCEfurE0ekwZ4YfS6wCvPSNqxClwg_kc73cPql5Y,323
 rfd3/configs/trainer/loss/losses/sequence_loss.yaml,sha256=kezbQcqwAZ0VKQPUBr2MsNr9DcDL3ENIP1i-j7h-6Co,64
 rfd3/configs/trainer/metrics/design_metrics.yaml,sha256=xVDpClhHqSHvsf-8StL26z51Vn-iuWMDG9KMB-kqOI0,719
-rc_foundry-0.1.6.dist-info/METADATA,sha256=EEOkAi2nABzo70kEP-n9t5aXZ8a4Gqr5wYZ2mjIBqp4,10585
-rc_foundry-0.1.6.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-rc_foundry-0.1.6.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
-rc_foundry-0.1.6.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
-rc_foundry-0.1.6.dist-info/RECORD,,
+rc_foundry-0.1.7.dist-info/METADATA,sha256=zlvCxfZ5-Ow7WuGKskfW6P1DGhZB9OfLIIBUBGncFeQ,11309
+rc_foundry-0.1.7.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+rc_foundry-0.1.7.dist-info/entry_points.txt,sha256=BmiWCbWGtrd_lSOFMuCLBXyo84B7Nco-alj7hB0Yw9A,130
+rc_foundry-0.1.7.dist-info/licenses/LICENSE.md,sha256=NKtPCJ7QMysFmzeDg56ZfUStvgzbq5sOvRQv7_ddZOs,1533
+rc_foundry-0.1.7.dist-info/RECORD,,

rfd3/configs/datasets/train/pdb/af3_train_interface.yaml

@@ -7,7 +7,7 @@ dataset:
     base_dir: ${paths.data.pdb_data_dir}
   dataset:
     name: interface
-    data: ${paths.data.pdb_parquet_dir}/interfaces_df_train.parquet
+    data: ${paths.data.pdb_parquet_dir}/interfaces_df.parquet
     filters:
       # filters common across all PDB datasets
       - "deposition_date < '2021-09-30'"

rfd3/configs/inference_engine/rfdiffusion3.yaml

@@ -7,7 +7,7 @@ _target_: rfd3.engine.RFD3InferenceEngine
 
 out_dir: ???
 inputs: ???  # null, json, pdb or 
-ckpt_path: /projects/ml/aa_design/models/rfd3_latest_cleaned.ckpt
+ckpt_path: rfd3
 json_keys_subset: null
 skip_existing: True
 
@@ -61,5 +61,5 @@ global_prefix: null
 dump_prediction_metadata_json: True
 dump_trajectories: False
 align_trajectory_structures: False
-prevalidate_inputs: True
+prevalidate_inputs: False
 low_memory_mode: False # False for standard mode, True for memory efficient tokenization mode

rfd3/configs/model/samplers/symmetry.yaml

@@ -4,7 +4,7 @@ defaults:
 
 kind: symmetry
 num_timesteps: 200
-gamma_0: 1.0  # 1.0 for SDE sampling
+gamma_0: 0.6  # 1.0 for SDE sampling
 gamma_min: 1.0
 gamma_min2: 0.0
 sym_step_frac: 0.9 # when 0.9, 90% of the trajectory from the start is symmetrized

rfd3/engine.py

@@ -23,7 +23,10 @@ from rfd3.inference.datasets import (
 )
 from rfd3.inference.input_parsing import DesignInputSpecification
 from rfd3.model.inference_sampler import SampleDiffusionConfig
-from rfd3.utils.inference import ensure_input_is_abspath
+from rfd3.utils.inference import (
+    ensure_inference_sampler_matches_design_spec,
+    ensure_input_is_abspath,
+)
 from rfd3.utils.io import (
     CIF_LIKE_EXTENSIONS,
     build_stack_from_atom_array_and_batched_coords,
@@ -171,6 +174,7 @@ class RFD3InferenceEngine(BaseInferenceEngine):
         )
         # save
         self.specification_overrides = dict(specification or {})
+        self.inference_sampler_overrides = dict(inference_sampler or {})
 
         # Setup output directories and args
         self.global_prefix = global_prefix
@@ -210,6 +214,9 @@ class RFD3InferenceEngine(BaseInferenceEngine):
             inputs=inputs,
             n_batches=n_batches,
         )
+        ensure_inference_sampler_matches_design_spec(
+            design_specifications, self.inference_sampler_overrides
+        )
         # init before
         self.initialize()
         outputs = self._run_multi(design_specifications)
@@ -383,6 +390,9 @@ class RFD3InferenceEngine(BaseInferenceEngine):
         # Based on inputs, construct the specifications to loop through
         design_specifications = {}
         for prefix, example_spec in inputs.items():
+            # Record task name in the specification
+            example_spec["extra"]["task_name"] = prefix
+
             # ... Create n_batches for example
             for batch_id in range((n_batches) if exists(n_batches) else 1):
                 # ... Example ID
@@ -524,21 +534,19 @@ def process_input(
 
 
 def _reshape_trajectory(traj, align_structures: bool):
-    traj = [traj[i] for i in range(len(traj))]
-    n_steps = len(traj)
+    traj = [traj[i] for i in range(len(traj))]  # make list of arrays
     max_frames = 100
-
+    if len(traj) > max_frames:
+        selected_indices = torch.linspace(0, len(traj) - 1, max_frames).long().tolist()
+        traj = [traj[i] for i in selected_indices]
     if align_structures:
         # ... align the trajectories on the last prediction
-        for step in range(n_steps - 1):
+        for step in range(len(traj) - 1):
             traj[step] = weighted_rigid_align(
-                X_L=traj[-1],
-                X_gt_L=traj[step],
-            )
+                X_L=traj[-1][None],
+                X_gt_L=traj[step][None],
+            ).squeeze(0)
     traj = traj[::-1]  # reverse to go from noised -> denoised
-    if n_steps > max_frames:
-        selected_indices = torch.linspace(0, n_steps - 1, max_frames).long().tolist()
-        traj = [traj[i] for i in selected_indices]
 
     traj = torch.stack(traj).cpu().numpy()
     return traj

rfd3/inference/input_parsing.py

@@ -696,7 +696,7 @@ class DesignInputSpecification(BaseModel):
             # Partial diffusion: use COM, keep all coordinates
             if exists(self.symmetry) and self.symmetry.id:
                 # For symmetric structures, avoid COM centering that would collapse chains
-                ranked_logger.info(
+                logger.info(
                     "Partial diffusion with symmetry: skipping COM centering to preserve chain spacing"
                 )
             else:

rfd3/inference/legacy_input_parsing.py

@@ -139,13 +139,18 @@ def fetch_motif_residue_(
         subarray, motif=True, unindexed=False, dtype=int
     )  # all values init to True (fix all)
 
+    to_unindex = f"{src_chain}{src_resid}" in unindexed_components
+    to_index = f"{src_chain}{src_resid}" in components
+
     # Assign is motif atom and sequence
     if exists(atoms := fixed_atoms.get(f"{src_chain}{src_resid}")):
+        # If specified, we set fixed atoms in the residue to be motif atoms
         atom_mask = get_name_mask(subarray.atom_name, atoms, res_name)
         subarray.set_annotation("is_motif_atom", atom_mask)
         # subarray.set_annotation("is_motif_atom_with_fixed_coord", atom_mask)  # BUGFIX: uncomment
 
     elif redesign_motif_sidechains and res_name in STANDARD_AA:
+        # If redesign_motif_sidechains is True, we only make the backbone atoms to be motif atoms
         n_atoms = subarray.shape[0]
         diffuse_oxygen = False
         if n_atoms < 3:
@@ -178,6 +183,18 @@ def fetch_motif_residue_(
         subarray.set_annotation(
             "is_motif_atom_with_fixed_seq", np.zeros(subarray.shape[0], dtype=int)
         )
+    elif to_index or to_unindex:
+        # If the residue is in the contig or unindexed components,
+        # we set all atoms in the residue to be motif atoms
+        subarray.set_annotation("is_motif_atom", np.ones(subarray.shape[0], dtype=int))
+    else:
+        if to_unindex and not (
+            unfix_all or f"{src_chain}{src_resid}" in unfix_residues
+        ):
+            raise ValueError(
+                f"{src_chain}{src_resid} is not found in fixed_atoms, contig or unindex contig."
+                "Please check your input and contig specification."
+            )
     if unfix_all or f"{src_chain}{src_resid}" in unfix_residues:
         subarray.set_annotation(
             "is_motif_atom_with_fixed_coord", np.zeros(subarray.shape[0], dtype=int)
@@ -197,7 +214,6 @@ def fetch_motif_residue_(
         subarray.set_annotation(
             "is_flexible_motif_atom", np.zeros(subarray.shape[0], dtype=bool)
         )
-    to_unindex = f"{src_chain}{src_resid}" in unindexed_components
     if to_unindex:
         subarray.set_annotation(
             "is_motif_atom_unindexed", subarray.is_motif_atom.copy()

rfd3/inference/parsing.py

@@ -117,6 +117,7 @@ def from_any_(v: Any, atom_array: AtomArray):
 
         # Split to atom names
         data_split[idx] = token.atom_name[comp_mask_subset].tolist()
+        # TODO: there is a bug where when you select specifc atoms within a ligand, output ligand is fragmented
 
         # Update mask & token dictionary
         mask[comp_mask] = comp_mask_subset

rfd3/inference/symmetry/atom_array.py

@@ -4,12 +4,8 @@ from rfd3.inference.symmetry.frames import (
     get_symmetry_frames_from_symmetry_id,
 )
 
-from foundry.utils.ddp import RankedLogger
-
 FIXED_TRANSFORM_ID = -1
 FIXED_ENTITY_ID = -1
-ranked_logger = RankedLogger(__name__, rank_zero_only=True)
-
 
 ########################################################
 # Symmetry annotations
@@ -28,7 +24,7 @@ def add_sym_annotations(atom_array, sym_conf):
     is_asu = np.full(n, True, dtype=np.bool_)
     atom_array.set_annotation("is_sym_asu", is_asu)
     # symmetry_id
-    symmetry_ids = np.full(n, sym_conf.get("id"), dtype="U6")
+    symmetry_ids = np.full(n, sym_conf.id, dtype="U6")
     atom_array.set_annotation("symmetry_id", symmetry_ids)
     return atom_array
 

rfd3/inference/symmetry/checks.py

@@ -1,10 +1,13 @@
 import numpy as np
-from rfd3.inference.symmetry.contigs import expand_contig_unsym_motif
+from rfd3.inference.symmetry.contigs import (
+    expand_contig_unsym_motif,
+    get_unsym_motif_mask,
+)
 from rfd3.transforms.conditioning_base import get_motif_features
 
 from foundry.utils.ddp import RankedLogger
 
-MIN_ATOMS_ALIGN = 100
+MIN_ATOMS_ALIGN = 30
 MAX_TRANSFORMS = 10
 RMSD_CUT = 1.0  # Angstroms
 
@@ -18,29 +21,33 @@ def check_symmetry_config(
     Check if the symmetry configuration is valid. Add all basic checks here.
     """
 
-    assert sym_conf.get("id"), "symmetry_id is required. e.g. {'id': 'C2'}"
+    assert sym_conf.id, "symmetry_id is required. e.g. {'id': 'C2'}"
     # if unsym motif is provided, check that each motif name is in the atom array
-    if sym_conf.get("is_unsym_motif"):
+
+    is_unsym_motif = np.zeros(atom_array.shape[0], dtype=bool)
+    if sym_conf.is_unsym_motif:
         assert (
             src_atom_array is not None
         ), "Source atom array must be provided for symmetric motifs"
-        unsym_motif_names = sym_conf["is_unsym_motif"].split(",")
+        unsym_motif_names = sym_conf.is_unsym_motif.split(",")
         unsym_motif_names = expand_contig_unsym_motif(unsym_motif_names)
+        is_unsym_motif = get_unsym_motif_mask(atom_array, unsym_motif_names)
         for n in unsym_motif_names:
             if (sm and n not in sm.split(",")) and (n not in atom_array.src_component):
                 raise ValueError(f"Unsym motif {n} not found in atom_array")
+
+    is_motif_token = get_motif_features(atom_array)["is_motif_token"]
     if (
-        get_motif_features(atom_array)["is_motif_token"].any()
-        and not sym_conf.get("is_symmetric_motif")
+        is_motif_token[~is_unsym_motif].any()
+        and not sym_conf.is_symmetric_motif
         and not has_dist_cond
     ):
         raise ValueError(
-            "Asymmetric motif inputs should be distance constrained. "
+            "Asymmetric motif inputs should be distance constrained."
             "Use atomwise_fixed_dist to constrain the distance between the motif atoms."
         )
-    # else: if unconditional symmetry, no need to have symmetric input motif
 
-    if partial and not sym_conf.get("is_symmetric_motif"):
+    if partial and not sym_conf.is_symmetric_motif:
         raise ValueError(
             "Partial diffusion with symmetry is only supported for symmetric inputs."
         )
@@ -54,9 +61,6 @@ def check_atom_array_is_symmetric(atom_array):
     Returns:
         bool: True if the atom array is symmetric, False otherwise
     """
-    # TODO: Implement something like this https://github.com/baker-laboratory/ipd/blob/main/ipd/sym/sym_detect.py#L303
-    #       and maybe this https://github.com/baker-laboratory/ipd/blob/main/ipd/sym/sym_detect.py#L231
-
     import biotite.structure as struc
     from rfd3.inference.symmetry.atom_array import (
         apply_symmetry_to_atomarray_coord,
@@ -68,8 +72,10 @@ def check_atom_array_is_symmetric(atom_array):
     # remove hetero atoms
     atom_array = atom_array[~atom_array.hetero]
     if len(atom_array) == 0:
-        ranked_logger.info("Atom array has no protein chains. Please check your input.")
-        return False
+        ranked_logger.warning(
+            "Atom array has no protein chains. Please check your input."
+        )
+        return True
 
     chains = np.unique(atom_array.chain_id)
     asu_mask = atom_array.chain_id == chains[0]
@@ -162,16 +168,22 @@ def find_optimal_rotation(coords1, coords2, max_points=1000):
         return None
 
 
-def check_input_frames_match_symmetry_frames(computed_frames, original_frames) -> None:
+def check_input_frames_match_symmetry_frames(
+    computed_frames, original_frames, nids_by_entity
+) -> None:
     """
     Check if the atom array matches the symmetry_id.
     Arguments:
         computed_frames: list of computed frames
         original_frames: list of original frames
     """
-    assert len(computed_frames) == len(
-        original_frames
-    ), "Number of computed frames does not match number of original frames"
+    assert len(computed_frames) == len(original_frames), (
+        "Number of computed frames does not match number of original frames.\n"
+        f"Computed Frames: {len(computed_frames)}. Original Frames: {len(original_frames)}.\n"
+        "If the computed frames are not as expected, please check if you have one-to-one mapping "
+        "(size, sequence, folding) of an entity across all chains.\n"
+        f"Computed Entity Mapping: {nids_by_entity}."
+    )
 
 
 def check_valid_multiplicity(nids_by_entity) -> None:
@@ -184,25 +196,35 @@ def check_valid_multiplicity(nids_by_entity) -> None:
     multiplicity = min([len(i) for i in nids_by_entity.values()])
     if multiplicity == 1:  # no possible symmetry
         raise ValueError(
-            "Input has no possible symmetry. If asymmetric motif, please use 2D conditioning inference instead."
+            "Input has no possible symmetry. If asymmetric motif, please use 2D conditioning inference instead.\n"
+            "Multiplicity: 1"
         )
 
     # Check that the input is not asymmetric
     multiplicity_good = [len(i) % multiplicity == 0 for i in nids_by_entity.values()]
     if not all(multiplicity_good):
-        raise ValueError("Invalid multiplicities of subunits. Please check your input.")
+        raise ValueError(
+            "Expected multiplicity does not match for some entities.\n"
+            "Please modify your input to have one-to-one mapping (size, sequence, folding) of an entity across all chains.\n"
+            f"Expected Multiplicity: {multiplicity}.\n"
+            f"Computed Entity Mapping: {nids_by_entity}."
+        )
 
 
 def check_valid_subunit_size(nids_by_entity, pn_unit_id) -> None:
     """
     Check that the subunits in the input are of the same size.
     Arguments:
-        nids_by_entity: dict mapping entity to ids
+        nids_by_entity: dict mapping entity to ids. e.g. {0: (['A_1', 'B_1', 'C_1']), 1: (['A_2', 'B_2', 'C_2'])}
+        pn_unit_id: array of ids. e.g. ['A_1', 'B_1', 'C_1', 'A_2', 'B_2', 'C_2']
     """
-    for i, js in nids_by_entity.items():
-        for j in js[1:]:
-            if (pn_unit_id == js[0]).sum() != (pn_unit_id == j).sum():
-                raise ValueError("Size mismatch in the input. Please check your file.")
+    for js in nids_by_entity.values():
+        for js_i in js[1:]:
+            if (pn_unit_id == js[0]).sum() != (pn_unit_id == js_i).sum():
+                raise ValueError(
+                    f"Size mismatch between chain {js[0]} ({(pn_unit_id == js[0]).sum()} atoms) "
+                    f"and chain {js_i} ({(pn_unit_id == js_i).sum()} atoms). Please check your input file."
+                )
 
 
 def check_min_atoms_to_align(natm_per_unique, reference_entity) -> None:
@@ -212,7 +234,10 @@ def check_min_atoms_to_align(natm_per_unique, reference_entity) -> None:
         nids_by_entity: dict mapping entity to ids
     """
     if natm_per_unique[reference_entity] < MIN_ATOMS_ALIGN:
-        raise ValueError("Not enough atoms to align. Please check your input.")
+        raise ValueError(
+            f"Not enough atoms to align < {MIN_ATOMS_ALIGN} atoms."
+            f"Please provide a input with at least {MIN_ATOMS_ALIGN} atoms."
+        )
 
 
 def check_max_transforms(chains_to_consider) -> None:
@@ -224,7 +249,7 @@ def check_max_transforms(chains_to_consider) -> None:
     """
     if len(chains_to_consider) > MAX_TRANSFORMS:
         raise ValueError(
-            "Number of transforms exceeds the max number of transforms (10)"
+            f"Number of transforms exceeds the max number of transforms ({MAX_TRANSFORMS})."
         )
 
 

rfd3/inference/symmetry/frames.py

@@ -10,12 +10,13 @@ def get_symmetry_frames_from_symmetry_id(symmetry_id):
     Returns:
         frames: list of rotation matrices
     """
+    from rfd3.inference.symmetry.symmetry_utils import SymmetryConfig
 
     # Get frames from symmetry id
     sym_conf = {}
-    if isinstance(symmetry_id, dict):
+    if isinstance(symmetry_id, SymmetryConfig):
         sym_conf = symmetry_id
-        symmetry_id = symmetry_id.get("id")
+        symmetry_id = symmetry_id.id
 
     if symmetry_id.lower().startswith("c"):
         order = int(symmetry_id[1:])
@@ -25,9 +26,9 @@ def get_symmetry_frames_from_symmetry_id(symmetry_id):
         frames = get_dihedral_frames(order)
     elif symmetry_id.lower() == "input_defined":
         assert (
-            "symmetry_file" in sym_conf
+            sym_conf.symmetry_file is not None
         ), "symmetry_file is required for input_defined symmetry"
-        frames = get_frames_from_file(sym_conf.get("symmetry_file"))
+        frames = get_frames_from_file(sym_conf.symmetry_file)
     else:
         raise ValueError(f"Symmetry id {symmetry_id} not supported")
 
@@ -120,7 +121,9 @@ def get_symmetry_frames_from_atom_array(src_atom_array, input_frames):
     computed_frames = [(R, np.array([0, 0, 0])) for R in Rs]
 
     # check that the computed frames match the input frames
-    check_input_frames_match_symmetry_frames(computed_frames, input_frames)
+    check_input_frames_match_symmetry_frames(
+        computed_frames, input_frames, nids_by_entity
+    )
 
     return computed_frames
 

rfd3/inference/symmetry/symmetry_utils.py

@@ -39,18 +39,36 @@ ranked_logger = RankedLogger(__name__, rank_zero_only=True)
 
 
 class SymmetryConfig(BaseModel):
-    # AM / HE TODO: feel free to flesh this out and add validation as needed
     model_config = ConfigDict(
         arbitrary_types_allowed=True,
         extra="allow",
     )
-    id: Optional[str] = Field(None)
-    # is_unsym_motif: Optional[np.ndarray[bool]] = Field(...)
-    # is_symmetric_motif: bool = Field(...)
+    id: Optional[str] = Field(
+        None,
+        description="Symmetry group ID. e.g. 'C3', 'D2'. Only C and D symmetry types are supported currently.",
+    )
+    is_unsym_motif: Optional[str] = Field(
+        None,
+        description="Comma separated list of contig/ligand names that should not be symmetrized such as DNA strands. \
+         e.g. 'HEM' or 'Y1-11,Z16-25'",
+    )
+    is_symmetric_motif: bool = Field(
+        True,
+        description="If True, the input motifs are expected to be already symmetric and won't be symmetrized. \
+        If False, the all input motifs are expected to be ASU and will be symmetrized.",
+    )
+
+
+def convery_sym_conf_to_symmetry_config(sym_conf: dict):
+    return SymmetryConfig(**sym_conf)
 
 
 def make_symmetric_atom_array(
-    asu_atom_array, sym_conf: SymmetryConfig, sm=None, has_2d=False, src_atom_array=None
+    asu_atom_array,
+    sym_conf: SymmetryConfig | dict,
+    sm=None,
+    has_dist_cond=False,
+    src_atom_array=None,
 ):
     """
     apply symmetry to an atom array.
@@ -58,39 +76,33 @@ def make_symmetric_atom_array(
         asu_atom_array: atom array of the asymmetric unit
         sym_conf: symmetry configuration (dict, "id" key is required)
         sm: optional small molecule names (str, comma separated)
-        has_2d: whether to add 2d entity annotations
+        has_dist_cond: whether to add 2d entity annotations
     Returns:
         new_asu_atom_array: atom array with symmetry applied
     """
-    sym_conf = (
-        sym_conf.model_dump()
-    )  # TODO: JB: remove this line to keep as symmetry config for cleaner syntax(?)
-    ranked_logger.info(f"Symmetry Configs: {sym_conf}")
+    if not isinstance(sym_conf, SymmetryConfig):
+        sym_conf = convery_sym_conf_to_symmetry_config(sym_conf)
 
-    # Making sure that the symmetry config is valid
     check_symmetry_config(
-        asu_atom_array,
-        sym_conf,
-        sm,
-        has_dist_cond=has_2d,
-        src_atom_array=src_atom_array,
+        asu_atom_array, sym_conf, sm, has_dist_cond, src_atom_array=src_atom_array
     )
     # Adding utility annotations to the asu atom array
     asu_atom_array = _add_util_annotations(asu_atom_array, sym_conf, sm)
 
-    if has_2d:  # NB: this will only work for asymmetric motifs at the moment - need to add functionality for symmetric motifs
+    if has_dist_cond:  # NB: this will only work for asymmetric motifs at the moment - need to add functionality for symmetric motifs
         asu_atom_array = add_2d_entity_annotations(asu_atom_array)
 
     frames = get_symmetry_frames_from_symmetry_id(sym_conf)
 
     # If the motif is symmetric, we get the frames instead from the source atom array.
-    if sym_conf.get("is_symmetric_motif"):
+    if sym_conf.is_symmetric_motif:
         assert (
             src_atom_array is not None
         ), "Source atom array must be provided for symmetric motifs"
-        # if symmetric motif is provided, get the frames from the src atom array
+        # if symmetric motif is provided, get the frames from the src atom array.
         frames = get_symmetry_frames_from_atom_array(src_atom_array, frames)
-    else:
+    elif (asu_atom_array._is_motif[~asu_atom_array._is_unsym_motif]).any():
+        # if the motifs that's not unsym motifs are present.
         raise NotImplementedError(
             "Asymmetric motif inputs are not implemented yet. please symmetrize the motif."
         )
@@ -101,7 +113,7 @@ def make_symmetric_atom_array(
     # Extracting all things at this moment that we will not want to symmetrize.
     # This includes: 1) unsym motifs, 2) ligands
     unsym_atom_arrays = []
-    if sym_conf.get("is_unsym_motif"):
+    if sym_conf.is_unsym_motif:
         # unsym_motif_atom_array = get_unsym_motif(asu_atom_array, asu_atom_array._is_unsym_motif)
         # Now remove the unsym motifs from the asu atom array
         unsym_atom_arrays.append(asu_atom_array[asu_atom_array._is_unsym_motif])
@@ -128,7 +140,7 @@ def make_symmetric_atom_array(
     symmetrized_atom_array = struc.concatenate(symmetry_unit_list)
 
     # add 2D conditioning annotations
-    if has_2d:
+    if has_dist_cond:
         symmetrized_atom_array = reannotate_2d_conditions(symmetrized_atom_array)
 
     # set all motifs to not have any symmetrization applied to them
@@ -183,7 +195,7 @@ def make_symmetric_atom_array_for_partial_diffusion(atom_array, sym_conf):
     frames = get_symmetry_frames_from_symmetry_id(sym_conf)
 
     # Add symmetry ID
-    symmetry_ids = np.full(n, sym_conf.get("id"), dtype="U6")
+    symmetry_ids = np.full(n, sym_conf.id, dtype="U6")
     atom_array.set_annotation("symmetry_id", symmetry_ids)
 
     # Initialize transform annotations (use same format as original system)
@@ -244,7 +256,7 @@ def _add_util_annotations(asu_atom_array, sym_conf, sm):
     """
     n = asu_atom_array.shape[0]
     is_motif = get_motif_features(asu_atom_array)["is_motif_atom"].astype(np.bool_)
-    is_sm = np.zeros(asu_atom_array.shape[0], dtype=bool)
+    is_sm = np.zeros(n, dtype=bool)
     is_asu = np.ones(n, dtype=bool)
     is_unsym_motif = np.zeros(n, dtype=bool)
 
@@ -257,8 +269,8 @@ def _add_util_annotations(asu_atom_array, sym_conf, sm):
         )
 
     # assign unsym motifs
-    if sym_conf.get("is_unsym_motif"):
-        unsym_motif_names = sym_conf["is_unsym_motif"].split(",")
+    if sym_conf.is_unsym_motif:
+        unsym_motif_names = sym_conf.is_unsym_motif.split(",")
         unsym_motif_names = expand_contig_unsym_motif(unsym_motif_names)
         is_unsym_motif = get_unsym_motif_mask(asu_atom_array, unsym_motif_names)
 
@@ -361,38 +373,4 @@ def apply_symmetry_to_xyz_atomwise(X_L, sym_feats, partial_diffusion=False):
                 "blc,cd->bld", asu_xyz, sym_transforms[target_id][0].to(asu_xyz.dtype)
             ) + sym_transforms[target_id][1].to(asu_xyz.dtype)
 
-    # Log inter-chain distances for debugging - use actual chain annotations
-    if sym_X_L.shape[1] > 100:  # Only for large structures
-        # Use symmetry entity annotations to find different chains
-        sym_entity_id = sym_feats["sym_entity_id"]
-        unique_entities = torch.unique(sym_entity_id)
-
-        if len(unique_entities) >= 2:
-            # Get atoms from first two different entities
-            entity_0_mask = sym_entity_id == unique_entities[0]
-            entity_1_mask = sym_entity_id == unique_entities[1]
-
-            if entity_0_mask.sum() > 0 and entity_1_mask.sum() > 0:
-                entity_0_atoms = sym_X_L[0, entity_0_mask, :]
-                entity_1_atoms = sym_X_L[0, entity_1_mask, :]
-
-                # Sample subset to avoid memory issues
-                entity_0_sample = entity_0_atoms[: min(50, entity_0_atoms.shape[0]), :]
-                entity_1_sample = entity_1_atoms[: min(50, entity_1_atoms.shape[0]), :]
-
-                min_distance = (
-                    torch.cdist(entity_0_sample, entity_1_sample).min().item()
-                )
-                ranked_logger.info(
-                    f"Min inter-chain distance after symmetry: {min_distance:.2f} Å"
-                )
-
-                # Also log the centers of each entity
-                entity_0_center = entity_0_atoms.mean(dim=0)
-                entity_1_center = entity_1_atoms.mean(dim=0)
-                center_distance = torch.norm(entity_0_center - entity_1_center).item()
-                ranked_logger.info(
-                    f"Distance between chain centers: {center_distance:.2f} Å"
-                )
-
     return sym_X_L

rfd3/run_inference.py

@@ -12,7 +12,9 @@ load_dotenv(override=True)
 
 # For pip-installed package, configs should be relative to this file
 # Adjust this path based on where configs are bundled in the package
-_config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "configs")
+_config_path = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.dirname(__file__))), "configs"
+)
 
 
 @hydra.main(

rfd3/utils/inference.py

@@ -391,6 +391,29 @@ def ensure_input_is_abspath(args: dict, path: PathLike | None):
     return args
 
 
+def ensure_inference_sampler_matches_design_spec(
+    design_spec: dict, inference_sampler: dict | None = None
+):
+    """
+    Ensure the inference sampler is set to the correct sampler for the design specification.
+    Args:
+        design_spec: Design specification dictionary
+        inference_sampler: Inference sampler dictionary
+    """
+    has_symmetry_specification = [
+        True if "symmetry" in item.keys() else False for item in design_spec.values()
+    ]
+    if any(has_symmetry_specification):
+        if (
+            inference_sampler is None
+            or inference_sampler.get("kind", "default") != "symmetry"
+        ):
+            raise ValueError(
+                "You requested for symmetric designs, but inference sampler is not set to symmetry. "
+                "Please add inference_sampler.kind='symmetry' to your command."
+            )
+
+
 #################################################################################
 # Custom infer_ori functions
 #################################################################################