proteinmpnn-cli 0.1.0__py3-none-any.whl

proteinmpnn/__init__.py

@@ -0,0 +1,30 @@
+"""ProteinMPNN - Protein sequence design using message passing neural networks."""
+
+import logging
+from importlib.metadata import version
+
+# Prevent "No handler found" warnings for library users who don't configure logging
+logging.getLogger("proteinmpnn").addHandler(logging.NullHandler())
+
+__version__ = version("proteinmpnn-cli")
+
+# Re-export main classes for convenient imports
+from proteinmpnn.inference import (  # noqa: E402
+    DesignResult,
+    InferenceRunner,
+    SequenceResult,
+)
+
+__all__ = [
+    "__version__",
+    "InferenceRunner",
+    "DesignResult",
+    "SequenceResult",
+]
+
+
+def main() -> None:
+    """Entry point for the CLI."""
+    from proteinmpnn.cli import app
+
+    app()

proteinmpnn/cli/__init__.py

@@ -0,0 +1,40 @@
+"""CLI for ProteinMPNN."""
+
+import logging
+from typing import Annotated
+
+import typer
+
+from proteinmpnn.utils.logging import setup_logging
+
+app = typer.Typer(
+    no_args_is_help=True,
+    add_completion=False,
+)
+
+
+@app.callback()
+def main(
+    log_level: Annotated[
+        str,
+        typer.Option(
+            "--log-level",
+            "-l",
+            help="Logging level",
+            case_sensitive=False,
+        ),
+    ] = "INFO",
+) -> None:
+    """ProteinMPNN: Structure-conditioned protein sequence design."""
+    level = getattr(logging, log_level.upper(), logging.INFO)
+    setup_logging(level=level)
+
+
+# Import commands to register them
+from proteinmpnn.cli import (  # noqa: F401, E402
+    compute_probs,  # noqa: F401, E402
+    run_single,  # noqa: F401, E402
+)
+
+if __name__ == "__main__":
+    app()

proteinmpnn/cli/compute_probs.py

@@ -0,0 +1,107 @@
+"""Compute conditional probabilities CLI command."""
+
+from __future__ import annotations
+
+from pathlib import Path  # noqa: TC003
+from typing import Annotated, Literal
+
+import typer
+
+from proteinmpnn.cli import app
+from proteinmpnn.cli.display import display_probs_results
+from proteinmpnn.cli.output import write_probs_csv, write_probs_npz
+from proteinmpnn.inference import InferenceRunner
+from proteinmpnn.utils.logging import get_logger
+
+logger = get_logger("cli.compute_probs")
+
+
+@app.command()
+def compute_probs(
+    pdb_path: Annotated[Path, typer.Argument(help="Path to the input PDB file")],
+    model_name: Annotated[
+        Literal[
+            "v_48_002",
+            "v_48_010",
+            "v_48_020",
+            "v_48_030",  # vanilla models
+            "ca_48_002",
+            "ca_48_010",
+            "ca_48_020",  # CA models
+            "s_48_002",
+            "s_48_010",
+            "s_48_020",
+            "s_48_030",  # soluble models
+        ],
+        typer.Option("--model", "-m", help="ProteinMPNN model to use"),
+    ] = "v_48_020",
+    designable_residues: Annotated[
+        str,
+        typer.Option(
+            "--design",
+            "-d",
+            help="Residues to compute conditional probs for (e.g., 'A1-A68')",
+        ),
+    ] = "",
+    unconditional: Annotated[
+        bool,
+        typer.Option(
+            "--unconditional",
+            "-u",
+            help="Compute unconditional probabilities (no sequence context)",
+        ),
+    ] = False,
+    output: Annotated[
+        Path | None,
+        typer.Option(
+            "--output",
+            "-o",
+            help="Output directory (defaults to PDB directory)",
+        ),
+    ] = None,
+    seed: Annotated[
+        int | None,
+        typer.Option("--seed", help="Random seed for reproducibility"),
+    ] = None,
+) -> None:
+    """Compute amino acid conditional (or unconditional) probabilities per residue.
+
+    Outputs both CSV and NPZ files with log probabilities for each residue position.
+
+    Example usage:
+
+        proteinmpnn compute-probs 6MRR.pdb --design A1-A68
+
+        proteinmpnn compute-probs 6MRR.pdb --unconditional
+
+    """
+    # Determine output directory
+    if output is None:
+        output = pdb_path.parent
+    output.mkdir(parents=True, exist_ok=True)
+
+    # Create runner and compute probabilities
+    logger.info("Loading model %s...", model_name)
+    runner = InferenceRunner(model_name=model_name)
+
+    mode_str = "unconditional" if unconditional else "conditional"
+    logger.info("Computing %s probabilities for %s...", mode_str, pdb_path.name)
+
+    result = runner.compute_probs(
+        pdb_path=pdb_path,
+        designable_res=designable_residues,
+        unconditional=unconditional,
+        seed=seed,
+    )
+
+    # Write outputs
+    csv_path = output / f"{pdb_path.stem}_probs.csv"
+    write_probs_csv(result, csv_path)
+    logger.info("Wrote CSV to %s", csv_path)
+
+    npz_path = output / f"{pdb_path.stem}_probs.npz"
+    write_probs_npz(result, npz_path)
+    logger.info("Wrote NPZ to %s", npz_path)
+
+    # Display rich summary
+    display_probs_results(result)

