get_unique_species_in_reference test

lcmd-epfl · Jan 3, 2025 · 6a9905e · 6a9905e
1 parent bc12822
commit 6a9905e
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 22 deletions.
diff --git a/cell2mol/charge_assignment.py b/cell2mol/charge_assignment.py
@@ -28,11 +28,13 @@
 # print("RDKIT Version:", rdBase.rdkitVersion)
 rdBase.DisableLog("rdApp.*")
 
+fullerene = ["C60", "C72", "C80"]
 #######################################################
 def get_possible_charge_state(spec: object, debug: int=0): 
     if not hasattr(spec,"protonation_states"): spec.get_protonation_states(debug=debug)
     if spec.protonation_states is None: 
-        return None 
+        return None
+
     if spec.formula in ["O4-Cl", "N3", "I3"]:
         ch_state = get_charge_manual(spec, debug=debug)
         possible_cs = [ch_state]
@@ -186,7 +188,7 @@ def get_protonation_states_specie(specie: object, debug: int=0) -> list:
     if   specie.type != "specie":                                   return None
     if   specie.subtype == "group":                                 return None
     elif specie.subtype == "molecule" and specie.iscomplex == True: return None     
-    elif (specie.subtype == "molecule" and specie.iscomplex == False) or specie.formula in ["O4-Cl", "N3", "I3"]: 
+    elif (specie.subtype == "molecule" and specie.iscomplex == False) or specie.formula in ["O4-Cl", "N3", "I3"] or specie.formula in fullerene: 
         if debug >= 2: print(f"\nPOSCHARGE: doing empty PROTONATION for this specie {specie.formula} ({specie.subtype})")
         #empty_list = list([np.zeros((len(specie.labels)))])
         empty_list = []
@@ -679,6 +681,7 @@ def get_charge_manual(spec, debug: int=0):
 
         if debug >= 2: print(f"I3: new_order={new_order}")
 
+
     temp_mol = Chem.MolFromSmiles(smiles, sanitize=False)
     mol = Chem.RenumberAtoms(temp_mol, new_order)
     atom_charge = []

diff --git a/cell2mol/refcell.py b/cell2mol/refcell.py
@@ -8,6 +8,7 @@
 from cell2mol.new_cell_reconstruction import modify_cov_factor_due_to_H, modify_cov_factor_due_to_possible_charges
 from cell2mol.other import handle_error
 import time
+import numpy as np
 
 # Constants
 VERSION = "2.0"
@@ -21,29 +22,40 @@ def process_refcell(input_path, name, current_dir, debug=0):
     output_fname = os.path.join(current_dir, "cell2mol.out")
     summary_fname = os.path.join(current_dir, "summary.out")
 
+    if os.path.exists(ref_cell_fname):
+        with open(output_fname, "a") as output:
+            with redirect_stdout(output):
+                print("=====================================================")
+                print(f"cell2mol version {VERSION}")
+                print(f"Reference cell file {ref_cell_fname} already exists. Skipping reference cell generation.")
+                print(f"Debug level: {debug}")
+                refcell = np.load(ref_cell_fname, allow_pickle=True)
+                if refcell.error_case == 0:
+                    get_unique_species_in_reference(refcell, debug) 
+                else:
+                    print(f"Error occurred in processing reference cell: error case {refcell.error_case}")   
+                refcell.save(ref_cell_fname)
+    else:
+        with open(output_fname, "w") as output:
+            with redirect_stdout(output):
+                print(f"cell2mol version {VERSION}")
+                print(f"INITIATING cell object from input path: {input_path}")
+                print(f"Debug level: {debug}")
 
-    # Redirect stdout to the output file for logging
-    with open(output_fname, "w") as output:
-        with redirect_stdout(output):
-            print(f"cell2mol version {VERSION}")
-            print(f"INITIATING cell object from input path: {input_path}")
-            print(f"Debug level: {debug}")
+                # # Read .cif file
+                structure = read(input_path)
+                cell_vector = structure.cell.array
+                cell_param = structure.cell.cellpar()      
 
-            # # Read .cif file
-            structure = read(input_path)
-            cell_vector = structure.cell.array
-            cell_param = structure.cell.cellpar()      
+                # Create the reference cell
+                refcell = create_reference(input_path, name, cell_vector, cell_param, debug)
 
-            # Create the reference cell
-            refcell = create_reference(input_path, name, cell_vector, cell_param, debug)
-
-            # Finalize and save the reference cell object if no errors
-            if refcell.error_case == 0:
-                pass
-                # get_unique_species_in_reference(refcell, debug)
-            else:
-                print(f"Error occurred in processing reference cell: error case {refcell.error_case}")
-            refcell.save(ref_cell_fname)
+                # Finalize and save the reference cell object if no errors
+                if refcell.error_case == 0:
+                    pass
+                else:
+                    print(f"Error occurred in processing reference cell: error case {refcell.error_case}")
+                refcell.save(ref_cell_fname)
 
     error_fname = os.path.join(current_dir, f"reference_error_{refcell.error_case}.out")
     with open(error_fname, "w") as error_output: