### Install and Setup ProteinEBM

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Installs necessary packages, clones the ProteinEBM repository, and downloads pre-trained expert and base models from Hugging Face. Verifies that the models are loaded correctly.

```python
#@title Install ProteinEBM

import sys
import os

# Clone repository
if not os.path.exists('ProteinEBM'):
    os.system("pip install -q ml_collections biopython mdtraj py3Dmol")
    os.system("pip install -q git+https://github.com/sokrypton/py2Dmol.git")
    print("✓ Installed missing packages")
    os.system("git clone https://github.com/jproney/ProteinEBM.git")
    print("✓ Cloned ProteinEBM repository")
else:
    print("✓ ProteinEBM already cloned")

# Add to Python path (no pip install needed!)
sys.path.insert(0, '/content/ProteinEBM')

# Direct download from Hugging Face
expert_model_url = "https://huggingface.co/jproney/ProteinEBM/resolve/main/model_6_expert_frozen_1m_md.pt"
expert_model_path = "weights/model_6_expert_frozen_1m_md.pt"


if not os.path.exists(expert_model_path):
    os.makedirs('weights', exist_ok=True)
    print("Downloading model from Hugging Face...")
    !wget -q {expert_model_url} -O {expert_model_path}
    print(f"✓ Downloaded to {expert_model_path}")
else:
    print(f"✓ Model already exists")

# Direct download from Hugging Face
base_model_url = "https://huggingface.co/jproney/ProteinEBM/resolve/main/model_1_frozen_1m_md.pt"
base_model_path = "weights/model_1_frozen_1m_md.pt"


if not os.path.exists(base_model_path):
    os.makedirs('weights', exist_ok=True)
    print("Downloading model from Hugging Face...")
    !wget -q {base_model_url} -O {base_model_path}
    print(f"✓ Downloaded to {base_model_path}")
else:
    print(f"✓ Model already exists")

# Verify it works
try:
  if "model" not in dir():
      import torch
      import yaml
      import numpy as np
      from ml_collections import ConfigDict
      import py2Dmol
      from protein_ebm.model.r3_diffuser import R3Diffuser
      from protein_ebm.model.ebm import ProteinEBM
      from protein_ebm.model.boltz_utils import center_random_augmentation
      from protein_ebm.data.protein_utils import residues_to_features

      # Load config
      with open("ProteinEBM/protein_ebm/config/expert_model_config.yaml", 'r') as f:
          config = yaml.safe_load(f)
      config = ConfigDict(config)

      # Create model
      diffuser = R3Diffuser(config.diffuser)
      x_model = ProteinEBM(config.model, diffuser).cuda()

      # Load weights
      ckpt = torch.load(expert_model_path, weights_only=False, map_location='cuda')
      x_model.load_state_dict({k[len("model."):]: v for k, v in ckpt['state_dict'].items() if k.startswith('model')})
      x_model.eval()

      # Load config
      with open("ProteinEBM/protein_ebm/config/base_model_config.yaml", 'r') as f:
          config = yaml.safe_load(f)
      config = ConfigDict(config)


      base_model = ProteinEBM(config.model, diffuser).cuda()

      # Load weights
      ckpt = torch.load(base_model_path, weights_only=False, map_location='cuda')
      base_model.load_state_dict({k[len("model."):]: v for k, v in ckpt['state_dict'].items() if k.startswith('model')})
      base_model.eval()


  print("✓ ProteinEBM models loaded")
except ImportError as e:
    print(f"✗ Error: {e}")

```

--------------------------------

### Load Model and Score Structures

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Minimal example demonstrating how to load a ProteinEBM model, parse a protein from a PDB file, and compute its energy. Ensure you have the necessary configuration file.

```python
import torch
import yaml
from ml_collections import ConfigDict

from protein_ebm.model.r3_diffuser import R3Diffuser
from protein_ebm.data.protein_utils import residues_to_features, plot_protein_frame
from protein_ebm.model.ebm import ProteinEBM
from protein_ebm.model.boltz_utils import center_random_augmentation
import numpy as np

with open("protein_ebm/config/expert_model_config.yaml", 'r') as f:
    config = yaml.safe_load(f)
    
config = ConfigDict(config)

# Create models
diffuser = R3Diffuser(config.diffuser)
model = ProteinEBM(config.model, diffuser).cuda()
```

--------------------------------

### Load and Initialize R3Diffuser and ProteinEBM Models

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/diffusion_inference.ipynb

Loads a base model configuration from a YAML file and initializes the R3Diffuser and ProteinEBM models. Ensure the configuration file path is correct and the necessary libraries are installed.

```python
import torch
import yaml
from ml_collections import ConfigDict
from argparse import ArgumentParser
import os

from protein_ebm.model.r3_diffuser import R3Diffuser
from protein_ebm.data.protein_utils import residues_to_features, plot_protein_frame, restypes, restype_order, restype_num
from protein_ebm.model.ebm import ProteinEBM
from protein_ebm.model.boltz_utils import center_random_augmentation
import numpy as np



with open("../protein_ebm/config/base_model_config.yaml", 'r') as f:
    config = yaml.safe_load(f)
    
config = ConfigDict(config)

# Create models
diffuser = R3Diffuser(config.diffuser)
model = ProteinEBM(config.model, diffuser).cuda()

```

--------------------------------

### Diffusion Simulation Setup

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Sets up parameters for running a reverse diffusion simulation. This includes defining reverse time steps, calculating the time increment for reverse steps, and setting a small time increment for language-based steps. A frame counter is also initialized.