proteinmpnn/cli/display.py

@@ -0,0 +1,124 @@
+"""Rich display utilities for CLI output."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+if TYPE_CHECKING:
+    from proteinmpnn.inference.results import ConditionalProbsResult, DesignResult
+
+console = Console()
+
+
+def display_design_results(result: DesignResult) -> None:
+    """Display design results in a rich table."""
+    # Create a table for designed sequences
+    table = Table(title=f"Designed Sequences for {result.protein_name}")
+
+    table.add_column("#", style="dim", width=4)
+    table.add_column("Temperature", justify="right", style="cyan")
+    table.add_column("Score", justify="right", style="green")
+    table.add_column("Recovery", justify="right", style="yellow")
+    table.add_column("Sequence", style="dim", overflow="fold", max_width=50)
+
+    # Add native sequence first
+    table.add_row(
+        "WT",
+        "-",
+        f"{result.native.score:.4f}",
+        "1.0000",
+        _truncate_seq(result.native.sequence),
+        style="bold",
+    )
+
+    # Add designed sequences sorted by score
+    sorted_seqs = sorted(result.sequences, key=lambda s: s.score)
+    for i, seq in enumerate(sorted_seqs, 1):
+        is_best = i == 1
+        table.add_row(
+            str(i),
+            f"{seq.temperature:.2f}",
+            f"{seq.score:.4f}",
+            f"{seq.seq_recovery:.4f}",
+            _truncate_seq(seq.sequence),
+            style="bold green" if is_best else None,
+        )
+
+    console.print()
+    console.print(table)
+
+    # Summary panel
+    best = sorted_seqs[0]
+    worst = sorted_seqs[-1]
+    summary = (
+        f"[bold]Native score:[/bold] {result.native.score:.4f}\n"
+        f"[bold green]Best score:[/bold green] {best.score:.4f} "
+        f"(T={best.temperature}, recovery={best.seq_recovery:.2%})\n"
+        f"[bold]Worst score:[/bold] {worst.score:.4f}\n"
+        f"[bold]Total sequences:[/bold] {len(result.sequences)}"
+    )
+    console.print(Panel(summary, title="Summary", border_style="blue"))
+
+
+def display_probs_results(result: ConditionalProbsResult) -> None:
+    """Display probability computation results."""
+    import numpy as np
+
+    from proteinmpnn.model.utils import ALPHABET
+
+    # Summary panel
+    mode_style = "yellow" if result.mode == "unconditional" else "cyan"
+    summary = (
+        f"[bold]Protein:[/bold] {result.protein_name}\n"
+        f"[bold]Model:[/bold] {result.model_name}\n"
+        f"[bold]Mode:[/bold] [{mode_style}]{result.mode}[/{mode_style}]\n"
+        f"[bold]Residues:[/bold] {len(result.residue_info)}\n"
+        f"[bold]Matrix shape:[/bold] {result.log_probs.shape}"
+    )
+    console.print()
+    console.print(Panel(summary, title="Probability Computation", border_style="blue"))
+
+    # Show top predictions for first few residues
+    table = Table(title="Top 3 Predictions per Position (first 10 residues)")
+    table.add_column("Position", style="dim")
+    table.add_column("Chain", style="cyan")
+    table.add_column("1st", justify="center")
+    table.add_column("P(1st)", justify="right", style="green")
+    table.add_column("2nd", justify="center")
+    table.add_column("P(2nd)", justify="right")
+    table.add_column("3rd", justify="center")
+    table.add_column("P(3rd)", justify="right", style="dim")
+
+    # Show first 10 residues
+    for i, res_info in enumerate(result.residue_info[:10]):
+        probs = np.exp(result.log_probs[i])
+        top_indices = np.argsort(probs)[::-1][:3]
+
+        table.add_row(
+            str(res_info.residue_idx),
+            res_info.chain,
+            ALPHABET[top_indices[0]],
+            f"{probs[top_indices[0]]:.3f}",
+            ALPHABET[top_indices[1]],
+            f"{probs[top_indices[1]]:.3f}",
+            ALPHABET[top_indices[2]],
+            f"{probs[top_indices[2]]:.3f}",
+        )
+
+    console.print(table)
+
+    if len(result.residue_info) > 10:
+        console.print(
+            f"[dim]... and {len(result.residue_info) - 10} more residues[/dim]"
+        )
+
+
+def _truncate_seq(seq: str, max_len: int = 40) -> str:
+    """Truncate sequence for display."""
+    if len(seq) <= max_len:
+        return seq
+    return seq[: max_len - 3] + "..."

proteinmpnn/cli/output.py

@@ -0,0 +1,76 @@
+"""Output formatting utilities for CLI."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from proteinmpnn.inference.results import ConditionalProbsResult, DesignResult
+
+
+def write_fasta(result: DesignResult, output_path: Path) -> None:
+    """Write design result to a FASTA file.
+
+    Args:
+        result: DesignResult from inference.
+        output_path: Path to output FASTA file.
+    """
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(result.to_fasta() + "\n")
+
+
+def write_af2_csv(result: DesignResult, output_path: Path) -> None:
+    """Write design result to AlphaFold2-compatible CSV.
+
+    Args:
+        result: DesignResult from inference.
+        output_path: Path to output CSV file.
+    """
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(result.to_af2_csv() + "\n")
+
+
+def write_probs(result: DesignResult, output_path: Path) -> None:
+    """Write probability matrix to NPZ file.
+
+    Args:
+        result: DesignResult from inference (must have probs).
+        output_path: Path to output NPZ file.
+
+    Raises:
+        ValueError: If result has no probability matrix.
+    """
+    if result.probs is None:
+        raise ValueError("DesignResult has no probability matrix")
+
+    import numpy as np
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(output_path, probs=result.probs)
+
+
+def write_probs_csv(result: ConditionalProbsResult, output_path: Path) -> None:
+    """Write conditional probabilities result to CSV file.
+
+    Args:
+        result: ConditionalProbsResult from inference.
+        output_path: Path to output CSV file.
+    """
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    output_path.write_text(result.to_csv() + "\n")
+
+
+def write_probs_npz(result: ConditionalProbsResult, output_path: Path) -> None:
+    """Write conditional probabilities result to NPZ file.
+
+    Args:
+        result: ConditionalProbsResult from inference.
+        output_path: Path to output NPZ file.
+    """
+    import numpy as np
+
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    npz_dict = result.to_npz_dict()
+    np.savez(output_path, **npz_dict)

