Reproduction of metal context sequence recovery

[DeNeutoy]

Hi @dauparas,

Thanks for the great code - I'm trying to reproduce the metal context sequence recovery evaluation and struggling to match the numbers reported in the paper.

I am taking the average of 10 resampled sequences temp=0.1,, only looking at residues within 5A of the non-protein atoms present in the pdb ids reported in training/test_metal.json.

I am seeing a recovery rate of ~0.6915 using ligandmpnn_v_32_010_25.pt.

Using the same script, I am able to reproduce the numbers for the small molecule and nucleotide evaluations - is there anything specific to the metal context evaluation that is different?

[DeNeutoy]

Here is my script, for context:

import argparse
import os
import sys
import random
from pathlib import Path
from typing import Dict, List, Tuple, Optional

import numpy as np
import torch

from data_utils import (
    parse_PDB,
    get_seq_rec,
    featurize,
)
from model_utils import ProteinMPNN


def get_residues_near_ligands(
    protein_dict: Dict,
    other_atoms,
    cutoff: float = 5.0,
    device: str = "cpu"
) -> torch.Tensor:
    """
    Identify residues within cutoff distance (default 5Å) of any non-protein atom.
    
    Args:
        protein_dict: Dictionary from parse_PDB containing protein data
        other_atoms: ProDy AtomGroup with non-protein atoms
        cutoff: Distance cutoff in Angstroms
        device: Device for tensors
    
    Returns:
        mask: Boolean tensor of shape [L] indicating which residues are within cutoff
    """
    if other_atoms is None or len(other_atoms) == 0:
        # No non-protein atoms, return all zeros
        return torch.zeros(protein_dict["mask"].shape[0], dtype=torch.bool, device=device)
    
    # Get CA coordinates (or CB if available)
    X = protein_dict["X"]  # [L, 4, 3] - N, CA, C, O
    CA = X[:, 1, :]  # CA atoms [L, 3]
    mask = protein_dict["mask"]  # [L] - valid residues
    
    ligand_coords = other_atoms.getCoords()  # [N_ligand, 3]
    
    if len(ligand_coords) == 0:
        return torch.zeros(protein_dict["mask"].shape[0], dtype=torch.bool, device=device)
    
    CA_tensor = torch.tensor(CA, device=device, dtype=torch.float32)
    ligand_coords_tensor = torch.tensor(ligand_coords, device=device, dtype=torch.float32)
    
    # Compute distances: [L, N_ligand]
    distances = torch.cdist(CA_tensor, ligand_coords_tensor)  # [L, N_ligand]
    
    min_distances = torch.min(distances, dim=1)[0]  # [L]
    
    # Mask: residues within cutoff AND have valid backbone atoms
    near_ligand_mask = (min_distances < cutoff) & (mask.bool())
    
    return near_ligand_mask


def load_model(
    checkpoint_path: str,
    device: str = "cpu"
) -> Tuple[ProteinMPNN, int]:
    """
    Load LigandMPNN model from checkpoint.
    
    Args:
        checkpoint_path: Path to model checkpoint
        device: Device for computation
    
    Returns:
        model: Loaded model
        atom_context_num: Number of ligand atoms to use
    """
    checkpoint = torch.load(checkpoint_path, map_location=device)
    atom_context_num = checkpoint["atom_context_num"]
    k_neighbors = checkpoint["num_edges"]
    ligand_mpnn_use_side_chain_context = 0  # Don't use side chain context for evaluation
    
    model = ProteinMPNN(
        node_features=128,
        edge_features=128,
        hidden_dim=128,
        num_encoder_layers=3,
        num_decoder_layers=3,
        k_neighbors=k_neighbors,
        device=device,
        atom_context_num=atom_context_num,
        model_type="ligand_mpnn",
        ligand_mpnn_use_side_chain_context=ligand_mpnn_use_side_chain_context,
    )
    
    model.load_state_dict(checkpoint["model_state_dict"])
    model.to(device)
    model.eval()
    
    return model, atom_context_num


