rc-foundry 0.1.4__py3-none-any.whl → 0.1.6__py3-none-any.whl

rf3/configs/model/components/rf3_net.yaml

@@ -0,0 +1,177 @@
+# Model architecture
+_target_: rf3.model.RF3.RF3
+
+# +---------- Channel dimensions ----------+
+c_s: 384
+c_z: 128
+c_atom: 128
+c_atompair: 16
+c_s_inputs: 449
+
+# +---------- Feature embedding ----------+
+feature_initializer:
+  # InputFeatureEmbedder
+  input_feature_embedder:
+    features:
+      - restype
+      - profile
+      - deletion_mean
+    atom_attention_encoder:
+      c_token: 384
+      c_atom_1d_features: 389
+      c_tokenpair: ${model.net.c_z}
+      use_inv_dist_squared: true
+      atom_1d_features:
+        - ref_pos
+        - ref_charge
+        - ref_mask
+        - ref_element
+        - ref_atom_name_chars
+      atom_transformer:
+        n_queries: 32
+        n_keys: 128
+        diffusion_transformer:
+          n_block: 3
+          diffusion_transformer_block:
+            n_head: 4
+            no_residual_connection_between_attention_and_transition: true
+            kq_norm: true
+
+  # RelativePositionEncoding
+  relative_position_encoding:
+    r_max: 32
+    s_max: 2
+
+# +---------- Recycler ----------+
+recycler:
+  # Pairformer
+  n_pairformer_blocks: 48
+  pairformer_block:
+    p_drop: 0.25
+    triangle_multiplication:
+      d_hidden: 128
+    triangle_attention:
+      n_head: 4
+      d_hidden: 32
+    attention_pair_bias:
+      n_head: 16
+
+  # TemplateEmbedder
+  template_embedder:
+    n_block: 2
+    raw_template_dim: 108
+    c: 64
+    p_drop: 0.25
+
+  # MSA module
+  msa_module:
+    n_block: 4
+    c_m: 64
+    p_drop_msa: 0.15
+    p_drop_pair: 0.25
+    msa_subsample_embedder:
+      num_sequences: 1024
+      dim_raw_msa: 34
+      c_s_inputs: ${model.net.c_s_inputs}
+      c_msa_embed: ${model.net.recycler.msa_module.c_m}
+    outer_product: 
+      c_msa_embed: ${model.net.recycler.msa_module.c_m}
+      c_outer_product: 32
+      c_out: ${model.net.c_z}
+    msa_pair_weighted_averaging: 
+      n_heads: 8
+      c_weighted_average: 32
+      c_msa_embed: ${model.net.recycler.msa_module.c_m}
+      c_z: ${model.net.c_z}
+      separate_gate_for_every_channel: true
+    msa_transition: 
+      n: 4  
+      c: ${model.net.recycler.msa_module.c_m}
+    triangle_multiplication_outgoing: 
+      d_pair: ${model.net.c_z}
+      d_hidden: 128
+      bias: True
+    triangle_multiplication_incoming: 
+      d_pair: ${model.net.c_z}
+      d_hidden: 128
+      bias: True
+    triangle_attention_starting: 
+      d_pair: ${model.net.c_z}
+      n_head: 4
+      d_hidden: 32
+      p_drop: 0.0 # This does not do anything: TODO: Remove
+    triangle_attention_ending: 
+      d_pair: ${model.net.c_z}
+      n_head: 4
+      d_hidden: 32
+      p_drop: 0.0 # This does not do anything; TODO: Remove
+    pair_transition: 
+      n: 4
+      c: ${model.net.c_z}
+
+# +---------- Diffusion module ----------+
+diffusion_module:
+  sigma_data: 16
+  c_token: 768
+  f_pred: edm
+  diffusion_conditioning:
+    c_s_inputs: ${model.net.c_s_inputs}
+    c_t_embed: 256
+    relative_position_encoding:
+      r_max: 32
+      s_max: 2
+  atom_attention_encoder:
+    c_tokenpair: ${model.net.c_z}
+    c_atom_1d_features: 389
+    use_inv_dist_squared: true
+    atom_1d_features:
+        - ref_pos
+        - ref_charge
+        - ref_mask
+        - ref_element
+        - ref_atom_name_chars
+    atom_transformer:
+      n_queries: 32
+      n_keys: 128
+      diffusion_transformer:
+        n_block: 3
+        diffusion_transformer_block:
+          n_head: 4
+          no_residual_connection_between_attention_and_transition: true
+          kq_norm: true
+    broadcast_trunk_feats_on_1dim_old: false
+    use_chiral_features: true
+    no_grad_on_chiral_center: false
+  diffusion_transformer:
+    n_block: 24
+    diffusion_transformer_block:
+      n_head: 16
+      no_residual_connection_between_attention_and_transition: true
+      kq_norm: true
+  atom_attention_decoder:
+    atom_transformer:
+      n_queries: 32
+      n_keys: 128
+      diffusion_transformer:
+        n_block: 3
+        diffusion_transformer_block:
+          n_head: 4
+          no_residual_connection_between_attention_and_transition: true
+          kq_norm: true
+distogram_head:
+  bins: 65 
+
+# +---------- Inference sampler ----------+
+inference_sampler:
+  solver: "af3"
+  num_timesteps: 200
+  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

