dh-cli 0.2.11__tar.gz → 0.3.1__tar.gz

{dh_cli-0.2.11 → dh_cli-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dh-cli
-Version: 0.2.11
+Version: 0.3.1
 Summary: Dayhoff Labs developer CLI
 Author-email: Dayhoff Labs <dev@dayhofflabs.com>
 License: # PolyForm Noncommercial License 1.0.0

{dh_cli-0.2.11 → dh_cli-0.3.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "dh-cli"
-version = "0.2.11"
+version = "0.3.1"
 description = "Dayhoff Labs developer CLI"
 requires-python = ">=3.11"
 readme = "README.md"

{dh_cli-0.2.11 → dh_cli-0.3.1}/src/dh_cli/batch/__init__.py

@@ -18,6 +18,8 @@ from .commands.finalize import finalize
 from .commands.list_jobs import list_jobs
 from .commands.local import local
 from .commands.logs import logs
+from .commands.protmpnn import protmpnn
+from .commands.protmpnn_to_boltz import protmpnn_to_boltz
 from .commands.retry import retry
 from .commands.status import status
 from .commands.submit import submit
@@ -49,6 +51,11 @@ def batch_cli():
     Structure Prediction:
       boltz      Predict protein structures with Boltz
 
+    \b
+    Sequence Design:
+      protmpnn          Design sequences with ProtMPNN/LigandMPNN
+      protmpnn-to-boltz Convert top variants to Boltz validation YAMLs
+
     \b
     ML Training:
       train      Submit hatchery ML training jobs
@@ -91,6 +98,8 @@ batch_cli.add_command(wait_for, name="wait-for")
 # Register pipeline commands
 batch_cli.add_command(embed_t5, name="embed-t5")
 batch_cli.add_command(boltz)
+batch_cli.add_command(protmpnn)
+batch_cli.add_command(protmpnn_to_boltz, name="protmpnn-to-boltz")
 batch_cli.add_command(train)
 
 __all__ = ["batch_cli"]

{dh_cli-0.2.11 → dh_cli-0.3.1}/src/dh_cli/batch/commands/finalize.py

@@ -145,6 +145,8 @@ def finalize(
                 )
             )
         _finalize_boltz(output_dir, output_path, full_output=full_output)
+    elif manifest.pipeline == "protmpnn":
+        _finalize_protmpnn(output_dir, output_path)
     else:
         _finalize_generic(output_dir, output_path)
 
@@ -447,3 +449,55 @@ def _finalize_generic(output_dir: Path, output_path: Path):
     click.echo(f"Copying output directory to {output_path}...")
     shutil.copytree(output_dir, output_path)
     click.echo(click.style("✓ Output copied successfully", fg="green"))
+
+
+def _finalize_protmpnn(output_dir: Path, output_path: Path):
+    """Merge per-worker ProtMPNN results into final output.
+
+    1. Merge results_worker_*.csv into results.csv (sorted by confidence)
+    2. Flatten all seqs/ and pdbs/ into output seqs/ and pdbs/ dirs
+    3. Print summary
+    """
+    import pandas as pd
+
+    csv_files = sorted(output_dir.glob("results_worker_*.csv"))
+
+    if not csv_files:
+        click.echo("No results_worker_*.csv files found in output directory.", err=True)
+        raise SystemExit(1)
+
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    dfs = [pd.read_csv(f) for f in csv_files]
+    merged = pd.concat(dfs, ignore_index=True)
+    merged = merged.sort_values("overall_confidence", ascending=False)
+    merged.to_csv(output_path / "results.csv", index=False)
+
+    num_variants = len(merged)
+    num_configs = merged["config_name"].nunique()
+
+    seqs_dest = output_path / "seqs"
+    seqs_dest.mkdir(exist_ok=True)
+    for config_dir in output_dir.iterdir():
+        if config_dir.is_dir():
+            config_seqs = config_dir / "seqs"
+            if config_seqs.exists():
+                for fa_file in config_seqs.glob("*.fa"):
+                    shutil.copy2(fa_file, seqs_dest / fa_file.name)
+
+    pdbs_dest = output_path / "pdbs"
+    pdbs_dest.mkdir(exist_ok=True)
+    for config_dir in output_dir.iterdir():
+        if config_dir.is_dir():
+            config_pdbs = config_dir / "pdbs"
+            if config_pdbs.exists():
+                for pdb_file in config_pdbs.glob("*.pdb"):
+                    shutil.copy2(pdb_file, pdbs_dest / pdb_file.name)
+
+    top_conf = merged.iloc[0]["overall_confidence"] if num_variants > 0 else "N/A"
+
+    click.echo(
+        f"{num_variants} variants from {num_configs} config(s), "
+        f"top confidence: {top_conf}"
+    )
+    click.echo(click.style(f"Results: {output_path / 'results.csv'}", fg="green"))

dh_cli-0.3.1/src/dh_cli/batch/commands/protmpnn.py

@@ -0,0 +1,418 @@
+"""ProtMPNN/LigandMPNN sequence design pipeline command."""
+
+import math
+import os
+import shutil
+from pathlib import Path
+
+import click
+
+from ..aws_batch import BatchClient, BatchError
+from ..job_id import generate_job_id
+from ..manifest import (
+    BATCH_JOBS_BASE,
+    BatchConfig,
+    InputConfig,
+    JobManifest,
+    JobStatus,
+    OutputConfig,
+    create_job_directory,
+    get_job_dir,
+    save_manifest,
+)
+
+DEFAULT_QUEUE = "t4-1x-spot"
+MAX_WORKERS = 50
+FILES_PER_WORKER = 10
+DEFAULT_JOB_DEFINITION = "dayhoff-protmpnn"
+DEFAULT_IMAGE_URI = (
+    "074735440724.dkr.ecr.us-east-1.amazonaws.com/dayhoff:protmpnn-latest"
+)
+
+
+@click.command()
+@click.argument("input_dir", type=click.Path(exists=True))
+@click.option(
+    "--workers",
+    default=None,
+    type=int,
+    help="Number of parallel workers [default: ~1 per 10 files]",
+)
+@click.option(
+    "--queue",
+    default=DEFAULT_QUEUE,
+    help=f"Batch queue [default: {DEFAULT_QUEUE}]",
+)
+@click.option("--dry-run", is_flag=True, help="Show plan without submitting")
+@click.option(
+    "--local",
+    "run_local",
+    is_flag=True,
+    help="Force local execution via Docker",
+)
+@click.option(
+    "--remote",
+    "run_remote",
+    is_flag=True,
+    help="Force Batch submission (override auto-detect)",
+)
+@click.option(
+    "--shell",
+    "run_shell",
+    is_flag=True,
+    help="Drop into container shell for debugging",
+)
+@click.option("--base-path", default=BATCH_JOBS_BASE, help="Base path for job data")
+def protmpnn(input_dir, workers, queue, dry_run, run_local, run_remote, run_shell, base_path):
+    """Design protein sequences with ProtMPNN/LigandMPNN.
+
+    Processes a directory of YAML config files, each specifying a PDB
+    structure and design parameters. Generates variant sequences ranked
+    by confidence.
+
+    \b
+    Examples:
+      # Run on a GPU engine (auto-detects GPU, runs locally)
+      dh batch protmpnn input/
+
+      # Force remote Batch submission
+      dh batch protmpnn input/ --remote
+
+      # Preview what would run
+      dh batch protmpnn input/ --dry-run
+
+      # Run locally via Docker
+      dh batch protmpnn input/ --local
+
+    \b
+    After remote job completes:
+      dh batch status <job-id>
+      dh batch finalize <job-id> --output ./results/
+
+    \b
+    YAML config format:
+      version: 1
+      pdb_path: 6DHI.pdb
+      model_type: ligand_mpnn
+      num_sequences: 20
+      seed: 42
+      temperature: 0.1
+      fixed_residues: "A42 A181 A217 A218"
+    """
+    input_path = Path(input_dir).resolve()
+
+    if run_shell:
+        _run_shell_mode(input_path)
+        return
+
+    if run_local:
+        _run_local_mode(input_path)
+        return
+
+    # Auto-detect GPU for smart defaulting
+    if not run_remote and not dry_run:
+        if _has_local_gpu():
+            click.echo("GPU detected — running locally (use --remote to override)")
+            _run_local_mode(input_path)
+            return
+
+    _submit_batch_job(input_path, workers, queue, dry_run, base_path)
+
+
+def _has_local_gpu() -> bool:
+    """Check if a local NVIDIA GPU is available."""
+    import subprocess
+
+    try:
+        result = subprocess.run(
+            ["nvidia-smi"],
+            capture_output=True,
+            timeout=5,
+        )
+        return result.returncode == 0
+    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return False
+
+
+def _count_yaml_files(input_path: Path) -> int:
+    return len(list(input_path.glob("*.yaml")))
+
+
+def _copy_inputs_to_job_dir(input_path: Path, job_dir: Path) -> int:
+    """Copy input YAML and PDB files to job directory."""
+    input_dir = job_dir / "input"
+    input_dir.mkdir(parents=True, exist_ok=True)
+
+    count = 0
+    for yaml_file in sorted(input_path.glob("*.yaml")):
+        shutil.copy2(yaml_file, input_dir / yaml_file.name)
+        count += 1
+
+    # Copy PDB files alongside YAMLs
+    for pdb_file in sorted(input_path.glob("*.pdb")):
+        shutil.copy2(pdb_file, input_dir / pdb_file.name)
+
+    return count
+
+
+def _submit_batch_job(
+    input_path: Path,
+    workers: int | None,
+    queue: str,
+    dry_run: bool,
+    base_path: str,
+):
+    """Submit ProtMPNN job to AWS Batch."""
+    click.echo(f"Scanning {input_path} for YAML files...")
+    num_files = _count_yaml_files(input_path)
+
+    if num_files == 0:
+        click.echo(
+            click.style("Error: No YAML files found in input directory", fg="red"),
+            err=True,
+        )
+        raise SystemExit(1)
+
+    click.echo(f"Found {num_files} config(s) to process")
+
+    if workers is None:
+        workers = max(1, min(math.ceil(num_files / FILES_PER_WORKER), MAX_WORKERS))
+    array_size = min(num_files, workers)
+
+    job_id = generate_job_id("protmpnn")
+
+    click.echo()
+    click.echo(f"Job ID:           {job_id}")
+    click.echo(f"Input:            {input_path}")
+    click.echo(f"Configs:          {num_files}")
+    click.echo(f"Workers:          {array_size}")
+    files_per_worker = math.ceil(num_files / array_size)
+    click.echo(f"Files per worker: ~{files_per_worker}")
+    click.echo(f"Queue:            {queue}")
+    click.echo(f"Job definition:   {DEFAULT_JOB_DEFINITION}")
+
+    if dry_run:
+        click.echo()
+        click.echo(click.style("Dry run - job not submitted", fg="yellow"))
+        return
+
+    if not click.confirm("\nSubmit job?", default=True):
+        click.echo("Cancelled.")
+        raise SystemExit(0)
+    click.echo()
+
+    job_dir = create_job_directory(job_id, base_path)
+    click.echo(f"Created job directory: {job_dir}")
+
+    click.echo("Copying input files...")
+    copied = _copy_inputs_to_job_dir(input_path, job_dir)
+    click.echo(f"Copied {copied} YAML files")
+
+    manifest = JobManifest(
+        job_id=job_id,
+        user=job_id.split("-")[0],
+        pipeline="protmpnn",
+        status=JobStatus.PENDING,
+        image_uri=DEFAULT_IMAGE_URI,
+        input=InputConfig(
+            source=str(input_path),
+            num_sequences=num_files,
+            num_chunks=array_size,
+        ),
+        batch=BatchConfig(
+            queue=queue,
+            job_definition=DEFAULT_JOB_DEFINITION,
+            array_size=array_size,
+        ),
+        output=OutputConfig(
+            destination=None,
+            finalized=False,
+        ),
+    )
+
+    save_manifest(manifest, base_path)
+
+    try:
+        client = BatchClient()
+
+        environment = {
+            "JOB_DIR": str(job_dir),
+            "JOB_ID": job_id,
+            "BATCH_ARRAY_SIZE": str(array_size),
+            "BATCH_NUM_FILES": str(num_files),
+        }
+
+        batch_job_id = client.submit_job(
+            job_name=job_id,
+            job_definition=DEFAULT_JOB_DEFINITION,
+            job_queue=queue,
+            array_size=array_size,
+            environment=environment,
+            timeout_seconds=1 * 3600,  # 1 hour
+            retry_attempts=5,
+        )
+
+        manifest.status = JobStatus.SUBMITTED
+        manifest.batch.job_id = batch_job_id
+        save_manifest(manifest, base_path)
+
+        click.echo()
+        click.echo(click.style("Job submitted successfully!", fg="green"))
+        click.echo()
+        click.echo(f"AWS Batch Job ID: {batch_job_id}")
+        click.echo()
+        click.echo("Next steps:")
+        click.echo(f"  Check status:  dh batch status {job_id}")
+        click.echo(f"  View logs:     dh batch logs {job_id}")
+        click.echo(f"  Cancel:        dh batch cancel {job_id}")
+        click.echo()
+        click.echo("After completion:")
+        click.echo(
+            f"  Finalize:      dh batch finalize {job_id} --output ./results/"
+        )
+
+    except BatchError as e:
+        manifest.status = JobStatus.FAILED
+        manifest.error_message = str(e)
+        save_manifest(manifest, base_path)
+        click.echo(click.style(f"Failed to submit job: {e}", fg="red"), err=True)
+        raise SystemExit(1)
+
+
+def _run_local_mode(input_path: Path):
+    """Run ProtMPNN locally in a Docker container."""
+    import subprocess
+
+    click.echo("Running ProtMPNN locally in container...")
+    click.echo(f"Input directory: {input_path}")
+
+    yaml_files = list(input_path.glob("*.yaml"))
+    if not yaml_files:
+        click.echo(click.style("Error: No YAML files found", fg="red"), err=True)
+        raise SystemExit(1)
+
+    click.echo(f"Found {len(yaml_files)} config file(s)")
+
+    temp_job_dir = input_path / ".local_protmpnn_job"
+    temp_input_dir = temp_job_dir / "input"
+    temp_output_dir = temp_job_dir / "output"
+
+    if temp_job_dir.exists():
+        shutil.rmtree(temp_job_dir)
+
+    temp_input_dir.mkdir(parents=True)
+    temp_output_dir.mkdir(parents=True)
+
+    for yaml_file in yaml_files:
+        shutil.copy2(yaml_file, temp_input_dir / yaml_file.name)
+    for pdb_file in input_path.glob("*.pdb"):
+        shutil.copy2(pdb_file, temp_input_dir / pdb_file.name)
+
+    click.echo(f"Output will be at: {temp_output_dir}/")
+    click.echo()
+
+    cmd = [
+        "docker",
+        "run",
+        "--rm",
+        "--gpus",
+        "all",
+        "-v",
+        "/primordial:/primordial",
+        "-v",
+        f"{temp_job_dir}:{temp_job_dir}",
+        "-e",
+        f"JOB_DIR={temp_job_dir}",
+        "-e",
+        "AWS_BATCH_JOB_ARRAY_INDEX=0",
+        "-e",
+        "BATCH_ARRAY_SIZE=1",
+        "-e",
+        f"BATCH_NUM_FILES={len(yaml_files)}",
+        DEFAULT_IMAGE_URI,
+    ]
+
+    click.echo(f"Running: {' '.join(cmd)}")
+    click.echo()
+
+    try:
+        result = subprocess.run(cmd)
+        if result.returncode != 0:
+            click.echo(
+                click.style(
+                    f"Container exited with code {result.returncode}", fg="red"
+                ),
+                err=True,
+            )
+            raise SystemExit(result.returncode)
+
+        csv_files = list(temp_output_dir.glob("results_worker_*.csv"))
+        if csv_files:
+            # Merge worker CSVs into results.csv for local mode
+            import pandas as pd
+
+            dfs = [pd.read_csv(f) for f in csv_files]
+            merged = pd.concat(dfs, ignore_index=True)
+            merged = merged.sort_values("overall_confidence", ascending=False)
+            merged.to_csv(temp_output_dir / "results.csv", index=False)
+
+            click.echo()
+            click.echo(click.style("Design complete!", fg="green"))
+            click.echo(f"Results: {temp_output_dir / 'results.csv'}")
+            click.echo(f"  {len(merged)} variants generated")
+        else:
+            click.echo(click.style("Warning: No results CSV found", fg="yellow"))
+
+    except FileNotFoundError:
+        click.echo(
+            click.style(
+                "Error: Docker not found. Is Docker installed and running?",
+                fg="red",
+            ),
+            err=True,
+        )
+        raise SystemExit(1)
+
+
+def _run_shell_mode(input_path: Path):
+    """Drop into container shell for debugging."""
+    import subprocess
+
+    click.echo("Dropping into container shell...")
+    click.echo(f"Input will be available at: /input/")
+    click.echo()
+
+    cmd = [
+        "docker",
+        "run",
+        "--rm",
+        "-it",
+        "--gpus",
+        "all",
+        "-v",
+        "/primordial:/primordial",
+        "-v",
+        f"{input_path}:/input",
+        "-e",
+        "JOB_DIR=/input",
+        "-e",
+        "AWS_BATCH_JOB_ARRAY_INDEX=0",
+        "--entrypoint",
+        "/bin/bash",
+        DEFAULT_IMAGE_URI,
+    ]
+
+    click.echo(f"Running: {' '.join(cmd)}")
+    click.echo()
+
+    try:
+        subprocess.run(cmd)
+    except FileNotFoundError:
+        click.echo(
+            click.style(
+                "Error: Docker not found. Is Docker installed and running?",
+                fg="red",
+            ),
+            err=True,
+        )
+        raise SystemExit(1)

dh_cli-0.3.1/src/dh_cli/batch/commands/protmpnn_to_boltz.py

@@ -0,0 +1,249 @@
+"""Convert ProtMPNN results to Boltz input YAMLs for structural validation."""
+
+import shutil
+from pathlib import Path
+
+import click
+import pandas as pd
+import yaml
+
+
+@click.command("protmpnn-to-boltz")
+@click.argument("results_dir", type=click.Path(exists=True))
+@click.option("--top", default=10, type=int, help="Number of top variants to convert")
+@click.option(
+    "--output",
+    "-o",
+    default=None,
+    type=click.Path(),
+    help="Output directory for Boltz YAMLs [default: boltz_input/]",
+)
+@click.option(
+    "--config",
+    "config_dir",
+    default=None,
+    type=click.Path(exists=True),
+    help="Directory containing original ProtMPNN config YAMLs (for ligand_smiles)",
+)
+def protmpnn_to_boltz(results_dir, top, output, config_dir):
+    """Convert top ProtMPNN variants to Boltz YAML configs.
+
+    Takes a ProtMPNN results directory (containing results.csv) and generates
+    Boltz-format YAML files for structural validation of the top-N variants.
+
+    \b
+    Examples:
+      # Convert top 10 from local run
+      dh batch protmpnn-to-boltz input/.local_protmpnn_job/output/ --output boltz_in/
+
+      # Convert top 20, pull ligand SMILES from original configs
+      dh batch protmpnn-to-boltz results/ --top 20 --config input/ -o boltz_in/
+
+    \b
+    The generated Boltz YAMLs can be used directly:
+      dh batch boltz boltz_in/
+      dh batch boltz --local boltz_in/
+    """
+    results_path = Path(results_dir).resolve()
+    csv_path = results_path / "results.csv"
+
+    if not csv_path.exists():
+        # Try worker CSVs if results.csv not found (e.g. raw output before finalize)
+        worker_csvs = sorted(results_path.glob("results_worker_*.csv"))
+        if worker_csvs:
+            dfs = [pd.read_csv(f) for f in worker_csvs]
+            df = pd.concat(dfs, ignore_index=True)
+            df = df.sort_values("overall_confidence", ascending=False)
+        else:
+            click.echo(
+                click.style(
+                    "Error: No results.csv or results_worker_*.csv found", fg="red"
+                ),
+                err=True,
+            )
+            raise SystemExit(1)
+    else:
+        df = pd.read_csv(csv_path)
+
+    if len(df) == 0:
+        click.echo(click.style("Error: Results CSV is empty", fg="red"), err=True)
+        raise SystemExit(1)
+
+    top_n = min(top, len(df))
+    top_variants = df.head(top_n)
+
+    # Resolve ligand SMILES from original config YAMLs
+    ligand_map = _load_ligand_smiles(config_dir, results_path)
+
+    output_path = Path(output or "boltz_input").resolve()
+    output_path.mkdir(parents=True, exist_ok=True)
+
+    click.echo(f"Converting top {top_n} variants to Boltz format...")
+
+    generated = []
+    for idx, row in top_variants.iterrows():
+        config_name = row.get("config_name", "unknown")
+        variant_id = int(row.get("variant_id", idx))
+        sequence = row["sequence"]
+        confidence = row.get("overall_confidence", float("nan"))
+
+        boltz_yaml = _build_boltz_yaml(
+            sequence=sequence,
+            config_name=config_name,
+            variant_id=variant_id,
+            ligand_smiles=ligand_map.get(config_name),
+        )
+
+        filename = f"{config_name}_var{variant_id:03d}.yaml"
+        yaml_path = output_path / filename
+
+        with open(yaml_path, "w") as f:
+            yaml.dump(boltz_yaml, f, default_flow_style=False, sort_keys=False)
+
+        generated.append((filename, confidence))
+
+    # Copy PDB files for reference if available
+    pdbs_src = results_path / "pdbs"
+    if pdbs_src.exists():
+        pdbs_dest = output_path / "reference_pdbs"
+        pdbs_dest.mkdir(exist_ok=True)
+        for pdb in pdbs_src.glob("*.pdb"):
+            shutil.copy2(pdb, pdbs_dest / pdb.name)
+
+    # Generate PyMOL visualization script
+    _write_pymol_script(output_path, results_path, generated, ligand_map)
+
+    click.echo()
+    click.echo(click.style(f"Generated {len(generated)} Boltz configs", fg="green"))
+    click.echo(f"Output: {output_path}/")
+    click.echo()
+    click.echo("Next steps:")
+    click.echo(f"  dh batch boltz {output_path}/")
+    click.echo(f"  dh batch boltz --local {output_path}/")
+
+
+def _load_ligand_smiles(
+    config_dir: str | None, results_path: Path
+) -> dict[str, str | None]:
+    """Load ligand_smiles from original ProtMPNN config YAMLs.
+
+    Searches config_dir first, then falls back to the input/ sibling
+    of the results directory (common in local runs).
+    """
+    smiles_map: dict[str, str | None] = {}
+
+    search_dirs = []
+    if config_dir:
+        search_dirs.append(Path(config_dir))
+
+    # For local runs: results are at input/.local_protmpnn_job/output/
+    # Config YAMLs are at input/
+    if results_path.name == "output":
+        job_dir = results_path.parent
+        input_dir = job_dir / "input"
+        if input_dir.exists():
+            search_dirs.append(input_dir)
+
+    for search_dir in search_dirs:
+        for yaml_file in search_dir.glob("*.yaml"):
+            try:
+                with open(yaml_file) as f:
+                    data = yaml.safe_load(f)
+                if isinstance(data, dict) and data.get("ligand_smiles"):
+                    smiles_map[yaml_file.stem] = data["ligand_smiles"]
+            except Exception:
+                continue
+
+    return smiles_map
+
+
+def _build_boltz_yaml(
+    sequence: str,
+    config_name: str,
+    variant_id: int,
+    ligand_smiles: str | None = None,
+) -> dict:
+    """Build a Boltz-format YAML dict for a single variant."""
+    sequences = [
+        {
+            "protein": {
+                "id": "A",
+                "sequence": sequence,
+            }
+        }
+    ]
+
+    if ligand_smiles:
+        sequences.append(
+            {
+                "ligand": {
+                    "id": "B",
+                    "smiles": ligand_smiles,
+                }
+            }
+        )
+
+    return {
+        "version": 1,
+        "sequences": sequences,
+    }
+
+
+def _write_pymol_script(
+    output_path: Path,
+    results_path: Path,
+    generated: list[tuple[str, float]],
+    ligand_map: dict[str, str | None],
+):
+    """Generate a PyMOL script for visualizing WT + variant structures.
+
+    This script is designed to be run after Boltz validation completes,
+    loading the predicted structures and aligning them to the WT.
+    """
+    pdbs_dir = results_path / "pdbs"
+    wt_pdbs = sorted(pdbs_dir.glob("*.pdb")) if pdbs_dir.exists() else []
+
+    lines = [
+        "# PyMOL visualization script for ProtMPNN variants",
+        "# Generated by: dh batch protmpnn-to-boltz",
+        "#",
+        "# Usage: pymol view_variants.pml",
+        "#   or:  pymol -r view_variants.pml",
+        "",
+        "from pymol import cmd",
+        "",
+    ]
+
+    if wt_pdbs:
+        wt_pdb = wt_pdbs[0]
+        lines.append(f'cmd.load("reference_pdbs/{wt_pdb.name}", "wildtype")')
+        lines.append('cmd.color("gray80", "wildtype")')
+        lines.append("")
+
+    lines.append("# Load variant structures after Boltz validation")
+    lines.append("# Boltz outputs will be in the finalized results directory")
+    for filename, confidence in generated:
+        obj_name = filename.replace(".yaml", "")
+        lines.append(f"# {obj_name}: confidence={confidence:.3f}")
+
+    lines.extend([
+        "",
+        "# Align all objects to wildtype",
+        'for obj in cmd.get_object_list():',
+        '    if obj != "wildtype":',
+        '        cmd.align(obj, "wildtype")',
+        "",
+        "# Show cartoon representation",
+        "cmd.show('cartoon')",
+        "cmd.hide('lines')",
+        "",
+        "# Highlight mutations (after loading Boltz results)",
+        "# cmd.select('mutations', 'wildtype and not (same sequence as variant)')",
+        "",
+        "cmd.zoom()",
+        "print('Loaded variant structures. Align Boltz results manually.')",
+    ])
+
+    script_path = output_path / "view_variants.pml"
+    with open(script_path, "w") as f:
+        f.write("\n".join(lines) + "\n")