def sample_sequences(
    model: ProteinMPNN,
    protein_dict: Dict,
    feature_dict: Dict,
    num_samples: int = 10,
    temperature: float = 0.1,
    device: str = "cpu"
) -> torch.Tensor:
    """
    Sample sequences from the model.
    
    Args:
        model: LigandMPNN model
        protein_dict: Protein dictionary from parse_PDB
        feature_dict: Feature dictionary from featurize
        num_samples: Number of sequences to sample
        temperature: Sampling temperature
        device: Device for computation
    
    Returns:
        S_stack: Stacked sequences [num_samples, L]
    """
    feature_dict["batch_size"] = 1
    feature_dict["temperature"] = temperature
    
    # No bias
    B, L, _, _ = feature_dict["X"].shape
    bias_AA = torch.zeros([21], device=device, dtype=torch.float32)
    feature_dict["bias"] = bias_AA[None, None, :].repeat([1, L, 1])
    feature_dict["symmetry_residues"] = [[]]
    feature_dict["symmetry_weights"] = [[]]
    
    S_list = []
    for _ in range(num_samples):
        feature_dict["randn"] = torch.randn(
            [feature_dict["batch_size"], feature_dict["mask"].shape[1]],
            device=device,
        )
        output_dict = model.sample(feature_dict)
        S_list.append(output_dict["S"])
    
    S_stack = torch.cat(S_list, 0)  # [num_samples, L]
    return S_stack


def process_pdb_file(
    pdb_path: str,
    model: ProteinMPNN,
    atom_context_num: int,
    cutoff: float = 5.0,
    num_samples: int = 10,
    temperature: float = 0.1,
    device: str = "cpu"
) -> Dict:
    """
    Process a single PDB file and compute sequence recovery for residues near ligands.
    
    Args:
        pdb_path: Path to PDB file
        model: LigandMPNN model
        atom_context_num: Number of ligand atoms to use
        cutoff: Distance cutoff in Angstroms
        num_samples: Number of sequences to sample
        temperature: Sampling temperature
        device: Device for computation
    
    Returns:
        Dictionary with results
    """
    results = {
        "pdb_path": pdb_path,
        "pdb_name": os.path.basename(pdb_path),
        "cutoff": cutoff,
        "has_ligands": False,
        "num_residues_near_ligand": 0,
        "sequence_recovery": None,
        "sequence_recovery_std": None,
        "individual_recoveries": [],
        "error": None
    }
    
    # Parse PDB
    protein_dict, backbone, other_atoms, icodes, water_atoms = parse_PDB(
        pdb_path,
        device=device,
        chains=[],
        parse_all_atoms=False,
        parse_atoms_with_zero_occupancy=False
    )
    
    # Check if there are non-protein atoms
    if other_atoms is not None and len(other_atoms) > 0:
        results["has_ligands"] = True
    
    native_seq = protein_dict["S"]
    
    # Get mask for residues near ligands
    near_ligand_mask = get_residues_near_ligands(
        protein_dict, other_atoms, cutoff=cutoff, device=device
    )
    
    num_near_ligand = torch.sum(near_ligand_mask).item()
    results["num_residues_near_ligand"] = num_near_ligand
    
    if num_near_ligand == 0:
        results["error"] = "No residues within cutoff distance of non-protein atoms"
        return results
    
    # Set up chain_mask to design all residues
    protein_dict["chain_mask"] = torch.ones_like(protein_dict["mask"])
    
    feature_dict = featurize(
        protein_dict,
        cutoff_for_score=cutoff,
        use_atom_context=True,
        number_of_ligand_atoms=atom_context_num,
        model_type="ligand_mpnn",
    )
    
    S_stack = sample_sequences(
        model, protein_dict, feature_dict,
        num_samples=num_samples,
        temperature=temperature,
        device=device
    )
    
    native_seq_batch = native_seq[None, :]  # [1, L]
    rec_mask = near_ligand_mask.int()[None, :]  # [1, L]
    
    individual_recoveries = []
    for i in range(num_samples):
        S_sample = S_stack[i:i+1, :]  # [1, L]
        recovery = get_seq_rec(native_seq_batch, S_sample, rec_mask)
        recovery_val = recovery[0].item()
        individual_recoveries.append(recovery_val)
    
    results["individual_recoveries"] = individual_recoveries
    results["sequence_recovery"] = np.mean(individual_recoveries)
    results["sequence_recovery_std"] = np.std(individual_recoveries)
    

    return results


