rc-foundry 0.1.5__py3-none-any.whl → 0.1.7__py3-none-any.whl

rfd3/configs/hydra/default.yaml

@@ -0,0 +1,18 @@
+# https://hydra.cc/docs/configure_hydra/intro/
+
+# enable color logging (requires `colorlog` to be installed)
+# defaults:
+#   - override hydra_logging: colorlog
+#   - override job_logging: colorlog
+
+
+# output directory, generated dynamically on each run
+run:
+  dir: ${paths.log_dir}/${task_name}/${name}/${now:%Y-%m-%d}_${now:%H-%M}_JOB_${oc.env:SLURM_JOB_ID,default}
+
+# ... this is where the log file is written (i.e. the programs output)
+job_logging:
+  handlers:
+    file:
+      # Incorporates fix from https://github.com/facebookresearch/hydra/pull/2242
+      filename: ${hydra.runtime.output_dir}/experiment.log

rfd3/configs/hydra/no_logging.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - override job_logging: disabled
+  - override hydra_logging: disabled
+
+output_subdir: null  
+run:  
+  dir: .

rfd3/configs/inference.yaml

@@ -0,0 +1,9 @@
+# @package _global_
+
+hydra:
+  searchpath:
+    - pkg://configs
+
+defaults:
+  - inference_engine: rfdiffusion3
+  - _self_

rfd3/configs/inference_engine/base.yaml

@@ -0,0 +1,15 @@
+# @package _global_
+
+defaults:
+  - /hydra: no_logging
+
+# Parameters for RFD3InferenceEngine.__init__()
+ckpt_path: ???
+num_nodes: 1
+devices_per_node: 1
+verbose: false
+seed: null
+
+# Parameters for RFD3InferenceEngine.run()
+inputs: ???
+out_dir: ???

rfd3/configs/inference_engine/dev.yaml

@@ -0,0 +1,20 @@
+# @package _global_
+defaults:
+  - rfdiffusion3
+  - _self_
+
+diffusion_batch_size: 8
+n_batches: 1
+seed: 42
+
+dump_trajectories: True
+verbose: true
+skip_existing: False
+cleanup_guideposts: False
+cleanup_virtual_atoms: False
+output_full_json: True
+
+inference_sampler:
+  gamma_0: 0.0
+
+out_dir: ./logs/benchmark

rfd3/configs/inference_engine/rfdiffusion3.yaml

@@ -0,0 +1,65 @@
+# @package _global_
+defaults:
+  - base
+  - _self_
+
+_target_: rfd3.engine.RFD3InferenceEngine
+
+out_dir: ???
+inputs: ???  # null, json, pdb or 
+ckpt_path: rfd3
+json_keys_subset: null
+skip_existing: True
+
+#########################################################
+# Design spec args: overrides args from input json
+specification: {}
+#########################################################
+
+# Diffusion args
+diffusion_batch_size: 8
+n_batches: 1
+
+# Inference sampler args | set to None to use the default in the checkpoint's config
+inference_sampler:
+  kind: "default" # "default" or "symmetry" to choose the sampler
+  # Classifier-free guidance args:
+  cfg_features:  # set to 0 in the reference CFG step
+    - active_donor
+    - active_acceptor
+    - ref_atomwise_rasa
+
+  use_classifier_free_guidance: False
+  cfg_t_max: null # max t to apply cfg guidance
+  cfg_scale: 1.5
+  center_option: "all"  # Options are ["all", "motif", "diffuse"]
+  s_trans: 1.0  # Translational noise scale for augmentation during inference
+  inference_noise_scaling_factor: 1.0
+  allow_realignment: False
+
+  # Diffusion args:
+  num_timesteps: 200
+  step_scale: 1.5  # 1.5 - 1.0   |   Higher values lead to less diverse, more designable, structures
+  noise_scale: 1.003
+  p: 7
+  gamma_0: 0.6  # Previously 1.0  | 0.0 for ODE sampling
+  gamma_min: 1.0
+  s_jitter_origin: 0.0  # Sigma of gaussian noise to jitter the motif offset (equivalent to ORI token Jitter)
+
+# Saving args
+cleanup_guideposts: True
+cleanup_virtual_atoms: True
+read_sequence_from_sequence_head: True
+output_full_json: True
+
+# Prefix to add to all output samples 
+# Default: None      -> f'{jsonfilebasename}_{jsonkey}_{batch}_{model}'
+# Otherwise: string  -> f'{string}{jsonkey}_{batch}_{model}'
+# e.g. Empty string  -> f'{jsonkey}_{batch}_{model}'
+# e.g. Chunk string  -> f'{chunkprefix_}{jsonkey}_{batch}_{model}' (pipelines usage)
+global_prefix: null
+dump_prediction_metadata_json: True
+dump_trajectories: False
+align_trajectory_structures: False
+prevalidate_inputs: False
+low_memory_mode: False # False for standard mode, True for memory efficient tokenization mode