```python
reverse_times = np.linspace(T_MIN, T_MAX, REVERSE_STEPS)[::-1] # Updated to use T_START and T_END
dt_rev = (T_MAX - T_MIN) / REVERSE_STEPS # Updated to use T_START and T_END
dt_lang = 0.001
frame_count = 0
```

--------------------------------

### Install ProteinEBM

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Installs ProteinEBM in a new conda environment. Ensure you activate the environment before proceeding.

```bash
conda create --name protebm python=3.11
conda activate protebm

git clone https://github.com/jproney/ProteinEBM
cd ProteinEBM-public
pip install .
```

--------------------------------

### Train ProteinEBM Model

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Pretrain the ProteinEBM model using the downloaded data. Ensure the expert model configuration file is updated for your GPU setup before running.

```bash
cd protein_ebm/scripts
python train.py ../config/expert_model_config.yaml
```

--------------------------------

### Load Model Checkpoint and State Dict

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Loads a PyTorch model checkpoint and selectively loads the state dictionary, filtering keys that start with 'model.'. Ensure the model architecture matches the checkpoint.

```python
ckpt = torch.load("weights/model_6_expert_frozen_1m_md.pt", weights_only=False)
model.load_state_dict({k[len("model."):]: v for k, v in ckpt['state_dict'].items() if k.startswith('model')})
```

--------------------------------

### Initialize Previous Step Coordinates (prev0)

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Initializes `prev0`, which typically represents the denoised coordinates from the previous step in the diffusion process. It is set to `pos_t.clone()` if diffusion starts from a base structure (`T_MAX < 1.0` and `initial_coords_no_noise` is available), otherwise it's initialized to zeros.

```python
# Initialize prev0 based on instructions
if initial_coords_no_noise is not None and T_MAX < 1.0:
    prev0 = pos_t.clone() # Instruction: Initialize prev0 to pos_t.clone() if T_START < 1.0 and initial_coords_no_noise was used
else:
    prev0 = torch.zeros_like(pos_t) # Otherwise, keep prev0 = torch.zeros_like(pos_t)
```

--------------------------------

### Score Decoys with ProteinEBM Script

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Executes the `score_decoys.py` script to score Rosetta decoys using a specified configuration file, model weights, and data lists. This command requires prior setup including downloading evaluation data and weights.

```bash
cd protein_ebm/scripts
python score_decoys.py ../config/expert_model_config.yaml ../../weights/model_6_expert_frozen_1m_md.pt ../data/data_lists/validation_decoys.txt ../../eval_data/model_6_val_decoy_scores.pt --decoy_dir ../../eval_data/decoys/ --n_samples 16 --t_max 0.15 --bsize 64 --template_self_condition
```

--------------------------------

### Run Fast-Folder Simulation

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Execute a fast-folder simulation using a combination of base and expert models. Ensure all specified paths and configurations are correctly set.

```bash
python run_dynamics.py --pdb_file ../../eval_data/fastfolders/experimental_structures/chignolin_cln025.pdb \
--config ../config/base_model_config.yaml --checkpoint ../../weights/model_1_frozen_1m_md.pt \
--lo_time_config ../config/expert_model_config.yaml --lo_time_checkpoint ../../weights/model_6_expert_frozen_1m_md.pt \
--lo_time_threshold 0.1 --steps 100 --min_steps 0 \
--t_min .01 --t_max 1.0 --ramp_start 0.5 --step_function_ramp --dt .001 --reverse_steps 200 --total_samples 400 \
--temp_scaling 0.85714 --temp_scaling_after_lo_time 1.0 --scoring_time 0.05 --use_aux_score_initial --batch_size 400 --log_dir ../../dynamics/ \
--experiment_name chignolin_dynamics
```

--------------------------------

### Initialize R3Diffuser and ProteinEBM Models

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/confbiasing.ipynb

Initializes the R3Diffuser and ProteinEBM models using the loaded configuration. The ProteinEBM model is moved to the CUDA-enabled GPU for accelerated computation. Ensure CUDA is available.

```python
# Create models
diffuser = R3Diffuser(config.diffuser)
model = ProteinEBM(config.model, diffuser).cuda()
```

--------------------------------

### Download Pretraining Data

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Download pretraining data for ProteinEBM. This script populates the weights/training_data directory with necessary data files.

```bash
cd download_scripts
./download_data.sh
```

--------------------------------

### Load Protein EBM Model Configuration and Weights

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/ddg_prediction.ipynb

Loads the configuration from a YAML file and then loads the pre-trained weights for the ProteinEBM model. Ensure the configuration file and checkpoint path are correct.

```python
import torch
import yaml
from ml_collections import ConfigDict
from argparse import ArgumentParser
import os

from protein_ebm.model.r3_diffuser import R3Diffuser
from protein_ebm.data.protein_utils import residues_to_features, plot_protein_frame, restypes, restype_order, restype_num
from protein_ebm.model.ebm import ProteinEBM
from protein_ebm.model.boltz_utils import center_random_augmentation
import numpy as np



with open("../protein_ebm/config/expert_model_config.yaml", 'r') as f:
    config = yaml.safe_load(f)
    
    
config = ConfigDict(config)

# Create models
diffuser = R3Diffuser(config.diffuser)
model = ProteinEBM(config.model, diffuser).cuda()


ckpt = torch.load("../weights/model_6_expert_frozen_1m_md.pt", weights_only=False)


model.load_state_dict({k[len("model."):]: v for k, v in ckpt['state_dict'].items() if k.startswith('model')})
```

--------------------------------