proteinmpnn/cli/run_single.py

@@ -0,0 +1,152 @@
+"""Run single protein design CLI command."""
+
+from __future__ import annotations
+
+from pathlib import Path  # noqa: TC003
+from typing import Annotated, Literal
+
+import typer
+
+from proteinmpnn.cli import app
+from proteinmpnn.cli.display import display_design_results
+from proteinmpnn.cli.output import write_af2_csv, write_fasta
+from proteinmpnn.inference import InferenceRunner
+from proteinmpnn.utils.logging import get_logger
+
+logger = get_logger("cli.run_single")
+
+
+@app.command()
+def run_single(
+    pdb_path: Annotated[Path, typer.Argument(help="Path to the input PDB file")],
+    model_name: Annotated[
+        Literal[
+            "v_48_002",
+            "v_48_010",
+            "v_48_020",
+            "v_48_030",  # vanilla models
+            "ca_48_002",
+            "ca_48_010",
+            "ca_48_020",  # CA models
+            "s_48_002",
+            "s_48_010",
+            "s_48_020",
+            "s_48_030",  # soluble models
+        ],
+        typer.Option("--model", "-m", help="ProteinMPNN model to use"),
+    ] = "v_48_020",
+    designable_residues: Annotated[
+        str,
+        typer.Option(
+            "--design",
+            "-d",
+            help="Designable residues (e.g., 'A1-A68' or 'A10,A12-A15')",
+        ),
+    ] = "",
+    symmetric_residues: Annotated[
+        str,
+        typer.Option(
+            "--symmetric",
+            "-s",
+            help="Symmetric residue pairs (e.g., 'A10:B10,A11:B11')",
+        ),
+    ] = "",
+    cluster_center: Annotated[
+        str,
+        typer.Option(
+            "--cluster",
+            "-c",
+            help="Cluster center residue(s) for radius-based selection",
+        ),
+    ] = "",
+    cluster_radius: Annotated[
+        float,
+        typer.Option("--radius", "-r", help="Cluster radius in Angstroms"),
+    ] = 10.0,
+    backbone_noise: Annotated[
+        float,
+        typer.Option("--noise", help="Backbone noise standard deviation"),
+    ] = 0.0,
+    num_seq_per_target: Annotated[
+        int,
+        typer.Option("-n", "--num-seqs", help="Number of sequences to generate"),
+    ] = 5,
+    batch_size: Annotated[
+        int,
+        typer.Option("--batch", "-b", help="Batch size for generation"),
+    ] = 1,
+    temperature: Annotated[
+        str,
+        typer.Option(
+            "--temp",
+            "-t",
+            help="Sampling temperature(s), space-separated (e.g., '0.1 0.2')",
+        ),
+    ] = "0.1",
+    output: Annotated[
+        Path | None,
+        typer.Option(
+            "--output",
+            "-o",
+            help="Output directory (defaults to PDB directory)",
+        ),
+    ] = None,
+    af2: Annotated[
+        bool,
+        typer.Option("--af2", help="Also output AlphaFold2-format CSV"),
+    ] = False,
+    seed: Annotated[
+        int | None,
+        typer.Option("--seed", help="Random seed for reproducibility"),
+    ] = None,
+) -> None:
+    """Design sequences for a single protein structure.
+
+    Example usage:
+
+        proteinmpnn run-single 6MRR.pdb --design A1-A68 -n 10
+
+        proteinmpnn run-single 4GYT.pdb --design A7-A183,B7-B183 \\
+            --symmetric A7-A183:B7-B183
+
+    """
+    # Parse temperatures
+    temperatures = [float(t) for t in temperature.split()]
+
+    # Determine output directory
+    if output is None:
+        output = pdb_path.parent
+    output.mkdir(parents=True, exist_ok=True)
+
+    # Create runner and generate sequences
+    logger.info("Loading model %s...", model_name)
+    runner = InferenceRunner(
+        model_name=model_name,
+        backbone_noise=backbone_noise,
+    )
+
+    logger.info("Designing sequences for %s...", pdb_path.name)
+    result = runner.design_single(
+        pdb_path=pdb_path,
+        designable_res=designable_residues,
+        symmetric_res=symmetric_residues,
+        cluster_center=cluster_center,
+        cluster_radius=cluster_radius,
+        num_sequences=num_seq_per_target,
+        batch_size=batch_size,
+        temperatures=temperatures,
+        seed=seed,
+    )
+
+    # Write outputs
+    fasta_path = output / f"{pdb_path.stem}.fasta"
+    write_fasta(result, fasta_path)
+    logger.info("Wrote %d sequences to %s", len(result.sequences) + 1, fasta_path)
+
+    if af2:
+        csv_path = output / f"{pdb_path.stem}.csv"
+        write_af2_csv(result, csv_path)
+        logger.info("Wrote AlphaFold2 CSV to %s", csv_path)
+
+    # Display rich summary
+    display_design_results(result)