rfd3/configs/logger/csv.yaml

@@ -0,0 +1,6 @@
+# https://lightning.ai/docs/fabric/latest/api/generated/lightning.fabric.loggers.CSVLogger.html#lightning.fabric.loggers.CSVLogger
+
+csv:
+  _target_: lightning.fabric.loggers.CSVLogger
+  root_dir: ${paths.output_dir}
+  flush_logs_every_n_steps: 1

rfd3/configs/logger/default.yaml

@@ -0,0 +1,2 @@
+defaults:
+  - csv

rfd3/configs/logger/wandb.yaml

@@ -0,0 +1,15 @@
+# https://wandb.ai
+
+wandb:
+  _target_: wandb.integration.lightning.fabric.WandbLogger
+  save_dir: ${paths.output_dir}
+  name: ${name}
+  offline: False
+  id: null # pass correct id (along with checkpoint path, and resume='allow' or 'must') to resume a run
+  anonymous: null # enable anonymous logging
+  project: ${project}
+  prefix: "" # a string to put at the beginning of metric keys
+  log_model: False # do not upload model checkpoints
+  tags: ${tags}
+  # (Default resume to "never" to avoid accidentally resuming runs; we want to be explicit about resuming)
+  resume: never # never, allow, or must (see: https://docs.wandb.ai/guides/runs/resuming/)

rfd3/configs/model/components/ema.yaml

@@ -0,0 +1 @@
+decay: 0.999 # From AF-3

rfd3/configs/model/components/rfd3_net.yaml

@@ -0,0 +1,131 @@
+_target_: rfd3.model.RFD3.RFD3
+
+c_s: 384
+c_z: 128
+c_atom: 128
+c_atompair: 16
+
+token_initializer: # formerly known as the trunk
+  relative_position_encoding:
+    r_max: 32
+    s_max: 2
+  
+  # Attention pair biases without batch dimensions
+  n_pairformer_blocks: 2
+  pairformer_block:
+    use_triangle_attn: false
+    use_triangle_mult: false
+    attention_pair_bias:
+      n_head: 16
+      kq_norm: True
+
+  token_1d_features:
+    ref_motif_token_type: 3
+    restype: 32
+    ref_plddt: 1
+    is_non_loopy: 1
+
+  downcast: ${model.net.diffusion_module.downcast}
+  atom_1d_features:
+    ref_atom_name_chars: 256
+    ref_element: 128
+    ref_charge: 1
+    ref_mask: 1
+    ref_is_motif_atom_with_fixed_coord: 1
+    ref_is_motif_atom_unindexed: 1
+    has_zero_occupancy: 1
+    ref_pos: 3
+    
+    # Guided features
+    ref_atomwise_rasa: 3
+    active_donor: 1
+    active_acceptor: 1
+    is_atom_level_hotspot: 1
+
+  atom_transformer:
+    n_blocks: 0
+    atom_transformer_block:
+      n_head: 4
+      kq_norm: True
+      no_residual_connection_between_attention_and_transition: False
+      dropout: 0.0
+      n_attn_seq_neighbours: 4
+      n_attn_keys: 128
+
+diffusion_module:
+  _target_: rfd3.model.RFD3_diffusion_module.RFD3DiffusionModule
+  c_token: 768
+  c_t_embed: 256  # Time embedding dimension
+  sigma_data: 16
+  f_pred: edm
+  n_attn_seq_neighbours: 2  # include self + n flanking neighbours
+  n_attn_keys: 128
+  n_recycle: 2
+  use_local_token_attention: false
+  
+  # Upcast/downcast mechanisms
+  upcast:
+    method: cross_attention
+    n_split: 3
+    cross_attention_block:
+      n_head: 4
+      c_model: 128
+      dropout: 0.0
+      kq_norm: True
+
+  downcast:
+    method: cross_attention
+    cross_attention_block:
+      n_head: 4
+      c_model: 128
+      dropout: 0.0
+      kq_norm: True
+
+  ########################################################################
+  # UNet level processing
+  ########################################################################
+  atom_attention_encoder:
+    n_blocks: 3
+    atom_transformer_block:
+      n_head: 4
+      kq_norm: True
+      no_residual_connection_between_attention_and_transition: False
+      dropout: 0.0
+
+  diffusion_token_encoder:  # encodes self conditioning information and distogram
+    use_distogram: True
+    use_self: True
+    use_sinusoidal_distogram_embedder: False
+    sigma_data: ${model.net.diffusion_module.sigma_data}
+
+    n_pairformer_blocks: 2
+    pairformer_block:
+      use_triangle_attn: false
+      use_triangle_mult: false
+      attention_pair_bias:
+        n_head: 16
+        kq_norm: True
+
+  diffusion_transformer:
+    n_block: 18
+    n_registers: 0  # 8 Idk if they do anything tbh
+    diffusion_transformer_block:
+      n_head: 16
+      kq_norm: True
+      no_residual_connection_between_attention_and_transition: False
+      dropout: 0.10
+
+  atom_attention_decoder:
+    n_blocks: 3
+    upcast: ${model.net.diffusion_module.upcast}
+    downcast: ${model.net.diffusion_module.downcast}
+
+    atom_transformer_block:
+      n_head: 4
+      kq_norm: True
+      no_residual_connection_between_attention_and_transition: False
+      dropout: 0.10
+
+  ########################################################################
+  # 
+  ########################################################################