### Run Direct Folding Simulation

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Execute a direct folding simulation. This simulation is configured for extended steps and specific time parameters to observe the folding process.

```bash
python run_dynamics.py --experiment_name nug2_folding_sim --pdb_file ../../eval_data/fastfolders/experimental_structures/proteing_1mi0.pdb \
--config ../config/expert_model_config.yaml --checkpoint ../../weights/model_6_expert_frozen_1m_md.pt \
--dt .001 --steps 10000 --reverse_steps 1 --t_min 0.05 --t_max 0.05 --scoring_time 0.05 --save_stride 10 \
 --total_samples 400 --batch_size 400 --use_aux_score_initial --start_unfolded --num_resample_rounds 1 --log_dir ../../dynamics/ 
```

--------------------------------

### Load and Parse Configuration

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/confbiasing.ipynb

Loads a YAML configuration file and parses it into an ML_collections ConfigDict for easy access to model parameters. Ensure the config file path is correct.

```python
import torch
import yaml
from ml_collections import ConfigDict
from argparse import ArgumentParser
import os

from protein_ebm.model.r3_diffuser import R3Diffuser
from protein_ebm.data.protein_utils import residues_to_features, plot_protein_frame
from protein_ebm.model.ebm import ProteinEBM
from protein_ebm.model.boltz_utils import center_random_augmentation
import numpy as np



with open("../protein_ebm/config/expert_model_config.yaml", 'r') as f:
    config = yaml.safe_load(f)
    
config = ConfigDict(config)
```

--------------------------------

### Prepare Input Features for ProteinEBM

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Prepares a dictionary of input features for the ProteinEBM model, including centered coordinates, amino acid types, masks, residue indices, and diffusion time. Coordinates are moved to CUDA.

```python
nres = atom_positions.shape[0]
ca_coords = center_random_augmentation(atom_positions[...,1,:].unsqueeze(0), torch.ones([1, nres])).view([1,nres,3])

t=0.05
input_feats = {
    'r_noisy': ca_coords.cuda(), # coordinates
    'aatype': aatype.unsqueeze(0).cuda(), # amino acid types
    'mask': torch.ones([1, nres]).cuda(), # amino acid mask (for multiple different-length proteins)
    'residue_idx': residue_idx.unsqueeze(0).cuda(), # residue indices
    't': torch.tensor([t], dtype=torch.float).cuda(), # diffusion time (set to 0.1 for scoring)
    'selfcond_coords' : ca_coords.cuda() # optional self-conditioning coordinate channel
}
```

--------------------------------

### Download Model Weights

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Downloads pre-trained model weights for ProteinEBM. The script downloads parameters for 6 different models, each suited for specific applications.

```bash
cd download_scripts
./download_weights.sh
```

--------------------------------

### ProteinEBM Configuration Parameters

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Set these parameters to control the simulation type, number of steps, and temperature range for trajectory generation. Adjust `LANGEVIN_PER_LEVEL` for direct folding and `REVERSE_STEPS` for reverse diffusion.

```python
SEQUENCE = "MQIFVKTLTGKTITLEVEPSDTIENVKAKIQDKEGIPPDQQRLIFAGKQLEDGRTLSDYNIQKESTLHLVLRLRGG"  # @param {type:"string"}
PDB_ID = "1MI0" # @param {type:"string"}
REVERSE_STEPS = 1 # @param ["1","50","100","200","400","800"] {"type":"raw"}
LANGEVIN_PER_LEVEL = 20000 # @param {"type":"raw"}
start_unfolded = True # @param {"type":"boolean"}
use_aux_score = True # @param {"type":"boolean"}
T_MAX = .05 # @param {"type":"raw"}
T_MIN = .05 # @param {"type":"raw"}
```

--------------------------------

### Initialize 2Dmol Viewer

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Initializes a 2Dmol viewer for scatter plot visualization. It is configured to detect cyclic properties and manage persistence. The viewer is set up with labels for the x and y axes representing 'Step' and 'Energy' respectively.

```python
viewer = py2Dmol.view(scatter=True,
                      detect_cyclic=False,
                      persistence=False)
viewer.new_obj(scatter_config={"xlabel": "Step", "ylabel": "Energy"})
viewer.show()
```

--------------------------------

### Initialize Diffusion Coordinates (pos_t)

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Sets the initial coordinates for the diffusion process (`pos_t`). If `T_MAX` is 1.0, it uses random noise. If `T_MAX` is less than 1.0, it applies forward diffusion to the `initial_coords_no_noise`. Handles `T_MAX` > 1.0 as a warning and defaults to random noise.

```python
# --- Conditional logic for pos_t initialization based on T_MAX ---

if T_MAX == 1.0:
    pos_t = torch.randn(1, nres, 3).cuda() * diffuser.config.coordinate_scaling
    # Ensure pos_t is centered even for random noise start
    pos_t = center_random_augmentation(
        pos_t.reshape([1, -1, 3]),
        torch.ones([1, nres], device=pos_t.device),
        rotate=False
    ).view([1, -1, 3])

elif T_MAX < 1.0:
    if initial_coords_no_noise is None:
        raise ValueError("A base structure (from PDB or unfolded sequence) must be provided when T_START < 1.0.")

    # Apply forward marginal diffusion to the base structure to reach T_MAX
    # Convert initial_coords_no_noise to numpy array on CPU before passing to diffuser
    pos_t_numpy = diffuser.forward_marginal(initial_coords_no_noise.cpu().numpy(), T_MAX)
    pos_t = torch.from_numpy(pos_t_numpy[0]).float().cuda() # Access the first element of the tuple and cast to float32

    # Ensure pos_t is centered after applying noise
    pos_t = center_random_augmentation(
        pos_t.reshape([1, -1, 3]),
        torch.ones([1, nres], device=pos_t.device),
        rotate=False
    ).view([1, -1, 3])
else:
    # Fallback for unexpected T_START > 1.0, treat as random start
    print("Warning: T_MAX > 1.0 is unusual. Starting with random noise.")
    pos_t = torch.randn(1, nres, 3).cuda() * diffuser.config.coordinate_scaling
    pos_t = center_random_augmentation(
        pos_t.reshape([1, -1, 3]),
        torch.ones([1, nres], device=pos_t.device),
        rotate=False
    ).view([1, -1, 3])
```

