Freeze backbone during NeuralPLexer NPT confidence runs

BioinfoMachineLearning · Aug 11, 2024 · 23f4c92 · 23f4c92
1 parent b0544cf
commit 23f4c92
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/forks/NeuralPLexer/neuralplexer/model/wrappers.py b/forks/NeuralPLexer/neuralplexer/model/wrappers.py
@@ -759,7 +759,9 @@ def _run_score_head(self, batch, embedding_iter_id, **kwargs):
 
     def _run_confidence_head(self, batch, **kwargs):
         batch = self._init_esdm_inputs(batch, "confidence")
-        return self.confidence_head(batch, frozen_prot=False, **kwargs)
+        return self.confidence_head(
+            batch, frozen_prot=self.global_config.frozen_backbone, **kwargs
+        )
 
     def run_confidence_estimation(self, batch, struct, return_avg_stats=False):
         batch_size = batch["metadata"]["num_structid"]
@@ -3164,6 +3166,8 @@ def sample_pl_complex_structures(
 
             if self.global_config.frozen_backbone:
                 frozen_input_protein_coords = batch["features"]["input_protein_coords"].clone()
+                frozen_input_protein_ca_centroid_coords = frozen_input_protein_coords[:, 1].mean(0)[None, ...]
+                frozen_input_protein_coords[batch["features"]["res_atom_mask"]] -= frozen_input_protein_ca_centroid_coords
 
             if return_all_states:
                 all_frames = []