diff --git a/src/moldrug/compat.py b/src/moldrug/compat.py new file mode 100644 index 0000000..eda3f14 --- /dev/null +++ b/src/moldrug/compat.py @@ -0,0 +1,28 @@ +""" +This module provides backward compatibility with moldrug <= 4.0. + +It allows pickle objects created with older versions of moldrug +to be successfully loaded after internal refactoring (e.g. moved classes). +""" +import dill +import importlib +from typing import Dict, Tuple + +# Mapping from old (module, class) → new (module, class) +_CLASS_RENAMES: Dict[Tuple[str, str], Tuple[str, str]] = { + ("moldrug.utils", "GA"): ("moldrug.opt", "GA"), + ("moldrug.utils", "Local"): ("moldrug.opt", "Local"), + # Add more mappings here as needed +} + + +class BackCompatUnpickler(dill.Unpickler): + """Custom unpickler that redirects renamed or moved classes.""" + + def find_class(self, module: str, name: str): + key = (module, name) + if key in _CLASS_RENAMES: + new_module, new_name = _CLASS_RENAMES[key] + mod = importlib.import_module(new_module) + return getattr(mod, new_name) + return super().find_class(module, name) diff --git a/src/moldrug/opt.py b/src/moldrug/opt.py index 53c3da3..88732e0 100644 --- a/src/moldrug/opt.py +++ b/src/moldrug/opt.py @@ -219,7 +219,7 @@ class GA: randomseed : Union[None, int] The random seed to use with random module. __moldrug_version__ : str - The molDrug version. + The moldrug version. costfunc : object The cost function set by the user. crem_db_path : str diff --git a/src/moldrug/utils.py b/src/moldrug/utils.py index 334dba2..8a6ff45 100644 --- a/src/moldrug/utils.py +++ b/src/moldrug/utils.py @@ -19,7 +19,7 @@ RDKitMolCreate) from rdkit import Chem, RDLogger from rdkit.Chem import AllChem, DataStructs, Descriptors, Lipinski, rdFMCS - +from moldrug.compat import BackCompatUnpickler logger = logging.getLogger(__name__) RDLogger.DisableLog('rdApp.*') @@ -465,6 +465,18 @@ def confgen(mol: Chem.rdchem.Mol, return_mol: bool = False, randomseed: Union[in else: return pdbqt_string + AllChem.EmbedMolecule(mol, randomSeed=randomSeed) + # The optimization introduce some sort of non-reproducible results. + # For that reason is not used when randomseed is set + if not randomseed: + AllChem.MMFFOptimizeMolecule(mol, maxIters=500) + preparator = MoleculePreparation() + mol_setups = preparator.prepare(mol) + pdbqt_string = PDBQTWriterLegacy.write_string(mol_setups[0])[0] + if return_mol: + return (pdbqt_string, mol) + else: + return pdbqt_string def update_reactant_zone(parent: Chem.rdchem.Mol, offspring: Chem.rdchem.Mol, parent_replace_ids: List[int] = None, parent_protected_ids: List[int] = None): @@ -841,9 +853,8 @@ def decompress_pickle(file: str): object The python object. """ - data = bz2.BZ2File(file, 'rb') - data = pickle.load(data) - return data + with bz2.BZ2File(file, 'rb') as f: + return BackCompatUnpickler(f).load() def is_iter(obj): diff --git a/streamlit/moldrug-dashboard.py b/streamlit/moldrug-dashboard.py index 1c265fb..9139ad6 100644 --- a/streamlit/moldrug-dashboard.py +++ b/streamlit/moldrug-dashboard.py @@ -21,7 +21,7 @@ import streamlit as st import streamlit.components.v1 as components -from moldrug import utils +from moldrug import utils, opt # TODO # add SyGma for metabolic prediction @@ -333,10 +333,10 @@ def lig_prot_overview(_pop, protein_pdb_string): def load_pbz2(pbz2): moldrug_result = utils.decompress_pickle(pbz2) is_GA = False - if isinstance(moldrug_result, utils.GA): + if isinstance(moldrug_result, opt.GA): gen, pop = moldrug_result.NumGens, moldrug_result.pop is_GA = True - elif isinstance(moldrug_result, utils.Local): + elif isinstance(moldrug_result, opt.Local): gen, pop = 0, moldrug_result.pop elif isinstance(moldrug_result, tuple): if isinstance(moldrug_result[0], int) and isinstance(moldrug_result[1][0], utils.Individual): @@ -608,7 +608,7 @@ def get_pubchem_dataframe(df: pd.DataFrame) -> pd.DataFrame: if 'kept_gens' in grid.dataframe.columns: props_to_drop.append('kept_gens') - prop_df = grid.dataframe.drop(['mol', 'pdbqt', 'kept_gens', 'img', 'mols2grid-id'], axis=1).set_index('idx') + prop_df = grid.dataframe.drop(props_to_drop, axis=1).set_index('idx') st.download_button( "Press to Download", convert_df(prop_df),