From 9f28b34df9b65ea6a665967c830202e7c3c2423e Mon Sep 17 00:00:00 2001
From: Alex Morehead <alex.morehead@gmail.com>
Date: Sat, 23 Nov 2024 17:22:40 -0600
Subject: [PATCH] Handle an edge case in relaxation

---
 posebench/models/minimize_energy.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/posebench/models/minimize_energy.py b/posebench/models/minimize_energy.py
index c17605b..f5e7223 100644
--- a/posebench/models/minimize_energy.py
+++ b/posebench/models/minimize_energy.py
@@ -16,12 +16,12 @@
 import shutil
 import sys
 from pathlib import Path
-from typing import Any
 
 import hydra
 import numpy as np
 import rootutils
 import timeout_decorator
+from beartype.typing import Any, Dict, Optional
 from Bio.PDB import MMCIFIO, PDBIO, MMCIFParser, PDBParser, Select
 from omegaconf import DictConfig
 from openff.toolkit.topology import Molecule
@@ -1208,7 +1208,7 @@ def optimize_ligand_in_pocket(
     config_path="../../configs/model",
     config_name="minimize_energy.yaml",
 )
-def minimize_energy(cfg: DictConfig):
+def minimize_energy(cfg: DictConfig) -> Optional[Dict[str, Any]]:
     """Minimize the energy of a ligand in a protein pocket using OpenMM."""
     logger.setLevel(cfg.log_level)
 
@@ -1272,6 +1272,30 @@ def minimize_energy(cfg: DictConfig):
         except OSError:
             pass
 
+        logger.info(f"Finalizing complex relaxation with relax_protein={cfg.relax_protein}")
+
+        # clean up temporary files
+        try:
+            os.remove(temp_protein_file_path)
+        except OSError:
+            pass
+        try:
+            os.remove(temp_ligand_file_path)
+        except OSError:
+            pass
+
+        temp_name = temp_protein_file_path.stem if cfg.name is None else cfg.name
+        try:
+            os.remove(temp_directory / f"{temp_name}_prepped_protein.pdb")
+        except OSError:
+            pass
+        try:
+            os.remove(temp_directory / f"{temp_name}_prepped_ligand.sdf")
+        except OSError:
+            pass
+
+        return None
+
     if prep_only:
         sys.exit(0)