rf3/configs/model/components/rf3_net_with_confidence_head.yaml

@@ -0,0 +1,45 @@
+defaults:
+  - rf3_net
+
+# Model architecture
+_target_: rf3.model.RF3.RF3WithConfidence
+
+# +---------- Mini rollout sampler ----------+
+# From the AF-3 main text:
+# > ... To remedy this, we developed a diffusion ‘rollout’ procedure for the full-structure prediction generation during training (using a larger step size than normal)
+# They do not further elaborate on how they adjusted the step size during diffusion rollout, but this may be a fruitful area of exploration moving forwards
+mini_rollout_sampler:
+  solver: "af3"
+  num_timesteps: 20 # 20 timesteps for the mini-rollout (vs. 200 for the full rollout during inference)
+  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
+
+# +---------- Confidence head architecture ----------+
+confidence_head:
+  c_s: ${model.net.c_s}
+  c_z: ${model.net.c_z}
+  n_pairformer_layers: 4
+  pairformer:
+    p_drop: 0.25
+    triangle_multiplication:
+      d_hidden: 128
+    triangle_attention:
+      n_head: 4
+      d_hidden: 32
+    attention_pair_bias:
+      n_head: 16 
+  n_bins_pae: 64
+  n_bins_pde: 64
+  n_bins_plddt: 50
+  n_bins_exp_resolved: 2
+  use_Cb_distances: False
+  use_af3_style_binning_and_final_layer_norms: True
+  symmetrize_Cb_logits: True

rf3/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

rf3/configs/model/rf3.yaml

@@ -0,0 +1,43 @@
+defaults:
+  - optimizers/adam@optimizer
+  - schedulers/af3@lr_scheduler
+  - components/ema@ema
+  - components/rf3_net@net
+  - _self_
+
+net:
+  feature_initializer:
+    input_feature_embedder:
+      atom_attention_encoder:
+        c_atom_1d_features: 393 # 392 + 1 has_atom_level_embedding = 393
+        atom_1d_features: 
+            - ref_pos                                                                             
+            - ref_charge                                                                          
+            - ref_mask                                                                            
+            - ref_element                                                                         
+            - ref_atom_name_chars                                                                 
+            - ref_pos_ground_truth 
+            - has_atom_level_embedding
+        use_atom_level_embedding: true
+        atom_level_embedding_dim: 384
+
+  recycler:
+    msa_module:
+        msa_subsample_embedder:
+          dim_raw_msa: 35
+    template_embedder:
+      raw_template_dim: 66
+
+  diffusion_module:
+    atom_attention_encoder:
+      c_atom_1d_features: 393 # 392 + 1 has_atom_level_embedding = 393
+      atom_1d_features:
+        - ref_pos                                                                             
+        - ref_charge                                                                          
+        - ref_mask                                                                            
+        - ref_element                                                                         
+        - ref_atom_name_chars                                                                 
+        - ref_pos_ground_truth 
+        - has_atom_level_embedding
+      use_atom_level_embedding: true
+      atom_level_embedding_dim: 384

rf3/configs/model/rf3_with_confidence.yaml

@@ -0,0 +1,7 @@
+defaults:
+  - components/rf3_net_with_confidence_head@net
+  - rf3
+  - _self_
+
+net:
+  _target_: rf3.model.RF3.RF3WithConfidence

rf3/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

rf3/configs/paths/data/default.yaml