proteinmpnn/data/__init__.py

@@ -0,0 +1,23 @@
+"""Data processing modules for ProteinMPNN."""
+
+from proteinmpnn.data.config import (
+    DesignableResidue,
+    MultiStateConfig,
+    SingleStateConfig,
+)
+from proteinmpnn.data.input import (
+    ProteinDesignInputFormatter,
+    create_design_input,
+)
+from proteinmpnn.data.multi_state import MultiStateDesignInput
+from proteinmpnn.data.single_state import SingleStateDesignInput
+
+__all__ = [
+    "create_design_input",
+    "DesignableResidue",
+    "MultiStateConfig",
+    "MultiStateDesignInput",
+    "ProteinDesignInputFormatter",
+    "SingleStateConfig",
+    "SingleStateDesignInput",
+]

proteinmpnn/data/config.py

@@ -0,0 +1,54 @@
+"""Pydantic models for ProteinMPNN design output configurations."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+from pydantic import BaseModel, Field
+
+
+class DesignableResidue(BaseModel):
+    """A single residue that can be mutated."""
+
+    chain: str = Field(description="Chain identifier")
+    resid: int = Field(description="Residue index (1-based, sequential within chain)")
+    WTAA: str = Field(description="Wild-type amino acid (1-letter code)")
+    MutTo: str = Field(default="all", description="Allowed amino acids for mutation")
+
+
+class SingleStateConfig(BaseModel):
+    """Output config for single-state protein design."""
+
+    sequence: dict[str, str] = Field(
+        default_factory=dict, description="Chain ID -> amino acid sequence"
+    )
+    designable: list[DesignableResidue] = Field(
+        default_factory=list, description="List of designable residues"
+    )
+    symmetric: list[list[str]] = Field(
+        default_factory=list, description="List of tied position groups"
+    )
+
+    def to_json(self, path: str | Path, indent: int = 2) -> None:
+        """Write config to JSON file.
+
+        Args:
+            path: Output file path.
+            indent: JSON indentation level.
+        """
+        Path(path).write_text(self.model_dump_json(indent=indent))
+
+
+class MultiStateConfig(SingleStateConfig):
+    """Output config for multi-state protein design.
+
+    Extends SingleStateConfig with additional MSD-specific fields.
+    """
+
+    tied_betas: dict[str, float] = Field(
+        default_factory=dict, description="Chain ID -> beta weight for MSD"
+    )
+    chain_key: dict[str, dict[str, str]] = Field(
+        default_factory=dict,
+        description="PDB name -> original chain -> remapped chain",
+    )