rfd3/configs/model/optimizers/adam.yaml

@@ -0,0 +1,5 @@
+# Optimizer
+_target_: torch.optim.Adam
+lr: 0 # Will be set by the scheduler (starts at 0, increasing to `base_lr`)
+betas: [0.9, 0.95]
+eps: 1.0e-8

rfd3/configs/model/rfd3_base.yaml

@@ -0,0 +1,8 @@
+
+defaults:
+  - optimizers/adam@optimizer
+  - schedulers/af3@lr_scheduler
+  - samplers/edm@net.inference_sampler
+  - components/ema@ema
+  - components/rfd3_net@net
+  - _self_

rfd3/configs/model/samplers/edm.yaml

@@ -0,0 +1,21 @@
+
+kind: "default" # "default", "symmetry", or "partial" to choose the sampler
+solver: "af3"
+num_timesteps: 100
+min_t: 0
+max_t: 1
+sigma_data: ${model.net.diffusion_module.sigma_data}
+s_min: 4e-4
+s_max: 160
+p: 7
+gamma_0: 0.8
+gamma_min: 1.0
+noise_scale: 1.003
+step_scale: 1.5
+allow_realignment: False
+use_classifier_free_guidance: False
+cfg_scale: 1.5
+cfg_features: # CFG_features will be set to 0 in the unconditional CFG step
+  - ref_atomwise_rasa
+  - active_donor
+  - active_acceptor

rfd3/configs/model/samplers/symmetry.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - edm
+  - _self_
+
+kind: symmetry
+num_timesteps: 200
+gamma_0: 0.6  # 1.0 for SDE sampling
+gamma_min: 1.0
+gamma_min2: 0.0
+sym_step_frac: 0.9 # when 0.9, 90% of the trajectory from the start is symmetrized

rfd3/configs/model/schedulers/af3.yaml

@@ -0,0 +1,6 @@
+# Learning rate scheduler
+_target_: foundry.training.schedulers.AF3Scheduler
+base_lr: 1.8e-3
+warmup_steps: 1000
+decay_factor: 0.95
+decay_steps: 50000

rfd3/configs/paths/data/default.yaml

@@ -0,0 +1,18 @@
+# path to directory with training splits
+pdb_data_dir: /projects/ml/frozen_pdb_copies/2025_07_13_pdb
+pdb_parquet_dir: /projects/ml/datahub/dfs/af3_splits/2024_12_16/  # TODO: uncomment
+
+# monomer distillation dataset
+monomer_distillation_data_dir: /squash/af2_distillation_facebook/
+monomer_distillation_parquet_dir: /projects/ml/datahub/dfs/distillation/af2_distillation_facebook
+
+# path to save examples that fail during the Transform pipeline (null = do not save)
+failed_examples_dir: null
+
+design_benchmark_data_dir: /projects/ml/aa_design/benchmarks
+design_model_weight_dir: /projects/ml/aa_design/models
+
+# path to directory with cached residue data
+residue_cache_dir: /net/tukwila/lschaaf/datahub/MACE-Egret-3-noH/mace_embeddings
+
+cif_cache_dir: /net/tukwila/ncorley/cifutils/cache

rfd3/configs/paths/default.yaml