--------------------------------

### Load Checkpoint and Visualize Protein

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/diffusion_inference.ipynb

Loads a diffusion model checkpoint and visualizes the resulting protein structure. Requires PyTorch and Py3Dmol. The code processes coordinates and applies an atom mask for visualization.

```python
%matplotlib notebook
non_ca = torch.tensor(all_pos[-1],dtype=torch.float).unsqueeze(-2) +  out['sidechain_coords'].squeeze().cpu()
pos_allatom = torch.cat([non_ca[...,:1,:], torch.tensor(all_pos[-1], dtype=torch.float).unsqueeze(-2), non_ca[...,1:,:]], dim=-2)

plot_protein_frame(pos_allatom[0], atom_mask, cartoon=True) # remove batch dim
```

--------------------------------

### Prepare Input Features for CUDA

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Ensures all input features are on CUDA and correctly shaped for model input. Raises an error if feature generation fails.

```python
if aatype_gen is not None:
    aatype_gen_gpu = aatype_gen.unsqueeze(0).cuda()
    atom_mask_gen_gpu = atom_mask_gen.unsqueeze(0).cuda()
    residue_idx_gen_gpu = residue_idx_gen.unsqueeze(0).cuda()
    residue_mask_gen_gpu = residue_mask_gen.unsqueeze(0).cuda()
    chain_encoding_gen_gpu = chain_encoding_gen.unsqueeze(0).cuda()
else:
    # This block should ideally not be reached if PDB/sequence parsing is successful
    raise RuntimeError("Feature generation failed, aatype_gen is None.")
```

--------------------------------

### Build Decoy Datasets

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Parses the decoy dataset into the tensor data format used by ProteinEBM. This populates the 'eval_data/decoys' directory.

```python
cd protein_ebm/data/data_scripts
python build_decoy_sets.py
```

--------------------------------

### Load Protein Dynamics Trajectory

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/analyze_dynamics.ipynb

Loads a protein dynamics trajectory from a .pt file using PyTorch. This is the initial step to access the trajectory data.

```python
import torch

#saved = torch.load("../dynamics/chignolin_dynamics_20260305_154302/dynamics_trajectory.pt")
#saved = torch.load("../dynamics/2chf_structure_prediction_20260305_215519/dynamics_trajectory.pt")
saved = torch.load("../dynamics/nug2_folding_sim_20260305_165239/dynamics_trajectory.pt")

print(saved.keys())
```

--------------------------------

### Load Model Checkpoint

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/diffusion_inference.ipynb

Loads a PyTorch model checkpoint, filtering and loading the state dictionary for the 'model' component.

```python
ckpt = torch.load("../weights/model_1_frozen_1m_md.pt", weights_only=False)


model.load_state_dict({k[len("model."):]: v for k, v in ckpt['state_dict'].items() if k.startswith('model')})
```

--------------------------------

### Initialize Protein Features from Sequence

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Converts an amino acid sequence into protein features when no PDB ID is provided or PDB loading fails. If `start_unfolded` is true, it generates unfolded backbone coordinates.

```python
if not PDB_ID: # This block will execute if PDB_ID was empty initially or if PDB loading failed
    print(f"Using sequence-based generation for protein: {SEQUENCE}.")
    aatype_gen, atom_mask_gen, residue_idx_gen, residue_mask_gen, chain_encoding_gen = seq_to_features(SEQUENCE)
    nres = len(aatype_gen)

    if start_unfolded:
        unfolded_coords, log_prob = generate_random_backbone_coords(SEQUENCE, uniform_sampling=True)
        unfolded_ca = unfolded_coords[:, 1, :].clone().unsqueeze(0)  # CA atoms only, [1, N, 3]
        # Center the structure and store as initial_coords_no_noise
        initial_coords_no_noise = center_random_augmentation(
            unfolded_ca,
            torch.ones([1, nres], device=unfolded_ca.device),
            rotate=False
        ).cuda().view([1, -1, 3])
```

--------------------------------

### Download Evaluation Data

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Downloads evaluation data for various protein structure analysis tasks. This script should be run from the 'download_scripts' directory.

```bash
cd download_scripts
./download_eval_data.sh
```

--------------------------------

### Compare Mean Top-1 Accuracy of Different Methods

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/rank_decoys.ipynb

Prints the mean top-1 accuracy for Rosetta and RG methods.

```python
print(torch.tensor([x for x in rosetta_top1s.values()]).mean())
print(torch.tensor([x for x in rg_top1s.values()]).mean())
```

--------------------------------

### Initialize Protein Features from PDB

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Loads protein features (amino acid types, atom masks, residue indices) from a PDB file. Handles potential errors during PDB loading and falls back to sequence-based generation if necessary. Optionally generates unfolded backbone coordinates if `start_unfolded` is true.

