Add DISCO motif guidance CodeOcean sweep

configs/base.yaml

@@ -10,6 +10,7 @@ defaults:
   - inference: default
   - structure_encoder: default
   - sequence_sampling_strategy: default
+  - cogen_guidance: default
   - experiment: null
   - effort: max
   - _self_

configs/cogen_guidance/default.yaml

@@ -0,0 +1,126 @@
+enabled: false
+
+# Comma-separated motifs or a YAML list. When null, the historical
+# ProteinGenerator motif list is used for compatibility experiments.
+motifs: null
+motif_weights: null
+
+logit_guidance:
+  enabled: false
+  gradient_enabled: true
+  gradient_scale: 1.0
+  motif_weight_all: 0.5
+  motif_weight_many: 0.5
+  motif_top_k: 3
+  motif_temperature: 0.5
+  local_rescue_enabled: false
+  local_rescue_bonus: 1.0
+  local_rescue_top_k: 3
+  local_rescue_min_probability: 0.03
+  local_rescue_max_motifs_per_step: 4
+  local_rescue_max_changes_per_step: 8
+  local_rescue_only_missing: true
+  hard_commit_enabled: false
+  hard_commit_spread: true
+  hard_commit_position_offset: 0
+  hard_commit_all_motifs: false
+  hard_commit_max_motifs_per_step: 4
+  hard_commit_max_changes_per_step: 16
+
+scoring:
+  motif_all_weight: 0.8
+  motif_many_weight: 0.2
+  motif_many_cap: 0
+
+save_best:
+  enabled: false
+  select_best_output: false
+  motif_score_weight: 1.0
+  confidence_weight: 0.0
+
+# Optional full coordinate trajectory tensor persistence. This is disabled by
+# default because the tensors can be large. Lightweight JSONL/CSV diagnostics are
+# written separately when guidance produces diagnostic records.
+trajectory_dump:
+  enabled: false
+  save_x1_denoised: true
+  save_xt_trajectory: true
+  save_final_coordinate: false
+  dtype: float32
+  max_steps: 0
+
+alternating:
+  enabled: false
+  motif_list1: null
+  motif_list2: null
+  expected_start_group: list1
+  combined_weight: 0.0
+
+  # Slot-based soft guidance adapted from ProteinGenerator's
+  # AlternatingMotifsBias. These values affect logits before DISCO's existing
+  # path-planning sampler runs; they are separate from the hard scoring weights
+  # below.
+  logit_guidance_enabled: false
+  logit_gradient_scale: 1.0
+  logit_temperature: 0.5
+  logit_max_slots: 0
+  logit_slot_gap: 3.0
+  logit_slot_width: 1.35
+  logit_tail_weight: 0.25
+  logit_coverage_weight: 0.35
+  logit_transition_weight: 0.40
+
+  # Slot-free comparison guidance adapted from ProteinGenerator's
+  # soft_transition_alternating_motifs_bias. Enable this instead of the slot
+  # guidance above to test transition-only pressure.
+  transition_logit_guidance_enabled: false
+  transition_logit_gradient_scale: 1.0
+  transition_logit_temperature: 0.5
+  transition_logit_tail_weight: 0.25
+  transition_logit_coverage_weight: 0.35
+  transition_logit_transition_weight: 0.40
+  transition_same_group_penalty_weight: 0.25
+  transition_soft_start_weight: 0.20
+  transition_soft_end_weight: 0.20
+  transition_distance_temperature: 8.0
+  transition_min_separation: 1
+
+  # Hard scoring weights used for diagnostics and save-best ranking.
+  alternation_weight: 0.45
+  balance_weight: 0.20
+  coverage_weight: 0.20
+  terminal_weight: 0.15
+  start_weight: 0.15
+
+immunogen:
+  enabled: false
+  motifs: null
+  motif_weights: null
+  fixed_fragments: []
+  combined_weight: 0.0
+  protect_fixed_fragments: true
+
+  # PG immunogen_mutable_motifs_bias equivalent. Motif windows are scored only
+  # when every residue in the window is mutable, so guidance pressure stays away
+  # from protected B-cell/immunogen fragments.
+  mutable_logit_guidance:
+    enabled: false
+    gradient_enabled: true
+    gradient_scale: 1.0
+    motif_weight_all: 0.5
+    motif_weight_many: 0.5
+    motif_top_k: 3
+    motif_temperature: 0.5
+    local_rescue_enabled: false
+    local_rescue_bonus: 1.0
+    local_rescue_top_k: 3
+    local_rescue_min_probability: 0.03
+    local_rescue_max_motifs_per_step: 4
+    local_rescue_max_changes_per_step: 8
+    local_rescue_only_missing: true
+    coverage_guided: false
+
+  # Optional PG immunogen spacing-aware save-score tie-breaker.
+  spacing_aware: false
+  spacing_weight: 0.05
+  spacing_target_distance: 0.0

configs/fabric/default.yaml

@@ -1,3 +1,5 @@
 # @package _global_
 fabric:
   num_nodes: 1
+  accelerator: auto
+  strategy: ddp

disco/model/cogen_inference/cogen_generator.py