@@ -0,0 +1,22 @@
+# NOTE: order of defaults determines the order in which configs override each other (higher up items are overridden by lower items)
+defaults:
+  - _self_
+  - data: default
+
+# path to root directory (requires the `PROJECT_ROOT` environment variable to be set)
+#  NOTE: This variable is auto-set upon loading via `rootutils`
+root_dir: ${oc.env:PROJECT_ROOT}
+
+# where to store data (checkpoints, logs, etc.) of all experiments in general 
+# (this influences the output_dir in the hydra/default.yaml config)
+# change this to e.g. /scratch if you are running larger experiments with lots lof logs, checkpoints, etc.
+# log_dir: ${.root_dir}/logs
+log_dir: /net/scratch/${oc.env:USER}/training/logs
+
+# path to output directory for this specific run, created dynamically by hydra
+# path generation pattern is specified in `configs/hydra/default.yaml`
+# use it to store all files generated during the run, like ckpts and metrics
+output_dir: ${hydra:runtime.output_dir}
+
+# path to working directory (auto-generated by hydra)
+work_dir: ${hydra:runtime.cwd}

rfd3/configs/train.yaml

@@ -0,0 +1,28 @@
+# @package _global_
+hydra:
+  searchpath:
+    - pkg://configs
+
+defaults:
+  - model: rfd3_base
+  - trainer: rfd3_base
+  - datasets: design_base
+  - callbacks: design_callbacks
+  - dataloader: fast
+  - paths: default
+  - hydra: default
+  - logger: default
+  - _self_
+  # Required overrides:
+  - experiment: ???
+  - debug: null
+
+# Definitions:
+task_name: train
+project: aa_design
+seed: 42
+ckpt_path: null
+
+# Placeholders
+name: aa_design
+tags: [aa_design]

rfd3/configs/trainer/cpu.yaml

@@ -0,0 +1,6 @@
+defaults:
+  - aa_design
+
+accelerator: cpu
+devices_per_node: 1
+num_nodes: 1

rfd3/configs/trainer/ddp.yaml

@@ -0,0 +1,5 @@
+strategy: ddp
+
+accelerator: gpu
+devices_per_node: 1
+num_nodes: 1

rfd3/configs/trainer/loss/losses/diffusion_loss.yaml

@@ -0,0 +1,12 @@
+_target_: rfd3.metrics.losses.DiffusionLoss
+sigma_data: ${model.net.diffusion_module.sigma_data}
+weight: 4.0
+lddt_weight: 0.25
+alpha_virtual_atom: 1.0
+alpha_polar_residues: 1.0
+lp_weight: 0.0
+unindexed_norm_p: 1.0
+alpha_unindexed_diffused: 1.0
+unindexed_t_alpha: 0.75
+normalize_virtual_atom_weight: False
+alpha_ligand: 10.0

rfd3/configs/trainer/loss/losses/sequence_loss.yaml

@@ -0,0 +1,3 @@
+_target_: rfd3.metrics.losses.SequenceLoss
+weight: 0.1
+max_t: 1

rfd3/configs/trainer/metrics/design_metrics.yaml

@@ -0,0 +1,22 @@
+general_metrics:
+  _target_: rfd3.metrics.design_metrics.AtomArrayMetrics
+  compute_for_diffused_region_only: True
+  compute_ss_adherence_if_possible: False
+
+backbone_metrics:
+  _target_: rfd3.metrics.design_metrics.BackboneMetrics
+
+sidechain_metrics:
+  _target_: rfd3.metrics.sidechain_metrics.SidechainMetrics
+  central_atom: CB
+  dist_threshold_min: 1.0 # min distance for identifying a bond
+  dist_threshold_max: 2.0 # max distance for identifying a bond
+  already_removed_virtual_atoms: ${trainer.cleanup_virtual_atoms}
+
+metadata_metrics:
+  _target_: rfd3.metrics.design_metrics.MetadataMetrics
+
+hbond_metrics:
+ _target_: rfd3.metrics.hbonds_hbplus_metrics.HbondMetrics
+ cutoff_HA_dist: 3
+ cutoff_DA_distance: 3.5

rfd3/configs/trainer/rfd3_base.yaml

@@ -0,0 +1,35 @@
+defaults:
+  - ddp
+  - loss/losses/diffusion_loss@loss.diffusion_loss
+  - loss/losses/sequence_loss@loss.sequence_loss
+  - metrics: design_metrics
+  - _self_
+
+_target_: rfd3.trainer.rfd3.AADesignTrainer
+
+# AADesign specific (atom-array related):
+output_full_json: False  # saves additional metadata in the output json
+allow_sequence_outputs: True
+cleanup_guideposts: False
+cleanup_virtual_atoms: False
+read_sequence_from_sequence_head: True
+compute_non_clash_metrics_for_diffused_region_only: False
+association_scheme: ${datasets.global_transform_args.association_scheme}
+
+# Other:
+n_examples_per_epoch: 2400 # 24000 # 10x as many epochs
+checkpoint_every_n_epochs: 10 #  Less often checkpointing for fewer epochs
+validate_every_n_epochs: 4  # Validate often
+
+max_epochs: 100_000
+prevalidate: False
+
+clip_grad_max_norm: 10.0
+output_dir: ${paths.output_dir}
+n_recycles_train: 2
+grad_accum_steps: 3  # overridden by launch.sh
+skip_optimizer_loading: True
+
+# Precision
+error_if_grad_nonfinite: False
+precision: bf16-mixed

