From aae325ae2808f7070fcf8ac23bbfbdf5d1f0f6b0 Mon Sep 17 00:00:00 2001
From: Ran Lu <ranl@princeton.edu>
Date: Wed, 29 Jan 2025 14:06:18 -0500
Subject: [PATCH] Replace imp module with importlib

imp was removed since python 3.12
---
 deepem/test/utils.py     |  3 +--
 deepem/train/data.py     |  9 +++++----
 deepem/train/utils.py    |  6 ++++--
 deepem/utils/py_utils.py | 10 ++++++++++
 4 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/deepem/test/utils.py b/deepem/test/utils.py
index 58d773b..893f40e 100644
--- a/deepem/test/utils.py
+++ b/deepem/test/utils.py
@@ -1,4 +1,3 @@
-import imp
 import numpy as np
 import os
 from types import SimpleNamespace
@@ -11,7 +10,7 @@
 
 def load_model(opt):
     # Create a model.
-    mod = imp.load_source('model', opt.model)
+    mod = py_utils.load_module('model', opt.model)
     if opt.onnx:
         model = OnnxModel(mod.create_model(opt), opt)
     else:
diff --git a/deepem/train/data.py b/deepem/train/data.py
index e320411..df05c3a 100644
--- a/deepem/train/data.py
+++ b/deepem/train/data.py
@@ -1,4 +1,5 @@
-import imp
+from deepem.utils.py_utils import load_module
+
 import numpy as np
 
 import torch
@@ -45,20 +46,20 @@ def requires_grad(self, key):
     def build(self, opt, data, is_train, prob):
         # Data augmentation
         if opt.augment:
-            mod = imp.load_source('augment', opt.augment)
+            mod = load_module('augment', opt.augment)
             aug = mod.get_augmentation(is_train, **opt.aug_params)
         else:
             aug = None
 
         # Data sampler
-        mod = imp.load_source('sampler', opt.sampler)
+        mod = load_module('sampler', opt.sampler)
         spec = mod.get_spec(opt.in_spec, opt.out_spec)
         zspecs = opt.zettaset_specs
         sampler = mod.Sampler(data, spec, is_train, aug, prob, zspecs)
 
         # Sample modifier
         if opt.modifier:
-            mod = imp.load_source('modifier', opt.modifier)
+            mod = load_module('modifier', opt.modifier)
             self.modifier = mod.Modifier(**opt.modifier_kwargs)
         else:
             def default_modifier(x, **kwargs):
diff --git a/deepem/train/utils.py b/deepem/train/utils.py
index b09cf47..fb6bf4f 100644
--- a/deepem/train/utils.py
+++ b/deepem/train/utils.py
@@ -1,4 +1,3 @@
-import imp
 import os
 import glob
 
@@ -9,6 +8,7 @@
 from deepem.train.data import Data
 from deepem.train.model import Model, AmpModel
 from deepem.loss.utils import BinaryWeightBalancer
+from deepem.utils.py_utils import load_module
 
 
 def get_criteria(opt):
@@ -59,7 +59,9 @@ def get_criteria(opt):
 
 def load_model(opt):
     # Create a model.
-    mod = imp.load_source('model', opt.model)
+
+    mod = load_module("model", opt.model)
+
     if opt.mixed_precision:
         model = AmpModel(mod.create_model(opt), get_criteria(opt), opt)
     else:
diff --git a/deepem/utils/py_utils.py b/deepem/utils/py_utils.py
index 5f676e2..c7f5a37 100644
--- a/deepem/utils/py_utils.py
+++ b/deepem/utils/py_utils.py
@@ -4,6 +4,16 @@
 
 from sklearn.decomposition import PCA
 
+import importlib
+import importlib.util
+
+
+def load_module(name, path):
+    spec = importlib.util.spec_from_file_location(name, path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
 
 def dict2tuple(d):
     return namedtuple('GenericDict', d.keys())(**d)