WIP towards writing relaxed structure as modelcif

openfold/np/relax/amber_minimize.py

@@ -524,9 +524,6 @@ def run_pipeline(
     _check_residues_are_well_defined(prot)
     pdb_string = clean_protein(prot, checks=checks)
 
-    # We keep the input around to restore metadata deleted by the relaxer
-    input_prot = prot
-
     exclude_residues = exclude_residues or []
     exclude_residues = set(exclude_residues)
     violations = np.inf

openfold/np/relax/relax.py

@@ -57,7 +57,7 @@ class AmberRelaxation(object):
         self._use_gpu = use_gpu
 
     def process(
-        self, *, prot: protein.Protein
+        self, *, prot: protein.Protein, cif_output: bool
     ) -> Tuple[str, Dict[str, Any], np.ndarray]:
         """Runs Amber relax on a prediction, adds hydrogens, returns PDB string."""
         out = amber_minimize.run_pipeline(
@@ -78,6 +78,8 @@ class AmberRelaxation(object):
             "attempts": out["min_attempts"],
             "rmsd": rmsd,
         }
+        # TODO write all this as ModelCIF if param is true. Should be simply proteint.to_modelcif(prot)
+        #  and then add the other pieces, except that clean_protein() does quite some additional stuff...
         pdb_str = amber_minimize.clean_protein(prot)
         min_pdb = utils.overwrite_pdb_coordinates(pdb_str, min_pos)
         min_pdb = utils.overwrite_b_factors(min_pdb, prot.b_factors)

openfold/utils/script_utils.py

@@ -228,7 +228,7 @@ def prep_output(out, batch, feature_dict, feature_processor, config_preset, mult
     return unrelaxed_protein
 
 
-def relax_protein(config, model_device, unrelaxed_protein, output_directory, output_name):
+def relax_protein(config, model_device, unrelaxed_protein, output_directory, output_name, cif_output):
     amber_relaxer = relax.AmberRelaxation(
         use_gpu=(model_device != "cpu"),
         **config.relax,
@@ -239,7 +239,8 @@ def relax_protein(config, model_device, unrelaxed_protein, output_directory, out
     if "cuda" in model_device:
         device_no = model_device.split(":")[-1]
         os.environ["CUDA_VISIBLE_DEVICES"] = device_no
-    relaxed_pdb_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein)
+    # the struct_str will contain either a PDB-format or a ModelCIF format string
+    struct_str, _, _ = amber_relaxer.process(prot=unrelaxed_protein, cif_output=cif_output)
     os.environ["CUDA_VISIBLE_DEVICES"] = visible_devices
     relaxation_time = time.perf_counter() - t
 
@@ -247,10 +248,13 @@ def relax_protein(config, model_device, unrelaxed_protein, output_directory, out
     update_timings({"relaxation": relaxation_time}, os.path.join(output_directory, "timings.json"))
 
     # Save the relaxed PDB.
+    suffix = "_relaxed.pdb"
+    if cif_output:
+        suffix = "_relaxed.cif"
     relaxed_output_path = os.path.join(
-        output_directory, f'{output_name}_relaxed.pdb'
+        output_directory, f'{output_name}{suffix}'
     )
     with open(relaxed_output_path, 'w') as fp:
-        fp.write(relaxed_pdb_str)
+        fp.write(struct_str)
 
     logger.info(f"Relaxed output written to {relaxed_output_path}...")
\ No newline at end of file

run_pretrained_openfold.py

@@ -288,7 +288,7 @@ def main(args):
             if not args.skip_relaxation:
                 # Relax the prediction.
                 logger.info(f"Running relaxation on {unrelaxed_output_path}...")
-                relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name)
+                relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name, args.cif_output)
 
             if args.save_outputs:
                 output_dict_path = os.path.join(

thread_sequence.py

@@ -106,7 +106,7 @@ def main(args):
         logger.info(f"Output written to {unrelaxed_output_path}...")
 
         logger.info(f"Running relaxation on {unrelaxed_output_path}...")
-        relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name)
+        relax_protein(config, args.model_device, unrelaxed_protein, output_directory, output_name, False)
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()