def main():
    parser = argparse.ArgumentParser(
        description="Evaluate sequence recovery for residues near ligands using LigandMPNN",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    
    parser.add_argument( "--pdb_dir", type=str, required=True, help="Directory containing PDB files")
    parser.add_argument( "--checkpoint", type=str, default="./model_params/ligandmpnn_v_32_010_25.pt", help="Path to LigandMPNN checkpoint")
    parser.add_argument( "--cutoff", type=float, default=5.0, help="Distance cutoff in Angstroms for considering residues near ligands")
    parser.add_argument( "--num_samples", type=int, default=10, help="Number of sequences to sample per PDB")
    parser.add_argument( "--temperature", type=float, default=0.1, help="Sampling temperature")
    parser.add_argument( "--output", type=str, default=None, help="Output file path for results (JSON format). If not provided, prints to stdout")
    parser.add_argument( "--device", type=str, default=None, help="Device for computation (cpu or cuda). Default: auto-detect")
    parser.add_argument( "--seed", type=int, default=None, help="Random seed for reproducibility")
    
    args = parser.parse_args()
    
    # Set device
    if args.device is None:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    else:
        device = torch.device(args.device)
    
    # Set random seed
    if args.seed is not None:
        torch.manual_seed(args.seed)
        random.seed(args.seed)
        np.random.seed(args.seed)
    
    # Validate inputs
    if not os.path.exists(args.pdb_dir):
        print(f"Error: PDB directory does not exist: {args.pdb_dir}", file=sys.stderr)
        sys.exit(1)
    
    if not os.path.exists(args.checkpoint):
        print(f"Error: Checkpoint does not exist: {args.checkpoint}", file=sys.stderr)
        sys.exit(1)
    
    print(f"Loading model from {args.checkpoint}...", file=sys.stderr)
    model, atom_context_num = load_model(args.checkpoint, device=device)
    print(f"Model loaded. Using device: {device}", file=sys.stderr)
    
    pdb_files = list(Path(args.pdb_dir).glob("*.pdb"))
    if len(pdb_files) == 0:
        print(f"Error: No PDB files found in {args.pdb_dir}", file=sys.stderr)
        sys.exit(1)
    
    print(f"Found {len(pdb_files)} PDB files", file=sys.stderr)
    
    # Process each PDB file
    all_results = []
    for pdb_path in sorted(pdb_files):
        pdb_path_str = str(pdb_path)
        print(f"\nProcessing {os.path.basename(pdb_path_str)}...", file=sys.stderr)
        
        results = process_pdb_file(
            pdb_path_str,
            model=model,
            atom_context_num=atom_context_num,
            cutoff=args.cutoff,
            num_samples=args.num_samples,
            temperature=args.temperature,
            device=device
        )
        
        all_results.append(results)
        
        # Print summary
        if results["error"]:
            print(f"  Error: {results['error']}", file=sys.stderr)
        else:
            print(f"  Residues near ligands: {results['num_residues_near_ligand']}", file=sys.stderr)
            if results["sequence_recovery"] is not None:
                print(f"  Sequence recovery: {results['sequence_recovery']:.4f} ± {results['sequence_recovery_std']:.4f}", file=sys.stderr)
    
    # Compute aggregate statistics
    valid_results = [r for r in all_results if r["sequence_recovery"] is not None]
    
    if len(valid_results) > 0:
        recovery_rates = [r["sequence_recovery"] for r in valid_results]
        mean_recovery = np.mean(recovery_rates)
        std_recovery = np.std(recovery_rates)
        
        print("\n" + "="*60, file=sys.stderr)
        print("SUMMARY", file=sys.stderr)
        print("="*60, file=sys.stderr)
        print(f"PDB files processed: {len(all_results)}", file=sys.stderr)
        print(f"PDB files with valid recovery: {len(valid_results)}", file=sys.stderr)
        print(f"Mean sequence recovery (per PDB): {mean_recovery:.4f} ± {std_recovery:.4f}", file=sys.stderr)
        print(f"Cutoff distance: {args.cutoff} Å", file=sys.stderr)
        print(f"Number of samples per PDB: {args.num_samples}", file=sys.stderr)
        print("="*60, file=sys.stderr)
    
    # Write output if requested
    if args.output:
        import json
        output_data = {
            "cutoff": args.cutoff,
            "num_samples": args.num_samples,
            "temperature": args.temperature,
            "checkpoint": args.checkpoint,
            "results": all_results,
            "summary": {
                "total_pdbs": len(all_results),
                "valid_recovery": len(valid_results),
                "mean_recovery": float(mean_recovery) if len(valid_results) > 0 else None,
                "std_recovery": float(std_recovery) if len(valid_results) > 0 else None,
            }
        }
        with open(args.output, 'w') as f:
            json.dump(output_data, f, indent=2)
        print(f"\nResults written to {args.output}", file=sys.stderr)


if __name__ == "__main__":
    main()