@@ -0,0 +1,43 @@
+########################
+# Datasets
+########################
+
+# path to directory with training splits
+pdb_data_dir: /projects/ml/datahub/dfs/af3_splits/2025_07_13
+
+# fb monomer distillation dataset
+monomer_distillation_data_dir: /squash/af2_distillation_facebook
+monomer_distillation_parquet_dir: /projects/ml/datahub/dfs/distillation/af2_distillation_facebook
+
+mgnify_distillation_data_dir: /squash/mgnify_distill_rf3/
+mgnify_distillation_parquet_dir: /home/dimaio/MGnify/
+
+# na complex distill set
+na_complex_distillation_data_dir: /projects/ml/prot_dna/rf3_newDL
+na_complex_distillation_parquet_dir: /projects/ml/prot_dna
+
+# disorder distill set
+disorder_distill_parquet_dir: /projects/ml/disorder_distill
+
+########################
+# MSAs
+########################
+
+# path(s) to search for protein MSAs (for PDB datasets)
+# e.g., {"dir": "/path/to/msas", "extension": ".a3m.gz", "directory_depth": 2}
+protein_msa_dirs:
+  - {"dir": "/projects/msa/hhblits", "extension": ".a3m.gz", "directory_depth": 2}
+  - {"dir": "/projects/msa/mmseqs_gpu", "extension": ".a3m.gz", "directory_depth": 2}
+  - {"dir": "/projects/msa/lab", "extension": ".a3m.gz", "directory_depth": 1}
+
+# path(s) to search for RNA MSAs
+# e.g., {"dir": "/path/to/msas", "extension": ".afa", "directory_depth": 0}
+rna_msa_dirs:
+  - {"dir": "/projects/msa/rna", "extension": ".afa", "directory_depth": 0}
+
+########################
+# Misc
+########################
+
+# path to save examples that fail during the Transform pipeline (null = do not save)
+failed_examples_dir: null

rf3/configs/paths/default.yaml

@@ -0,0 +1,21 @@
+# 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: /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}

rf3/configs/train.yaml

@@ -0,0 +1,42 @@
+# @package _global_
+# ^ The "package" determines where the content of the config is placed in the output config
+# For more information about overriding configs, see: https://hydra.cc/docs/advanced/overriding_packages/#overriding-packages-using-the-defaults-list
+
+# NOTE: order of defaults determines the order in which configs override each other (higher up items are overridden by lower items)
+defaults:
+  - callbacks: default
+  - logger: csv
+  - trainer: ???
+  - paths: default
+  - datasets: ???
+  - dataloader: default
+  - hydra: default
+  - model: ???
+  # We must keep _self_ before experiment and debug to ensure that the experiment and debug configs can override
+  - _self_
+
+  # experiment configs allow for version control of specific hyperparameters
+  # e.g. best hyperparameters for given model and datamodule
+  - experiment: ???
+
+  # debug configs to add onto any experiment for quickly testing or debugging code
+  - debug: null
+
+
+# DO NOT set these here. Set them in the relevant experiment config file.
+# ... these are just here to ensure users always specify these fields in their experiment configs.
+name: ???
+tags: ???
+
+# NOTE: These values will be overwritten by the experiment config if they are set there. They are just provided as defaults
+#       here.
+# ... task name (determines the output directory path)
+task_name: "train"
+
+project: ??? # required for W&B logging
+
+seed: 1
+
+# Provide checkpoint path to resume training from a checkpoint
+# NOTE: If using W&B, must also set the `id` and `resume` fields in the `logger/wandb` config
+ckpt_path: null

rf3/configs/trainer/cpu.yaml

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

rf3/configs/trainer/ddp.yaml

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

rf3/configs/trainer/loss/losses/confidence_loss.yaml

@@ -0,0 +1,29 @@
+_target_: rf3.loss.af3_confidence_loss.ConfidenceLoss
+weight: 1.0
+
+plddt:
+  weight: 1.0
+  n_bins: 50
+  max_value: 1.0
+
+pae:
+  weight: 1.0
+  n_bins: 64
+  max_value: 32
+
+pde:  
+  weight: 1.0
+  n_bins: 64
+  max_value: 32
+
+exp_resolved:
+  weight: 1.0
+  n_bins: 2
+  max_value: 1
+
+# Adds to loss_dict true and predicted average plddt, pae, and pde per batch, also info about the spread and correlation of those values within a batch
+log_statistics: True 
+
+rank_loss:
+  use_listnet_loss: False
+  weight: 0.0

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

@@ -0,0 +1,9 @@
+_target_: rf3.loss.af3_losses.DiffusionLoss
+weight: 4.0
+sigma_data: ${model.net.diffusion_module.sigma_data}
+alpha_dna: 5
+alpha_rna: 5
+alpha_ligand: 10
+edm_lambda: True
+se3_invariant_loss: True
+clamp_diffusion_loss: False