rfd3/configs/validate.yaml

@@ -0,0 +1,34 @@
+# @package _global_
+
+hydra:
+  searchpath:
+    - pkg://configs
+
+defaults:
+  - model: rfd3_base
+  - trainer: rfd3_base
+  - datasets: design_base
+  - callbacks: design_callbacks
+  - dataloader: fast
+  - paths: default
+  - hydra: default
+  - logger: csv
+  - _self_
+  - experiment: ???
+  - debug: null
+
+name: ???
+tags: ???
+project: aa_design
+# Names;
+task_name: "validate"
+seed: 42
+
+callbacks:
+  dump_validation_structures_callback:
+    dump_predictions: True
+    one_model_per_file: True
+    dump_trajectories: False
+
+# Args:
+ckpt_path: ???

rfd3/engine.py

@@ -23,7 +23,10 @@ from rfd3.inference.datasets import (
 )
 from rfd3.inference.input_parsing import DesignInputSpecification
 from rfd3.model.inference_sampler import SampleDiffusionConfig
-from rfd3.utils.inference import ensure_input_is_abspath
+from rfd3.utils.inference import (
+    ensure_inference_sampler_matches_design_spec,
+    ensure_input_is_abspath,
+)
 from rfd3.utils.io import (
     CIF_LIKE_EXTENSIONS,
     build_stack_from_atom_array_and_batched_coords,
@@ -171,6 +174,7 @@ class RFD3InferenceEngine(BaseInferenceEngine):
         )
         # save
         self.specification_overrides = dict(specification or {})
+        self.inference_sampler_overrides = dict(inference_sampler or {})
 
         # Setup output directories and args
         self.global_prefix = global_prefix
@@ -210,6 +214,9 @@ class RFD3InferenceEngine(BaseInferenceEngine):
             inputs=inputs,
             n_batches=n_batches,
         )
+        ensure_inference_sampler_matches_design_spec(
+            design_specifications, self.inference_sampler_overrides
+        )
         # init before
         self.initialize()
         outputs = self._run_multi(design_specifications)
@@ -383,6 +390,9 @@ class RFD3InferenceEngine(BaseInferenceEngine):
         # Based on inputs, construct the specifications to loop through
         design_specifications = {}
         for prefix, example_spec in inputs.items():
+            # Record task name in the specification
+            example_spec["extra"]["task_name"] = prefix
+
             # ... Create n_batches for example
             for batch_id in range((n_batches) if exists(n_batches) else 1):
                 # ... Example ID
@@ -524,21 +534,19 @@ def process_input(
 
 
 def _reshape_trajectory(traj, align_structures: bool):
-    traj = [traj[i] for i in range(len(traj))]
-    n_steps = len(traj)
+    traj = [traj[i] for i in range(len(traj))]  # make list of arrays
     max_frames = 100
-
+    if len(traj) > max_frames:
+        selected_indices = torch.linspace(0, len(traj) - 1, max_frames).long().tolist()
+        traj = [traj[i] for i in selected_indices]
     if align_structures:
         # ... align the trajectories on the last prediction
-        for step in range(n_steps - 1):
+        for step in range(len(traj) - 1):
             traj[step] = weighted_rigid_align(
-                X_L=traj[-1],
-                X_gt_L=traj[step],
-            )
+                X_L=traj[-1][None],
+                X_gt_L=traj[step][None],
+            ).squeeze(0)
     traj = traj[::-1]  # reverse to go from noised -> denoised
-    if n_steps > max_frames:
-        selected_indices = torch.linspace(0, n_steps - 1, max_frames).long().tolist()
-        traj = [traj[i] for i in selected_indices]
 
     traj = torch.stack(traj).cpu().numpy()
     return traj

rfd3/inference/input_parsing.py

@@ -696,7 +696,7 @@ class DesignInputSpecification(BaseModel):
             # Partial diffusion: use COM, keep all coordinates
             if exists(self.symmetry) and self.symmetry.id:
                 # For symmetric structures, avoid COM centering that would collapse chains
-                ranked_logger.info(
+                logger.info(
                     "Partial diffusion with symmetry: skipping COM centering to preserve chain spacing"
                 )
             else: