Removed unused imports and hoisted all nested imports to the global module scope

Author: alinelena

src/process_omol25/cli.py

@@ -5,6 +5,7 @@
 from pathlib import Path
 
 from .process_omol25 import S3DataProcessor, setup_logging
+from mpi4py import MPI as mpi
 
 
 def parse_args():
@@ -92,7 +93,6 @@ def main():
     args = parse_args()
 
     if args.mpi:
-        from mpi4py import MPI as mpi
         comm = mpi.COMM_WORLD
         rank = comm.Get_rank()
         size = comm.Get_size()

src/process_omol25/download_omol25.py

@@ -10,9 +10,10 @@
 from botocore.config import Config
 from tqdm import tqdm
 import time
-import psutil
 
-from .process_omol25 import setup_logging, get_ranges
+from .process_omol25 import setup_logging
+from mpi4py import MPI
+from mpi4py import MPI as mpi
 
 logger = logging.getLogger(__name__)
 
@@ -121,7 +122,6 @@ def extract_buffer(buffer: BytesIO, x: str, k: str):
 
 def manager_loop(keys, data, restart_file, comm, size):
     """Rank 0 Dispatcher for downloads (Synchronous)."""
-    from mpi4py import MPI
     start_time = time.time()
     logger.info(f"Download Manager starting with {size - 1} workers.")
 
@@ -158,7 +158,6 @@ def manager_loop(keys, data, restart_file, comm, size):
 
 def worker_loop(data, args, comm):
     """Rank > 0 Downloader (Synchronous)."""
-    from mpi4py import MPI
     s3_client = None
     if not args.local_dir:
         with open(args.login_file, "r") as f:
@@ -211,7 +210,6 @@ def main():
     args = parse_args()
 
     if args.mpi:
-        from mpi4py import MPI as mpi
         comm = mpi.COMM_WORLD
         rank = comm.Get_rank()
         size = comm.Get_size()

src/process_omol25/process_omol25.py

@@ -14,15 +14,14 @@
 ase.parallel.world = DummyMPI()
 import boto3
 from botocore.config import Config
-from botocore.exceptions import ClientError
 from tarfile import open as tar_open
 from zstandard import ZstdDecompressor
-from typing import Optional, Dict, Any, List, Tuple
+from typing import Optional, Dict, Any, Tuple
 import pandas as pd
 import numpy as np
-import psutil
 from tqdm import tqdm
 import time
+from mpi4py import MPI
 
 MiB = 1024 ** 2
 GiB = 1024 * MiB
@@ -491,7 +490,6 @@ def process_single(self, idx: int, s3_client=None):
 
     def _manager_loop(self):
         """Rank 0 Result Collector loop (Hybrid RMA)."""
-        from mpi4py import MPI
         start_time = time.time()
         num_tasks = len(self.indices_to_process)
         logger.info(f"Manager starting. Collecting results for {num_tasks} tasks.")
@@ -533,7 +531,6 @@ def _manager_loop(self):
 
     def _worker_loop(self):
         """Rank > 0 Processing loop (Hybrid RMA)."""
-        from mpi4py import MPI
         s3_client = None
         if not self.local_dir:
             s3_client = boto3.client(
@@ -642,7 +639,6 @@ def _final_merge(self, elapsed_time):
 
     def run_mpi(self):
         """Main entry point for MPI runs (Hybrid RMA)."""
-        from mpi4py import MPI
         self.comm = MPI.COMM_WORLD
         self.rank = self.comm.Get_rank()
         self.size = self.comm.Get_size()

src/process_omol25/verify_processed_omol25.py

@@ -10,6 +10,7 @@
 
 try:
     import pandas as pd
+    import numpy as np
     import ase.io
 except ImportError as e:
     print(f"Error: Missing required dependency. {e}")
@@ -122,7 +123,6 @@ def main():
                         mismatches += 1
                 elif pq_val != xyz_val:
                     try:
-                        import numpy as np
                         if isinstance(pq_val, np.ndarray) or isinstance(xyz_val, np.ndarray):
                             if not np.allclose(pq_val, xyz_val, equal_nan=True):
                                 print(f"  [Mismatch] sha={sha}, key={key}: NDArray mismatch")

tests/test_verify_processed_omol25.py

@@ -1,6 +1,5 @@
 import subprocess
 import pandas as pd
-from pathlib import Path
 from ase import Atoms
 from ase.io import write