rf3/configs/trainer/loss/losses/distogram_loss.yaml

@@ -0,0 +1,2 @@
+_target_: rf3.loss.af3_losses.DistogramLoss
+weight: 3e-2

rf3/configs/trainer/loss/structure_prediction.yaml

@@ -0,0 +1,4 @@
+defaults:
+  # Note that the SmoothedLDDTLoss is included within the DiffusionLoss
+  - losses/diffusion_loss@diffusion_loss
+  - losses/distogram_loss@distogram_loss

rf3/configs/trainer/loss/structure_prediction_with_confidence.yaml

@@ -0,0 +1,2 @@
+defaults:
+  - losses/confidence_loss@confidence_loss

rf3/configs/trainer/metrics/structure_prediction.yaml

@@ -0,0 +1,14 @@
+by_type_lddt:
+  _target_: rf3.metrics.lddt.ByTypeLDDT
+all_atom_lddt:
+  _target_: rf3.metrics.lddt.AllAtomLDDT
+distogram:
+  _target_: rf3.metrics.distogram.DistogramLoss
+distogram_comparisons:
+  _target_: rf3.metrics.distogram.DistogramComparisons
+distogram_entropy:
+  _target_: rf3.metrics.distogram.DistogramEntropy
+chiral_loss:
+  _target_: rf3.metrics.chiral.ChiralLoss
+unresolved_rasa:
+  _target_: rf3.metrics.rasa.UnresolvedRegionRASA

rf3/configs/trainer/rf3.yaml

@@ -0,0 +1,20 @@
+defaults:
+  - ddp
+  - loss: structure_prediction
+  - metrics: structure_prediction
+
+_target_: rf3.trainers.rf3.RF3Trainer
+validate_every_n_epochs: 1
+max_epochs: 10_000
+n_examples_per_epoch: 24000
+prevalidate: True
+
+# We must pre-specify the number of recycles during training so we can pre-sample recycles per batch consistently for each GPU
+n_recycles_train: ${datasets.n_recycles_train}
+
+clip_grad_max_norm: 10.0
+
+output_dir: ${paths.output_dir}
+checkpoint_every_n_epochs: 1
+
+precision: bf16-mixed

rf3/configs/trainer/rf3_with_confidence.yaml

@@ -0,0 +1,13 @@
+defaults:
+  - rf3
+  - override loss: structure_prediction_with_confidence
+
+_target_: rf3.trainers.rf3.RF3TrainerWithConfidence
+
+metrics:
+  ptm:
+    _target_: rf3.metrics.predicted_error.ComputePTM
+  iptm:
+    _target_: rf3.metrics.predicted_error.ComputeIPTM
+  count_clashing_chains:
+    _target_: rf3.metrics.clashing_chains.CountClashingChains

rf3/configs/validate.yaml

@@ -0,0 +1,45 @@
+# @package _global_
+# ^ The "package" determines where the content of the config is placed in the output config
+# For more information about overriding configs, see: https://hydra.cc/docs/advanced/overriding_packages/#overriding-packages-using-the-defaults-list
+
+# NOTE: order of defaults determines the order in which configs override each other (higher up items are overridden by lower items)
+defaults:
+  - callbacks: default
+  - logger: csv
+  - trainer: ???
+  - paths: default
+  - datasets: ???
+  - dataloader: default
+  - hydra: default
+  - model: ???
+  # We must keep _self_ before experiment and debug to ensure that the experiment and debug configs can override
+  - _self_
+
+  # experiment configs allow for version control of specific hyperparameters
+  # e.g. best hyperparameters for given model and datamodule
+  - experiment: ???
+
+  # debug configs to add onto any experiment for quickly testing or debugging code
+  - debug: null
+
+
+# DO NOT set these here. Set them in the relevant experiment config file.
+# ... these are just here to ensure users always specify these fields in their experiment configs.
+name: ???
+tags: ???
+
+# NOTE: These values will be overwritten by the experiment config if they are set there. They are just provided as defaults
+#       here.
+# ... task name (determines the output directory path)
+task_name: "validate"
+
+project: ??? # required for W&B logging
+
+seed: 1
+
+# Dump CIF files for validation structures
+callbacks:
+  dump_validation_structures_callback:
+    dump_predictions: True
+    one_model_per_file: False
+    dump_trajectories: False

rfd3/cli.py

@@ -12,10 +12,16 @@ app = typer.Typer()
 def design(ctx: typer.Context):
     """Run design using hydra config overrides and input files."""
     # Find the RFD3 configs directory relative to this file
-    # This file is at: models/rfd3/src/rfd3/cli.py
-    # Configs are at: models/rfd3/configs/
-    rfd3_package_dir = Path(__file__).parent.parent.parent  # Go up to models/rfd3/
-    config_path = str(rfd3_package_dir / "configs")
+    # Development: models/rfd3/src/rfd3/cli.py -> models/rfd3/configs/
+    # Installed: site-packages/rfd3/cli.py -> site-packages/rfd3/configs/
+
+    # Try development location first
+    dev_config_path = Path(__file__).parent.parent.parent / "configs"
+    if dev_config_path.exists():
+        config_path = str(dev_config_path)
+    else:
+        # Fall back to installed package location
+        config_path = str(Path(__file__).parent / "configs")
 
     # Get all arguments
     args = ctx.params.get("args", []) + ctx.args

rfd3/configs/callbacks/design_callbacks.yaml

@@ -0,0 +1,10 @@
+defaults:
+  - train_logging
+  - _self_
+
+log_learning_rate_callback:
+  log_every_n: 25 # default 10
+
+log_af3_training_losses_callback:
+  log_full_batch_losses: False
+  log_every_n: 25 # default 10

rfd3/configs/callbacks/metrics_logging.yaml

@@ -0,0 +1,20 @@
+store_validation_metrics_in_df_callback:
+    _target_: foundry.callbacks.metrics_logging.StoreValidationMetricsInDFCallback
+    save_dir: ${paths.output_dir}/val_metrics
+    metrics_to_save: "all"
+
+dump_validation_structures_callback:
+  _target_: rfd3.trainer.dump_validation_structures.DumpValidationStructuresCallback
+  save_dir: ${paths.output_dir}/val_structures
+  dump_predictions: True
+  dump_prediction_metadata_json: True
+  dump_trajectories: False
+  dump_denoised_trajectories_only: False
+  
+  one_model_per_file: True
+  dump_every_n: 4
+  align_trajectories: False
+  verbose: False
+
+log_design_validation_metrics_callback:
+  _target_: rfd3.callbacks.LogDesignValidationMetricsCallback

rfd3/configs/callbacks/train_logging.yaml

@@ -0,0 +1,24 @@
+log_af3_training_losses_callback:
+    _target_: foundry.callbacks.train_logging.LogAF3TrainingLossesCallback
+    log_every_n: 10
+    log_full_batch_losses: true
+
+log_learning_rate_callback:
+    _target_: foundry.callbacks.train_logging.LogLearningRateCallback
+    log_every_n: 10
+
+log_model_parameters_callback:
+    _target_: foundry.callbacks.train_logging.LogModelParametersCallback
+
+log_dataset_sampling_ratios_callback:
+    _target_: foundry.callbacks.train_logging.LogDatasetSamplingRatiosCallback
+
+# Optional health logging
+#   activations_tracking_callback:
+#     _target_: foundry.callbacks.health_logging.ActivationsGradientsWeightsTracker
+# log_freq: 100
+# keep_cache: True  # --> WARNING: Do not run this in a production run, this will lead to a memory leak! Meant for debugging.
+# activations_tracking_callback:
+#   _target_: foundry.callbacks.health_logging.ActivationsGradientsWeightsTracker
+#   log_freq: 100
+# keep_cache: True  # --> WARNING: Do not run this in a production run, this will lead to a memory leak! Meant for debugging.

rfd3/configs/dataloader/default.yaml

@@ -0,0 +1,15 @@
+train: 
+  dataloader_params:
+    # These parameters will be unpacked as kwargs for the DataLoader
+    batch_size: 1
+    num_workers: 2
+    prefetch_factor: 3
+  n_fallback_retries: 4
+
+val: 
+  dataloader_params:
+    # These parameters will be unpacked as kwargs for the DataLoader
+    batch_size: 1
+    num_workers: 2
+    prefetch_factor: 3
+  n_fallback_retries: 0 # Disable fallback retries for validation

rfd3/configs/dataloader/fast.yaml

@@ -0,0 +1,11 @@
+defaults:
+  - default
+
+train:
+  dataloader_params:
+    num_workers: 2
+    prefetch_factor: 6
+val: 
+  dataloader_params:
+    num_workers: 2
+    prefetch_factor: 6