```python
if PDB_ID:
    print(f'PDB ID {PDB_ID} provided. Loading from PDB.')
    try:
        pdb_file_path = download_pdb(PDB_ID)
        aatype_gen, atom_mask_gen, residue_idx_gen, residue_mask_gen, chain_encoding_gen, ca_coords_initial = parse_pdb_and_get_features(pdb_file_path)
        nres = len(aatype_gen)

        # Create a reverse mapping from index to amino acid symbol for sequence generation
        rev_restype_order = {v: k for k, v in restype_order.items()}
        pdb_sequence = "".join([rev_restype_order.get(idx.item(), 'X') for idx in aatype_gen])

        if start_unfolded:
            # If start_unfolded is True, generate unfolded coordinates based on PDB sequence
            print(f"Using unfolded initialization from PDB sequence (derived from {PDB_ID}).")
            unfolded_coords, log_prob = generate_random_backbone_coords(pdb_sequence, uniform_sampling=True)
            unfolded_ca = unfolded_coords[:, 1, :].clone().unsqueeze(0)  # CA atoms only, [1, N, 3]
            initial_coords_no_noise = center_random_augmentation(
                unfolded_ca,
                torch.ones([1, nres], device=unfolded_ca.device),
                rotate=False
            ).cuda().view([1, -1, 3])
        else:
            # Otherwise, use coordinates directly from PDB
            ca_coords_initial_tensor = ca_coords_initial.unsqueeze(0).cuda()
            initial_coords_no_noise = center_random_augmentation(
                ca_coords_initial_tensor,
                torch.ones([1, nres], device=ca_coords_initial_tensor.device),
                rotate=False
            ).view([1, -1, 3])
            print(f"Loaded PDB {PDB_ID} with {nres} residues.")

    except Exception as e:
        print(f"Error loading PDB {PDB_ID}: {e}")
        print("Falling back to sequence-based generation.")
        PDB_ID = "" # Reset PDB_ID to trigger sequence logic below
```

--------------------------------

### Visualize Protein Structure at Minimum Energy

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/analyze_dynamics.ipynb

Renders a 3D visualization of the protein structure at a specific time point corresponding to the minimum energy trajectory. Requires py3Dmol for visualization.

```python
plot_protein_frame(tseries[700,min_nrg_traj], atom_mask, cartoon=True)
```

--------------------------------

### Download PDB File

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Downloads a PDB file from the RCSB PDB database if it does not already exist locally. Ensures the PDB file is available for parsing.

```python
import os

def download_pdb(pdb_id, pdb_dir="pdb_files"):
    os.makedirs(pdb_dir, exist_ok=True)
    pdb_file = os.path.join(pdb_dir, f'{pdb_id.lower()}.pdb')
    if not os.path.exists(pdb_file):
        print(f"Downloading PDB {pdb_id}...")
        url = f"https://files.rcsb.org/download/{pdb_id.upper()}.pdb"
        response = requests.get(url)
        response.raise_for_status()
        with open(pdb_file, 'w') as f:
            f.write(response.text)
        print(f"Downloaded {pdb_id} to {pdb_file}")
    else:
        print(f"PDB {pdb_id} already exists at {pdb_file}")
    return pdb_file
```

--------------------------------

### Load Decoy Data

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/rank_decoys.ipynb

Loads decoy scores from a PyTorch file. Handles ensemble analysis by loading a second dataset or using the first if the second is not provided.

```python
import torch
from scipy.stats import spearmanr
import matplotlib.pyplot as plt

decoy_data_path = "../eval_data/model_6_val_decoy_scores.pt"
decoy_data_path_2 = None # For ensemble analysis

decoy_data = torch.load(decoy_data_path, weights_only=False, map_location="cpu")

if decoy_data_path_2 is not None:
    decoy_data2 = torch.load(decoy_data_path_2, weights_only=False, map_location="cpu")
else:
    decoy_data2 = decoy_data
```

--------------------------------

### Langevin Annealing and Reverse Diffusion Loop

Source: https://github.com/sokrypton/proteinebm/blob/main/proteinebm_diffusion.ipynb

Iterates through reverse diffusion time steps, applying Langevin annealing and reverse diffusion updates. Selects the appropriate model based on time.

