Merge pull request #1 from hjhk258/main

Fix

mace-bench/util/env.sh

-#!/bin/bash
 
 export CUDA_HOME=/usr/local/cuda
 export PATH=$CUDA_HOME/bin:$PATH

mace-bench/util/mps_clean.sh

-#!/bin/bash
 
 echo quit | nvidia-cuda-mps-control
 nvidia-smi -i 0 -c DEFAULT

mace-bench/util/mps_start.sh

-#!/bin/bash
+
 nvidia-smi -i 0 -c EXCLUSIVE_PROCESS # Set GPU 0 to exclusive mode.  
 nvidia-smi -i 1 -c EXCLUSIVE_PROCESS # Set GPU 1 to exclusive mode.  
 nvidia-smi -i 2 -c EXCLUSIVE_PROCESS # Set GPU 2 to exclusive mode.  

main.py

@@ -23,6 +23,7 @@ if __name__ == '__main__':
     parser.add_argument('--space_group_list', type=str, nargs='+', default=["2,14"], help='Space group list for structure generation, spilt by comma to add mutiple groups, split by space for multiple packings')
     parser.add_argument('--add_name', type=str, nargs='+', default=["CRYSTAL"], help='Add name for the generated structures, split by space for multiple packings')
     parser.add_argument('--max_workers', type=int, default=8, help='Maximum number of workers for parallel processing')
+    parser.add_argument('--mode', type=str, default=8, choices=['all', 'conformer_only', 'structure_only'], help='choose the jobs to do')
     args = parser.parse_args()
 
     target_folder = args.path
@@ -34,6 +35,7 @@ if __name__ == '__main__':
     space_group_list = [list(map(int, group.split(','))) for group in args.space_group_list]
     add_name = args.add_name
     max_workers = args.max_workers
+    mode = args.mode
 
     num_molecules = len(smiles_list)
     num_packings = max(len(molecule_num_in_cell), len(space_group_list))
@@ -63,15 +65,21 @@ if __name__ == '__main__':
     for i in range(num_molecules):
         molecule_folder = os.path.join(target_folder, f"molecule_{i+1}")
         molecule_data.append([])
-        if generate_conformers > 0:
+        if generate_conformers > 0 and mode != "structure_only":
             conformer_search.conformer_search(smiles_list[i], molecule_folder, num_conformers=generate_conformers, max_attempts=10000, rms_thresh=0.1)
             with open(os.path.join(molecule_folder, "info.txt"), "w") as smiles_file:
                 smiles_file.write(f"SMILES: {smiles_list[i]}")
-        for j in range(use_conformers):
+        file_num = len(os.listdir(os.path.join(molecule_folder, "conformers")))
+        cnt = 0
+        for j in range(file_num):
+            if cnt >= use_conformers:
+                break
             temp_path = os.path.join(molecule_folder, "conformers", f"conformer_{j}.xyz")
             if not os.path.exists(temp_path):
                 break
             molecule_data[i].append(format_parser.read_xyz_file(temp_path))
+            cnt += 1
+            
         if len(molecule_data[i]) <= 0:
             print(f"No conformer loaded for molecule_{i+1}. Check configurations!")
             break
@@ -82,6 +90,7 @@ if __name__ == '__main__':
 
     # step2: structure generation
     ##############################################################################################
+    if mode != "conformer_only":
         for i in range(num_packings):
             for combination in combinations:
                 molecule_list = []

post_process/run_remove.sh

-#!/bin/sh -l
 #An example for GPU job.
 #SBATCH -D ./
 #SBATCH --export=ALL