@@ -69,6 +69,7 @@ def sample_diffusion_cogen(
     integrator: str = "euler",
     gamma_anneal: str = "none",
     noisy_guidance: DictConfig | None = None,
+    cogen_guidance: DictConfig | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, AtomArray]:
     """Implements Algorithm 18 in AF3.
     It performs denoising steps from time 0 to time T.
@@ -195,6 +196,7 @@ def sample_diffusion_cogen(
             integrator=integrator,
             gamma_anneal=gamma_anneal,
             noisy_guidance=noisy_guidance,
+            cogen_guidance=cogen_guidance,
         )
 
         inf_loop_strat = InferenceLoopStrategyFactory.create_loop_strategy(

disco/model/cogen_inference/cogen_inference_loop_body_impl.py

@@ -28,6 +28,7 @@
     BaseSequenceSamplingStrategy,
     get_seq_t,
 )
+from disco.model.cogen_inference.guidance import build_guidance_controller
 from disco.model.utils import centre_random_augmentation, InferenceNoiseScheduler
 
 logger = logging.getLogger(__name__)
@@ -246,6 +247,7 @@ def __init__(
         integrator: str = "euler",
         gamma_anneal: str = "none",
         noisy_guidance: DictConfig | None = None,
+        cogen_guidance: DictConfig | None = None,
     ):
         self.model = model
         self.sequence_sampling_strategy = sequence_sampling_strategy
@@ -277,6 +279,7 @@ def __init__(
         self.seq_backbone_noise_sigma_min = seq_backbone_noise_sigma_min
         self.seq_backbone_noise_anneal = seq_backbone_noise_anneal
         self.seq_backbone_noise_start_t = seq_backbone_noise_start_t
+        self.guidance_controller = build_guidance_controller(cogen_guidance)
 
         assert integrator in ("euler", "heun"), f"Unknown integrator: {integrator}"
         self.integrator = integrator
@@ -781,6 +784,8 @@ def execute(
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         if self.is_inverse_folding:
             x_l = self.inverse_folding_struct
+        if self.guidance_controller is not None:
+            can_change_idx = self.guidance_controller.modify_change_mask(can_change_idx)
 
         # [..., N_sample, N_atom, 3]
         x_l = (
@@ -1078,6 +1083,16 @@ def execute(
             **seq_logit_inf_input
         )
 
+        if self.guidance_controller is not None:
+            seq_decoder_logits = self.guidance_controller.modify_logits(
+                logits=seq_decoder_logits,
+                xt_seq=xt_seq,
+                can_change_idx=can_change_idx,
+                step=i,
+                x_denoised=x_denoised,
+                x_l=x_l,
+            )
+
         xt_seq_new = self.sequence_sampling_strategy.step(
             seq_decoder_logits,
             x_l,
@@ -1092,6 +1107,11 @@ def execute(
 
         if freeze_seq:
             xt_seq_new = xt_seq
+        elif self.guidance_controller is not None:
+            xt_seq_new = self.guidance_controller.modify_sampled_sequence(
+                xt_seq=xt_seq_new,
+                can_change_idx=can_change_idx,
+            )
 
         xt_seq_changed_idx = xt_seq_new != xt_seq
         xt_seq_changed = xt_seq_changed_idx.any()
@@ -1134,6 +1154,16 @@ def execute(
             step_diag["noisy_guidance_omega_seq"] = self.omega_seq
         if hasattr(self.sequence_sampling_strategy, "last_step_diagnostics"):
             step_diag.update(self.sequence_sampling_strategy.last_step_diagnostics)
+        if self.guidance_controller is not None:
+            step_diag.update(
+                self.guidance_controller.observe_step(
+                    xt_seq=xt_seq,
+                    structure=x_l,
+                    logits=seq_decoder_logits,
+                    can_change_idx=can_change_idx,
+                    step=i,
+                )
+            )
         self.diagnostic_records.append(step_diag)
 
         return x_l, xt_seq, input_feature_dict, xt_seq_changed

disco/model/cogen_inference/cogen_inference_loop_strategies.py

@@ -143,6 +143,34 @@ def run_loop(
                 can_change_idx,
             )
 
+        guidance_controller = self.inner_loop_executor.guidance_controller
+        if (
+            guidance_controller is not None
+            and guidance_controller.select_best_output
+        ):
+            best_output = guidance_controller.select_best_outputs(
+                xt_seq=xt_seq,
+                structure=x_l,
+            )
+            if best_output is not None:
+                best_seq, best_structure, best_details = best_output
+                changed = best_seq != xt_seq
+                if changed.any():
+                    updater = self.inner_loop_executor.ftr_dict_updater
+                    input_feature_dict = updater.update(
+                        input_feature_dict,
+                        self.inner_loop_executor.sample2feat,
+                        best_seq,
+                        changed,
+                        self.inner_loop_executor.bb_only,
+                        self.inner_loop_executor.random_transform_ref_pos,
+                        self.inner_loop_executor.random_transform_msk_res,
+                        self.inner_loop_executor.ref_pos_augment,
+                    )
+                xt_seq = best_seq
+                x_l = best_structure
+                self.inner_loop_executor.diagnostic_records.append(best_details)
+
         return (
             x_l,
             xt_seq,