```python
with torch.no_grad():

    for time_idx, t in enumerate(reverse_times):

        if t < .1:
          model = x_model
        else:
          model = base_model

        # LANGEVIN ANNEALING at this time level
        for lang_step in range(LANGEVIN_PER_LEVEL):

            input_feats = {
                'r_noisy': pos_t,
                'aatype': aatype_gen_gpu,
                'mask': residue_mask_gen_gpu,
                'residue_idx': residue_idx_gen_gpu,
                't': torch.tensor([t], dtype=torch.float).cuda(),
                'chain_encoding': chain_encoding_gen_gpu,
            }

            if use_aux_score:
              out = model.compute_energy(input_feats)
              score = out['r_update_aux']
              energy = out['energy']
              pred_coords = out['pred_coords_aux']
            else:
              out = model.compute_score(input_feats)
              score = out['trans_score']
              energy = out['energy']
              pred_coords = out['pred_coords']

            # Backward
            drift = diffuser.drift_coef(pos_t, t)
            diffusion = diffuser.diffusion_coef(t)
            pos_t = pos_t - (drift - diffusion**2 * score  / diffuser.config.coordinate_scaling
                            ) * dt_lang + diffusion * np.sqrt(dt_lang) * torch.randn_like(pos_t) / diffuser.config.coordinate_scaling

            # Center
            pos_t, prev0 = center_random_augmentation(
                pos_t.reshape([1, -1, 3]), torch.ones([1, nres], device=pos_t.device),
                rotate=False, second_coords=pred_coords.reshape([1, -1, 3]),
                return_second_coords=True
            )
            pos_t, prev0 = pos_t.view([1, -1, 3]), prev0.view([1, -1, 3])

            # Forward
            pos_t = pos_t + diffuser.drift_coef(pos_t, t - dt_lang) * dt_lang + \
                    diffuser.diffusion_coef(t - dt_lang) * np.sqrt(dt_lang) * \
                    torch.randn_like(pos_t) / diffuser.config.coordinate_scaling

            # ADD FRAME for every Langevin step
            viewer.add(prev0.cpu().squeeze(0).numpy(), scatter=[frame_count, energy.cpu().item()])
            frame_count += 1

        # REVERSE DIFFUSION STEP
        input_feats = {
            'r_noisy': pos_t,
            'aatype': aatype_gen_gpu,
            'mask': residue_mask_gen_gpu,
            'residue_idx': residue_idx_gen_gpu,
            't': torch.tensor([t], dtype=torch.float).cuda(),
            'chain_encoding': chain_encoding_gen_gpu,
            'atom_mask': atom_mask_gen_gpu
        }

        if use_aux_score:
          out = model.compute_energy(input_feats)
          score = out['r_update_aux']
          energy = out['energy']
        else:
          out = model.compute_score(input_feats)
          score = out['trans_score']
          energy = out['energy']

        pos_t = pos_t - (diffuser.drift_coef(pos_t, t) -
                        diffuser.diffusion_coef(t)**2 * score / diffuser.config.coordinate_scaling
                       ) * dt_rev + diffuser.diffusion_coef(t) * np.sqrt(dt_rev) * torch.randn_like(pos_t) / diffuser.config.coordinate_scaling

        pos_t, prev0 = center_random_augmentation(
            pos_t.reshape([1, -1, 3]), torch.ones([1, nres], device=pos_t.device),
            rotate=False, second_coords=pred_coords.reshape([1, -1, 3]),
            return_second_coords=True
        )
        pos_t, prev0 = pos_t.view([1, -1, 3]), prev0.view([1, -1, 3])

        # ADD FRAME for every reverse step
        viewer.add(prev0.cpu().squeeze(0).numpy(), scatter=[frame_count, energy.cpu().item()], align=False)
        frame_count += 1
```

--------------------------------

### Load and Display CSV Data

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/confbiasing.ipynb

Reads protein evaluation data from a CSV file into a pandas DataFrame and displays the first few rows.

```python
import pandas as pd

lpla_dat = pd.read_csv('../eval_data/lpla.csv')
lpla_dat
```

--------------------------------

### Run Two-Stage Structure Prediction

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Perform a two-stage structure prediction run. This involves resampling dynamics and requires specific configuration for both the base and resampling models.

```bash
cd protein_ebm/scripts
python run_dynamics.py --pdb_file ../../eval_data/decoys/natives/2chf.pdb --config ../config/base_model_config.yaml \
 --checkpoint ../../weights/model_4_unfrozen_3m_md.pt --resample_dynamics_config ../config/expert_model_config.yaml\
 --resample_dynamics_checkpoint ../../weights/model_5_expert_pretrained.pt \
 --steps 100 --resample_steps 10 --min_steps 0 --ramp_start 0.5 --step_function_ramp --dt .001 --reverse_steps 200 \
 --resample_reverse_steps 20 --t_min .01 --t_max 1.0 --total_samples 400 --resample_total_samples 800 --temp_scaling 0.85714 \
 --resample_temp_scaling 1.0 --resample_noise_time 0.1 --scoring_time 0.05 --num_resample_rounds 4 --use_aux_score_initial \
 --batch_size 100 --resample_batch_size 10 --log_dir ../../dynamics/ --experiment_name 2chf_structure_prediction
```

--------------------------------

### Parse PDB Structure

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/confbiasing.ipynb

Initializes a PDB parser and loads protein structures from specified PDB files. It then extracts the first chain from each structure.

```python
from Bio.PDB import PDBParser
from Bio.PDB.Polypeptide import is_aa

parser = PDBParser(QUIET=True)

open_lpla = '../eval_data/confbiasing/3a7r.pdb'
closed_lpla = '../eval_data/confbiasing/1x2g.pdb'

structs = []
t = 0.05
for pdb_path in [open_lpla, closed_lpla]:
    structure = parser.get_structure("my_structure", pdb_path) 

    
    chain = list(structure.get_chains())[0]
```

--------------------------------

### Find Threshold Transitions in Time Series Data

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/analyze_dynamics.ipynb

Identifies transitions between high and low states in time-series data. Use this function to detect when a signal crosses predefined thresholds.

