Mdc/add protpardelle

.actlignore

@@ -3,7 +3,9 @@
 .pixi/
 grid_search_results/
 outputs/
-/data/
+data/
+!/src/sampleworks/data/
+!/src/sampleworks/data/**
 initial_dataset_40*/
 checkpoints/
 release_data/

pyproject.toml

@@ -24,16 +24,51 @@ pypi-dependencies = {"boltz" = "*", "cuequivariance-torch" = "*", "cuequivarianc
 platforms = ["linux-64"]
 
 [tool.pixi.feature.boltz-osx]
-pypi-dependencies = {"boltz" = "*", "cuequivariance-torch" = "*", "rdkit" = ">=2025.3.6"}
+pypi-dependencies = {"boltz" = "*", "cuequivariance-torch" = "*", "rdkit" = ">=2025.3.6", "joblib" = "*", "sfcalculator-torch" = ">=0.3.2"}
 platforms = ["osx-arm64"]
 
 
 [tool.pixi.feature.protenix]
 pypi-dependencies = {"protenix" = "*", "einx" = "*", "triton" = "*"}
 platforms = ["linux-64"]
 
+[tool.pixi.feature.protpardelle]
+platforms = ["linux-64", "osx-arm64"]
+
+[tool.pixi.feature.protpardelle.dependencies]
+biopython = "*"
+einops = "*"
+huggingface_hub = "*"
+hydra-core = "*"
+jupyter = "*"
+numpy = ">=1.25.0,<2.0"
+omegaconf = "*"
+pandas = "==2.3.1"
+prody = "*"
+pyyaml = "*"
+scipy = "*"
+pytorch = ">=2.6.0,<2.8"
+tqdm = "*"
+typer = "*"
+
+[tool.pixi.feature.protpardelle.pypi-dependencies]
+protpardelle = {git = "https://github.com/ProteinDesignLab/protpardelle-1c.git"}
+biotite = "*"
+dm-tree = "*"
+jaxtyping = "*"
+ml_collections = "*"
+modelcif = "*"
+transformers = "*"
+wandb = "*"
+
+[tool.pixi.feature.protpardelle.pypi-options.dependency-overrides]
+pandas = "==2.3.1"
+
 [project]
-authors = [{email = "karson.chrispens@ucsf.edu", name = "Karson Chrispens"}]
+authors = [
+    {email = "karson.chrispens@ucsf.edu", name = "Karson Chrispens"},
+    {email = "marcus.collins@astera.org", name = "Marcus D. Collins"}
+]
 classifiers = [
   "Development Status :: 3 - Alpha",
   "Intended Audience :: Science/Research",
@@ -132,6 +167,8 @@ boltz-dev = {features = ["boltz", "dev"]}
 boltz-osx = {features = ["boltz-osx", "dev", "analysis"]}
 protenix = {features = ["protenix"]}
 protenix-dev = {features = ["protenix", "dev"]}
+protpardelle = {features = ["protpardelle"]}
+protpardelle-dev = {features = ["protpardelle", "dev"]}
 rf3 = {features = ["rf3"]}
 rf3-dev = {features = ["rf3", "dev"]}
 
@@ -141,14 +178,16 @@ rdkit = ">=2024.3.5,<2025.9"
 [tool.pixi.feature.rf3.pypi-dependencies]
 rc-foundry = {editable = true, extras = ["rf3"], git = "https://github.com/k-chrispens/foundry.git"}
 
-[tool.pixi.pypi-dependencies]
-sampleworks = {editable = true, path = "."}
-
 # Workspace-level override: reciprocalspaceship (sfcalculator-torch transitive dep)
 # caps pandas<=2.2.3 in its metadata, but runs fine on 2.3.1; protenix hard-pins
 # pandas==2.3.1. Standardize all envs on 2.3.1 to avoid silent downgrades.
 [tool.pixi.pypi-options.dependency-overrides]
 pandas = "==2.3.1"
+gemmi = "==0.6.7"
+
+
+[tool.pixi.pypi-dependencies]
+sampleworks = {editable = true, path = "."}
 
 [tool.pixi.system-requirements]
 cuda = "12"

src/sampleworks/cli/guidance.py

@@ -3,18 +3,25 @@
 from __future__ import annotations
 
 import sys
+from loguru import logger
 
 from sampleworks.utils.guidance_script_arguments import GuidanceConfig
-from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance
 
 
 def main(argv: list[str] | None = None) -> int:
     config = GuidanceConfig.from_cli(argv)
+
+    from loguru import logger
+
+    from sampleworks.utils.guidance_script_utils import get_model_and_device, run_guidance
+
+    logger.info(f"Running guidance with config: {config}")
     device, model_wrapper = get_model_and_device(
         config.device,
         getattr(config, "model_checkpoint", None),
         config.model,
         method=getattr(config, "method", None),
+        protpardelle_config_path=getattr(config, "protpardelle_config_path", None),
     )
     result = run_guidance(config, config.guidance_type, model_wrapper, device)
     return result.exit_code

src/sampleworks/core/samplers/edm.py

@@ -415,15 +415,17 @@ def step(
 
         # Store eps separately for proper frame transformation
         # eps_scale will be float if check_context didn't raise
-        eps = torch.randn_like(maybe_augmented_state) * eps_scale  # ty: ignore[unsupported-operator]
+        eps = torch.randn_like(maybe_augmented_state) * eps_scale
         noisy_state = maybe_augmented_state + eps
         noisy_state = torch.as_tensor(noisy_state).detach().requires_grad_(allow_gradients)
 
         # t_hat will be float if check_context didn't raise
         # Use no_grad when gradients aren't needed to avoid memory overhead from
         # gradient checkpointing holding intermediate activations
+        # TODO testing adding eps to signature for use with Protpardelle-1c, if successful,
+        #   I need to modify the Protocol itself. @Michael Anzuoni
         with torch.set_grad_enabled(allow_gradients):
-            x_hat_0 = model_wrapper.step(noisy_state, t_hat, features=features)
+            x_hat_0 = model_wrapper.step(noisy_state, t_hat, eps, features=features)
 
         reconciler = (
             context.reconciler.to(torch.as_tensor(x_hat_0).device)

src/sampleworks/data/cc89_epoch415.yaml

@@ -0,0 +1,101 @@
+data:
+  auto_calc_sigma_data: true
+  chain_residx_gap: 200
+  dummy_fill_mode: zero
+  fixed_size: 512
+  mixing_ratios:
+  - 1.0
+  n_aatype_tokens: 21
+  n_examples_for_sigma_data: 500
+  pdb_paths:
+  - /scratch/users/tianyulu/augmented_ingraham_cath_bugfree
+  se3_data_augment: true
+  sigma_data: 10.3
+  subset:
+  - designable_
+  translation_scale: 1.0
+diffusion:
+  sampling:
+    function: uniform
+    s_max: 80
+    s_min: 0.001
+  training:
+    function: lognormal
+    psigma_mean: -0.5
+    psigma_std: 1.5
+model:
+  compute_loss_on_all_atoms: false
+  conditioning_style: concat
+  crop_conditional: true
+  dummy_fill_masked_atoms: false
+  full_mpnn_model_path: /scratch/users/tianyulu/farfalle/ProteinMPNN/vanilla_model_weights
+  mpnn_model:
+    label_smoothing: 0.1
+    n_channel: 128
+    n_layers: 3
+    n_neighbors: 32
+    noise_cond_mult: 4
+    use_self_conditioning: true
+  mpnn_model_checkpoint: ''
+  pretrained_modules: []
+  struct_model:
+    arch: dit
+    n_atoms: 37
+    n_channel: 256
+    noise_cond_mult: 4
+    uvit:
+      cat_pwd_to_conv: false
+      conv_skip_connection: false
+      dim_head: 32
+      n_blocks_per_layer: 2
+      n_filt_per_layer: []
+      n_heads: 8
+      n_layers: 10
+      patch_size: 1
+      position_embedding_max: 32
+      position_embedding_type: rotary
+  struct_model_checkpoint: ''
+  task: ai-allatom
+train:
+  batch_size: 32
+  checkpoint_freq: 1
+  checkpoints: []
+  ckpt_path: /scratch/users/tianyulu/farfalle/out_dir/farfalle/cc89/checkpoints/epoch206_training_state.pth
+  clip_grad_norm: true
+  crop_cond:
+    contiguous_prob: 0.05
+    discontiguous_prob: 0.9
+    dist_threshold: 45.0
+    max_discontiguous_res: 24
+    max_span_len: 12
+    recenter_coords: true
+    sidechain_only_prob: 0.0
+    sidechain_prob: 0.9
+    terms_prob: 0.5
+  crop_conditional: true
+  decay_steps: 2000000
+  eval_freq: 8000000
+  eval_loss_t:
+  - 0.1
+  - 0.3
+  - 0.5
+  - 0.7
+  - 0.9
+  fpd_length_ranges_per_chain:
+  - - 50
+    - 256
+  grad_clip_val: 1.0
+  length_ranges_per_chain:
+  - - 166
+    - 188
+  lr: 0.0001
+  max_epochs: 10000
+  n_eval_samples: 10
+  n_fpd_samples: 0
+  sc_num_seqs: 4
+  seed: 0
+  self_cond_train_prob: 0.9
+  shapes_path: /scratch/users/tianyulu/protein_shapes
+  subsample_eval_set: 0.05
+  warmup_steps: 1000
+  weight_decay: 0.0

src/sampleworks/models/protpardelle/__init__.py

@@ -0,0 +1 @@
+"""Protpardelle-1c model wrapper."""