```python
import numpy as np

HIGH = 10
LOW = 4

def find_threshold_transitions(X):
    T, N = X.shape

    # Regime masks
    high = X > HIGH
    low  = X < LOW

    # Detect entries and exits
    enter_high = np.vstack([high[0:1],  high[1:] & ~high[:-1]])
    exit_high  = np.vstack([np.zeros((1,high.shape[1]), dtype=bool), ~high[1:] & high[:-1]])

    enter_low  = np.vstack([low[0:1],   low[1:] & ~low[:-1]])
    exit_low   = np.vstack([np.zeros((1,high.shape[1]), dtype=bool),  ~low[1:] & low[:-1]])

    hi_to_lo = []
    lo_to_hi = []

    for col in range(N):

        # Collect event times and labels
        # type codes:
        # 0 = exit_high
        # 1 = enter_low
        # 2 = exit_low
        # 3 = enter_high

        times = np.concatenate([
            np.where(exit_high[:, col])[0],
            np.where(enter_low[:, col])[0],
            np.where(exit_low[:, col])[0],
            np.where(enter_high[:, col])[0],
        ])

        types = np.concatenate([
            np.zeros(np.sum(exit_high[:, col]), dtype=int),
            np.ones(np.sum(enter_low[:, col]), dtype=int),
            np.full(np.sum(exit_low[:, col]), 2, dtype=int),
            np.full(np.sum(enter_high[:, col]), 3, dtype=int),
        ])

        if len(times) == 0:
            continue

        # Sort events by time
        order = np.argsort(times)
        times = times[order]
        types = types[order]

        # Look for consecutive valid pairs
        for i in range(len(times) - 1):

            t1, t2 = times[i], times[i + 1]
            e1, e2 = types[i], types[i + 1]

            # exit_high -> enter_low
            if e1 == 0 and e2 == 1:
                hi_to_lo.append([t1, t2, col])

            # exit_low -> enter_high
            if e1 == 2 and e2 == 3:
                lo_to_hi.append([t1, t2, col])


    return (
        np.array(hi_to_lo, dtype=int),
        np.array(lo_to_hi, dtype=int)
    )


hi_to_lo, lo_to_hi = find_threshold_transitions(rms_bybatch.numpy())
```

--------------------------------

### Compute Energy with ProteinEBM

Source: https://github.com/sokrypton/proteinebm/blob/main/README.md

Computes the energy of the input features using the ProteinEBM model within a `torch.no_grad()` context to disable gradient calculations. Prints the computed energy.

```python
with torch.no_grad():
    out = model.compute_energy(input_feats)

print(out['energy'])
```

--------------------------------

### Load ProteinEBM Energy Data

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/confbiasing.ipynb

Assigns calculated open and closed energy values to new columns in a filtered DataFrame. Prints the resulting DataFrame to display the loaded data.

```python
lpla_dat_filtered["ProteinEBM Open Energy"] = open_nrg
lpla_dat_filtered["ProteinEBM Closed Energy"] = closed_nrg
print(lpla_dat)
```

--------------------------------

### Analyze Decoy Scores and Calculate Metrics

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/rank_decoys.ipynb

Iterates through decoy data to calculate Spearman correlations between different scoring functions (energy, Rosetta, radius of gyration) and TM-scores. It also identifies top-ranked models based on TM-score and stores various metrics for further analysis.

```python
import numpy as np

summary = {}
all_spearmans = {}
all_tms = {}
rosetta_spearmans = {}
rg_spearmans = {}
rgs = {}
all_top1s = {}
rosetta_top1s = {}
rg_top1s = {}


for pdb in decoy_data.keys():
    print(pdb)
    decoys = torch.load(f"../eval_data/decoys/{pdb}.pt", weights_only=False)

    spearmans = []
    top1s = []
    for i,(nrg1, nrg2) in enumerate(zip(decoy_data[pdb], decoy_data2[pdb])):

        t = i/20
        nrg = nrg1 + nrg2 
            
        spearmans.append(spearmanr(nrg.squeeze().numpy(), -decoys['tmscore'].numpy())[0])
        
        top1s.append(decoys['tmscore'][nrg.argmin()])

    rg = torch.sqrt((decoys['atom37'][...,:3,:].reshape([decoys['atom37'].shape[0], -1, 3]) - decoys['atom37'][...,:3,:].reshape([decoys['atom37'].shape[0], -1, 3]).mean(dim=-2, keepdim=True)).pow(2).sum(dim=-1).mean(dim=-1))
    rgs[pdb] = rg
    rg_spearmans[pdb] = spearmanr(rg, -decoys['tmscore'].numpy())[0]
    rosetta_spearmans[pdb] = spearmanr(decoys['rosetta'].numpy(), -decoys['tmscore'].numpy())[0]

    rosetta_top1s[pdb] = decoys['tmscore'][decoys['rosetta'].argmin()]
    rg_top1s[pdb] = decoys['tmscore'][rg.argmin()]


    all_tms[pdb] = decoys['tmscore']
    all_spearmans[pdb] = spearmans
    all_top1s[pdb] = top1s

    
    eval_idx = 2 
    print(spearmans[eval_idx])
    print(rosetta_spearmans[pdb])
    print(rg_spearmans[pdb])

    summary[pdb] = spearmans[eval_idx]
```

--------------------------------

### Load and Filter Megascale Stability Datasets

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/ddg_prediction.ipynb

Loads megascale stability datasets from a specified path, filtering for files containing 'Tsuboyama' and excluding them from the 'stab_files' list. It then reads these filtered CSV files into a list of pandas DataFrames.

```python
megascale_files = list(reversed([x for x in os.listdir(protgym_path) if x in stab_files and 'Tsuboyama' in x]))
dats = [pd.read_csv(protgym_path + x) for x in megascale_files]
print(len(megascale_files))
print(megascale_files)
```

--------------------------------

### Analyze RMSD and Index for Minimum Energy Trajectory

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/analyze_dynamics.ipynb

Extracts and prints the minimum RMSD value and its corresponding index for a given trajectory (identified by minimum energy). Plots the RMSD over time for this trajectory.

```python
i=min_nrg_traj
print(rms_bybatch[:,i].min())
print(rms_bybatch[:,i].argmin())
plt.plot(rms_bybatch[:,i])
```

--------------------------------

### Load and Filter ProteinGym DMS Data

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/ddg_prediction.ipynb

Loads protein substitution data from a CSV file and filters it to select entries related to protein stability. This prepares the data for subsequent analysis or evaluation.

```python
import os
import pandas as pd

gymdat = pd.read_csv("../eval_data/proteingym/DMS_substitutions.csv")

stab_files = set(gymdat[gymdat['coarse_selection_type'] == 'Stability']['DMS_filename'])


protgym_path = "../eval_data/proteingym/DMS_ProteinGym_substitutions/"
```

--------------------------------

### Diffusion Inference Loop

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/diffusion_inference.ipynb

Performs the reverse diffusion process to generate protein structures. This involves iterating through reverse time steps, computing score predictions, and applying the reverse diffusion step. Requires PyTorch and NumPy.

```python
num_t = 300
t_max = 1.0
min_t = .01
reverse_steps = np.linspace(min_t, t_max, num_t)[::-1]
dt = (t_max - min_t)/num_t


r_noisy, trans_score = diffuser.forward_marginal(
    ca_coords.numpy(),
    t_max)

r_noisy = torch.tensor(r_noisy, dtype=torch.float)

pos_t = r_noisy.cpu().numpy()
all_pos = [pos_t]

align_steps=False
self_condition=True
aux_score=False
prev0 = torch.zeros(pos_t.shape).cuda()

with torch.no_grad():
    for t in reverse_steps:
        if t > min_t:
            print(t)
            # Set timestep features

            input_feats = {
                'r_noisy': torch.tensor(pos_t, dtype=torch.float).cuda(),
                'aatype': aatype.unsqueeze(0).cuda(),
                'mask': residue_mask.unsqueeze(0).cuda(),
                'residue_idx': residue_idx.unsqueeze(0).cuda(),
                't': torch.tensor([t], dtype=torch.float).cuda(),
                'selfcond_coords' : prev0,
                'atom_mask' : atom_mask.unsqueeze(0).cuda()
            }
            
            #Get model predictions
            out = model.compute_score(input_feats)

            if aux_score:
                score = out['r_update_aux'].cpu().numpy()
            else:
                score = out['trans_score'].cpu().numpy()

            # Reverse diffusion step
            pos_t = diffuser.reverse(
                x_t=pos_t,
                score_t=score,
                mask=residue_mask.unsqueeze(0).numpy(),
                t=t,
                dt=dt,
                center=False,
                noise_scale=1.0,
            )

            # rotate
            pos_t, prev0_new = center_random_augmentation(torch.tensor(pos_t, dtype=torch.float).cuda(), torch.ones([1, nres]).cuda(), rotate=False, second_coords=out['pred_coords'].reshape([1,-1,3]),return_second_coords=True)
            pos_t = pos_t.cpu().numpy()
            
            if self_condition:
                prev0 = prev0_new

            
            all_pos.append(pos_t)

all_pos.append(out['pred_coords'].detach().cpu().numpy())
```

--------------------------------

### Calculate Mutation Energies

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/confbiasing.ipynb

Iterates through a filtered dataset of protein mutations, applies each mutation to a set of pre-computed features, and calculates the energy difference between open and closed states. Requires pre-loaded structures and a trained model.

```python
from protein_ebm.data.protein_utils import restypes, restype_order

open_nrg = []
closed_nrg = []
native_nrgs = []

lpla_dat_filtered = lpla_dat[lpla_dat['Promiscuous Activity'] == lpla_dat['Promiscuous Activity']]

for i,r in lpla_dat_filtered.iterrows():
    mut = r['Mutant (WT Context)']
    aa1 = mut[0]
    aa2 = mut[-1]
    idx = int(mut[1:-1])
    print(mut)

    nrgs = []

    for input_feats in structs:
        aatype_mut = input_feats['aatype'].clone()

        if not torch.any(input_feats['residue_idx'] == idx):
            nrgs.append(float('nan'))
            print("Index not present in one or more structures!")
            continue
            
        
        seq_idx = (input_feats['residue_idx'] == idx).nonzero()[0,-1]


        assert restypes[aatype_mut[0,seq_idx]] == aa1

        aatype_mut[:,seq_idx] = restype_order[aa2]

            
        mut_feats = {**input_feats, 'aatype' : aatype_mut} 

        with torch.no_grad():
            out = model.compute_energy(mut_feats)
            
        nrgs.append(out['energy'].mean().item())

    open_nrg.append(nrgs[0])
    closed_nrg.append(nrgs[1])

    print(f"Row {i} Open {open_nrg[-1]}, Closed {closed_nrg[-1]}")
```

--------------------------------

### Load and Process Protein Structure

Source: https://github.com/sokrypton/proteinebm/blob/main/notebooks/diffusion_inference.ipynb

Loads a PDB file using Bio.PDB, extracts residue features, and prepares them for model input. Requires Bio.PDB and PyTorch.

```python
from Bio.PDB import PDBParser
from Bio.PDB.Polypeptide import is_aa

#Load the structure
parser = PDBParser(QUIET=True)
structure = parser.get_structure("my_structure", "../eval_data/decoys/natives/2chf.pdb")

chain = [c for c in structure.get_chains()][0]
atom_positions, atom_mask, aatype, residue_idx = residues_to_features([r for r in chain.get_residues() if is_aa(r)])
nres = atom_positions.shape[0]

residue_mask = torch.ones(nres) 
aatype = torch.tensor(aatype, dtype=torch.long)
ca_coords = center_random_augmentation(atom_positions[...,1,:].unsqueeze(0), torch.ones([1, nres])).view([1,-1,3])
```