Format fix. More options in readme

basic_function/format_parser.py

-import re
-from basic_function import operation
-from basic_function import data_classes
-from basic_function import chemical_knowledge
-import copy
-
-
-def read_xyz_file(file_path):
-    input_file = open(file_path, 'r')
-    lines = input_file.readlines()
-    number_of_atoms = int(lines[0])
-    name = str(lines[1][:-1])
-    atoms = []
-    for index,line in enumerate(lines):
-        split_line = list(filter(lambda x: x != '', re.split("\\s+", line)))
-        if len(split_line)==4 and operation.is_number(split_line[1]) and \
-                operation.is_number(split_line[2]) and operation.is_number(split_line[3]):
-            atoms.append(data_classes.Atom(element=split_line[0],
-                                            cart_xyz=[float(split_line[1]), float(split_line[2]), float(split_line[3])],
-                                           atom_id=index-2))
-
-    if number_of_atoms!=len(atoms):
-        print("Warning! The length of atoms don't match the number of atoms given")
-
-    molecule = data_classes.Molecule(atoms=atoms, name=name, system_name=name)
-
-    return molecule
-
-def write_cif_file(crystal, sym=False, name="zcx"):
-    """
-    Accept crystal class, give the cif file out
-    :param crystal: crystal class
-    :param coordinates: frac or cart
-    :param sym: False:give all atoms out; True:with symmetry
-    :param name: file name
-    :return: cif_out
-    cif file in list format should be print using the following function:
-                                        target=open("D:\\zcx.cif",'w')
-                                        target.writelines(cif_out)
-                                        target.close()
-    """
-    
-    if crystal.system_name!="unknown":
-        name = crystal.system_name
-    cif_file = []
-
-    cif_file.append("data_"+str(name)+"\n")
-    if sym==False:
-        if crystal.space_group==1:
-            crystal_temp = crystal
-        else:
-            crystal_temp = copy.deepcopy(crystal)
-            crystal_temp.make_p1()
-        cif_file.append("_symmetry_space_group_name_H-M    \'P1\'"+"\n")
-        cif_file.append("_symmetry_Int_Tables_number       1"+"\n")
-
-        cif_file.append("loop_"+"\n")
-        cif_file.append("_symmetry_equiv_pos_site_id"+"\n")
-        cif_file.append("_symmetry_equiv_pos_as_xyz"+"\n")
-        cif_file.append("1 x,y,z"+"\n")
-        cif_file.append("_cell_length_a                   "+str(crystal_temp.cell_para[0][0])+"\n")
-        cif_file.append("_cell_length_b                   "+str(crystal_temp.cell_para[0][1])+"\n")
-        cif_file.append("_cell_length_c                   "+str(crystal_temp.cell_para[0][2])+"\n")
-        cif_file.append("_cell_angle_alpha                "+str(crystal_temp.cell_para[1][0])+"\n")
-        cif_file.append("_cell_angle_beta                 "+str(crystal_temp.cell_para[1][1])+"\n")
-        cif_file.append("_cell_angle_gamma                "+str(crystal_temp.cell_para[1][2])+"\n")
-
-        cif_file.append("loop_"+"\n")
-        cif_file.append("_atom_site_label"+"\n")
-        cif_file.append("_atom_site_type_symbol"+"\n")
-        cif_file.append("_atom_site_fract_x"+"\n")
-        cif_file.append("_atom_site_fract_y"+"\n")
-        cif_file.append("_atom_site_fract_z"+"\n")
-        for i in range(0,len(crystal_temp.atoms)):
-            cif_file.append("{:6}  {:4} {:16.8f} {:16.8f} {:16.8f}\n"
-                             .format(i+1,crystal_temp.atoms[i].element,crystal_temp.atoms[i].frac_xyz[0],
-                                     crystal_temp.atoms[i].frac_xyz[1],crystal_temp.atoms[i].frac_xyz[2]))
-
-        return cif_file
-
-    elif sym==True:
-        cif_file.append("_symmetry_space_group_name_H-M    \'{}\'".format(chemical_knowledge.space_group[crystal.space_group][1])+"\n")
-        cif_file.append("_symmetry_Int_Tables_number       {}".format(crystal.space_group)+"\n")
-
-        cif_file.append("loop_"+"\n")
-        cif_file.append("_symmetry_equiv_pos_site_id"+"\n")
-        cif_file.append("_symmetry_equiv_pos_as_xyz"+"\n")
-        for idx, SYMM in enumerate(crystal.SYMM):
-            cif_file.append("{} {}".format(idx+1,SYMM)+"\n")
-        cif_file.append("_cell_length_a                   "+str(crystal.cell_para[0][0])+"\n")
-        cif_file.append("_cell_length_b                   "+str(crystal.cell_para[0][1])+"\n")
-        cif_file.append("_cell_length_c                   "+str(crystal.cell_para[0][2])+"\n")
-        cif_file.append("_cell_angle_alpha                "+str(crystal.cell_para[1][0])+"\n")
-        cif_file.append("_cell_angle_beta                 "+str(crystal.cell_para[1][1])+"\n")
-        cif_file.append("_cell_angle_gamma                "+str(crystal.cell_para[1][2])+"\n")
-
-        cif_file.append("loop_"+"\n")
-        cif_file.append("_atom_site_label"+"\n")
-        cif_file.append("_atom_site_type_symbol"+"\n")
-        cif_file.append("_atom_site_fract_x"+"\n")
-        cif_file.append("_atom_site_fract_y"+"\n")
-        cif_file.append("_atom_site_fract_z"+"\n")
-        for i in range(0,len(crystal.atoms)):
-            cif_file.append("{:6}  {:4} {:16.8f} {:16.8f} {:16.8f}\n"
-                             .format(i+1,crystal.atoms[i].element,crystal.atoms[i].frac_xyz[0],
-                                     crystal.atoms[i].frac_xyz[1],crystal.atoms[i].frac_xyz[2]))
-
-        return cif_file
-
-
-def write_cifs_file(crystals, sym=False, name="zcx"):
-    cifs_file = []
-    for crystal in crystals:
-        single_cif = write_cif_file(crystal,sym=sym, name=name)
-        cifs_file.extend(single_cif)
-    return cifs_file
-
-
-def read_cif_file(file_path,on_sym_check=False,shut_up=False,system_name="unknown",comment_name="unknown"):
-    input_file = open(file_path, 'r')
-    lines = input_file.readlines()
-    step_pickle = []
-    crystal_all = []
-    if system_name=="unknown":
-        no_name = True
-    else:
-        no_name = False
-    # first time scan
-    for index,line in enumerate(lines):
-        # find out all the step pickle
-        if line.startswith("data_"):
-            step_pickle.append(index)
-    step_pickle.append(len(lines))
-
-    # treat every step and return a crystal
-    for m in range(0,len(step_pickle)-1):
-        atoms = []
-        atoms_P1 = []
-        SYMM = []
-        cell_para = [["unknown","unknown","unknown"],["unknown","unknown","unknown"]]
-
-        for index, line in enumerate(lines[step_pickle[m]:step_pickle[m+1]]):
-            split_line = list(filter(lambda x: x != '', re.split("\\s+", line)))
-            if line.startswith("#"):
-                pass
-            elif len(split_line)==0:
-                pass
-            # read the loop of symmetry
-            # elif split_line[0]=="loop_" and lines[step_pickle[m]+index+1]=="_symmetry_equiv_pos_as_xyz\n":
-            elif split_line[0] == "loop_" and lines[step_pickle[m] + index + 1] == "_symmetry_equiv_pos_as_xyz\n":
-                temp_number = 1
-                while "_" not in lines[step_pickle[m]+index+1+temp_number]:
-                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))
-                    temp_number+=1
-                    if not operation.is_number(split_line_temp[0]):
-                        SYMM.append(split_line_temp[0])
-                    else:
-                        SYMM.append(split_line_temp[1])
-            elif split_line[0] == "loop_" and lines[step_pickle[m]+index+2].strip(" ")=="_symmetry_equiv_pos_as_xyz\n":
-                temp_number = 1
-                while "_" not in lines[step_pickle[m]+index+2+temp_number]:
-                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+2+temp_number])))
-                    temp_number+=1
-                    SYMM.append("".join(split_line_temp[1:]))
-            elif split_line[0] == "loop_" and "_space_group_symop_operation_xyz\n" in lines[step_pickle[m] + index + 1]:
-                # ase format
-                temp_number = 1
-                while "_" not in lines[step_pickle[m]+index+1+temp_number]:
-                    if lines[step_pickle[m]+index+1+temp_number]=="\n":
-                        temp_number += 1
-                        continue
-                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))
-                    temp_number+=1
-                    if not operation.is_number(split_line_temp[0]):
-                        SYMM.append("".join(split_line_temp))
-
-            # read the loop of atoms:
-            elif (split_line[0] == "loop_" and lines[step_pickle[m] + index + 1].strip(" ") == "_atom_site_label\n") or \
-                    (split_line[0] == "loop_" and lines[step_pickle[m] + index + 2].strip(" ") == "_atom_site_label\n"):
-                temp_number = 0
-                while "_" in lines[step_pickle[m]+index+1+temp_number]:
-                    if lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_type_symbol\n":
-                        ele_pos = temp_number
-                    elif lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_fract_x\n":
-                        x_pos = temp_number
-                    elif lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_fract_y\n":
-                        y_pos = temp_number
-                    elif lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_fract_z\n":
-                        z_pos = temp_number
-                    temp_number+=1
-                how_long = temp_number
-
-                while len(list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))) == how_long:
-                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))
-                    atoms.append(data_classes.Atom(element=split_line_temp[ele_pos],
-                                           frac_xyz=[float(split_line_temp[x_pos]),float(split_line_temp[y_pos]),
-                                                     float(split_line_temp[z_pos])]))
-                    temp_number += 1
-                    if step_pickle[m]+index+1+temp_number==len(lines):
-                        break
-
-            elif split_line[0] == "_cell_length_a":
-                cell_para[0][0] = float(split_line[1])
-            elif split_line[0] == "_cell_length_b":
-                cell_para[0][1] = float(split_line[1])
-            elif split_line[0] == "_cell_length_c":
-                cell_para[0][2] = float(split_line[1])
-            elif split_line[0] == "_cell_angle_alpha":
-                cell_para[1][0] = float(split_line[1])
-            elif split_line[0] == "_cell_angle_beta":
-                cell_para[1][1] = float(split_line[1])
-            elif split_line[0] == "_cell_angle_gamma":
-                cell_para[1][2] = float(split_line[1])
-            elif "data_" in line:
-                if no_name == True:
-                    system_name = line[5:]
-                    system_name = system_name.replace(" ","_")
-                    system_name = system_name.replace("\\", "_")
-                    system_name = system_name.replace("\n", "")
-        for atom in atoms:
-            all_reflect_position = operation.space_group_transfer_for_single_atom(atom.frac_xyz, SYMM)
-            for new_position in all_reflect_position:
-                atoms_P1.append(data_classes.Atom(element=atom.element,
-                                            frac_xyz=[new_position[0], new_position[1], new_position[2]]))
-        crystal_all.append(data_classes.Crystal(cell_para=cell_para, atoms=atoms_P1, comment=comment_name, system_name=system_name))
-        if on_sym_check == True:
-            raise Exception("Not finished part, TODO in code")
-        if shut_up==False:
-            if m%100 == 0:
-                print("{} structures have been treated".format(m))
-
-    return crystal_all
-
-
-def write_poscar_file(crystal, coordinates = 'frac', name = "parser_zcx_create"):
-
-    vasp_file = []
-
-    vasp_file.append('{}\n'.format(name))
-    vasp_file.append('1.0\n')
-    cell_vect = crystal.cell_vect
-    for vect in cell_vect:
-        vasp_file.append("{:16.8f} {:16.8f} {:16.8f}\n".format(vect[0],vect[1],vect[2]))
-    crystal.sort_by_element()
-    vasp_file.append("".join("{:>6s}".format(x) for x in crystal.get_element()) + "\n")
-    vasp_file.append("".join("{:>6.0f}".format(x) for x in crystal.get_element_amount()) + "\n")
-    if coordinates == 'frac':
-        vasp_file.append('Direct\n')
-        for ELEMENT in crystal.get_element():
-            for ATOM in crystal.atoms:
-                if ATOM.element == ELEMENT:
-                    vasp_file.append(
-                        "{:16.8f} {:16.8f} {:16.8f}\n".format(ATOM.frac_xyz[0], ATOM.frac_xyz[1], ATOM.frac_xyz[2]))
-    elif coordinates == 'cart':
-        vasp_file.append('Cartesian\n')
-        for ELEMENT in crystal.get_element():
-            for ATOM in crystal.atoms:
-                if ATOM.element == ELEMENT:
-                    vasp_file.append(
-                        "{:16.8f} {:16.8f} {:16.8f}\n".format(ATOM.cart_xyz[0], ATOM.cart_xyz[1], ATOM.cart_xyz[2]))
-    else:
-        raise Exception("Wrong coordinates type: {}".format(coordinates))
-
-    return vasp_file
-    
-    
-def read_ase_pbc_file(file_path,shut_up=False):
-    input_file = open(file_path, 'r')
-    lines = input_file.readlines()[2:]
-    step_pickle = []
-    crystal_all = []
-
-    # first time scan
-    for index,line in enumerate(lines):
-        # find out all the step pickle
-        if line.startswith("Step "):
-            step_pickle.append(index)
-    step_pickle.append(len(lines))
-
-    # treat every step and return a crystal
-    for m in range(0,len(step_pickle)-1):
-        atoms_P1 = []
-        force_matrix = []
-        position_matrix = []
-        in_forces = False
-        in_positions = False
-
-
-        for index, line in enumerate(lines[step_pickle[m]:step_pickle[m+1]]):
-            # split_line = list(filter(lambda x: x != '', re.split("\\s+", line)))
-            line = line.strip()
-            # check Forces part
-            if line.startswith("Forces:"):
-                in_forces = True
-                in_positions = False
-                continue
-
-            # check Positions part
-            if line.startswith("Positions:"):
-                in_positions = True
-                in_forces = False
-                continue
-
-            if in_forces and line.startswith("[") and line.endswith("]"):
-                line = line.replace("[", "").replace("]", "")
-                force_matrix.append([float(x) for x in line.split()])
-
-            # analyse Positions part
-            if in_positions and line.startswith("[") and line.endswith("]"):
-                line = line.replace("[", "").replace("]", "")
-                position_matrix.append([float(x) for x in line.split()])
-
-            if line.startswith("Elements:"):
-                elements_string = line.strip().split(":", 1)[-1].strip()
-                elements_string = elements_string[1:-1]  
-
-                elements = [elem.strip().strip("'") for elem in elements_string.split(",")]
-
-            if line.startswith("cell:"):
-                matrix_string = line[len("cell: Cell("):-1]
-                rows = matrix_string.split("], [")
-                cell_vect = [
-                    [float(value) for value in row.replace('[', '').replace(']', '').replace(')', '').split(", ")]
-                    for row in rows
-                ]
-
-        for i in range(0,len(elements)):
-            atoms_P1.append(data_classes.Atom(element=elements[i], cart_xyz=[position_matrix[i][0], position_matrix[i][1], position_matrix[i][2]]))
-
-        crystal_all.append(data_classes.Crystal(cell_vect=cell_vect, atoms=atoms_P1))
-
-    return crystal_all    
+import re
+from basic_function import operation
+from basic_function import data_classes
+from basic_function import chemical_knowledge
+import copy
+
+
+def read_xyz_file(file_path):
+    input_file = open(file_path, 'r')
+    lines = input_file.readlines()
+    number_of_atoms = int(lines[0])
+    name = str(lines[1][:-1])
+    atoms = []
+    for index,line in enumerate(lines):
+        split_line = list(filter(lambda x: x != '', re.split("\\s+", line)))
+        if len(split_line)==4 and operation.is_number(split_line[1]) and \
+                operation.is_number(split_line[2]) and operation.is_number(split_line[3]):
+            atoms.append(data_classes.Atom(element=split_line[0],
+                                            cart_xyz=[float(split_line[1]), float(split_line[2]), float(split_line[3])],
+                                           atom_id=index-2))
+
+    if number_of_atoms!=len(atoms):
+        print("Warning! The length of atoms don't match the number of atoms given")
+
+    molecule = data_classes.Molecule(atoms=atoms, name=name, system_name=name)
+
+    return molecule
+
+def write_cif_file(crystal, sym=False, name="zcx"):
+    """
+    Accept crystal class, give the cif file out
+    :param crystal: crystal class
+    :param coordinates: frac or cart
+    :param sym: False:give all atoms out; True:with symmetry
+    :param name: file name
+    :return: cif_out
+    cif file in list format should be print using the following function:
+                                        target=open("D:\\zcx.cif",'w')
+                                        target.writelines(cif_out)
+                                        target.close()
+    """
+    
+    if crystal.system_name!="unknown":
+        name = crystal.system_name
+    cif_file = []
+
+    cif_file.append("data_"+str(name)+"\n")
+    if sym==False:
+        if crystal.space_group==1:
+            crystal_temp = crystal
+        else:
+            crystal_temp = copy.deepcopy(crystal)
+            crystal_temp.make_p1()
+        cif_file.append("_symmetry_space_group_name_H-M    \'P1\'"+"\n")
+        cif_file.append("_symmetry_Int_Tables_number       1"+"\n")
+
+        cif_file.append("loop_"+"\n")
+        cif_file.append("_symmetry_equiv_pos_site_id"+"\n")
+        cif_file.append("_symmetry_equiv_pos_as_xyz"+"\n")
+        cif_file.append("1 x,y,z"+"\n")
+        cif_file.append("_cell_length_a                   "+str(crystal_temp.cell_para[0][0])+"\n")
+        cif_file.append("_cell_length_b                   "+str(crystal_temp.cell_para[0][1])+"\n")
+        cif_file.append("_cell_length_c                   "+str(crystal_temp.cell_para[0][2])+"\n")
+        cif_file.append("_cell_angle_alpha                "+str(crystal_temp.cell_para[1][0])+"\n")
+        cif_file.append("_cell_angle_beta                 "+str(crystal_temp.cell_para[1][1])+"\n")
+        cif_file.append("_cell_angle_gamma                "+str(crystal_temp.cell_para[1][2])+"\n")
+
+        cif_file.append("loop_"+"\n")
+        cif_file.append("_atom_site_label"+"\n")
+        cif_file.append("_atom_site_type_symbol"+"\n")
+        cif_file.append("_atom_site_fract_x"+"\n")
+        cif_file.append("_atom_site_fract_y"+"\n")
+        cif_file.append("_atom_site_fract_z"+"\n")
+        for i in range(0,len(crystal_temp.atoms)):
+            cif_file.append("{:6}  {:4} {:16.8f} {:16.8f} {:16.8f}\n"
+                             .format(i+1,crystal_temp.atoms[i].element,crystal_temp.atoms[i].frac_xyz[0],
+                                     crystal_temp.atoms[i].frac_xyz[1],crystal_temp.atoms[i].frac_xyz[2]))
+
+        return cif_file
+
+    elif sym==True:
+        cif_file.append("_symmetry_space_group_name_H-M    \'{}\'".format(chemical_knowledge.space_group[crystal.space_group][1])+"\n")
+        cif_file.append("_symmetry_Int_Tables_number       {}".format(crystal.space_group)+"\n")
+
+        cif_file.append("loop_"+"\n")
+        cif_file.append("_symmetry_equiv_pos_site_id"+"\n")
+        cif_file.append("_symmetry_equiv_pos_as_xyz"+"\n")
+        for idx, SYMM in enumerate(crystal.SYMM):
+            cif_file.append("{} {}".format(idx+1,SYMM)+"\n")
+        cif_file.append("_cell_length_a                   "+str(crystal.cell_para[0][0])+"\n")
+        cif_file.append("_cell_length_b                   "+str(crystal.cell_para[0][1])+"\n")
+        cif_file.append("_cell_length_c                   "+str(crystal.cell_para[0][2])+"\n")
+        cif_file.append("_cell_angle_alpha                "+str(crystal.cell_para[1][0])+"\n")
+        cif_file.append("_cell_angle_beta                 "+str(crystal.cell_para[1][1])+"\n")
+        cif_file.append("_cell_angle_gamma                "+str(crystal.cell_para[1][2])+"\n")
+
+        cif_file.append("loop_"+"\n")
+        cif_file.append("_atom_site_label"+"\n")
+        cif_file.append("_atom_site_type_symbol"+"\n")
+        cif_file.append("_atom_site_fract_x"+"\n")
+        cif_file.append("_atom_site_fract_y"+"\n")
+        cif_file.append("_atom_site_fract_z"+"\n")
+        for i in range(0,len(crystal.atoms)):
+            cif_file.append("{:6}  {:4} {:16.8f} {:16.8f} {:16.8f}\n"
+                             .format(i+1,crystal.atoms[i].element,crystal.atoms[i].frac_xyz[0],
+                                     crystal.atoms[i].frac_xyz[1],crystal.atoms[i].frac_xyz[2]))
+
+        return cif_file
+
+
+def write_cifs_file(crystals, sym=False, name="zcx"):
+    cifs_file = []
+    for crystal in crystals:
+        single_cif = write_cif_file(crystal,sym=sym, name=name)
+        cifs_file.extend(single_cif)
+    return cifs_file
+
+
+def read_cif_file(file_path,on_sym_check=False,shut_up=False,system_name="unknown",comment_name="unknown"):
+    input_file = open(file_path, 'r')
+    lines = input_file.readlines()
+    step_pickle = []
+    crystal_all = []
+    if system_name=="unknown":
+        no_name = True
+    else:
+        no_name = False
+    # first time scan
+    for index,line in enumerate(lines):
+        # find out all the step pickle
+        if line.startswith("data_"):
+            step_pickle.append(index)
+    step_pickle.append(len(lines))
+
+    # treat every step and return a crystal
+    for m in range(0,len(step_pickle)-1):
+        atoms = []
+        atoms_P1 = []
+        SYMM = []
+        cell_para = [["unknown","unknown","unknown"],["unknown","unknown","unknown"]]
+
+        for index, line in enumerate(lines[step_pickle[m]:step_pickle[m+1]]):
+            split_line = list(filter(lambda x: x != '', re.split("\\s+", line)))
+            if line.startswith("#"):
+                pass
+            elif len(split_line)==0:
+                pass
+            # read the loop of symmetry
+            # elif split_line[0]=="loop_" and lines[step_pickle[m]+index+1]=="_symmetry_equiv_pos_as_xyz\n":
+            elif split_line[0] == "loop_" and lines[step_pickle[m] + index + 1] == "_symmetry_equiv_pos_as_xyz\n":
+                temp_number = 1
+                while "_" not in lines[step_pickle[m]+index+1+temp_number]:
+                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))
+                    temp_number+=1
+                    if not operation.is_number(split_line_temp[0]):
+                        SYMM.append(split_line_temp[0])
+                    else:
+                        SYMM.append(split_line_temp[1])
+            elif split_line[0] == "loop_" and lines[step_pickle[m]+index+2].strip(" ")=="_symmetry_equiv_pos_as_xyz\n":
+                temp_number = 1
+                while "_" not in lines[step_pickle[m]+index+2+temp_number]:
+                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+2+temp_number])))
+                    temp_number+=1
+                    SYMM.append("".join(split_line_temp[1:]))
+            elif split_line[0] == "loop_" and "_space_group_symop_operation_xyz\n" in lines[step_pickle[m] + index + 1]:
+                # ase format
+                temp_number = 1
+                while "_" not in lines[step_pickle[m]+index+1+temp_number]:
+                    if lines[step_pickle[m]+index+1+temp_number]=="\n":
+                        temp_number += 1
+                        continue
+                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))
+                    temp_number+=1
+                    if not operation.is_number(split_line_temp[0]):
+                        SYMM.append("".join(split_line_temp))
+
+            # read the loop of atoms:
+            elif (split_line[0] == "loop_" and lines[step_pickle[m] + index + 1].strip(" ") == "_atom_site_label\n") or \
+                    (split_line[0] == "loop_" and lines[step_pickle[m] + index + 2].strip(" ") == "_atom_site_label\n"):
+                temp_number = 0
+                while "_" in lines[step_pickle[m]+index+1+temp_number]:
+                    if lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_type_symbol\n":
+                        ele_pos = temp_number
+                    elif lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_fract_x\n":
+                        x_pos = temp_number
+                    elif lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_fract_y\n":
+                        y_pos = temp_number
+                    elif lines[step_pickle[m] + index + 1 + temp_number].strip(" ") == "_atom_site_fract_z\n":
+                        z_pos = temp_number
+                    temp_number+=1
+                how_long = temp_number
+
+                while len(list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))) == how_long:
+                    split_line_temp = list(filter(lambda x: x != '', re.split("\\s+", lines[step_pickle[m]+index+1+temp_number])))
+                    atoms.append(data_classes.Atom(element=split_line_temp[ele_pos],
+                                           frac_xyz=[float(split_line_temp[x_pos]),float(split_line_temp[y_pos]),
+                                                     float(split_line_temp[z_pos])]))
+                    temp_number += 1
+                    if step_pickle[m]+index+1+temp_number==len(lines):
+                        break
+
+            elif split_line[0] == "_cell_length_a":
+                cell_para[0][0] = float(split_line[1])
+            elif split_line[0] == "_cell_length_b":
+                cell_para[0][1] = float(split_line[1])
+            elif split_line[0] == "_cell_length_c":
+                cell_para[0][2] = float(split_line[1])
+            elif split_line[0] == "_cell_angle_alpha":
+                cell_para[1][0] = float(split_line[1])
+            elif split_line[0] == "_cell_angle_beta":
+                cell_para[1][1] = float(split_line[1])
+            elif split_line[0] == "_cell_angle_gamma":
+                cell_para[1][2] = float(split_line[1])
+            elif "data_" in line:
+                if no_name == True:
+                    system_name = line[5:]
+                    system_name = system_name.replace(" ","_")
+                    system_name = system_name.replace("\\", "_")
+                    system_name = system_name.replace("\n", "")
+        for atom in atoms:
+            all_reflect_position = operation.space_group_transfer_for_single_atom(atom.frac_xyz, SYMM)
+            for new_position in all_reflect_position:
+                atoms_P1.append(data_classes.Atom(element=atom.element,
+                                            frac_xyz=[new_position[0], new_position[1], new_position[2]]))
+        crystal_all.append(data_classes.Crystal(cell_para=cell_para, atoms=atoms_P1, comment=comment_name, system_name=system_name))
+        if on_sym_check == True:
+            raise Exception("Not finished part, TODO in code")
+        if shut_up==False:
+            if m%100 == 0:
+                print("{} structures have been treated".format(m))
+
+    return crystal_all
+
+
+def write_poscar_file(crystal, coordinates = 'frac', name = "parser_zcx_create"):
+
+    vasp_file = []
+
+    vasp_file.append('{}\n'.format(name))
+    vasp_file.append('1.0\n')
+    cell_vect = crystal.cell_vect
+    for vect in cell_vect:
+        vasp_file.append("{:16.8f} {:16.8f} {:16.8f}\n".format(vect[0],vect[1],vect[2]))
+    crystal.sort_by_element()
+    vasp_file.append("".join("{:>6s}".format(x) for x in crystal.get_element()) + "\n")
+    vasp_file.append("".join("{:>6.0f}".format(x) for x in crystal.get_element_amount()) + "\n")
+    if coordinates == 'frac':
+        vasp_file.append('Direct\n')
+        for ELEMENT in crystal.get_element():
+            for ATOM in crystal.atoms:
+                if ATOM.element == ELEMENT:
+                    vasp_file.append(
+                        "{:16.8f} {:16.8f} {:16.8f}\n".format(ATOM.frac_xyz[0], ATOM.frac_xyz[1], ATOM.frac_xyz[2]))
+    elif coordinates == 'cart':
+        vasp_file.append('Cartesian\n')
+        for ELEMENT in crystal.get_element():
+            for ATOM in crystal.atoms:
+                if ATOM.element == ELEMENT:
+                    vasp_file.append(
+                        "{:16.8f} {:16.8f} {:16.8f}\n".format(ATOM.cart_xyz[0], ATOM.cart_xyz[1], ATOM.cart_xyz[2]))
+    else:
+        raise Exception("Wrong coordinates type: {}".format(coordinates))
+
+    return vasp_file
+    
+    
+def read_ase_pbc_file(file_path,shut_up=False):
+    input_file = open(file_path, 'r')
+    lines = input_file.readlines()[2:]
+    step_pickle = []
+    crystal_all = []
+
+    # first time scan
+    for index,line in enumerate(lines):
+        # find out all the step pickle
+        if line.startswith("Step "):
+            step_pickle.append(index)
+    step_pickle.append(len(lines))
+
+    # treat every step and return a crystal
+    for m in range(0,len(step_pickle)-1):
+        atoms_P1 = []
+        force_matrix = []
+        position_matrix = []
+        in_forces = False
+        in_positions = False
+
+
+        for index, line in enumerate(lines[step_pickle[m]:step_pickle[m+1]]):
+            # split_line = list(filter(lambda x: x != '', re.split("\\s+", line)))
+            line = line.strip()
+            # check Forces part
+            if line.startswith("Forces:"):
+                in_forces = True
+                in_positions = False
+                continue
+
+            # check Positions part
+            if line.startswith("Positions:"):
+                in_positions = True
+                in_forces = False
+                continue
+
+            if in_forces and line.startswith("[") and line.endswith("]"):
+                line = line.replace("[", "").replace("]", "")
+                force_matrix.append([float(x) for x in line.split()])
+
+            # analyse Positions part
+            if in_positions and line.startswith("[") and line.endswith("]"):
+                line = line.replace("[", "").replace("]", "")
+                position_matrix.append([float(x) for x in line.split()])
+
+            if line.startswith("Elements:"):
+                elements_string = line.strip().split(":", 1)[-1].strip()
+                elements_string = elements_string[1:-1]  
+
+                elements = [elem.strip().strip("'") for elem in elements_string.split(",")]
+
+            if line.startswith("cell:"):
+                matrix_string = line[len("cell: Cell("):-1]
+                rows = matrix_string.split("], [")
+                cell_vect = [
+                    [float(value) for value in row.replace('[', '').replace(']', '').replace(')', '').split(", ")]
+                    for row in rows
+                ]
+
+        for i in range(0,len(elements)):
+            atoms_P1.append(data_classes.Atom(element=elements[i], cart_xyz=[position_matrix[i][0], position_matrix[i][1], position_matrix[i][2]]))
+
+        crystal_all.append(data_classes.Crystal(cell_vect=cell_vect, atoms=atoms_P1))
+
+    return crystal_all    
     
\ No newline at end of file

basic_function/operation.py

-# -*- coding: utf-8 -*-
-"""
-A collection of functions for performing crystallographic and molecular operations,
-such as symmetry application, supercell generation, and geometric analysis.
-"""
-
-# --- Standard Library Imports ---
-import copy
-import fractions
-import re
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-# --- Third-Party Imports ---
-import networkx as nx
-import numpy as np
-import numpy.typing as npt
-from scipy.spatial import cKDTree as KDTree
-
-# --- Local Application Imports ---
-from basic_function import chemical_knowledge, data_classes
-
-# Type aliases for clarity
-NDArrayFloat = npt.NDArray[np.float64]
-CellVectors = List[List[float]]
-SymmetryOperations = List[str]
-
-
-def is_number(s: str) -> bool:
-    """Checks if a string can be interpreted as a number (float or fraction).
-
-    Args:
-        s: The input string.
-
-    Returns:
-        True if the string represents a number, False otherwise.
-    """
-    try:
-        float(s)
-        return True
-    except ValueError:
-        pass
-
-    try:
-        # Check for fractional representations like "1/2"
-        float(fractions.Fraction(s))
-        return True
-    except ValueError:
-        return False
-
-
-def _parse_symmetry_operations(
-    sym_ops: SymmetryOperations,
-) -> Tuple[List[NDArrayFloat], List[NDArrayFloat]]:
-    """Parses a list of symmetry operation strings into matrices.
-
-    This is an internal helper function to avoid code duplication in public functions.
-
-    Args:
-        sym_ops: A list of symmetry operation strings (e.g., ['x, y, z+1/2']).
-
-    Returns:
-        A tuple containing two lists:
-        - A list of 3x3 rotation/reflection matrices (M).
-        - A list of 1x3 translation vectors (C).
-
-    Raises:
-        ValueError: If a symmetry operation string is malformed.
-    """
-    rotation_matrices = []
-    translation_vectors = []
-
-    for sym_op_str in sym_ops:
-        sym_op_parts = sym_op_str.lower().replace(" ", "").split(",")
-        if len(sym_op_parts) != 3:
-            raise ValueError(f"Symmetry operation '{sym_op_str}' is invalid.")
-
-        matrix_m = np.zeros((3, 3))
-        matrix_c = np.zeros((1, 3))
-
-        for i, part in enumerate(sym_op_parts):
-            # Regex to find elements like '+x', '-y', 'z', '1/2', '-0.5'
-            tokens = re.findall(r"([+-]?[xyz0-9./]+)", part)
-            for token in tokens:
-                token = token.strip()
-                if not token:
-                    continue
-
-                if "x" in token:
-                    matrix_m[0, i] = -1.0 if token.startswith("-") else 1.0
-                elif "y" in token:
-                    matrix_m[1, i] = -1.0 if token.startswith("-") else 1.0
-                elif "z" in token:
-                    matrix_m[2, i] = -1.0 if token.startswith("-") else 1.0
-                elif is_number(token):
-                    matrix_c[0, i] += float(fractions.Fraction(token))
-                else:
-                    raise ValueError(f"Invalid fragment '{token}' in symmetry operation.")
-
-        rotation_matrices.append(matrix_m)
-        translation_vectors.append(matrix_c)
-
-    return rotation_matrices, translation_vectors
-
-
-def space_group_transfer_for_single_atom(
-    frac_xyz: List[float], space_group_ops: SymmetryOperations
-) -> List[List[float]]:
-    """Applies space group symmetry operations to a single atomic coordinate.
-
-    Args:
-        frac_xyz: The fractional coordinates [x, y, z] of a single atom.
-        space_group_ops: A list of space group symmetry operation strings.
-
-    Returns:
-        A list of all symmetrically equivalent fractional coordinates.
-    """
-    rot_matrices, trans_vectors = _parse_symmetry_operations(space_group_ops)
-    
-    equivalent_positions = []
-    atom_pos = np.array(frac_xyz)
-
-    for rot, trans in zip(rot_matrices, trans_vectors):
-        new_pos = np.dot(atom_pos, rot.T) + trans.squeeze()
-        equivalent_positions.append(new_pos.tolist())
-
-    return equivalent_positions
-
-
-def super_cell(
-    crystal: "data_classes.Crystal",
-    cell_range: Optional[List[List[int]]] = None,
-) -> "data_classes.Crystal":
-    """Constructs a supercell from a unit cell.
-
-    Args:
-        crystal: The input Crystal object.
-        cell_range: A list of ranges for each lattice vector, e.g.,
-                    [[-1, 1], [-1, 1], [-1, 1]] creates a 3x3x3 supercell.
-                    If None, defaults to [[-1, 1], [-1, 1], [-1, 1]].
-
-    Returns:
-        A new Crystal object representing the supercell.
-    """
-    if cell_range is None:
-        cell_range = [[-1, 1], [-1, 1], [-1, 1]]
-
-    dims = [r[1] - r[0] + 1 for r in cell_range]
-    
-    new_lattice = [
-        [dim * val for val in crystal.cell_vect[i]] for i, dim in enumerate(dims)
-    ]
-
-    translation_vectors = []
-    for h in range(cell_range[0][0], cell_range[0][1] + 1):
-        for k in range(cell_range[1][0], cell_range[1][1] + 1):
-            for l in range(cell_range[2][0], cell_range[2][1] + 1):
-                translation_vectors.append([h, k, l])
-
-    new_atoms = []
-    for atom in crystal.atoms:
-        for trans_vec in translation_vectors:
-            new_frac_xyz = [
-                (atom.frac_xyz[i] + trans_vec[i]) / dims[i] for i in range(3)
-            ]
-            new_atoms.append(
-                data_classes.Atom(element=atom.element, frac_xyz=new_frac_xyz)
-            )
-
-    if crystal.energy != "unknown":
-        total_cells = dims[0] * dims[1] * dims[2]
-        new_energy = crystal.energy * total_cells
-    else:
-        new_energy = "unknown"
-
-    return data_classes.Crystal(
-        cell_vect=new_lattice, energy=new_energy, atoms=new_atoms
-    )
-
-
-def orient_molecule(molecule: "data_classes.Molecule") -> "data_classes.Molecule":
-    """Orients a molecule along its principal axes of inertia.
-
-    The method uses the Moment of Inertia tensor to define a canonical orientation.
-    The molecule's coordinates are modified in-place. For more details, see:
-    http://sobereva.com/426
-
-    Args:
-        molecule: The Molecule object to be oriented.
-
-    Returns:
-        The same Molecule object with its atoms reoriented.
-    """
-    all_ele, all_cart = molecule.get_ele_and_cart()
-
-    if len(all_cart) <= 1:
-        return molecule  # No orientation needed for single atoms or empty molecules.
-
-    masses = np.array([chemical_knowledge.element_masses[el] for el in all_ele])
-    relative_position = all_cart - molecule.get_center_of_mass()
-
-    # Calculate the moment of inertia tensor
-    I_xx = np.sum(masses * (relative_position[:, 1] ** 2 + relative_position[:, 2] ** 2))
-    I_yy = np.sum(masses * (relative_position[:, 0] ** 2 + relative_position[:, 2] ** 2))
-    I_zz = np.sum(masses * (relative_position[:, 0] ** 2 + relative_position[:, 1] ** 2))
-    I_xy = -np.sum(masses * relative_position[:, 0] * relative_position[:, 1])
-    I_xz = -np.sum(masses * relative_position[:, 0] * relative_position[:, 2])
-    I_yz = -np.sum(masses * relative_position[:, 1] * relative_position[:, 2])
-
-    I_matrix = np.array([[I_xx, I_xy, I_xz], [I_xy, I_yy, I_yz], [I_xz, I_yz, I_zz]])
-
-    # Eigenvectors of the inertia tensor are the principal axes.
-    # np.linalg.eigh is used for symmetric matrices.
-    eigenvalues, eigenvectors = np.linalg.eigh(I_matrix)
-    principal_axes = eigenvectors.T
-    
-    # Project the relative positions onto the new axes system.
-    new_positions = np.dot(relative_position, principal_axes.T)
-    
-    molecule.put_ele_cart_back(all_ele, new_positions)
-    return molecule
-
-
-def get_rotate_matrix(v: NDArrayFloat) -> NDArrayFloat:
-    """Generates a 3x3 rotation matrix from a 3D vector `v`.
-
-    This function uses a mapping from a 3D vector to a quaternion, which is then
-    used to construct the rotation matrix. This method avoids gimbal lock. A
-    left-handed coordinate system is assumed.
-
-    Args:
-        v: A 3-element NumPy array used to generate the quaternion.
-
-    Returns:
-        A 3x3 rotation matrix.
-    """
-    # Ensure v elements are within valid ranges if necessary, though the
-    # formulas handle most inputs gracefully.
-    v0_sqrt = np.sqrt(max(v[0], 0))
-    v0_1_sqrt = np.sqrt(max(1.0 - v[0], 0))
-    
-    angle1 = 2.0 * np.pi * v[1]
-    angle2 = 2.0 * np.pi * v[2]
-
-    # Quaternion components (x, y, z, w)
-    qx = v0_1_sqrt * np.sin(angle1)
-    qy = v0_1_sqrt * np.cos(angle1)
-    qz = v0_sqrt * np.sin(angle2)
-    qw = v0_sqrt * np.cos(angle2)
-
-    return np.array([
-        [1 - 2*qy**2 - 2*qz**2, 2*qx*qy + 2*qw*qz, 2*qx*qz - 2*qw*qy],
-        [2*qx*qy - 2*qw*qz, 1 - 2*qx**2 - 2*qz**2, 2*qy*qz + 2*qw*qx],
-        [2*qx*qz + 2*qw*qy, 2*qy*qz - 2*qw*qx, 1 - 2*qx**2 - 2*qy**2]
-    ])
-
-
-def f2c_matrix(
-    cell_params: Tuple[List[float], List[float]]
-) -> Optional[NDArrayFloat]:
-    """Calculates the fractional-to-Cartesian transformation matrix.
-
-    Args:
-        cell_params: A tuple containing [[a, b, c], [alpha, beta, gamma]],
-                     where lengths are in Angstroms and angles are in degrees.
-
-    Returns:
-        The 3x3 transformation matrix, or None if cell parameters are invalid.
-    """
-    lengths, angles = cell_params
-    a, b, c = lengths
-    alpha, beta, gamma = np.deg2rad(angles)
-
-    cos_a, cos_b, cos_g = np.cos([alpha, beta, gamma])
-    sin_g = np.sin(gamma)
-
-    # Volume calculation term
-    volume_term_sq = (
-        1.0 - cos_a**2 - cos_b**2 - cos_g**2 + 2.0 * cos_a * cos_b * cos_g
-    )
-    if volume_term_sq < 0:
-        return None
-    
-    volume = a * b * c * np.sqrt(volume_term_sq)
-
-    matrix = np.zeros((3, 3))
-    matrix[0, 0] = a
-    matrix[0, 1] = b * cos_g
-    matrix[0, 2] = c * cos_b
-    matrix[1, 1] = b * sin_g
-    matrix[1, 2] = c * (cos_a - cos_b * cos_g) / sin_g
-    matrix[2, 2] = volume / (a * b * sin_g)
-    
-    return matrix.T
-
-
-def c2f_matrix(
-    cell_params: Tuple[List[float], List[float]]
-) -> Optional[NDArrayFloat]:
-    """Calculates the Cartesian-to-fractional transformation matrix.
-
-    This is the inverse of the matrix generated by `f2c_matrix`.
-
-    Args:
-        cell_params: A tuple containing [[a, b, c], [alpha, beta, gamma]],
-                     where lengths are in Angstroms and angles are in degrees.
-
-    Returns:
-        The 3x3 transformation matrix, or None if cell parameters are invalid.
-    """
-    f2c = f2c_matrix(cell_params)
-    if f2c is None:
-        return None
-    
-    try:
-        return np.linalg.inv(f2c)
-    except np.linalg.LinAlgError:
-        return None
-
-
-def apply_SYMM(
-    frac_xyz: NDArrayFloat, symm_ops: SymmetryOperations
-) -> NDArrayFloat:
-    """Applies symmetry operations to a single set of fractional coordinates.
-
-    Args:
-        frac_xyz: A NumPy array of fractional coordinates [x, y, z].
-        symm_ops: A list of symmetry operation strings.
-
-    Returns:
-        A NumPy array of all symmetrically equivalent fractional coordinates.
-    """
-    rot_matrices, trans_vectors = _parse_symmetry_operations(symm_ops)
-
-    equivalent_positions = [
-        np.dot(frac_xyz, rot.T) + trans.squeeze()
-        for rot, trans in zip(rot_matrices, trans_vectors)
-    ]
-    
-    return np.array(equivalent_positions)
-
-
-def apply_SYMM_with_element(
-    elements: Union[str, List[str]],
-    frac_xyzs: NDArrayFloat,
-    symm_ops: SymmetryOperations,
-) -> Tuple[NDArrayFloat, NDArrayFloat]:
-    """Applies symmetry operations, returning new elements and coordinates.
-
-    Args:
-        elements: The element symbol(s) corresponding to the coordinates.
-        frac_xyzs: A NumPy array of fractional coordinates.
-        symm_ops: A list of symmetry operation strings.
-
-    Returns:
-        A tuple containing:
-        - A NumPy array of element symbols for each new position.
-        - A NumPy array of all symmetrically equivalent fractional coordinates.
-    """
-    equivalent_positions = apply_SYMM(frac_xyzs, symm_ops)
-    num_ops = len(equivalent_positions)
-    
-    replicated_elements = np.tile(np.array(elements).squeeze(), (num_ops, 1))
-    
-    return replicated_elements, equivalent_positions
-
-
-def calculate_longest_diagonal_length(cell_vect: CellVectors) -> float:
-    """Calculates the length of the longest space diagonal of a unit cell.
-
-    The longest diagonal connects the origin (0,0,0) to the opposite
-    corner (1,1,1) of the unit cell.
-
-    Args:
-        cell_vect: The three lattice vectors of the cell.
-
-    Returns:
-        The length of the longest diagonal in Angstroms.
-    """
-    cell_vect_np = np.array(cell_vect)
-    diagonal_vector = np.sum(cell_vect_np, axis=0)
-    return float(np.linalg.norm(diagonal_vector))
-
-
-def calculate_distance_of_parallel_plane_in_crystal(cell_vect: CellVectors) -> List[float]:
-    """Calculates inter-planar distances for primary crystallographic planes.
-
-    This computes the distances for the (100), (010), and (001) families of planes.
-
-    Args:
-        cell_vect: The three lattice vectors [a, b, c] of the cell.
-
-    Returns:
-        A list of three distances [d_a, d_b, d_c], where d_a is the distance
-        between planes parallel to the b-c plane, and so on.
-    """
-    distances = []
-    vectors = [np.array(v) for v in cell_vect]
-    
-    # Permutations to calculate distance for each primary plane
-    # (a to b-c plane, b to a-c plane, c to a-b plane)
-    indices = [(0, 1, 2), (1, 0, 2), (2, 0, 1)]
-
-    for i, j, k in indices:
-        point_p = vectors[i]
-        plane_v1 = vectors[j]
-        plane_v2 = vectors[k]
-        
-        # Normal vector to the plane defined by plane_v1 and plane_v2
-        normal_vector = np.cross(plane_v1, plane_v2)
-        
-        # Distance from point P to the plane is |N · P| / ||N||
-        distance = abs(np.dot(normal_vector, point_p)) / np.linalg.norm(normal_vector)
-        distances.append(distance)
-        
-    return distances
-
-
-def detect_is_frame_vdw_new(crystal: "data_classes.Crystal", tolerance: float = 1.2) -> bool:
-    """Detects if a crystal structure forms a connected framework via VdW radii.
-
-    The method involves:
-    1. Expanding the crystal to a P1 symmetry supercell.
-    2. Building a 3x3x3 supercell to ensure periodic connections are considered.
-    3. Constructing a graph where atoms are nodes and an edge exists if their
-       distance is within a scaled sum of their van der Waals radii.
-    4. Checking if the largest connected component in the graph is large enough
-       to be considered a single, percolating framework.
-
-    Args:
-        crystal: The Crystal object to analyze.
-        tolerance: A tolerance factor to scale the VdW radii sum.
-
-    Returns:
-        True if the structure is a connected framework, False otherwise.
-    """
-    crystal_temp = copy.deepcopy(crystal)
-    crystal_temp.make_p1()
-    crystal_temp.move_atom_into_cell()
-
-    # Create a 3x3x3 supercell to check for connectivity across boundaries
-    crystal_supercell = super_cell(crystal_temp, cell_range=[[-1, 1], [-1, 1], [-1, 1]])
-    
-    all_ele, all_carts = crystal_supercell.get_ele_and_cart()
-
-    vdw_radii_map = chemical_knowledge.element_vdw_radii
-    vdw_max = max(vdw_radii_map[el] for el in set(all_ele))
-    distance_threshold = vdw_max * tolerance * 2
-
-    # KDTree for efficient nearest-neighbor search
-    tree = KDTree(all_carts)
-    pairs = tree.query_pairs(r=distance_threshold)
-
-    # Build a graph to find connected components
-    graph = nx.Graph()
-    graph.add_nodes_from(range(len(all_carts)))
-    graph.add_edges_from(list(pairs))
-
-    if not graph.nodes:
-        return False
-
-    # Find the largest connected component
-    largest_cc = max(nx.connected_components(graph), key=len)
-
-    # A heuristic to check for a percolating framework. A connected framework
-    # should connect most atoms. The threshold '9' is empirical but robustly
-    # distinguishes between isolated molecules and a fully connected lattice.
-    # In a 3x3x3 supercell (27 unit cells), a connected framework should involve
-    # significantly more atoms than in a few unit cells.
+# -*- coding: utf-8 -*-
+"""
+A collection of functions for performing crystallographic and molecular operations,
+such as symmetry application, supercell generation, and geometric analysis.
+"""
+
+# --- Standard Library Imports ---
+import copy
+import fractions
+import re
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+# --- Third-Party Imports ---
+import networkx as nx
+import numpy as np
+import numpy.typing as npt
+from scipy.spatial import cKDTree as KDTree
+
+# --- Local Application Imports ---
+from basic_function import chemical_knowledge, data_classes
+
+# Type aliases for clarity
+NDArrayFloat = npt.NDArray[np.float64]
+CellVectors = List[List[float]]
+SymmetryOperations = List[str]
+
+
+def is_number(s: str) -> bool:
+    """Checks if a string can be interpreted as a number (float or fraction).
+
+    Args:
+        s: The input string.
+
+    Returns:
+        True if the string represents a number, False otherwise.
+    """
+    try:
+        float(s)
+        return True
+    except ValueError:
+        pass
+
+    try:
+        # Check for fractional representations like "1/2"
+        float(fractions.Fraction(s))
+        return True
+    except ValueError:
+        return False
+
+
+def _parse_symmetry_operations(
+    sym_ops: SymmetryOperations,
+) -> Tuple[List[NDArrayFloat], List[NDArrayFloat]]:
+    """Parses a list of symmetry operation strings into matrices.
+
+    This is an internal helper function to avoid code duplication in public functions.
+
+    Args:
+        sym_ops: A list of symmetry operation strings (e.g., ['x, y, z+1/2']).
+
+    Returns:
+        A tuple containing two lists:
+        - A list of 3x3 rotation/reflection matrices (M).
+        - A list of 1x3 translation vectors (C).
+
+    Raises:
+        ValueError: If a symmetry operation string is malformed.
+    """
+    rotation_matrices = []
+    translation_vectors = []
+
+    for sym_op_str in sym_ops:
+        sym_op_parts = sym_op_str.lower().replace(" ", "").split(",")
+        if len(sym_op_parts) != 3:
+            raise ValueError(f"Symmetry operation '{sym_op_str}' is invalid.")
+
+        matrix_m = np.zeros((3, 3))
+        matrix_c = np.zeros((1, 3))
+
+        for i, part in enumerate(sym_op_parts):
+            # Regex to find elements like '+x', '-y', 'z', '1/2', '-0.5'
+            tokens = re.findall(r"([+-]?[xyz0-9./]+)", part)
+            for token in tokens:
+                token = token.strip()
+                if not token:
+                    continue
+
+                if "x" in token:
+                    matrix_m[0, i] = -1.0 if token.startswith("-") else 1.0
+                elif "y" in token:
+                    matrix_m[1, i] = -1.0 if token.startswith("-") else 1.0
+                elif "z" in token:
+                    matrix_m[2, i] = -1.0 if token.startswith("-") else 1.0
+                elif is_number(token):
+                    matrix_c[0, i] += float(fractions.Fraction(token))
+                else:
+                    raise ValueError(f"Invalid fragment '{token}' in symmetry operation.")
+
+        rotation_matrices.append(matrix_m)
+        translation_vectors.append(matrix_c)
+
+    return rotation_matrices, translation_vectors
+
+
+def space_group_transfer_for_single_atom(
+    frac_xyz: List[float], space_group_ops: SymmetryOperations
+) -> List[List[float]]:
+    """Applies space group symmetry operations to a single atomic coordinate.
+
+    Args:
+        frac_xyz: The fractional coordinates [x, y, z] of a single atom.
+        space_group_ops: A list of space group symmetry operation strings.
+
+    Returns:
+        A list of all symmetrically equivalent fractional coordinates.
+    """
+    rot_matrices, trans_vectors = _parse_symmetry_operations(space_group_ops)
+    
+    equivalent_positions = []
+    atom_pos = np.array(frac_xyz)
+
+    for rot, trans in zip(rot_matrices, trans_vectors):
+        new_pos = np.dot(atom_pos, rot.T) + trans.squeeze()
+        equivalent_positions.append(new_pos.tolist())
+
+    return equivalent_positions
+
+
+def super_cell(
+    crystal: "data_classes.Crystal",
+    cell_range: Optional[List[List[int]]] = None,
+) -> "data_classes.Crystal":
+    """Constructs a supercell from a unit cell.
+
+    Args:
+        crystal: The input Crystal object.
+        cell_range: A list of ranges for each lattice vector, e.g.,
+                    [[-1, 1], [-1, 1], [-1, 1]] creates a 3x3x3 supercell.
+                    If None, defaults to [[-1, 1], [-1, 1], [-1, 1]].
+
+    Returns:
+        A new Crystal object representing the supercell.
+    """
+    if cell_range is None:
+        cell_range = [[-1, 1], [-1, 1], [-1, 1]]
+
+    dims = [r[1] - r[0] + 1 for r in cell_range]
+    
+    new_lattice = [
+        [dim * val for val in crystal.cell_vect[i]] for i, dim in enumerate(dims)
+    ]
+
+    translation_vectors = []
+    for h in range(cell_range[0][0], cell_range[0][1] + 1):
+        for k in range(cell_range[1][0], cell_range[1][1] + 1):
+            for l in range(cell_range[2][0], cell_range[2][1] + 1):
+                translation_vectors.append([h, k, l])
+
+    new_atoms = []
+    for atom in crystal.atoms:
+        for trans_vec in translation_vectors:
+            new_frac_xyz = [
+                (atom.frac_xyz[i] + trans_vec[i]) / dims[i] for i in range(3)
+            ]
+            new_atoms.append(
+                data_classes.Atom(element=atom.element, frac_xyz=new_frac_xyz)
+            )
+
+    if crystal.energy != "unknown":
+        total_cells = dims[0] * dims[1] * dims[2]
+        new_energy = crystal.energy * total_cells
+    else:
+        new_energy = "unknown"
+
+    return data_classes.Crystal(
+        cell_vect=new_lattice, energy=new_energy, atoms=new_atoms
+    )
+
+
+def orient_molecule(molecule: "data_classes.Molecule") -> "data_classes.Molecule":
+    """Orients a molecule along its principal axes of inertia.
+
+    The method uses the Moment of Inertia tensor to define a canonical orientation.
+    The molecule's coordinates are modified in-place. For more details, see:
+    http://sobereva.com/426
+
+    Args:
+        molecule: The Molecule object to be oriented.
+
+    Returns:
+        The same Molecule object with its atoms reoriented.
+    """
+    all_ele, all_cart = molecule.get_ele_and_cart()
+
+    if len(all_cart) <= 1:
+        return molecule  # No orientation needed for single atoms or empty molecules.
+
+    masses = np.array([chemical_knowledge.element_masses[el] for el in all_ele])
+    relative_position = all_cart - molecule.get_center_of_mass()
+
+    # Calculate the moment of inertia tensor
+    I_xx = np.sum(masses * (relative_position[:, 1] ** 2 + relative_position[:, 2] ** 2))
+    I_yy = np.sum(masses * (relative_position[:, 0] ** 2 + relative_position[:, 2] ** 2))
+    I_zz = np.sum(masses * (relative_position[:, 0] ** 2 + relative_position[:, 1] ** 2))
+    I_xy = -np.sum(masses * relative_position[:, 0] * relative_position[:, 1])
+    I_xz = -np.sum(masses * relative_position[:, 0] * relative_position[:, 2])
+    I_yz = -np.sum(masses * relative_position[:, 1] * relative_position[:, 2])
+
+    I_matrix = np.array([[I_xx, I_xy, I_xz], [I_xy, I_yy, I_yz], [I_xz, I_yz, I_zz]])
+
+    # Eigenvectors of the inertia tensor are the principal axes.
+    # np.linalg.eigh is used for symmetric matrices.
+    eigenvalues, eigenvectors = np.linalg.eigh(I_matrix)
+    principal_axes = eigenvectors.T
+    
+    # Project the relative positions onto the new axes system.
+    new_positions = np.dot(relative_position, principal_axes.T)
+    
+    molecule.put_ele_cart_back(all_ele, new_positions)
+    return molecule
+
+
+def get_rotate_matrix(v: NDArrayFloat) -> NDArrayFloat:
+    """Generates a 3x3 rotation matrix from a 3D vector `v`.
+
+    This function uses a mapping from a 3D vector to a quaternion, which is then
+    used to construct the rotation matrix. This method avoids gimbal lock. A
+    left-handed coordinate system is assumed.
+
+    Args:
+        v: A 3-element NumPy array used to generate the quaternion.
+
+    Returns:
+        A 3x3 rotation matrix.
+    """
+    # Ensure v elements are within valid ranges if necessary, though the
+    # formulas handle most inputs gracefully.
+    v0_sqrt = np.sqrt(max(v[0], 0))
+    v0_1_sqrt = np.sqrt(max(1.0 - v[0], 0))
+    
+    angle1 = 2.0 * np.pi * v[1]
+    angle2 = 2.0 * np.pi * v[2]
+
+    # Quaternion components (x, y, z, w)
+    qx = v0_1_sqrt * np.sin(angle1)
+    qy = v0_1_sqrt * np.cos(angle1)
+    qz = v0_sqrt * np.sin(angle2)
+    qw = v0_sqrt * np.cos(angle2)
+
+    return np.array([
+        [1 - 2*qy**2 - 2*qz**2, 2*qx*qy + 2*qw*qz, 2*qx*qz - 2*qw*qy],
+        [2*qx*qy - 2*qw*qz, 1 - 2*qx**2 - 2*qz**2, 2*qy*qz + 2*qw*qx],
+        [2*qx*qz + 2*qw*qy, 2*qy*qz - 2*qw*qx, 1 - 2*qx**2 - 2*qy**2]
+    ])
+
+
+def f2c_matrix(
+    cell_params: Tuple[List[float], List[float]]
+) -> Optional[NDArrayFloat]:
+    """Calculates the fractional-to-Cartesian transformation matrix.
+
+    Args:
+        cell_params: A tuple containing [[a, b, c], [alpha, beta, gamma]],
+                     where lengths are in Angstroms and angles are in degrees.
+
+    Returns:
+        The 3x3 transformation matrix, or None if cell parameters are invalid.
+    """
+    lengths, angles = cell_params
+    a, b, c = lengths
+    alpha, beta, gamma = np.deg2rad(angles)
+
+    cos_a, cos_b, cos_g = np.cos([alpha, beta, gamma])
+    sin_g = np.sin(gamma)
+
+    # Volume calculation term
+    volume_term_sq = (
+        1.0 - cos_a**2 - cos_b**2 - cos_g**2 + 2.0 * cos_a * cos_b * cos_g
+    )
+    if volume_term_sq < 0:
+        return None
+    
+    volume = a * b * c * np.sqrt(volume_term_sq)
+
+    matrix = np.zeros((3, 3))
+    matrix[0, 0] = a
+    matrix[0, 1] = b * cos_g
+    matrix[0, 2] = c * cos_b
+    matrix[1, 1] = b * sin_g
+    matrix[1, 2] = c * (cos_a - cos_b * cos_g) / sin_g
+    matrix[2, 2] = volume / (a * b * sin_g)
+    
+    return matrix.T
+
+
+def c2f_matrix(
+    cell_params: Tuple[List[float], List[float]]
+) -> Optional[NDArrayFloat]:
+    """Calculates the Cartesian-to-fractional transformation matrix.
+
+    This is the inverse of the matrix generated by `f2c_matrix`.
+
+    Args:
+        cell_params: A tuple containing [[a, b, c], [alpha, beta, gamma]],
+                     where lengths are in Angstroms and angles are in degrees.
+
+    Returns:
+        The 3x3 transformation matrix, or None if cell parameters are invalid.
+    """
+    f2c = f2c_matrix(cell_params)
+    if f2c is None:
+        return None
+    
+    try:
+        return np.linalg.inv(f2c)
+    except np.linalg.LinAlgError:
+        return None
+
+
+def apply_SYMM(
+    frac_xyz: NDArrayFloat, symm_ops: SymmetryOperations
+) -> NDArrayFloat:
+    """Applies symmetry operations to a single set of fractional coordinates.
+
+    Args:
+        frac_xyz: A NumPy array of fractional coordinates [x, y, z].
+        symm_ops: A list of symmetry operation strings.
+
+    Returns:
+        A NumPy array of all symmetrically equivalent fractional coordinates.
+    """
+    rot_matrices, trans_vectors = _parse_symmetry_operations(symm_ops)
+
+    equivalent_positions = [
+        np.dot(frac_xyz, rot.T) + trans.squeeze()
+        for rot, trans in zip(rot_matrices, trans_vectors)
+    ]
+    
+    return np.array(equivalent_positions)
+
+
+def apply_SYMM_with_element(
+    elements: Union[str, List[str]],
+    frac_xyzs: NDArrayFloat,
+    symm_ops: SymmetryOperations,
+) -> Tuple[NDArrayFloat, NDArrayFloat]:
+    """Applies symmetry operations, returning new elements and coordinates.
+
+    Args:
+        elements: The element symbol(s) corresponding to the coordinates.
+        frac_xyzs: A NumPy array of fractional coordinates.
+        symm_ops: A list of symmetry operation strings.
+
+    Returns:
+        A tuple containing:
+        - A NumPy array of element symbols for each new position.
+        - A NumPy array of all symmetrically equivalent fractional coordinates.
+    """
+    equivalent_positions = apply_SYMM(frac_xyzs, symm_ops)
+    num_ops = len(equivalent_positions)
+    
+    replicated_elements = np.tile(np.array(elements).squeeze(), (num_ops, 1))
+    
+    return replicated_elements, equivalent_positions
+
+
+def calculate_longest_diagonal_length(cell_vect: CellVectors) -> float:
+    """Calculates the length of the longest space diagonal of a unit cell.
+
+    The longest diagonal connects the origin (0,0,0) to the opposite
+    corner (1,1,1) of the unit cell.
+
+    Args:
+        cell_vect: The three lattice vectors of the cell.
+
+    Returns:
+        The length of the longest diagonal in Angstroms.
+    """
+    cell_vect_np = np.array(cell_vect)
+    diagonal_vector = np.sum(cell_vect_np, axis=0)
+    return float(np.linalg.norm(diagonal_vector))
+
+
+def calculate_distance_of_parallel_plane_in_crystal(cell_vect: CellVectors) -> List[float]:
+    """Calculates inter-planar distances for primary crystallographic planes.
+
+    This computes the distances for the (100), (010), and (001) families of planes.
+
+    Args:
+        cell_vect: The three lattice vectors [a, b, c] of the cell.
+
+    Returns:
+        A list of three distances [d_a, d_b, d_c], where d_a is the distance
+        between planes parallel to the b-c plane, and so on.
+    """
+    distances = []
+    vectors = [np.array(v) for v in cell_vect]
+    
+    # Permutations to calculate distance for each primary plane
+    # (a to b-c plane, b to a-c plane, c to a-b plane)
+    indices = [(0, 1, 2), (1, 0, 2), (2, 0, 1)]
+
+    for i, j, k in indices:
+        point_p = vectors[i]
+        plane_v1 = vectors[j]
+        plane_v2 = vectors[k]
+        
+        # Normal vector to the plane defined by plane_v1 and plane_v2
+        normal_vector = np.cross(plane_v1, plane_v2)
+        
+        # Distance from point P to the plane is |N · P| / ||N||
+        distance = abs(np.dot(normal_vector, point_p)) / np.linalg.norm(normal_vector)
+        distances.append(distance)
+        
+    return distances
+
+
+def detect_is_frame_vdw_new(crystal: "data_classes.Crystal", tolerance: float = 1.2) -> bool:
+    """Detects if a crystal structure forms a connected framework via VdW radii.
+
+    The method involves:
+    1. Expanding the crystal to a P1 symmetry supercell.
+    2. Building a 3x3x3 supercell to ensure periodic connections are considered.
+    3. Constructing a graph where atoms are nodes and an edge exists if their
+       distance is within a scaled sum of their van der Waals radii.
+    4. Checking if the largest connected component in the graph is large enough
+       to be considered a single, percolating framework.
+
+    Args:
+        crystal: The Crystal object to analyze.
+        tolerance: A tolerance factor to scale the VdW radii sum.
+
+    Returns:
+        True if the structure is a connected framework, False otherwise.
+    """
+    crystal_temp = copy.deepcopy(crystal)
+    crystal_temp.make_p1()
+    crystal_temp.move_atom_into_cell()
+
+    # Create a 3x3x3 supercell to check for connectivity across boundaries
+    crystal_supercell = super_cell(crystal_temp, cell_range=[[-1, 1], [-1, 1], [-1, 1]])
+    
+    all_ele, all_carts = crystal_supercell.get_ele_and_cart()
+
+    vdw_radii_map = chemical_knowledge.element_vdw_radii
+    vdw_max = max(vdw_radii_map[el] for el in set(all_ele))
+    distance_threshold = vdw_max * tolerance * 2
+
+    # KDTree for efficient nearest-neighbor search
+    tree = KDTree(all_carts)
+    pairs = tree.query_pairs(r=distance_threshold)
+
+    # Build a graph to find connected components
+    graph = nx.Graph()
+    graph.add_nodes_from(range(len(all_carts)))
+    graph.add_edges_from(list(pairs))
+
+    if not graph.nodes:
+        return False
+
+    # Find the largest connected component
+    largest_cc = max(nx.connected_components(graph), key=len)
+
+    # A heuristic to check for a percolating framework. A connected framework
+    # should connect most atoms. The threshold '9' is empirical but robustly
+    # distinguishes between isolated molecules and a fully connected lattice.
+    # In a 3x3x3 supercell (27 unit cells), a connected framework should involve
+    # significantly more atoms than in a few unit cells.
     return len(largest_cc) > 9 * len(crystal_temp.atoms)
\ No newline at end of file

basic_function/packaged_function.py

-from basic_function import format_parser
-from basic_function import CSP_generator_normal
-import os
-import concurrent.futures
-import sys
-
-
-
-def process_crystal(seed, sg, molecules,output_path,add_name):
-    aaa = CSP_generator_normal.CrystalGenerator(molecules, space_group=sg)
-    molecules_number = sum(len(molecule.atoms) for molecule in molecules)
-    new_crystal = aaa.generate(seed=seed)
-    sys.stdout.flush()
-    if new_crystal is not None:
-        cif_out = format_parser.write_cif_file(new_crystal)
-        with open(f"{output_path}/structures/{add_name}_{sg}_{seed}_z{len(molecules)}_{molecules_number}.cif", 'w') as target:
-            target.writelines(cif_out)
-        return True
-    return False
-
-def CSP_generater_parallel(molecules,output_path,need_structure = 100, space_group_list=[1],max_workers=8,add_name='',start_seed=1):
-    space_groups = space_group_list
-    accept_count = need_structure
-
-    try:
-        os.makedirs("{}/structures".format(output_path))
-    except:
-        print("Warning, these is already an structures folder in this path, skip mkdir")
-    for sg in space_groups:
-        accept = 0
-        seed = start_seed
-
-        with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-            futures = {}
-            while accept < accept_count:
-                # submit new task
-                while len(futures) < max_workers and accept + len(futures) < accept_count:
-                    future = executor.submit(process_crystal, seed, sg, molecules, output_path,add_name)
-                    futures[future] = seed
-                    seed += 1
-
-                # check the finished task
-                done, _ = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED)
-                for future in done:
-                    if future.result():
-                        accept += 1
-                    # remove it from list, no matter what result it is
-                    del futures[future]
-
-                # cancel all task if the number need is arrived.
-                if accept >= accept_count:
-                    for future in futures:
-                        future.cancel()
-                    break
-
-
-def CSP_generater_serial(molecules,output_path,need_structure = 100, densely_pack_method=False, space_group_list=[1]):
-    """
-    :param molecules: a list [molecule1, molecule2, ...]
-    :param output_path: a str indicate the path of output folder
-    :param need_structure: int
-    :param space_group_list:a list indicate the space group need to search
-    """
-    try:
-        os.makedirs("{}\\structures".format(output_path))
-    except:
-        print("Warning, these is already an structures folder in this path, skip mkdir")
-    for sg in space_group_list:
-        aaa = CSP_generator_normal.CrystalGenerator(molecules, space_group=sg)
-        accept=0
-        i=1
-        while accept<need_structure:
-            new_crystal = aaa.generate(seed=i,densely_pack_method=densely_pack_method)
-            if new_crystal==None:
-                i += 1
-                continue
-            cif_out = format_parser.write_cif_file(new_crystal)
-            target = open("{}\\structures\\{}_{}.cif".format(output_path,sg,i), 'w')
-            target.writelines(cif_out)
-            target.close()
-            accept+=1
-            i += 1
-
-
-def CSP_generater_one_test(molecules,output_path,space_group=1,seed=1):
-    """
-    used to test the generator for a given space group and seed
-    :param molecules: a list [molecule1, molecule2, ...]
-    :param output_path: a str indicate the path of output folder
-    :param space_group: a list indicate the space group need to search
-    :param seed: int
-    :return: write out a cif
-    """
-    aaa = CSP_generator_normal.CrystalGenerator(molecules, space_group=space_group)
-    new_crystal = aaa.generate(seed=seed, test=True)
-    if new_crystal==None:
-        print("Return failed generate")
-    else:
-        cif_out = format_parser.write_cif_file(new_crystal,sym=True)
-        target = open("{}\\{}_{}.cif".format(output_path,space_group,seed), 'w')
-        target.writelines(cif_out)
-        target.close()
+from basic_function import format_parser
+from basic_function import CSP_generator_normal
+import os
+import concurrent.futures
+import sys
+
+
+
+def process_crystal(seed, sg, molecules,output_path,add_name):
+    aaa = CSP_generator_normal.CrystalGenerator(molecules, space_group=sg)
+    molecules_number = sum(len(molecule.atoms) for molecule in molecules)
+    new_crystal = aaa.generate(seed=seed)
+    sys.stdout.flush()
+    if new_crystal is not None:
+        cif_out = format_parser.write_cif_file(new_crystal)
+        with open(f"{output_path}/structures/{add_name}_{sg}_{seed}_z{len(molecules)}_{molecules_number}.cif", 'w') as target:
+            target.writelines(cif_out)
+        return True
+    return False
+
+def CSP_generater_parallel(molecules,output_path,need_structure = 100, space_group_list=[1],max_workers=8,add_name='',start_seed=1):
+    space_groups = space_group_list
+    accept_count = need_structure
+
+    try:
+        os.makedirs("{}/structures".format(output_path))
+    except:
+        print("Warning, these is already an structures folder in this path, skip mkdir")
+    for sg in space_groups:
+        accept = 0
+        seed = start_seed
+
+        with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+            futures = {}
+            while accept < accept_count:
+                # submit new task
+                while len(futures) < max_workers and accept + len(futures) < accept_count:
+                    future = executor.submit(process_crystal, seed, sg, molecules, output_path,add_name)
+                    futures[future] = seed
+                    seed += 1
+
+                # check the finished task
+                done, _ = concurrent.futures.wait(futures, return_when=concurrent.futures.FIRST_COMPLETED)
+                for future in done:
+                    if future.result():
+                        accept += 1
+                    # remove it from list, no matter what result it is
+                    del futures[future]
+
+                # cancel all task if the number need is arrived.
+                if accept >= accept_count:
+                    for future in futures:
+                        future.cancel()
+                    break
+
+
+def CSP_generater_serial(molecules,output_path,need_structure = 100, densely_pack_method=False, space_group_list=[1]):
+    """
+    :param molecules: a list [molecule1, molecule2, ...]
+    :param output_path: a str indicate the path of output folder
+    :param need_structure: int
+    :param space_group_list:a list indicate the space group need to search
+    """
+    try:
+        os.makedirs("{}\\structures".format(output_path))
+    except:
+        print("Warning, these is already an structures folder in this path, skip mkdir")
+    for sg in space_group_list:
+        aaa = CSP_generator_normal.CrystalGenerator(molecules, space_group=sg)
+        accept=0
+        i=1
+        while accept<need_structure:
+            new_crystal = aaa.generate(seed=i,densely_pack_method=densely_pack_method)
+            if new_crystal==None:
+                i += 1
+                continue
+            cif_out = format_parser.write_cif_file(new_crystal)
+            target = open("{}\\structures\\{}_{}.cif".format(output_path,sg,i), 'w')
+            target.writelines(cif_out)
+            target.close()
+            accept+=1
+            i += 1
+
+
+def CSP_generater_one_test(molecules,output_path,space_group=1,seed=1):
+    """
+    used to test the generator for a given space group and seed
+    :param molecules: a list [molecule1, molecule2, ...]
+    :param output_path: a str indicate the path of output folder
+    :param space_group: a list indicate the space group need to search
+    :param seed: int
+    :return: write out a cif
+    """
+    aaa = CSP_generator_normal.CrystalGenerator(molecules, space_group=space_group)
+    new_crystal = aaa.generate(seed=seed, test=True)
+    if new_crystal==None:
+        print("Return failed generate")
+    else:
+        cif_out = format_parser.write_cif_file(new_crystal,sym=True)
+        target = open("{}\\{}_{}.cif".format(output_path,space_group,seed), 'w')
+        target.writelines(cif_out)
+        target.close()

basic_function/unit_cell_parser.py

-# -*- coding: utf-8 -*-
-"""
-Provides functions for converting between different representations of a
-crystallographic unit cell (cell parameters and lattice vectors) and for
-transforming atomic coordinates between fractional and Cartesian systems.
-"""
-
-# --- Standard Library Imports ---
-from typing import List, Tuple, Union
-
-# --- Third-Party Imports ---
-import numpy as np
-import numpy.typing as npt
-
-# --- Type Aliases for Clarity ---
-NDArrayFloat = npt.NDArray[np.float64]
-CellParameters = Tuple[List[float], List[float]]
-CellVectors = Union[List[List[float]], NDArrayFloat]
-Coordinates = Union[List[float], NDArrayFloat]
-
-
-def cell_para_to_vect(
-    cell_para: CellParameters, check: bool = False
-) -> CellVectors:
-    """Converts cell parameters to lattice vectors.
-
-    The lattice vector `a` is aligned with the x-axis. The vector `b` lies in
-    the xy-plane.
-
-    Args:
-        cell_para: A tuple containing [[a, b, c], [alpha, beta, gamma]],
-                   where lengths are in Angstroms and angles are in degrees.
-        check: If True, asserts the input shape is correct.
-
-    Returns:
-        A 3x3 list of lists representing the cell vectors [a, b, c].
-    """
-    if check:
-        shape_check = np.array(cell_para)
-        assert shape_check.shape == (2, 3), "Input `cell_para` must have shape (2, 3)."
-
-    lengths = cell_para[0]
-    angles_deg = cell_para[1]
-    
-    a, b, c = lengths
-    alpha, beta, gamma = np.deg2rad(angles_deg)
-
-    cos_a, cos_b, cos_g = np.cos([alpha, beta, gamma])
-    sin_g = np.sin(gamma)
-
-    # This term is related to the square of the cell volume.
-    # It ensures the cell parameters are physically valid.
-    volume_term_sq = (
-        1.0 - cos_a**2 - cos_b**2 - cos_g**2 + 2.0 * cos_a * cos_b * cos_g
-    )
-    
-    # Ensure the argument for sqrt is non-negative
-    volume_term = np.sqrt(max(0, volume_term_sq))
-
-    cell_vect = np.zeros((3, 3))
-    cell_vect[0, 0] = a
-    cell_vect[1, 0] = b * cos_g
-    cell_vect[1, 1] = b * sin_g
-    cell_vect[2, 0] = c * cos_b
-    cell_vect[2, 1] = c * (cos_a - cos_b * cos_g) / sin_g
-    cell_vect[2, 2] = c * volume_term / sin_g
-
-    return cell_vect.tolist()
-
-
-def cell_vect_to_para(cell_vect: CellVectors, check: bool = False) -> CellParameters:
-    """Converts lattice vectors to cell parameters.
-
-    Args:
-        cell_vect: A 3x3 array-like object representing the lattice vectors.
-        check: If True, asserts the input shape is correct.
-
-    Returns:
-        A tuple containing [[a, b, c], [alpha, beta, gamma]].
-    """
-    cell_vect_np = np.array(cell_vect)
-    if check:
-        assert cell_vect_np.shape == (3, 3), "Input `cell_vect` must have shape (3, 3)."
-
-    vec_a, vec_b, vec_c = cell_vect_np
-    
-    len_a = np.linalg.norm(vec_a)
-    len_b = np.linalg.norm(vec_b)
-    len_c = np.linalg.norm(vec_c)
-    
-    lengths = [len_a, len_b, len_c]
-    
-    # Calculate angles using the dot product formula; handle potential floating point inaccuracies.
-    def _calculate_angle(v1, v2, norm1, norm2):
-        cosine_angle = np.dot(v1, v2) / (norm1 * norm2)
-        # Clip to handle values slightly outside [-1, 1] due to precision issues
-        return np.arccos(np.clip(cosine_angle, -1.0, 1.0))
-
-    alpha_rad = _calculate_angle(vec_b, vec_c, len_b, len_c)
-    beta_rad = _calculate_angle(vec_a, vec_c, len_a, len_c)
-    gamma_rad = _calculate_angle(vec_a, vec_b, len_a, len_b)
-    
-    angles_deg = np.rad2deg([alpha_rad, beta_rad, gamma_rad]).tolist()
-    
-    return (lengths, angles_deg)
-
-
-def atom_frac_to_cart_by_cell_vect(
-    atom_frac: Coordinates, cell_vect: CellVectors, check: bool = False
-) -> List[float]:
-    """Converts fractional coordinates to Cartesian coordinates using cell vectors.
-
-    Args:
-        atom_frac: A 3-element list or array of fractional coordinates.
-        cell_vect: A 3x3 matrix of lattice vectors.
-        check: If True, asserts input shapes are correct.
-
-    Returns:
-        A list of 3 Cartesian coordinates.
-    """
-    atom_frac_np = np.array(atom_frac)
-    cell_vect_np = np.array(cell_vect)
-
-    if check:
-        assert cell_vect_np.shape == (3, 3), "Input `cell_vect` must have shape (3, 3)."
-        assert atom_frac_np.shape == (3,), "Input `atom_frac` must have 3 elements."
-
-    # The transformation is a linear combination of the basis vectors.
-    # atom_cart = frac_x * vec_a + frac_y * vec_b + frac_z * vec_c
-    # This is equivalent to a dot product: [fx, fy, fz] @ [[ax,ay,az],[bx,by,bz],[cx,cy,cz]]
-    atom_cart = np.dot(atom_frac_np, cell_vect_np)
-    return atom_cart.tolist()
-
-
-def atom_frac_to_cart_by_cell_para(
-    atom_frac: Coordinates, cell_para: CellParameters, check: bool = False
-) -> List[float]:
-    """Converts fractional coordinates to Cartesian using cell parameters.
-
-    Args:
-        atom_frac: A 3-element list or array of fractional coordinates.
-        cell_para: The cell parameters [[a, b, c], [alpha, beta, gamma]].
-        check: If True, performs validation checks in underlying functions.
-
-    Returns:
-        A list of 3 Cartesian coordinates.
-    """
-    cell_vect = cell_para_to_vect(cell_para, check=check)
-    return atom_frac_to_cart_by_cell_vect(atom_frac, cell_vect, check=check)
-
-
-def atom_cart_to_frac_by_cell_vect(
-    atom_cart: Coordinates, cell_vect: CellVectors, check: bool = False
-) -> List[float]:
-    """Converts Cartesian coordinates to fractional coordinates using cell vectors.
-
-    Args:
-        atom_cart: A 3-element list or array of Cartesian coordinates.
-        cell_vect: A 3x3 matrix of lattice vectors.
-        check: If True, asserts input shapes are correct.
-
-    Returns:
-        A list of 3 fractional coordinates.
-    """
-    atom_cart_np = np.array(atom_cart)
-    cell_vect_np = np.array(cell_vect)
-
-    if check:
-        assert cell_vect_np.shape == (3, 3), "Input `cell_vect` must have shape (3, 3)."
-        assert atom_cart_np.shape == (3,), "Input `atom_cart` must have 3 elements."
-
-    # The transformation is atom_frac = atom_cart @ inverse(cell_vect)
-    inv_cell_vect = np.linalg.inv(cell_vect_np)
-    atom_frac = np.dot(atom_cart_np, inv_cell_vect)
-    return atom_frac.tolist()
-
-
-def atom_cart_to_frac_by_cell_para(
-    atom_cart: Coordinates, cell_para: CellParameters, check: bool = False
-) -> List[float]:
-    """Converts Cartesian coordinates to fractional using cell parameters.
-
-    Args:
-        atom_cart: A 3-element list or array of Cartesian coordinates.
-        cell_para: The cell parameters [[a, b, c], [alpha, beta, gamma]].
-        check: If True, performs validation checks in underlying functions.
-
-    Returns:
-        A list of 3 fractional coordinates.
-    """
-    cell_vect = cell_para_to_vect(cell_para, check=check)
-    return atom_cart_to_frac_by_cell_vect(atom_cart, cell_vect, check=check)
-
-
-def calculate_volume(cell_info: Union[CellParameters, CellVectors]) -> float:
-    """Calculates the volume of the unit cell.
-
-    Args:
-        cell_info: Can be either cell parameters [[a,b,c], [al,be,ga]] or
-                   a 3x3 matrix of cell vectors.
-
-    Returns:
-        The volume of the cell in cubic Angstroms.
-
-    Raises:
-        ValueError: If the shape of `cell_info` is not (2, 3) or (3, 3).
-    """
-    cell_info_np = np.array(cell_info)
-
-    if cell_info_np.shape == (3, 3):
-        # Input is cell vectors, calculate volume using the scalar triple product.
-        return float(np.abs(np.dot(cell_info_np[0], np.cross(cell_info_np[1], cell_info_np[2]))))
-    
-    elif cell_info_np.shape == (2, 3):
-        # Input is cell parameters.
-        lengths, angles_deg = cell_info_np
-        a, b, c = lengths
-        alpha, beta, gamma = np.deg2rad(angles_deg)
-
-        cos_a, cos_b, cos_g = np.cos([alpha, beta, gamma])
-
-        # Standard formula for volume from cell parameters
-        volume_sq = (
-            a**2 * b**2 * c**2 * (1 - cos_a**2 - cos_b**2 - cos_g**2 + 2 * cos_a * cos_b * cos_g)
-        )
-        return float(np.sqrt(max(0, volume_sq)))
-
-    else:
+# -*- coding: utf-8 -*-
+"""
+Provides functions for converting between different representations of a
+crystallographic unit cell (cell parameters and lattice vectors) and for
+transforming atomic coordinates between fractional and Cartesian systems.
+"""
+
+# --- Standard Library Imports ---
+from typing import List, Tuple, Union
+
+# --- Third-Party Imports ---
+import numpy as np
+import numpy.typing as npt
+
+# --- Type Aliases for Clarity ---
+NDArrayFloat = npt.NDArray[np.float64]
+CellParameters = Tuple[List[float], List[float]]
+CellVectors = Union[List[List[float]], NDArrayFloat]
+Coordinates = Union[List[float], NDArrayFloat]
+
+
+def cell_para_to_vect(
+    cell_para: CellParameters, check: bool = False
+) -> CellVectors:
+    """Converts cell parameters to lattice vectors.
+
+    The lattice vector `a` is aligned with the x-axis. The vector `b` lies in
+    the xy-plane.
+
+    Args:
+        cell_para: A tuple containing [[a, b, c], [alpha, beta, gamma]],
+                   where lengths are in Angstroms and angles are in degrees.
+        check: If True, asserts the input shape is correct.
+
+    Returns:
+        A 3x3 list of lists representing the cell vectors [a, b, c].
+    """
+    if check:
+        shape_check = np.array(cell_para)
+        assert shape_check.shape == (2, 3), "Input `cell_para` must have shape (2, 3)."
+
+    lengths = cell_para[0]
+    angles_deg = cell_para[1]
+    
+    a, b, c = lengths
+    alpha, beta, gamma = np.deg2rad(angles_deg)
+
+    cos_a, cos_b, cos_g = np.cos([alpha, beta, gamma])
+    sin_g = np.sin(gamma)
+
+    # This term is related to the square of the cell volume.
+    # It ensures the cell parameters are physically valid.
+    volume_term_sq = (
+        1.0 - cos_a**2 - cos_b**2 - cos_g**2 + 2.0 * cos_a * cos_b * cos_g
+    )
+    
+    # Ensure the argument for sqrt is non-negative
+    volume_term = np.sqrt(max(0, volume_term_sq))
+
+    cell_vect = np.zeros((3, 3))
+    cell_vect[0, 0] = a
+    cell_vect[1, 0] = b * cos_g
+    cell_vect[1, 1] = b * sin_g
+    cell_vect[2, 0] = c * cos_b
+    cell_vect[2, 1] = c * (cos_a - cos_b * cos_g) / sin_g
+    cell_vect[2, 2] = c * volume_term / sin_g
+
+    return cell_vect.tolist()
+
+
+def cell_vect_to_para(cell_vect: CellVectors, check: bool = False) -> CellParameters:
+    """Converts lattice vectors to cell parameters.
+
+    Args:
+        cell_vect: A 3x3 array-like object representing the lattice vectors.
+        check: If True, asserts the input shape is correct.
+
+    Returns:
+        A tuple containing [[a, b, c], [alpha, beta, gamma]].
+    """
+    cell_vect_np = np.array(cell_vect)
+    if check:
+        assert cell_vect_np.shape == (3, 3), "Input `cell_vect` must have shape (3, 3)."
+
+    vec_a, vec_b, vec_c = cell_vect_np
+    
+    len_a = np.linalg.norm(vec_a)
+    len_b = np.linalg.norm(vec_b)
+    len_c = np.linalg.norm(vec_c)
+    
+    lengths = [len_a, len_b, len_c]
+    
+    # Calculate angles using the dot product formula; handle potential floating point inaccuracies.
+    def _calculate_angle(v1, v2, norm1, norm2):
+        cosine_angle = np.dot(v1, v2) / (norm1 * norm2)
+        # Clip to handle values slightly outside [-1, 1] due to precision issues
+        return np.arccos(np.clip(cosine_angle, -1.0, 1.0))
+
+    alpha_rad = _calculate_angle(vec_b, vec_c, len_b, len_c)
+    beta_rad = _calculate_angle(vec_a, vec_c, len_a, len_c)
+    gamma_rad = _calculate_angle(vec_a, vec_b, len_a, len_b)
+    
+    angles_deg = np.rad2deg([alpha_rad, beta_rad, gamma_rad]).tolist()
+    
+    return (lengths, angles_deg)
+
+
+def atom_frac_to_cart_by_cell_vect(
+    atom_frac: Coordinates, cell_vect: CellVectors, check: bool = False
+) -> List[float]:
+    """Converts fractional coordinates to Cartesian coordinates using cell vectors.
+
+    Args:
+        atom_frac: A 3-element list or array of fractional coordinates.
+        cell_vect: A 3x3 matrix of lattice vectors.
+        check: If True, asserts input shapes are correct.
+
+    Returns:
+        A list of 3 Cartesian coordinates.
+    """
+    atom_frac_np = np.array(atom_frac)
+    cell_vect_np = np.array(cell_vect)
+
+    if check:
+        assert cell_vect_np.shape == (3, 3), "Input `cell_vect` must have shape (3, 3)."
+        assert atom_frac_np.shape == (3,), "Input `atom_frac` must have 3 elements."
+
+    # The transformation is a linear combination of the basis vectors.
+    # atom_cart = frac_x * vec_a + frac_y * vec_b + frac_z * vec_c
+    # This is equivalent to a dot product: [fx, fy, fz] @ [[ax,ay,az],[bx,by,bz],[cx,cy,cz]]
+    atom_cart = np.dot(atom_frac_np, cell_vect_np)
+    return atom_cart.tolist()
+
+
+def atom_frac_to_cart_by_cell_para(
+    atom_frac: Coordinates, cell_para: CellParameters, check: bool = False
+) -> List[float]:
+    """Converts fractional coordinates to Cartesian using cell parameters.
+
+    Args:
+        atom_frac: A 3-element list or array of fractional coordinates.
+        cell_para: The cell parameters [[a, b, c], [alpha, beta, gamma]].
+        check: If True, performs validation checks in underlying functions.
+
+    Returns:
+        A list of 3 Cartesian coordinates.
+    """
+    cell_vect = cell_para_to_vect(cell_para, check=check)
+    return atom_frac_to_cart_by_cell_vect(atom_frac, cell_vect, check=check)
+
+
+def atom_cart_to_frac_by_cell_vect(
+    atom_cart: Coordinates, cell_vect: CellVectors, check: bool = False
+) -> List[float]:
+    """Converts Cartesian coordinates to fractional coordinates using cell vectors.
+
+    Args:
+        atom_cart: A 3-element list or array of Cartesian coordinates.
+        cell_vect: A 3x3 matrix of lattice vectors.
+        check: If True, asserts input shapes are correct.
+
+    Returns:
+        A list of 3 fractional coordinates.
+    """
+    atom_cart_np = np.array(atom_cart)
+    cell_vect_np = np.array(cell_vect)
+
+    if check:
+        assert cell_vect_np.shape == (3, 3), "Input `cell_vect` must have shape (3, 3)."
+        assert atom_cart_np.shape == (3,), "Input `atom_cart` must have 3 elements."
+
+    # The transformation is atom_frac = atom_cart @ inverse(cell_vect)
+    inv_cell_vect = np.linalg.inv(cell_vect_np)
+    atom_frac = np.dot(atom_cart_np, inv_cell_vect)
+    return atom_frac.tolist()
+
+
+def atom_cart_to_frac_by_cell_para(
+    atom_cart: Coordinates, cell_para: CellParameters, check: bool = False
+) -> List[float]:
+    """Converts Cartesian coordinates to fractional using cell parameters.
+
+    Args:
+        atom_cart: A 3-element list or array of Cartesian coordinates.
+        cell_para: The cell parameters [[a, b, c], [alpha, beta, gamma]].
+        check: If True, performs validation checks in underlying functions.
+
+    Returns:
+        A list of 3 fractional coordinates.
+    """
+    cell_vect = cell_para_to_vect(cell_para, check=check)
+    return atom_cart_to_frac_by_cell_vect(atom_cart, cell_vect, check=check)
+
+
+def calculate_volume(cell_info: Union[CellParameters, CellVectors]) -> float:
+    """Calculates the volume of the unit cell.
+
+    Args:
+        cell_info: Can be either cell parameters [[a,b,c], [al,be,ga]] or
+                   a 3x3 matrix of cell vectors.
+
+    Returns:
+        The volume of the cell in cubic Angstroms.
+
+    Raises:
+        ValueError: If the shape of `cell_info` is not (2, 3) or (3, 3).
+    """
+    cell_info_np = np.array(cell_info)
+
+    if cell_info_np.shape == (3, 3):
+        # Input is cell vectors, calculate volume using the scalar triple product.
+        return float(np.abs(np.dot(cell_info_np[0], np.cross(cell_info_np[1], cell_info_np[2]))))
+    
+    elif cell_info_np.shape == (2, 3):
+        # Input is cell parameters.
+        lengths, angles_deg = cell_info_np
+        a, b, c = lengths
+        alpha, beta, gamma = np.deg2rad(angles_deg)
+
+        cos_a, cos_b, cos_g = np.cos([alpha, beta, gamma])
+
+        # Standard formula for volume from cell parameters
+        volume_sq = (
+            a**2 * b**2 * c**2 * (1 - cos_a**2 - cos_b**2 - cos_g**2 + 2 * cos_a * cos_b * cos_g)
+        )
+        return float(np.sqrt(max(0, volume_sq)))
+
+    else:
         raise ValueError(f"Cannot understand input shape {cell_info_np.shape} for `cell_info`.")
\ No newline at end of file

csp.sh

-#!/bin/bash
-TOP_DIR=$(pwd)
-TAR_DIR="${TOP_DIR}/test"
-
-mkdir -p "${TAR_DIR}"
-cd ${TAR_DIR}
-
-# generate structures
-python "${TOP_DIR}/main.py" --path ${TAR_DIR} --smiles "OC(=O)c1cc(O)c(O)c(O)c1.O" \
-    --molecule_num_in_cell 1,1 --space_group_list 13,14 --add_name KONTIQ --max_workers 16\
-     --num_generation 100 --generate_conformers 20 --use_conformers 4 > generate.log 2>&1
-
-# opt structures using mace, --batch_size 0 means auto batch size only for mace
-mkdir -p "${TAR_DIR}/mace_opt"
-cd "${TAR_DIR}/mace_opt"
-python "${TOP_DIR}/mace-bench/scripts/mace_opt_batch.py" --target_folder "${TAR_DIR}/structures" \
-    --molecule_single 21 --gpu_offset 0 --n_gpus 8 --num_workers 80 --batch_size 0 \
-    --max_steps 3000 --filter1 UnitCellFilter --filter2 UnitCellFilter \
-    --optimizer1 BFGSFusedLS --optimizer2 BFGS --num_threads 1 --cueq true \
-    --use_ordered_files true --model mace > opt.log 2>&1
-
-# opt structures using 7net
-# mkdir -p "${TAR_DIR}/7net_opt"
-# cd "${TAR_DIR}/7net_opt"
-# python "${TOP_DIR}/mace-bench/scripts/mace_opt_batch.py" --target_folder "${TAR_DIR}/structures" \
-#     --molecule_single 21 --gpu_offset 0 --n_gpus 8 --num_workers 48 --batch_size 2 \
-#     --max_steps 3000 --filter1 UnitCellFilter --filter2 UnitCellFilter \
-#     --optimizer1 BFGSFusedLS --optimizer2 BFGS --num_threads 2 --cueq true \
-#     --use_ordered_files true --model sevennet > opt.log 2>&1
-
-# Postprocess the opt structures
-python "${TOP_DIR}/post_process/clean_table.py"
-## Make sure you have installed csd-python-api in current env before execuing following commands
-# conda activate ccdc
-# python "${TOP_DIR}/post_process/check_match.py" --workers 80 --timeout 20 --ref_path "${TAR_DIR}/refs"
-# python "${TOP_DIR}/post_process/duplicate_remove.py" --workers 80
+
+TOP_DIR=$(pwd)
+TAR_DIR="${TOP_DIR}/test"
+
+mkdir -p "${TAR_DIR}"
+cd ${TAR_DIR}
+
+# conformer search and structure generation
+# change --mode to conformer_only or structure_only to seperate the process. 
+python "${TOP_DIR}/main.py" --path ${TAR_DIR} --smiles "OC(=O)c1cc(O)c(O)c(O)c1.O" \
+    --molecule_num_in_cell 1,1 --space_group_list 13,14 --add_name KONTIQ --max_workers 16\
+     --num_generation 100 --generate_conformers 20 --use_conformers 4 --mode all > generate.log 2>&1
+
+# opt structures using mace, --batch_size 0 means auto batch size only for mace
+mkdir -p "${TAR_DIR}/mace_opt"
+cd "${TAR_DIR}/mace_opt"
+python "${TOP_DIR}/mace-bench/scripts/mace_opt_batch.py" --target_folder "${TAR_DIR}/structures" \
+    --molecule_single 21 --gpu_offset 0 --n_gpus 8 --num_workers 80 --batch_size 0 \
+    --max_steps 3000 --filter1 UnitCellFilter --filter2 UnitCellFilter \
+    --optimizer1 BFGSFusedLS --optimizer2 BFGS --num_threads 1 --cueq true \
+    --use_ordered_files true --model mace > opt.log 2>&1
+
+# opt structures using 7net
+# mkdir -p "${TAR_DIR}/7net_opt"
+# cd "${TAR_DIR}/7net_opt"
+# python "${TOP_DIR}/mace-bench/scripts/mace_opt_batch.py" --target_folder "${TAR_DIR}/structures" \
+#     --molecule_single 21 --gpu_offset 0 --n_gpus 8 --num_workers 48 --batch_size 2 \
+#     --max_steps 3000 --filter1 UnitCellFilter --filter2 UnitCellFilter \
+#     --optimizer1 BFGSFusedLS --optimizer2 BFGS --num_threads 2 --cueq true \
+#     --use_ordered_files true --model sevennet > opt.log 2>&1
+
+# Postprocess the opt structures
+python "${TOP_DIR}/post_process/clean_table.py"
+## Make sure you have installed csd-python-api in current env before execuing following commands
+# conda activate ccdc
+# python "${TOP_DIR}/post_process/check_match.py" --workers 80 --timeout 20 --ref_path "${TAR_DIR}/refs"
+# python "${TOP_DIR}/post_process/duplicate_remove.py" --workers 80

mace-bench/3rdparty/SevenNet/sevenn/__init__.py

-from importlib.metadata import version
-
-from packaging.version import Version
-
-__version__ = version('sevenn')
-
-from e3nn import __version__ as e3nn_ver
-
-if Version(e3nn_ver) < Version('0.5.0'):
-    raise ValueError(
-        'The e3nn version MUST be 0.5.0 or later due to changes in CG coefficient '
-        'convention.'
-    )
+from importlib.metadata import version
+
+from packaging.version import Version
+
+__version__ = version('sevenn')
+
+from e3nn import __version__ as e3nn_ver
+
+if Version(e3nn_ver) < Version('0.5.0'):
+    raise ValueError(
+        'The e3nn version MUST be 0.5.0 or later due to changes in CG coefficient '
+        'convention.'
+    )

mace-bench/3rdparty/SevenNet/sevenn/_const.py

-import os
-from enum import Enum
-from typing import Dict
-
-import torch
-
-import sevenn._keys as KEY
-from sevenn.nn.activation import ShiftedSoftPlus
-
-NUM_UNIV_ELEMENT = 119  # Z = 0 ~ 118
-
-IMPLEMENTED_RADIAL_BASIS = ['bessel']
-IMPLEMENTED_CUTOFF_FUNCTION = ['poly_cut', 'XPLOR']
-# TODO: support None. This became difficult because of parallel model
-IMPLEMENTED_SELF_CONNECTION_TYPE = ['nequip', 'linear']
-IMPLEMENTED_INTERACTION_TYPE = ['nequip']
-
-IMPLEMENTED_SHIFT = ['per_atom_energy_mean', 'elemwise_reference_energies']
-IMPLEMENTED_SCALE = ['force_rms', 'per_atom_energy_std', 'elemwise_force_rms']
-
-SUPPORTING_METRICS = ['RMSE', 'ComponentRMSE', 'MAE', 'Loss']
-SUPPORTING_ERROR_TYPES = [
-    'TotalEnergy',
-    'Energy',
-    'Force',
-    'Stress',
-    'Stress_GPa',
-    'TotalLoss',
-]
-
-IMPLEMENTED_MODEL = ['E3_equivariant_model']
-
-# string input to real torch function
-ACTIVATION = {
-    'relu': torch.nn.functional.relu,
-    'silu': torch.nn.functional.silu,
-    'tanh': torch.tanh,
-    'abs': torch.abs,
-    'ssp': ShiftedSoftPlus,
-    'sigmoid': torch.sigmoid,
-    'elu': torch.nn.functional.elu,
-}
-ACTIVATION_FOR_EVEN = {
-    'ssp': ShiftedSoftPlus,
-    'silu': torch.nn.functional.silu,
-}
-ACTIVATION_FOR_ODD = {'tanh': torch.tanh, 'abs': torch.abs}
-ACTIVATION_DICT = {'e': ACTIVATION_FOR_EVEN, 'o': ACTIVATION_FOR_ODD}
-
-_prefix = os.path.abspath(f'{os.path.dirname(__file__)}/pretrained_potentials')
-SEVENNET_0_11Jul2024 = f'{_prefix}/SevenNet_0__11Jul2024/checkpoint_sevennet_0.pth'
-SEVENNET_0_22May2024 = f'{_prefix}/SevenNet_0__22May2024/checkpoint_sevennet_0.pth'
-SEVENNET_l3i5 = f'{_prefix}/SevenNet_l3i5/checkpoint_l3i5.pth'
-SEVENNET_MF_0 = f'{_prefix}/SevenNet_MF_0/checkpoint_sevennet_mf_0.pth'
-SEVENNET_MF_ompa = f'{_prefix}/SevenNet_MF_ompa/checkpoint_sevennet_mf_ompa.pth'
-SEVENNET_omat = f'{_prefix}/SevenNet_omat/checkpoint_sevennet_omat.pth'
-
-_git_prefix = 'https://github.com/MDIL-SNU/SevenNet/releases/download'
-CHECKPOINT_DOWNLOAD_LINKS = {
-    SEVENNET_MF_ompa: f'{_git_prefix}/v0.11.0.cp/checkpoint_sevennet_mf_ompa.pth',
-    SEVENNET_omat: f'{_git_prefix}/v0.11.0.cp/checkpoint_sevennet_omat.pth',
-}
-# to avoid torch script to compile torch_geometry.data
-AtomGraphDataType = Dict[str, torch.Tensor]
-
-
-class LossType(Enum):  # only used for train_v1, do not use it afterwards
-    ENERGY = 'energy'  # eV or eV/atom
-    FORCE = 'force'  # eV/A
-    STRESS = 'stress'  # kB
-
-
-def error_record_condition(x):
-    if type(x) is not list:
-        return False
-    for v in x:
-        if type(v) is not list or len(v) != 2:
-            return False
-        if v[0] not in SUPPORTING_ERROR_TYPES:
-            return False
-        if v[0] == 'TotalLoss':
-            continue
-        if v[1] not in SUPPORTING_METRICS:
-            return False
-    return True
-
-
-DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG = {
-    KEY.CUTOFF: 4.5,
-    KEY.NODE_FEATURE_MULTIPLICITY: 32,
-    KEY.IRREPS_MANUAL: False,
-    KEY.LMAX: 1,
-    KEY.LMAX_EDGE: -1,  # -1 means lmax_edge = lmax
-    KEY.LMAX_NODE: -1,  # -1 means lmax_node = lmax
-    KEY.IS_PARITY: True,
-    KEY.NUM_CONVOLUTION: 3,
-    KEY.RADIAL_BASIS: {
-        KEY.RADIAL_BASIS_NAME: 'bessel',
-    },
-    KEY.CUTOFF_FUNCTION: {
-        KEY.CUTOFF_FUNCTION_NAME: 'poly_cut',
-    },
-    KEY.ACTIVATION_RADIAL: 'silu',
-    KEY.ACTIVATION_SCARLAR: {'e': 'silu', 'o': 'tanh'},
-    KEY.ACTIVATION_GATE: {'e': 'silu', 'o': 'tanh'},
-    KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: [64, 64],
-    # KEY.AVG_NUM_NEIGH: True,  # deprecated
-    # KEY.TRAIN_AVG_NUM_NEIGH: False,  # deprecated
-    KEY.CONV_DENOMINATOR: 'avg_num_neigh',
-    KEY.TRAIN_DENOMINTAOR: False,
-    KEY.TRAIN_SHIFT_SCALE: False,
-    # KEY.OPTIMIZE_BY_REDUCE: True,  # deprecated, always True
-    KEY.USE_BIAS_IN_LINEAR: False,
-    KEY.USE_MODAL_NODE_EMBEDDING: False,
-    KEY.USE_MODAL_SELF_INTER_INTRO: False,
-    KEY.USE_MODAL_SELF_INTER_OUTRO: False,
-    KEY.USE_MODAL_OUTPUT_BLOCK: False,
-    KEY.READOUT_AS_FCN: False,
-    # Applied af readout as fcn is True
-    KEY.READOUT_FCN_HIDDEN_NEURONS: [30, 30],
-    KEY.READOUT_FCN_ACTIVATION: 'relu',
-    KEY.SELF_CONNECTION_TYPE: 'nequip',
-    KEY.INTERACTION_TYPE: 'nequip',
-    KEY._NORMALIZE_SPH: True,
-    KEY.CUEQUIVARIANCE_CONFIG: {},
-}
-
-
-# Basically, "If provided, it should be type of ..."
-MODEL_CONFIG_CONDITION = {
-    KEY.NODE_FEATURE_MULTIPLICITY: int,
-    KEY.LMAX: int,
-    KEY.LMAX_EDGE: int,
-    KEY.LMAX_NODE: int,
-    KEY.IS_PARITY: bool,
-    KEY.RADIAL_BASIS: {
-        KEY.RADIAL_BASIS_NAME: lambda x: x in IMPLEMENTED_RADIAL_BASIS,
-    },
-    KEY.CUTOFF_FUNCTION: {
-        KEY.CUTOFF_FUNCTION_NAME: lambda x: x in IMPLEMENTED_CUTOFF_FUNCTION,
-    },
-    KEY.CUTOFF: float,
-    KEY.NUM_CONVOLUTION: int,
-    KEY.CONV_DENOMINATOR: lambda x: isinstance(x, float)
-    or x
-    in [
-        'avg_num_neigh',
-        'sqrt_avg_num_neigh',
-    ],
-    KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: list,
-    KEY.TRAIN_SHIFT_SCALE: bool,
-    KEY.TRAIN_DENOMINTAOR: bool,
-    KEY.USE_BIAS_IN_LINEAR: bool,
-    KEY.USE_MODAL_NODE_EMBEDDING: bool,
-    KEY.USE_MODAL_SELF_INTER_INTRO: bool,
-    KEY.USE_MODAL_SELF_INTER_OUTRO: bool,
-    KEY.USE_MODAL_OUTPUT_BLOCK: bool,
-    KEY.READOUT_AS_FCN: bool,
-    KEY.READOUT_FCN_HIDDEN_NEURONS: list,
-    KEY.READOUT_FCN_ACTIVATION: str,
-    KEY.ACTIVATION_RADIAL: str,
-    KEY.SELF_CONNECTION_TYPE: lambda x: (
-        x in IMPLEMENTED_SELF_CONNECTION_TYPE
-        or (
-            isinstance(x, list)
-            and all(sc in IMPLEMENTED_SELF_CONNECTION_TYPE for sc in x)
-        )
-    ),
-    KEY.INTERACTION_TYPE: lambda x: x in IMPLEMENTED_INTERACTION_TYPE,
-    KEY._NORMALIZE_SPH: bool,
-    KEY.CUEQUIVARIANCE_CONFIG: dict,
-}
-
-
-def model_defaults(config):
-    defaults = DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG
-
-    if KEY.READOUT_AS_FCN not in config:
-        config[KEY.READOUT_AS_FCN] = defaults[KEY.READOUT_AS_FCN]
-    if config[KEY.READOUT_AS_FCN] is False:
-        defaults.pop(KEY.READOUT_FCN_ACTIVATION, None)
-        defaults.pop(KEY.READOUT_FCN_HIDDEN_NEURONS, None)
-
-    return defaults
-
-
-DEFAULT_DATA_CONFIG = {
-    KEY.DTYPE: 'single',
-    KEY.DATA_FORMAT: 'ase',
-    KEY.DATA_FORMAT_ARGS: {},
-    KEY.SAVE_DATASET: False,
-    KEY.SAVE_BY_LABEL: False,
-    KEY.SAVE_BY_TRAIN_VALID: False,
-    KEY.RATIO: 0.0,
-    KEY.BATCH_SIZE: 6,
-    KEY.PREPROCESS_NUM_CORES: 1,
-    KEY.COMPUTE_STATISTICS: True,
-    KEY.DATASET_TYPE: 'graph',
-    # KEY.USE_SPECIES_WISE_SHIFT_SCALE: False,
-    KEY.USE_MODAL_WISE_SHIFT: False,
-    KEY.USE_MODAL_WISE_SCALE: False,
-    KEY.SHIFT: 'per_atom_energy_mean',
-    KEY.SCALE: 'force_rms',
-    # KEY.DATA_SHUFFLE: True,
-    # KEY.DATA_WEIGHT: False,
-    # KEY.DATA_MODALITY: False,
-}
-
-DATA_CONFIG_CONDITION = {
-    KEY.DTYPE: str,
-    KEY.DATA_FORMAT: str,
-    KEY.DATA_FORMAT_ARGS: dict,
-    KEY.SAVE_DATASET: str,
-    KEY.SAVE_BY_LABEL: bool,
-    KEY.SAVE_BY_TRAIN_VALID: bool,
-    KEY.RATIO: float,
-    KEY.BATCH_SIZE: int,
-    KEY.PREPROCESS_NUM_CORES: int,
-    KEY.DATASET_TYPE: lambda x: x in ['graph', 'atoms'],
-    # KEY.USE_SPECIES_WISE_SHIFT_SCALE: bool,
-    KEY.SHIFT: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SHIFT,
-    KEY.SCALE: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SCALE,
-    KEY.USE_MODAL_WISE_SHIFT: bool,
-    KEY.USE_MODAL_WISE_SCALE: bool,
-    # KEY.DATA_SHUFFLE: bool,
-    KEY.COMPUTE_STATISTICS: bool,
-    # KEY.DATA_WEIGHT: bool,
-    # KEY.DATA_MODALITY: bool,
-}
-
-
-def data_defaults(config):
-    defaults = DEFAULT_DATA_CONFIG
-    if KEY.LOAD_VALIDSET in config:
-        defaults.pop(KEY.RATIO, None)
-    return defaults
-
-
-DEFAULT_TRAINING_CONFIG = {
-    KEY.RANDOM_SEED: 1,
-    KEY.EPOCH: 300,
-    KEY.LOSS: 'mse',
-    KEY.LOSS_PARAM: {},
-    KEY.OPTIMIZER: 'adam',
-    KEY.OPTIM_PARAM: {},
-    KEY.SCHEDULER: 'exponentiallr',
-    KEY.SCHEDULER_PARAM: {},
-    KEY.FORCE_WEIGHT: 0.1,
-    KEY.STRESS_WEIGHT: 1e-6,  # SIMPLE-NN default
-    KEY.PER_EPOCH: 5,
-    # KEY.USE_TESTSET: False,
-    KEY.CONTINUE: {
-        KEY.CHECKPOINT: False,
-        KEY.RESET_OPTIMIZER: False,
-        KEY.RESET_SCHEDULER: False,
-        KEY.RESET_EPOCH: False,
-        KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: True,
-        KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: True,
-    },
-    # KEY.DEFAULT_MODAL: 'common',
-    KEY.CSV_LOG: 'log.csv',
-    KEY.NUM_WORKERS: 0,
-    KEY.IS_TRAIN_STRESS: True,
-    KEY.TRAIN_SHUFFLE: True,
-    KEY.ERROR_RECORD: [
-        ['Energy', 'RMSE'],
-        ['Force', 'RMSE'],
-        ['Stress', 'RMSE'],
-        ['TotalLoss', 'None'],
-    ],
-    KEY.BEST_METRIC: 'TotalLoss',
-    KEY.USE_WEIGHT: False,
-    KEY.USE_MODALITY: False,
-}
-
-
-TRAINING_CONFIG_CONDITION = {
-    KEY.RANDOM_SEED: int,
-    KEY.EPOCH: int,
-    KEY.FORCE_WEIGHT: float,
-    KEY.STRESS_WEIGHT: float,
-    KEY.USE_TESTSET: None,  # Not used
-    KEY.NUM_WORKERS: int,
-    KEY.PER_EPOCH: int,
-    KEY.CONTINUE: {
-        KEY.CHECKPOINT: str,
-        KEY.RESET_OPTIMIZER: bool,
-        KEY.RESET_SCHEDULER: bool,
-        KEY.RESET_EPOCH: bool,
-        KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: bool,
-        KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: bool,
-    },
-    KEY.DEFAULT_MODAL: str,
-    KEY.IS_TRAIN_STRESS: bool,
-    KEY.TRAIN_SHUFFLE: bool,
-    KEY.ERROR_RECORD: error_record_condition,
-    KEY.BEST_METRIC: str,
-    KEY.CSV_LOG: str,
-    KEY.USE_MODALITY: bool,
-    KEY.USE_WEIGHT: bool,
-}
-
-
-def train_defaults(config):
-    defaults = DEFAULT_TRAINING_CONFIG
-    if KEY.IS_TRAIN_STRESS not in config:
-        config[KEY.IS_TRAIN_STRESS] = defaults[KEY.IS_TRAIN_STRESS]
-    if not config[KEY.IS_TRAIN_STRESS]:
-        defaults.pop(KEY.STRESS_WEIGHT, None)
-    return defaults
+import os
+from enum import Enum
+from typing import Dict
+
+import torch
+
+import sevenn._keys as KEY
+from sevenn.nn.activation import ShiftedSoftPlus
+
+NUM_UNIV_ELEMENT = 119  # Z = 0 ~ 118
+
+IMPLEMENTED_RADIAL_BASIS = ['bessel']
+IMPLEMENTED_CUTOFF_FUNCTION = ['poly_cut', 'XPLOR']
+# TODO: support None. This became difficult because of parallel model
+IMPLEMENTED_SELF_CONNECTION_TYPE = ['nequip', 'linear']
+IMPLEMENTED_INTERACTION_TYPE = ['nequip']
+
+IMPLEMENTED_SHIFT = ['per_atom_energy_mean', 'elemwise_reference_energies']
+IMPLEMENTED_SCALE = ['force_rms', 'per_atom_energy_std', 'elemwise_force_rms']
+
+SUPPORTING_METRICS = ['RMSE', 'ComponentRMSE', 'MAE', 'Loss']
+SUPPORTING_ERROR_TYPES = [
+    'TotalEnergy',
+    'Energy',
+    'Force',
+    'Stress',
+    'Stress_GPa',
+    'TotalLoss',
+]
+
+IMPLEMENTED_MODEL = ['E3_equivariant_model']
+
+# string input to real torch function
+ACTIVATION = {
+    'relu': torch.nn.functional.relu,
+    'silu': torch.nn.functional.silu,
+    'tanh': torch.tanh,
+    'abs': torch.abs,
+    'ssp': ShiftedSoftPlus,
+    'sigmoid': torch.sigmoid,
+    'elu': torch.nn.functional.elu,
+}
+ACTIVATION_FOR_EVEN = {
+    'ssp': ShiftedSoftPlus,
+    'silu': torch.nn.functional.silu,
+}
+ACTIVATION_FOR_ODD = {'tanh': torch.tanh, 'abs': torch.abs}
+ACTIVATION_DICT = {'e': ACTIVATION_FOR_EVEN, 'o': ACTIVATION_FOR_ODD}
+
+_prefix = os.path.abspath(f'{os.path.dirname(__file__)}/pretrained_potentials')
+SEVENNET_0_11Jul2024 = f'{_prefix}/SevenNet_0__11Jul2024/checkpoint_sevennet_0.pth'
+SEVENNET_0_22May2024 = f'{_prefix}/SevenNet_0__22May2024/checkpoint_sevennet_0.pth'
+SEVENNET_l3i5 = f'{_prefix}/SevenNet_l3i5/checkpoint_l3i5.pth'
+SEVENNET_MF_0 = f'{_prefix}/SevenNet_MF_0/checkpoint_sevennet_mf_0.pth'
+SEVENNET_MF_ompa = f'{_prefix}/SevenNet_MF_ompa/checkpoint_sevennet_mf_ompa.pth'
+SEVENNET_omat = f'{_prefix}/SevenNet_omat/checkpoint_sevennet_omat.pth'
+
+_git_prefix = 'https://github.com/MDIL-SNU/SevenNet/releases/download'
+CHECKPOINT_DOWNLOAD_LINKS = {
+    SEVENNET_MF_ompa: f'{_git_prefix}/v0.11.0.cp/checkpoint_sevennet_mf_ompa.pth',
+    SEVENNET_omat: f'{_git_prefix}/v0.11.0.cp/checkpoint_sevennet_omat.pth',
+}
+# to avoid torch script to compile torch_geometry.data
+AtomGraphDataType = Dict[str, torch.Tensor]
+
+
+class LossType(Enum):  # only used for train_v1, do not use it afterwards
+    ENERGY = 'energy'  # eV or eV/atom
+    FORCE = 'force'  # eV/A
+    STRESS = 'stress'  # kB
+
+
+def error_record_condition(x):
+    if type(x) is not list:
+        return False
+    for v in x:
+        if type(v) is not list or len(v) != 2:
+            return False
+        if v[0] not in SUPPORTING_ERROR_TYPES:
+            return False
+        if v[0] == 'TotalLoss':
+            continue
+        if v[1] not in SUPPORTING_METRICS:
+            return False
+    return True
+
+
+DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG = {
+    KEY.CUTOFF: 4.5,
+    KEY.NODE_FEATURE_MULTIPLICITY: 32,
+    KEY.IRREPS_MANUAL: False,
+    KEY.LMAX: 1,
+    KEY.LMAX_EDGE: -1,  # -1 means lmax_edge = lmax
+    KEY.LMAX_NODE: -1,  # -1 means lmax_node = lmax
+    KEY.IS_PARITY: True,
+    KEY.NUM_CONVOLUTION: 3,
+    KEY.RADIAL_BASIS: {
+        KEY.RADIAL_BASIS_NAME: 'bessel',
+    },
+    KEY.CUTOFF_FUNCTION: {
+        KEY.CUTOFF_FUNCTION_NAME: 'poly_cut',
+    },
+    KEY.ACTIVATION_RADIAL: 'silu',
+    KEY.ACTIVATION_SCARLAR: {'e': 'silu', 'o': 'tanh'},
+    KEY.ACTIVATION_GATE: {'e': 'silu', 'o': 'tanh'},
+    KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: [64, 64],
+    # KEY.AVG_NUM_NEIGH: True,  # deprecated
+    # KEY.TRAIN_AVG_NUM_NEIGH: False,  # deprecated
+    KEY.CONV_DENOMINATOR: 'avg_num_neigh',
+    KEY.TRAIN_DENOMINTAOR: False,
+    KEY.TRAIN_SHIFT_SCALE: False,
+    # KEY.OPTIMIZE_BY_REDUCE: True,  # deprecated, always True
+    KEY.USE_BIAS_IN_LINEAR: False,
+    KEY.USE_MODAL_NODE_EMBEDDING: False,
+    KEY.USE_MODAL_SELF_INTER_INTRO: False,
+    KEY.USE_MODAL_SELF_INTER_OUTRO: False,
+    KEY.USE_MODAL_OUTPUT_BLOCK: False,
+    KEY.READOUT_AS_FCN: False,
+    # Applied af readout as fcn is True
+    KEY.READOUT_FCN_HIDDEN_NEURONS: [30, 30],
+    KEY.READOUT_FCN_ACTIVATION: 'relu',
+    KEY.SELF_CONNECTION_TYPE: 'nequip',
+    KEY.INTERACTION_TYPE: 'nequip',
+    KEY._NORMALIZE_SPH: True,
+    KEY.CUEQUIVARIANCE_CONFIG: {},
+}
+
+
+# Basically, "If provided, it should be type of ..."
+MODEL_CONFIG_CONDITION = {
+    KEY.NODE_FEATURE_MULTIPLICITY: int,
+    KEY.LMAX: int,
+    KEY.LMAX_EDGE: int,
+    KEY.LMAX_NODE: int,
+    KEY.IS_PARITY: bool,
+    KEY.RADIAL_BASIS: {
+        KEY.RADIAL_BASIS_NAME: lambda x: x in IMPLEMENTED_RADIAL_BASIS,
+    },
+    KEY.CUTOFF_FUNCTION: {
+        KEY.CUTOFF_FUNCTION_NAME: lambda x: x in IMPLEMENTED_CUTOFF_FUNCTION,
+    },
+    KEY.CUTOFF: float,
+    KEY.NUM_CONVOLUTION: int,
+    KEY.CONV_DENOMINATOR: lambda x: isinstance(x, float)
+    or x
+    in [
+        'avg_num_neigh',
+        'sqrt_avg_num_neigh',
+    ],
+    KEY.CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS: list,
+    KEY.TRAIN_SHIFT_SCALE: bool,
+    KEY.TRAIN_DENOMINTAOR: bool,
+    KEY.USE_BIAS_IN_LINEAR: bool,
+    KEY.USE_MODAL_NODE_EMBEDDING: bool,
+    KEY.USE_MODAL_SELF_INTER_INTRO: bool,
+    KEY.USE_MODAL_SELF_INTER_OUTRO: bool,
+    KEY.USE_MODAL_OUTPUT_BLOCK: bool,
+    KEY.READOUT_AS_FCN: bool,
+    KEY.READOUT_FCN_HIDDEN_NEURONS: list,
+    KEY.READOUT_FCN_ACTIVATION: str,
+    KEY.ACTIVATION_RADIAL: str,
+    KEY.SELF_CONNECTION_TYPE: lambda x: (
+        x in IMPLEMENTED_SELF_CONNECTION_TYPE
+        or (
+            isinstance(x, list)
+            and all(sc in IMPLEMENTED_SELF_CONNECTION_TYPE for sc in x)
+        )
+    ),
+    KEY.INTERACTION_TYPE: lambda x: x in IMPLEMENTED_INTERACTION_TYPE,
+    KEY._NORMALIZE_SPH: bool,
+    KEY.CUEQUIVARIANCE_CONFIG: dict,
+}
+
+
+def model_defaults(config):
+    defaults = DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG
+
+    if KEY.READOUT_AS_FCN not in config:
+        config[KEY.READOUT_AS_FCN] = defaults[KEY.READOUT_AS_FCN]
+    if config[KEY.READOUT_AS_FCN] is False:
+        defaults.pop(KEY.READOUT_FCN_ACTIVATION, None)
+        defaults.pop(KEY.READOUT_FCN_HIDDEN_NEURONS, None)
+
+    return defaults
+
+
+DEFAULT_DATA_CONFIG = {
+    KEY.DTYPE: 'single',
+    KEY.DATA_FORMAT: 'ase',
+    KEY.DATA_FORMAT_ARGS: {},
+    KEY.SAVE_DATASET: False,
+    KEY.SAVE_BY_LABEL: False,
+    KEY.SAVE_BY_TRAIN_VALID: False,
+    KEY.RATIO: 0.0,
+    KEY.BATCH_SIZE: 6,
+    KEY.PREPROCESS_NUM_CORES: 1,
+    KEY.COMPUTE_STATISTICS: True,
+    KEY.DATASET_TYPE: 'graph',
+    # KEY.USE_SPECIES_WISE_SHIFT_SCALE: False,
+    KEY.USE_MODAL_WISE_SHIFT: False,
+    KEY.USE_MODAL_WISE_SCALE: False,
+    KEY.SHIFT: 'per_atom_energy_mean',
+    KEY.SCALE: 'force_rms',
+    # KEY.DATA_SHUFFLE: True,
+    # KEY.DATA_WEIGHT: False,
+    # KEY.DATA_MODALITY: False,
+}
+
+DATA_CONFIG_CONDITION = {
+    KEY.DTYPE: str,
+    KEY.DATA_FORMAT: str,
+    KEY.DATA_FORMAT_ARGS: dict,
+    KEY.SAVE_DATASET: str,
+    KEY.SAVE_BY_LABEL: bool,
+    KEY.SAVE_BY_TRAIN_VALID: bool,
+    KEY.RATIO: float,
+    KEY.BATCH_SIZE: int,
+    KEY.PREPROCESS_NUM_CORES: int,
+    KEY.DATASET_TYPE: lambda x: x in ['graph', 'atoms'],
+    # KEY.USE_SPECIES_WISE_SHIFT_SCALE: bool,
+    KEY.SHIFT: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SHIFT,
+    KEY.SCALE: lambda x: type(x) in [float, list] or x in IMPLEMENTED_SCALE,
+    KEY.USE_MODAL_WISE_SHIFT: bool,
+    KEY.USE_MODAL_WISE_SCALE: bool,
+    # KEY.DATA_SHUFFLE: bool,
+    KEY.COMPUTE_STATISTICS: bool,
+    # KEY.DATA_WEIGHT: bool,
+    # KEY.DATA_MODALITY: bool,
+}
+
+
+def data_defaults(config):
+    defaults = DEFAULT_DATA_CONFIG
+    if KEY.LOAD_VALIDSET in config:
+        defaults.pop(KEY.RATIO, None)
+    return defaults
+
+
+DEFAULT_TRAINING_CONFIG = {
+    KEY.RANDOM_SEED: 1,
+    KEY.EPOCH: 300,
+    KEY.LOSS: 'mse',
+    KEY.LOSS_PARAM: {},
+    KEY.OPTIMIZER: 'adam',
+    KEY.OPTIM_PARAM: {},
+    KEY.SCHEDULER: 'exponentiallr',
+    KEY.SCHEDULER_PARAM: {},
+    KEY.FORCE_WEIGHT: 0.1,
+    KEY.STRESS_WEIGHT: 1e-6,  # SIMPLE-NN default
+    KEY.PER_EPOCH: 5,
+    # KEY.USE_TESTSET: False,
+    KEY.CONTINUE: {
+        KEY.CHECKPOINT: False,
+        KEY.RESET_OPTIMIZER: False,
+        KEY.RESET_SCHEDULER: False,
+        KEY.RESET_EPOCH: False,
+        KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: True,
+        KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: True,
+    },
+    # KEY.DEFAULT_MODAL: 'common',
+    KEY.CSV_LOG: 'log.csv',
+    KEY.NUM_WORKERS: 0,
+    KEY.IS_TRAIN_STRESS: True,
+    KEY.TRAIN_SHUFFLE: True,
+    KEY.ERROR_RECORD: [
+        ['Energy', 'RMSE'],
+        ['Force', 'RMSE'],
+        ['Stress', 'RMSE'],
+        ['TotalLoss', 'None'],
+    ],
+    KEY.BEST_METRIC: 'TotalLoss',
+    KEY.USE_WEIGHT: False,
+    KEY.USE_MODALITY: False,
+}
+
+
+TRAINING_CONFIG_CONDITION = {
+    KEY.RANDOM_SEED: int,
+    KEY.EPOCH: int,
+    KEY.FORCE_WEIGHT: float,
+    KEY.STRESS_WEIGHT: float,
+    KEY.USE_TESTSET: None,  # Not used
+    KEY.NUM_WORKERS: int,
+    KEY.PER_EPOCH: int,
+    KEY.CONTINUE: {
+        KEY.CHECKPOINT: str,
+        KEY.RESET_OPTIMIZER: bool,
+        KEY.RESET_SCHEDULER: bool,
+        KEY.RESET_EPOCH: bool,
+        KEY.USE_STATISTIC_VALUES_OF_CHECKPOINT: bool,
+        KEY.USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY: bool,
+    },
+    KEY.DEFAULT_MODAL: str,
+    KEY.IS_TRAIN_STRESS: bool,
+    KEY.TRAIN_SHUFFLE: bool,
+    KEY.ERROR_RECORD: error_record_condition,
+    KEY.BEST_METRIC: str,
+    KEY.CSV_LOG: str,
+    KEY.USE_MODALITY: bool,
+    KEY.USE_WEIGHT: bool,
+}
+
+
+def train_defaults(config):
+    defaults = DEFAULT_TRAINING_CONFIG
+    if KEY.IS_TRAIN_STRESS not in config:
+        config[KEY.IS_TRAIN_STRESS] = defaults[KEY.IS_TRAIN_STRESS]
+    if not config[KEY.IS_TRAIN_STRESS]:
+        defaults.pop(KEY.STRESS_WEIGHT, None)
+    return defaults

mace-bench/3rdparty/SevenNet/sevenn/_keys.py

-"""
-How to add new feature?
-
-1. Add new key to this file.
-2. Add new key to _const.py
-2.1. if the type of input is consistent,
-    write adequate condition and default to _const.py.
-2.2. if the type of input is not consistent,
-    you must add your own input validation code to
-    parse_input.py
-"""
-
-from typing import Final
-
-# see
-# https://github.com/pytorch/pytorch/issues/52312
-# for FYI
-
-# ~~ keys ~~ #
-# PyG : primitive key of torch_geometric.data.Data type
-
-# ==================================================#
-# ~~~~~~~~~~~~~~~~~ KEY for data ~~~~~~~~~~~~~~~~~~ #
-# ==================================================#
-# some raw properties of graph
-ATOMIC_NUMBERS: Final[str] = 'atomic_numbers'  # (N)
-POS: Final[str] = 'pos'  # (N, 3) PyG
-CELL: Final[str] = 'cell_lattice_vectors'  # (3, 3)
-CELL_SHIFT: Final[str] = 'pbc_shift'  # (N, 3)
-CELL_VOLUME: Final[str] = 'cell_volume'
-
-EDGE_VEC: Final[str] = 'edge_vec'  # (N_edge, 3)
-EDGE_LENGTH: Final[str] = 'edge_length'  # (N_edge, 1)
-
-# some primary data of graph
-EDGE_IDX: Final[str] = 'edge_index'  # (2, N_edge) PyG
-ATOM_TYPE: Final[str] = 'atom_type'  # (N) one-hot index of nodes
-NODE_FEATURE: Final[str] = 'x'  # (N, ?) PyG
-NODE_FEATURE_GHOST: Final[str] = 'x_ghost'
-NODE_ATTR: Final[str] = 'node_attr'  # (N, N_species) from one_hot
-MODAL_ATTR: Final[str] = (
-    'modal_attr'  # (1, N_modalities) for handling multi-modal
-)
-MODAL_TYPE: Final[str] = 'modal_type'  # (1) one-hot index of modal
-EDGE_ATTR: Final[str] = 'edge_attr'  # (from spherical harmonics)
-EDGE_EMBEDDING: Final[str] = 'edge_embedding'  # (from edge embedding)
-
-# inputs of loss function
-ENERGY: Final[str] = 'total_energy'  # (1)
-FORCE: Final[str] = 'force_of_atoms'  # (N, 3)
-STRESS: Final[str] = 'stress'  # (6)
-
-# This is for training, per atom scale.
-SCALED_ENERGY: Final[str] = 'scaled_total_energy'
-
-# general outputs of models
-SCALED_ATOMIC_ENERGY: Final[str] = 'scaled_atomic_energy'
-ATOMIC_ENERGY: Final[str] = 'atomic_energy'
-PRED_TOTAL_ENERGY: Final[str] = 'inferred_total_energy'
-
-PRED_PER_ATOM_ENERGY: Final[str] = 'inferred_per_atom_energy'
-PER_ATOM_ENERGY: Final[str] = 'per_atom_energy'
-
-PRED_FORCE: Final[str] = 'inferred_force'
-SCALED_FORCE: Final[str] = 'scaled_force'
-
-PRED_STRESS: Final[str] = 'inferred_stress'
-SCALED_STRESS: Final[str] = 'scaled_stress'
-
-# very general data property for AtomGraphData
-NUM_ATOMS: Final[str] = 'num_atoms'  # int
-NUM_GHOSTS: Final[str] = 'num_ghosts'
-NLOCAL: Final[str] = 'nlocal'  # only for lammps parallel, must be on cpu
-USER_LABEL: Final[str] = 'user_label'
-DATA_WEIGHT: Final[str] = 'data_weight'  # weight for given data
-DATA_MODALITY: Final[str] = (
-    'data_modality'  # modality of given data. e.g. PBE and SCAN
-)
-BATCH: Final[str] = 'batch'
-
-TAG = 'tag'  # replace USER_LABEL
-
-# etc
-SELF_CONNECTION_TEMP: Final[str] = 'self_cont_tmp'
-BATCH_SIZE: Final[str] = 'batch_size'
-INFO: Final[str] = 'data_info'
-
-# something special
-LABEL_NONE: Final[str] = 'No_label'
-
-# ==================================================#
-# ~~~~~~ KEY for train/data configuration ~~~~~~~~ #
-# ==================================================#
-PREPROCESS_NUM_CORES = 'preprocess_num_cores'
-SAVE_DATASET = 'save_dataset_path'
-SAVE_BY_LABEL = 'save_by_label'
-SAVE_BY_TRAIN_VALID = 'save_by_train_valid'
-DATA_FORMAT = 'data_format'
-DATA_FORMAT_ARGS = 'data_format_args'
-STRUCTURE_LIST = 'structure_list'
-LOAD_DATASET = 'load_dataset_path'  # not used in v2
-LOAD_TRAINSET = 'load_trainset_path'
-LOAD_VALIDSET = 'load_validset_path'
-LOAD_TESTSET = 'load_testset_path'
-FORMAT_OUTPUTS = 'format_outputs_for_ase'
-COMPUTE_STATISTICS = 'compute_statistics'
-DATASET_TYPE = 'dataset_type'
-
-RANDOM_SEED = 'random_seed'
-RATIO = 'data_divide_ratio'
-USE_TESTSET = 'use_testset'
-EPOCH = 'epoch'
-LOSS = 'loss'
-LOSS_PARAM = 'loss_param'
-OPTIMIZER = 'optimizer'
-OPTIM_PARAM = 'optim_param'
-SCHEDULER = 'scheduler'
-SCHEDULER_PARAM = 'scheduler_param'
-FORCE_WEIGHT = 'force_loss_weight'
-STRESS_WEIGHT = 'stress_loss_weight'
-DEVICE = 'device'
-DTYPE = 'dtype'
-
-TRAIN_SHUFFLE = 'train_shuffle'
-
-IS_TRAIN_STRESS = 'is_train_stress'
-
-CONTINUE = 'continue'
-CHECKPOINT = 'checkpoint'
-RESET_OPTIMIZER = 'reset_optimizer'
-RESET_SCHEDULER = 'reset_scheduler'
-RESET_EPOCH = 'reset_epoch'
-USE_STATISTIC_VALUES_OF_CHECKPOINT = 'use_statistic_values_of_checkpoint'
-USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY = (
-    'use_statistic_values_for_cp_modal_only'
-)
-
-CSV_LOG = 'csv_log'
-
-ERROR_RECORD = 'error_record'
-BEST_METRIC = 'best_metric'
-
-NUM_WORKERS = 'num_workers'  # not work
-
-RANK = 'rank'
-LOCAL_RANK = 'local_rank'
-WORLD_SIZE = 'world_size'
-IS_DDP = 'is_ddp'
-DDP_BACKEND = 'ddp_backend'
-PER_EPOCH = 'per_epoch'
-
-USE_WEIGHT = 'use_weight'
-USE_MODALITY = 'use_modality'
-DEFAULT_MODAL = 'default_modal'
-
-
-# ==================================================#
-# ~~~~~~~~ KEY for model configuration ~~~~~~~~~~~ #
-# ==================================================#
-# ~~ global model configuration ~~ #
-# note that these names are directly used for input.yaml for user input
-MODEL_TYPE = '_model_type'
-CUTOFF = 'cutoff'
-CHEMICAL_SPECIES = 'chemical_species'
-MODAL_LIST = 'modal_list'
-CHEMICAL_SPECIES_BY_ATOMIC_NUMBER = '_chemical_species_by_atomic_number'
-NUM_SPECIES = '_number_of_species'
-NUM_MODALITIES = '_number_of_modalities'
-TYPE_MAP = '_type_map'
-MODAL_MAP = '_modal_map'
-
-# ~~ E3 equivariant model build configuration keys ~~ #
-# see model_build default_config for type
-IRREPS_MANUAL = 'irreps_manual'
-NODE_FEATURE_MULTIPLICITY = 'channel'
-
-RADIAL_BASIS = 'radial_basis'
-BESSEL_BASIS_NUM = 'bessel_basis_num'
-
-CUTOFF_FUNCTION = 'cutoff_function'
-POLY_CUT_P = 'poly_cut_p_value'
-
-LMAX = 'lmax'
-LMAX_EDGE = 'lmax_edge'
-LMAX_NODE = 'lmax_node'
-IS_PARITY = 'is_parity'
-CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS = 'weight_nn_hidden_neurons'
-NUM_CONVOLUTION = 'num_convolution_layer'
-ACTIVATION_SCARLAR = 'act_scalar'
-ACTIVATION_GATE = 'act_gate'
-ACTIVATION_RADIAL = 'act_radial'
-
-SELF_CONNECTION_TYPE = 'self_connection_type'
-
-RADIAL_BASIS_NAME = 'radial_basis_name'
-CUTOFF_FUNCTION_NAME = 'cutoff_function_name'
-
-USE_BIAS_IN_LINEAR = 'use_bias_in_linear'
-
-USE_MODAL_NODE_EMBEDDING = 'use_modal_node_embedding'
-USE_MODAL_SELF_INTER_INTRO = 'use_modal_self_inter_intro'
-USE_MODAL_SELF_INTER_OUTRO = 'use_modal_self_inter_outro'
-USE_MODAL_OUTPUT_BLOCK = 'use_modal_output_block'
-
-READOUT_AS_FCN = 'readout_as_fcn'
-READOUT_FCN_HIDDEN_NEURONS = 'readout_fcn_hidden_neurons'
-READOUT_FCN_ACTIVATION = 'readout_fcn_activation'
-
-AVG_NUM_NEIGH = 'avg_num_neigh'
-CONV_DENOMINATOR = 'conv_denominator'
-SHIFT = 'shift'
-SCALE = 'scale'
-
-USE_SPECIES_WISE_SHIFT_SCALE = 'use_species_wise_shift_scale'
-USE_MODAL_WISE_SHIFT = 'use_modal_wise_shift'
-USE_MODAL_WISE_SCALE = 'use_modal_wise_scale'
-
-TRAIN_SHIFT_SCALE = 'train_shift_scale'
-TRAIN_DENOMINTAOR = 'train_denominator'
-INTERACTION_TYPE = 'interaction_type'
-TRAIN_AVG_NUM_NEIGH = 'train_avg_num_neigh'  # deprecated
-
-CUEQUIVARIANCE_CONFIG = 'cuequivariance_config'
-
-_NORMALIZE_SPH = '_normalize_sph'
-OPTIMIZE_BY_REDUCE = 'optimize_by_reduce'
+"""
+How to add new feature?
+
+1. Add new key to this file.
+2. Add new key to _const.py
+2.1. if the type of input is consistent,
+    write adequate condition and default to _const.py.
+2.2. if the type of input is not consistent,
+    you must add your own input validation code to
+    parse_input.py
+"""
+
+from typing import Final
+
+# see
+# https://github.com/pytorch/pytorch/issues/52312
+# for FYI
+
+# ~~ keys ~~ #
+# PyG : primitive key of torch_geometric.data.Data type
+
+# ==================================================#
+# ~~~~~~~~~~~~~~~~~ KEY for data ~~~~~~~~~~~~~~~~~~ #
+# ==================================================#
+# some raw properties of graph
+ATOMIC_NUMBERS: Final[str] = 'atomic_numbers'  # (N)
+POS: Final[str] = 'pos'  # (N, 3) PyG
+CELL: Final[str] = 'cell_lattice_vectors'  # (3, 3)
+CELL_SHIFT: Final[str] = 'pbc_shift'  # (N, 3)
+CELL_VOLUME: Final[str] = 'cell_volume'
+
+EDGE_VEC: Final[str] = 'edge_vec'  # (N_edge, 3)
+EDGE_LENGTH: Final[str] = 'edge_length'  # (N_edge, 1)
+
+# some primary data of graph
+EDGE_IDX: Final[str] = 'edge_index'  # (2, N_edge) PyG
+ATOM_TYPE: Final[str] = 'atom_type'  # (N) one-hot index of nodes
+NODE_FEATURE: Final[str] = 'x'  # (N, ?) PyG
+NODE_FEATURE_GHOST: Final[str] = 'x_ghost'
+NODE_ATTR: Final[str] = 'node_attr'  # (N, N_species) from one_hot
+MODAL_ATTR: Final[str] = (
+    'modal_attr'  # (1, N_modalities) for handling multi-modal
+)
+MODAL_TYPE: Final[str] = 'modal_type'  # (1) one-hot index of modal
+EDGE_ATTR: Final[str] = 'edge_attr'  # (from spherical harmonics)
+EDGE_EMBEDDING: Final[str] = 'edge_embedding'  # (from edge embedding)
+
+# inputs of loss function
+ENERGY: Final[str] = 'total_energy'  # (1)
+FORCE: Final[str] = 'force_of_atoms'  # (N, 3)
+STRESS: Final[str] = 'stress'  # (6)
+
+# This is for training, per atom scale.
+SCALED_ENERGY: Final[str] = 'scaled_total_energy'
+
+# general outputs of models
+SCALED_ATOMIC_ENERGY: Final[str] = 'scaled_atomic_energy'
+ATOMIC_ENERGY: Final[str] = 'atomic_energy'
+PRED_TOTAL_ENERGY: Final[str] = 'inferred_total_energy'
+
+PRED_PER_ATOM_ENERGY: Final[str] = 'inferred_per_atom_energy'
+PER_ATOM_ENERGY: Final[str] = 'per_atom_energy'
+
+PRED_FORCE: Final[str] = 'inferred_force'
+SCALED_FORCE: Final[str] = 'scaled_force'
+
+PRED_STRESS: Final[str] = 'inferred_stress'
+SCALED_STRESS: Final[str] = 'scaled_stress'
+
+# very general data property for AtomGraphData
+NUM_ATOMS: Final[str] = 'num_atoms'  # int
+NUM_GHOSTS: Final[str] = 'num_ghosts'
+NLOCAL: Final[str] = 'nlocal'  # only for lammps parallel, must be on cpu
+USER_LABEL: Final[str] = 'user_label'
+DATA_WEIGHT: Final[str] = 'data_weight'  # weight for given data
+DATA_MODALITY: Final[str] = (
+    'data_modality'  # modality of given data. e.g. PBE and SCAN
+)
+BATCH: Final[str] = 'batch'
+
+TAG = 'tag'  # replace USER_LABEL
+
+# etc
+SELF_CONNECTION_TEMP: Final[str] = 'self_cont_tmp'
+BATCH_SIZE: Final[str] = 'batch_size'
+INFO: Final[str] = 'data_info'
+
+# something special
+LABEL_NONE: Final[str] = 'No_label'
+
+# ==================================================#
+# ~~~~~~ KEY for train/data configuration ~~~~~~~~ #
+# ==================================================#
+PREPROCESS_NUM_CORES = 'preprocess_num_cores'
+SAVE_DATASET = 'save_dataset_path'
+SAVE_BY_LABEL = 'save_by_label'
+SAVE_BY_TRAIN_VALID = 'save_by_train_valid'
+DATA_FORMAT = 'data_format'
+DATA_FORMAT_ARGS = 'data_format_args'
+STRUCTURE_LIST = 'structure_list'
+LOAD_DATASET = 'load_dataset_path'  # not used in v2
+LOAD_TRAINSET = 'load_trainset_path'
+LOAD_VALIDSET = 'load_validset_path'
+LOAD_TESTSET = 'load_testset_path'
+FORMAT_OUTPUTS = 'format_outputs_for_ase'
+COMPUTE_STATISTICS = 'compute_statistics'
+DATASET_TYPE = 'dataset_type'
+
+RANDOM_SEED = 'random_seed'
+RATIO = 'data_divide_ratio'
+USE_TESTSET = 'use_testset'
+EPOCH = 'epoch'
+LOSS = 'loss'
+LOSS_PARAM = 'loss_param'
+OPTIMIZER = 'optimizer'
+OPTIM_PARAM = 'optim_param'
+SCHEDULER = 'scheduler'
+SCHEDULER_PARAM = 'scheduler_param'
+FORCE_WEIGHT = 'force_loss_weight'
+STRESS_WEIGHT = 'stress_loss_weight'
+DEVICE = 'device'
+DTYPE = 'dtype'
+
+TRAIN_SHUFFLE = 'train_shuffle'
+
+IS_TRAIN_STRESS = 'is_train_stress'
+
+CONTINUE = 'continue'
+CHECKPOINT = 'checkpoint'
+RESET_OPTIMIZER = 'reset_optimizer'
+RESET_SCHEDULER = 'reset_scheduler'
+RESET_EPOCH = 'reset_epoch'
+USE_STATISTIC_VALUES_OF_CHECKPOINT = 'use_statistic_values_of_checkpoint'
+USE_STATISTIC_VALUES_FOR_CP_MODAL_ONLY = (
+    'use_statistic_values_for_cp_modal_only'
+)
+
+CSV_LOG = 'csv_log'
+
+ERROR_RECORD = 'error_record'
+BEST_METRIC = 'best_metric'
+
+NUM_WORKERS = 'num_workers'  # not work
+
+RANK = 'rank'
+LOCAL_RANK = 'local_rank'
+WORLD_SIZE = 'world_size'
+IS_DDP = 'is_ddp'
+DDP_BACKEND = 'ddp_backend'
+PER_EPOCH = 'per_epoch'
+
+USE_WEIGHT = 'use_weight'
+USE_MODALITY = 'use_modality'
+DEFAULT_MODAL = 'default_modal'
+
+
+# ==================================================#
+# ~~~~~~~~ KEY for model configuration ~~~~~~~~~~~ #
+# ==================================================#
+# ~~ global model configuration ~~ #
+# note that these names are directly used for input.yaml for user input
+MODEL_TYPE = '_model_type'
+CUTOFF = 'cutoff'
+CHEMICAL_SPECIES = 'chemical_species'
+MODAL_LIST = 'modal_list'
+CHEMICAL_SPECIES_BY_ATOMIC_NUMBER = '_chemical_species_by_atomic_number'
+NUM_SPECIES = '_number_of_species'
+NUM_MODALITIES = '_number_of_modalities'
+TYPE_MAP = '_type_map'
+MODAL_MAP = '_modal_map'
+
+# ~~ E3 equivariant model build configuration keys ~~ #
+# see model_build default_config for type
+IRREPS_MANUAL = 'irreps_manual'
+NODE_FEATURE_MULTIPLICITY = 'channel'
+
+RADIAL_BASIS = 'radial_basis'
+BESSEL_BASIS_NUM = 'bessel_basis_num'
+
+CUTOFF_FUNCTION = 'cutoff_function'
+POLY_CUT_P = 'poly_cut_p_value'
+
+LMAX = 'lmax'
+LMAX_EDGE = 'lmax_edge'
+LMAX_NODE = 'lmax_node'
+IS_PARITY = 'is_parity'
+CONVOLUTION_WEIGHT_NN_HIDDEN_NEURONS = 'weight_nn_hidden_neurons'
+NUM_CONVOLUTION = 'num_convolution_layer'
+ACTIVATION_SCARLAR = 'act_scalar'
+ACTIVATION_GATE = 'act_gate'
+ACTIVATION_RADIAL = 'act_radial'
+
+SELF_CONNECTION_TYPE = 'self_connection_type'
+
+RADIAL_BASIS_NAME = 'radial_basis_name'
+CUTOFF_FUNCTION_NAME = 'cutoff_function_name'
+
+USE_BIAS_IN_LINEAR = 'use_bias_in_linear'
+
+USE_MODAL_NODE_EMBEDDING = 'use_modal_node_embedding'
+USE_MODAL_SELF_INTER_INTRO = 'use_modal_self_inter_intro'
+USE_MODAL_SELF_INTER_OUTRO = 'use_modal_self_inter_outro'
+USE_MODAL_OUTPUT_BLOCK = 'use_modal_output_block'
+
+READOUT_AS_FCN = 'readout_as_fcn'
+READOUT_FCN_HIDDEN_NEURONS = 'readout_fcn_hidden_neurons'
+READOUT_FCN_ACTIVATION = 'readout_fcn_activation'
+
+AVG_NUM_NEIGH = 'avg_num_neigh'
+CONV_DENOMINATOR = 'conv_denominator'
+SHIFT = 'shift'
+SCALE = 'scale'
+
+USE_SPECIES_WISE_SHIFT_SCALE = 'use_species_wise_shift_scale'
+USE_MODAL_WISE_SHIFT = 'use_modal_wise_shift'
+USE_MODAL_WISE_SCALE = 'use_modal_wise_scale'
+
+TRAIN_SHIFT_SCALE = 'train_shift_scale'
+TRAIN_DENOMINTAOR = 'train_denominator'
+INTERACTION_TYPE = 'interaction_type'
+TRAIN_AVG_NUM_NEIGH = 'train_avg_num_neigh'  # deprecated
+
+CUEQUIVARIANCE_CONFIG = 'cuequivariance_config'
+
+_NORMALIZE_SPH = '_normalize_sph'
+OPTIMIZE_BY_REDUCE = 'optimize_by_reduce'

mace-bench/3rdparty/SevenNet/sevenn/atom_graph_data.py

-from typing import Optional
-
-import torch
-import torch_geometric.data
-
-import sevenn._keys as KEY
-import sevenn.util
-
-
-class AtomGraphData(torch_geometric.data.Data):
-    """
-    Args:
-        x (Tensor, optional): atomic numbers with shape :obj:`[num_nodes,
-            atomic_numbers]`. (default: :obj:`None`)
-        edge_index (LongTensor, optional): Graph connectivity in coordinate
-            format with shape :obj:`[2, num_edges]`. (default: :obj:`None`)
-        edge_attr (Tensor, optional): Edge feature matrix with shape
-            :obj:`[num_edges, num_edge_features]`. (default: :obj:`None`)
-        y_energy: scalar # unit of eV (VASP raw)
-        y_force: [num_nodes, 3] # unit of eV/A (VASP raw)
-        y_stress: [6]  # [xx, yy, zz, xy, yz, zx] # unit of eV/A^3 (VASP raw)
-        pos (Tensor, optional): Node position matrix with shape
-            :obj:`[num_nodes, num_dimensions]`. (default: :obj:`None`)
-        **kwargs (optional): Additional attributes.
-
-    x, y_force, pos should be aligned with each other.
-    """
-
-    def __init__(
-        self,
-        x: Optional[torch.Tensor] = None,
-        edge_index: Optional[torch.Tensor] = None,
-        pos: Optional[torch.Tensor] = None,
-        edge_attr: Optional[torch.Tensor] = None,
-        **kwargs
-    ):
-        super(AtomGraphData, self).__init__(x, edge_index, edge_attr, pos=pos)
-        self[KEY.NODE_ATTR] = x  # ?
-        for k, v in kwargs.items():
-            self[k] = v
-
-    def to_numpy_dict(self):
-        # This is not debugged yet!
-        dct = {
-            k: v.detach().cpu().numpy() if type(v) is torch.Tensor else v
-            for k, v in self.items()
-        }
-        return dct
-
-    def fit_dimension(self):
-        per_atom_keys = [
-            KEY.ATOMIC_NUMBERS,
-            KEY.ATOMIC_ENERGY,
-            KEY.POS,
-            KEY.FORCE,
-            KEY.PRED_FORCE,
-        ]
-        natoms = self.num_atoms.item()
-        for k, v in self.items():
-            if not isinstance(v, torch.Tensor):
-                continue
-            if natoms == 1 and k in per_atom_keys:
-                self[k] = v.squeeze().unsqueeze(0)
-            else:
-                self[k] = v.squeeze()
-        return self
-
-    @staticmethod
-    def from_numpy_dict(dct):
-        for k, v in dct.items():
-            if k == KEY.CELL_SHIFT:
-                dct[k] = torch.Tensor(v)  # this is special
-            else:
-                dct[k] = sevenn.util.dtype_correct(v)
-        return AtomGraphData(**dct)
+from typing import Optional
+
+import torch
+import torch_geometric.data
+
+import sevenn._keys as KEY
+import sevenn.util
+
+
+class AtomGraphData(torch_geometric.data.Data):
+    """
+    Args:
+        x (Tensor, optional): atomic numbers with shape :obj:`[num_nodes,
+            atomic_numbers]`. (default: :obj:`None`)
+        edge_index (LongTensor, optional): Graph connectivity in coordinate
+            format with shape :obj:`[2, num_edges]`. (default: :obj:`None`)
+        edge_attr (Tensor, optional): Edge feature matrix with shape
+            :obj:`[num_edges, num_edge_features]`. (default: :obj:`None`)
+        y_energy: scalar # unit of eV (VASP raw)
+        y_force: [num_nodes, 3] # unit of eV/A (VASP raw)
+        y_stress: [6]  # [xx, yy, zz, xy, yz, zx] # unit of eV/A^3 (VASP raw)
+        pos (Tensor, optional): Node position matrix with shape
+            :obj:`[num_nodes, num_dimensions]`. (default: :obj:`None`)
+        **kwargs (optional): Additional attributes.
+
+    x, y_force, pos should be aligned with each other.
+    """
+
+    def __init__(
+        self,
+        x: Optional[torch.Tensor] = None,
+        edge_index: Optional[torch.Tensor] = None,
+        pos: Optional[torch.Tensor] = None,
+        edge_attr: Optional[torch.Tensor] = None,
+        **kwargs
+    ):
+        super(AtomGraphData, self).__init__(x, edge_index, edge_attr, pos=pos)
+        self[KEY.NODE_ATTR] = x  # ?
+        for k, v in kwargs.items():
+            self[k] = v
+
+    def to_numpy_dict(self):
+        # This is not debugged yet!
+        dct = {
+            k: v.detach().cpu().numpy() if type(v) is torch.Tensor else v
+            for k, v in self.items()
+        }
+        return dct
+
+    def fit_dimension(self):
+        per_atom_keys = [
+            KEY.ATOMIC_NUMBERS,
+            KEY.ATOMIC_ENERGY,
+            KEY.POS,
+            KEY.FORCE,
+            KEY.PRED_FORCE,
+        ]
+        natoms = self.num_atoms.item()
+        for k, v in self.items():
+            if not isinstance(v, torch.Tensor):
+                continue
+            if natoms == 1 and k in per_atom_keys:
+                self[k] = v.squeeze().unsqueeze(0)
+            else:
+                self[k] = v.squeeze()
+        return self
+
+    @staticmethod
+    def from_numpy_dict(dct):
+        for k, v in dct.items():
+            if k == KEY.CELL_SHIFT:
+                dct[k] = torch.Tensor(v)  # this is special
+            else:
+                dct[k] = sevenn.util.dtype_correct(v)
+        return AtomGraphData(**dct)

mace-bench/3rdparty/SevenNet/sevenn/calculator.py

-import ctypes
-import os
-import pathlib
-import warnings
-from typing import Any, Dict, Optional, Union
-
-import numpy as np
-import torch
-import torch.jit
-import torch.jit._script
-from ase.calculators.calculator import Calculator, all_changes
-from ase.calculators.mixing import SumCalculator
-from ase.data import chemical_symbols
-
-import sevenn._keys as KEY
-import sevenn.util as util
-from sevenn.atom_graph_data import AtomGraphData
-from sevenn.nn.sequential import AtomGraphSequential
-from sevenn.train.dataload import unlabeled_atoms_to_graph
-import logging
-
-torch_script_type = torch.jit._script.RecursiveScriptModule
-
-
-class SevenNetCalculator(Calculator):
-    """Supporting properties:
-    'free_energy', 'energy', 'forces', 'stress', 'energies'
-    free_energy equals energy. 'energies' stores atomic energy.
-
-    Multi-GPU acceleration is not supported with ASE calculator.
-    You should use LAMMPS for the acceleration.
-    """
-
-    def __init__(
-        self,
-        model: Union[str, pathlib.PurePath, AtomGraphSequential] = '7net-0',
-        file_type: str = 'checkpoint',
-        device: Union[torch.device, str] = 'auto',
-        modal: Optional[str] = None,
-        enable_cueq: bool = False,
-        sevennet_config: Optional[Dict] = None,  # Not used in logic, just meta info
-        **kwargs,
-    ):
-        """Initialize SevenNetCalculator.
-
-        Parameters
-        ----------
-        model: str | Path | AtomGraphSequential, default='7net-0'
-            Name of pretrained models (7net-mf-ompa, 7net-omat, 7net-l3i5, 7net-0) or
-            path to the checkpoint, deployed model or the model itself
-        file_type: str, default='checkpoint'
-            one of 'checkpoint' | 'torchscript' | 'model_instance'
-        device: str | torch.device, default='auto'
-            if not given, use CUDA if available
-        modal: str | None, default=None
-            modal (fidelity) if given model is multi-modal model. for 7net-mf-ompa,
-            it should be one of 'mpa' (MPtrj + sAlex) or 'omat24' (OMat24)
-            case insensitive
-        enable_cueq: bool, default=False
-            if True, use cuEquivariant to accelerate inference.
-        sevennet_config: dict | None, default=None
-            Not used, but can be used to carry meta information of this calculator
-        """
-        print("&&& Initializing SevenNetCalculator")
-        super().__init__(**kwargs)
-        self.sevennet_config = None
-
-        if isinstance(model, pathlib.PurePath):
-            model = str(model)
-
-        allowed_file_types = ['checkpoint', 'torchscript', 'model_instance']
-        file_type = file_type.lower()
-        if file_type not in allowed_file_types:
-            raise ValueError(f'file_type not in {allowed_file_types}')
-
-        if enable_cueq and file_type in ['model_instance', 'torchscript']:
-            warnings.warn(
-                'file_type should be checkpoint to enable cueq. cueq set to False'
-            )
-            enable_cueq = False
-
-        if isinstance(device, str):  # TODO: do we really need this?
-            if device == 'auto':
-                self.device = torch.device(
-                    'cuda' if torch.cuda.is_available() else 'cpu'
-                )
-            else:
-                self.device = torch.device(device)
-        else:
-            self.device = device
-
-        if file_type == 'checkpoint' and isinstance(model, str):
-            cp = util.load_checkpoint(model)
-
-            backend = 'e3nn' if not enable_cueq else 'cueq'
-            model_loaded = cp.build_model(backend)
-            model_loaded.set_is_batch_data(False)
-
-            self.type_map = cp.config[KEY.TYPE_MAP]
-            self.cutoff = cp.config[KEY.CUTOFF]
-            self.sevennet_config = cp.config
-
-        elif file_type == 'torchscript' and isinstance(model, str):
-            if modal:
-                raise NotImplementedError()
-            extra_dict = {
-                'chemical_symbols_to_index': b'',
-                'cutoff': b'',
-                'num_species': b'',
-                'model_type': b'',
-                'version': b'',
-                'dtype': b'',
-                'time': b'',
-            }
-            model_loaded = torch.jit.load(
-                model, _extra_files=extra_dict, map_location=self.device
-            )
-            chem_symbols = extra_dict['chemical_symbols_to_index'].decode('utf-8')
-            sym_to_num = {sym: n for n, sym in enumerate(chemical_symbols)}
-            self.type_map = {
-                sym_to_num[sym]: i for i, sym in enumerate(chem_symbols.split())
-            }
-            self.cutoff = float(extra_dict['cutoff'].decode('utf-8'))
-
-        elif isinstance(model, AtomGraphSequential):
-            if model.type_map is None:
-                raise ValueError(
-                    'Model must have the type_map to be used with calculator'
-                )
-            if model.cutoff == 0.0:
-                raise ValueError('Model cutoff seems not initialized')
-            model.eval_type_map = torch.tensor(True)  # ?
-            model.set_is_batch_data(False)
-            model_loaded = model
-            self.type_map = model.type_map
-            self.cutoff = model.cutoff
-        else:
-            raise ValueError('Unexpected input combinations')
-
-        if self.sevennet_config is None and sevennet_config is not None:
-            self.sevennet_config = sevennet_config
-
-        self.model = model_loaded
-
-        self.modal = None
-        if isinstance(self.model, AtomGraphSequential):
-            modal_map = self.model.modal_map
-            if modal_map:
-                modal_ava = list(modal_map.keys())
-                if not modal:
-                    raise ValueError(f'modal argument missing (avail: {modal_ava})')
-                elif modal not in modal_ava:
-                    raise ValueError(f'unknown modal {modal} (not in {modal_ava})')
-                self.modal = modal
-            elif not self.model.modal_map and modal:
-                warnings.warn(f'modal={modal} is ignored as model has no modal_map')
-
-        self.model.to(self.device)
-        self.model.eval()
-        self.implemented_properties = [
-            'free_energy',
-            'energy',
-            'forces',
-            'stress',
-            'energies',
-        ]
-
-    def set_atoms(self, atoms):
-        # called by ase, when atoms.calc = calc
-        zs = tuple(set(atoms.get_atomic_numbers()))
-        for z in zs:
-            if z not in self.type_map:
-                sp = list(self.type_map.keys())
-                raise ValueError(
-                    f'Model do not know atomic number: {z}, (knows: {sp})'
-                )
-
-    def output_to_results(self, output):
-        energy = output[KEY.PRED_TOTAL_ENERGY].detach().cpu().item()
-        num_atoms = output['num_atoms'].item()
-        atomic_energies = output[KEY.ATOMIC_ENERGY].detach().cpu().numpy().flatten()
-        forces = output[KEY.PRED_FORCE].detach().cpu().numpy()[:num_atoms, :]
-        stress = np.array(
-            (-output[KEY.PRED_STRESS])
-            .detach()
-            .cpu()
-            .numpy()[[0, 1, 2, 4, 5, 3]]  # as voigt notation
-        )
-        # Store results
-        return {
-            'free_energy': energy,
-            'energy': energy,
-            'energies': atomic_energies,
-            'forces': forces,
-            'stress': stress,
-            'num_edges': output[KEY.EDGE_IDX].shape[1],
-        }
-
-    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
-        # call parent class to set necessary atom attributes
-        Calculator.calculate(self, atoms, properties, system_changes)
-        if atoms is None:
-            raise ValueError('No atoms to evaluate')
-        data = AtomGraphData.from_numpy_dict(
-            unlabeled_atoms_to_graph(atoms, self.cutoff)
-        )
-        if self.modal:
-            data[KEY.DATA_MODALITY] = self.modal
-
-        data.to(self.device)  # type: ignore
-
-        if isinstance(self.model, torch_script_type):
-            data[KEY.NODE_FEATURE] = torch.tensor(
-                [self.type_map[z.item()] for z in data[KEY.NODE_FEATURE]],
-                dtype=torch.int64,
-                device=self.device,
-            )
-            data[KEY.POS].requires_grad_(True)  # backward compatibility
-            data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
-            data = data.to_dict()
-            del data['data_info']
-
-        import logging
-        logging.debug(f"data: {data}")
-        # logging.debug(f"data[pos]: {data['pos']}")
-        # logging.debug(f"data[x]: {data['x']}")
-        logging.debug(f"data[cell_lattice_vectors]: {data['cell_lattice_vectors']}")
-        logging.debug(f"data[cell_volume]: {data['cell_volume']}")
-        output = self.model(data)
-        # logging.info(f"input: {data}")
-        # logging.info(f"output[{KEY.PRED_TOTAL_ENERGY}] = {output[KEY.PRED_TOTAL_ENERGY]}")
-        # logging.info(f"output[{KEY.PRED_FORCE}] = {output[KEY.PRED_FORCE]}")
-        # logging.info(f"output[{KEY.PRED_STRESS}] = {output[KEY.PRED_STRESS]}")
-        self.results = self.output_to_results(output)
-        # logging.debug(f"results['energy'] = {self.results['energy']}")
-        # logging.debug(f"results['forces'] = {self.results['forces']}")
-        # logging.debug(f"results['stress'] = {self.results['stress']}")
-
-    def predict_one(self, atoms):
-        if atoms is None:
-            raise ValueError('No atoms to evaluate')
-        data = AtomGraphData.from_numpy_dict(
-            unlabeled_atoms_to_graph(atoms, self.cutoff)
-        )
-        if self.modal:
-            data[KEY.DATA_MODALITY] = self.modal
-
-        data.to(self.device)  # type: ignore
-
-        if isinstance(self.model, torch_script_type):
-            data[KEY.NODE_FEATURE] = torch.tensor(
-                [self.type_map[z.item()] for z in data[KEY.NODE_FEATURE]],
-                dtype=torch.int64,
-                device=self.device,
-            )
-            data[KEY.POS].requires_grad_(True)  # backward compatibility
-            data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
-            data = data.to_dict()
-            del data['data_info']
-
-        return self.model(data)
-
-
-
-    def predict(self, atoms_list, properties=None):
-        
-        # if len(atoms_list) == 1:
-        #     output = self.predict_one(atoms_list[0])
-        #     predictions = {}
-        #     predictions['energy'] = output[KEY.PRED_TOTAL_ENERGY].to(torch.float64).unsqueeze(0)
-        #     predictions['forces'] = output[KEY.PRED_FORCE].to(torch.float64).unsqueeze(0)
-        #     voigt = (-output[KEY.PRED_STRESS])[[0, 1, 2, 4, 5, 3]].to(torch.float64).unsqueeze(0)
-        #     stress_list = []
-        #     for i in range(voigt.shape[0]):
-        #         stress_list.append(self._stress2tensor(voigt[i,:]))
-        #     predictions['stress'] = torch.stack(stress_list, dim=0).view(-1,3,3)
-        #     return predictions
-
-
-        if not atoms_list:
-            raise ValueError("Empty atoms_list provided")
-            
-        if not isinstance(atoms_list, list):
-            atoms_list = [atoms_list]
-            
-        # Convert atoms to graph data
-        graph_list = []
-        for atoms in atoms_list:
-            data = AtomGraphData.from_numpy_dict(
-                unlabeled_atoms_to_graph(atoms, self.cutoff)
-            )
-            if self.modal:
-                data[KEY.DATA_MODALITY] = self.modal
-                
-            if isinstance(self.model, torch_script_type):
-                data[KEY.NODE_FEATURE] = torch.tensor(
-                    [self.type_map[z.item()] for z in data[KEY.NODE_FEATURE]],
-                    dtype=torch.int64,
-                    device=self.device,
-                )
-                data[KEY.POS].requires_grad_(True)  # backward compatibility
-                data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
-                
-            graph_list.append(data)
-            
-        # Process graphs based on model type
-        # was_batch_mode = True
-        if isinstance(self.model, AtomGraphSequential):
-            # was_batch_mode = self.model.is_batch_data
-            self.model.set_is_batch_data(True)
-            self.model.eval()
-            
-        # Batch the data if there are multiple atoms
-        from torch_geometric.loader.dataloader import Collater
-        batched_data = Collater(graph_list)(graph_list)
-        batched_data = batched_data.to(self.device)
-        
-        import logging
-        logging.debug(f"batched_data: {batched_data}")
-        # logging.debug(f"batched_data[pos]: {batched_data['pos']}")
-        # logging.debug(f"batched_data[x]: {batched_data['x']}")
-        logging.debug(f"batched_data[cell_lattice_vectors]: {batched_data['cell_lattice_vectors']}")
-        logging.debug(f"batched_data[cell_volume]: {batched_data['cell_volume']}")
-        # Run model on batched data
-        if isinstance(self.model, torch_script_type):
-            batched_dict = batched_data.to_dict()
-            if 'data_info' in batched_dict:
-                del batched_dict['data_info']
-            output = self.model(batched_dict)
-        else:
-            output = self.model(batched_data)
-            
-        # Convert to list of individual outputs using util.to_atom_graph_list
-        # logging.info(f"input: {batched_data}")
-        # logging.info(f"output[{KEY.PRED_TOTAL_ENERGY}] = {output[KEY.PRED_TOTAL_ENERGY]}")
-        # logging.info(f"output[{KEY.PRED_FORCE}] = {output[KEY.PRED_FORCE]}")
-        # logging.info(f"output[{KEY.PRED_STRESS}] = {output[KEY.PRED_STRESS]}")
-
-        predictions = {}
-        predictions['energy'] = output[KEY.PRED_TOTAL_ENERGY].to(torch.float64).detach()
-        predictions['forces'] = output[KEY.PRED_FORCE].to(torch.float64).detach()
-        voigt = (-output[KEY.PRED_STRESS])[:, [0, 1, 2, 4, 5, 3]].to(torch.float64).detach()
-        stress_list = []
-        for i in range(voigt.shape[0]):
-            stress_list.append(self._stress2tensor(voigt[i,:]))
-        predictions['stress'] = torch.stack(stress_list, dim=0).view(-1,3,3).detach()
-
-        # logging.debug(f"predictions['energy'] = {predictions['energy']}")
-        # logging.debug(f"predictions['forces'] = {predictions['forces']}")
-        # logging.debug(f"predictions['stress'] = {predictions['stress']}")
-        return predictions
-
-    def _stress2tensor(self, stress):
-        tensor = torch.tensor(
-            [
-                [stress[0], stress[5], stress[4]],
-                [stress[5], stress[1], stress[3]],
-                [stress[4], stress[3], stress[2]],
-            ], 
-            device=self.device
-        )
-        return tensor
-
-
-class SevenNetD3Calculator(SumCalculator):
-    def __init__(
-        self,
-        model: Union[str, pathlib.PurePath, AtomGraphSequential] = '7net-0',
-        file_type: str = 'checkpoint',
-        device: Union[torch.device, str] = 'auto',
-        sevennet_config: Optional[Any] = None,  # hold meta information
-        damping_type: str = 'damp_bj',
-        functional_name: str = 'pbe',
-        vdw_cutoff: float = 9000,  # au^2, 0.52917726 angstrom = 1 au
-        cn_cutoff: float = 1600,  # au^2, 0.52917726 angstrom = 1 au
-        batch_size=10,
-        **kwargs,
-    ):
-        """Initialize SevenNetD3Calculator. CUDA required.
-
-        Parameters
-        ----------
-        model: str | Path | AtomGraphSequential
-            Name of pretrained models (7net-mf-ompa, 7net-omat, 7net-l3i5, 7net-0) or
-            path to the checkpoint, deployed model or the model itself
-        file_type: str, default='checkpoint'
-            one of 'checkpoint' | 'torchscript' | 'model_instance'
-        device: str | torch.device, default='auto'
-            if not given, use CUDA if available
-        modal: str | None, default=None
-            modal (fidelity) if given model is multi-modal model. for 7net-mf-ompa,
-            it should be one of 'mpa' (MPtrj + sAlex) or 'omat24' (OMat24)
-        enable_cueq: bool, default=False
-            if True, use cuEquivariant to accelerate inference.
-        damping_type: str, default='damp_bj'
-            Damping type of D3, one of 'damp_bj' | 'damp_zero'
-        functional_name: str, default='pbe'
-            Target functional name of D3 parameters.
-        vdw_cutoff: float, default=9000
-            vdw cutoff of D3 calculator in au
-        cn_cutoff: float, default=1600
-            cn cutoff of D3 calculator in au
-        """
-        self.d3_calc = D3Calculator(
-            damping_type=damping_type,
-            functional_name=functional_name,
-            vdw_cutoff=vdw_cutoff,
-            cn_cutoff=cn_cutoff,
-            **kwargs,
-        )
-
-        self.sevennet_calc = SevenNetCalculator(
-            model=model,
-            file_type=file_type,
-            device=device,
-            sevennet_config=sevennet_config,
-            **kwargs,
-        )
-
-        super().__init__([self.sevennet_calc, self.d3_calc])
-
-        self.device = device
-        self.d3_calcs = []
-        for _ in range(batch_size):
-            self.d3_calcs.append(
-                D3Calculator(
-                    damping_type=damping_type,
-                    functional_name=functional_name,
-                    vdw_cutoff=vdw_cutoff,
-                    cn_cutoff=cn_cutoff,
-                    **kwargs,
-                )
-            )
-
-
-    def predict(self, atoms_list):
-        """Predict the energy and forces for a list of atoms.
-        """
-        # Call the predict method of the first calculator (SevenNetCalculator)
-        predictions = self.sevennet_calc.predict(atoms_list)
-        
-        energy_list = []
-        forces_list = []
-        stress_list = []
-        predictions3d = {}
-        for i, atoms in enumerate(atoms_list):
-            prediction = self.d3_calcs[i].predict_one(atoms)
-            energy_list.append(torch.tensor(prediction['energy']))
-            forces_list.append(torch.from_numpy(prediction['forces']).to(self.device))
-            stress_list.append(self._stress2tensor(torch.from_numpy(prediction['stress'])))
-
-        # Convert lists to tensors
-        predictions3d['energy'] = torch.stack(energy_list, dim=0).to(self.device)
-        predictions3d['forces'] = torch.cat(forces_list, dim=0).view(-1, 3)
-        predictions3d['stress'] = torch.stack(stress_list, dim=0).view(-1, 3, 3)
-
-        predictions['energy'] += predictions3d['energy'].detach()
-        predictions['forces'] += predictions3d['forces'].detach()
-        predictions['stress'] += predictions3d['stress'].detach()
-        
-        return predictions
-
-    def _stress2tensor(self, stress):
-        tensor = torch.tensor(
-            [
-                # [stress[0], stress[3], stress[4]],
-                # [stress[3], stress[1], stress[5]],
-                # [stress[4], stress[5], stress[2]],
-                [stress[0], stress[5], stress[4]],
-                [stress[5], stress[1], stress[3]],
-                [stress[4], stress[3], stress[2]],
-            ], 
-            device=self.device
-        )
-        return tensor
-
-
-
-def _load(name: str) -> ctypes.CDLL:
-    from torch.utils.cpp_extension import LIB_EXT, _get_build_directory, load
-
-    # Load the library from the candidate locations
-
-    package_dir = os.path.dirname(os.path.abspath(__file__))
-    try:
-        return ctypes.CDLL(os.path.join(package_dir, f'{name}{LIB_EXT}'))
-    except OSError:
-        pass
-
-    cache_dir = _get_build_directory(name, verbose=False)
-    try:
-        return ctypes.CDLL(os.path.join(cache_dir, f'{name}{LIB_EXT}'))
-    except OSError:
-        pass
-
-    # Compile the library if it is not found
-
-    if os.access(package_dir, os.W_OK):
-        compile_dir = package_dir
-    else:
-        print('Warning: package directory is not writable. Using cache directory.')
-        compile_dir = cache_dir
-
-    if 'TORCH_CUDA_ARCH_LIST' not in os.environ:
-        print('Warning: TORCH_CUDA_ARCH_LIST is not set.')
-        print('Warning: Use default CUDA architectures: 61, 70, 75, 80, 86, 89, 90')
-        os.environ['TORCH_CUDA_ARCH_LIST'] = '6.1;7.0;7.5;8.0;8.6;8.9;9.0'
-
-    load(
-        name=name,
-        sources=[os.path.join(package_dir, 'pair_e3gnn', 'pair_d3_for_ase.cu')],
-        extra_cuda_cflags=['-O3', '--expt-relaxed-constexpr', '-fmad=false'],
-        build_directory=compile_dir,
-        verbose=True,
-        is_python_module=False,
-    )
-
-    return ctypes.CDLL(os.path.join(compile_dir, f'{name}{LIB_EXT}'))
-
-
-class PairD3(ctypes.Structure):
-    pass  # Opaque structure; only used as a pointer
-
-
-class D3Calculator(Calculator):
-    """ASE calculator for accelerated D3 van der Waals (vdW) correction.
-
-    Example:
-        from ase.calculators.mixing import SumCalculator
-        calc_1 = SevenNetCalculator()
-        calc_2 = D3Calculator()
-        return SumCalculator([calc_1, calc_2])
-
-    This calculator interfaces with the `libpaird3.so` library,
-    which is compiled by nvcc during the package installation.
-    If you encounter any errors, please verify
-    the installation process and the compilation options in `setup.py`.
-    Note: Multi-GPU parallel MD is not supported in this mode.
-    Note: Cffi could be used, but it was avoided to reduce dependencies.
-    """
-
-    # Here, free_energy = energy
-    implemented_properties = ['free_energy', 'energy', 'forces', 'stress']
-
-    def __init__(
-        self,
-        damping_type: str = 'damp_bj',  # damp_bj, damp_zero
-        functional_name: str = 'pbe',  # check the source code
-        vdw_cutoff: float = 9000,  # au^2, 0.52917726 angstrom = 1 au
-        cn_cutoff: float = 1600,  # au^2, 0.52917726 angstrom = 1 au
-        **kwargs,
-    ):
-        super().__init__(**kwargs)
-
-        if not torch.cuda.is_available():
-            raise NotImplementedError('CPU + D3 is not implemented yet')
-
-        self.rthr = vdw_cutoff
-        self.cnthr = cn_cutoff
-        self.damp_name = damping_type.lower()
-        self.func_name = functional_name.lower()
-
-        if self.damp_name not in ['damp_bj', 'damp_zero']:
-            raise ValueError('Error: Invalid damping type.')
-
-        self._lib = _load('pair_d3')
-
-        self._lib.pair_init.restype = ctypes.POINTER(PairD3)
-        self.pair = self._lib.pair_init()
-
-        self._lib.pair_set_atom.argtypes = [
-            ctypes.POINTER(PairD3),  # PairD3* pair
-            ctypes.c_int,  # int natoms
-            ctypes.c_int,  # int ntypes
-            ctypes.POINTER(ctypes.c_int),  # int* types
-            ctypes.POINTER(ctypes.c_double),  # double* x
-        ]
-        self._lib.pair_set_atom.restype = None
-
-        self._lib.pair_set_domain.argtypes = [
-            ctypes.POINTER(PairD3),  # PairD3* pair
-            ctypes.c_int,  # int xperiodic
-            ctypes.c_int,  # int yperiodic
-            ctypes.c_int,  # int zperiodic
-            ctypes.POINTER(ctypes.c_double),  # double* boxlo
-            ctypes.POINTER(ctypes.c_double),  # double* boxhi
-            ctypes.c_double,  # double xy
-            ctypes.c_double,  # double xz
-            ctypes.c_double,  # double yz
-        ]
-        self._lib.pair_set_domain.restype = None
-
-        self._lib.pair_run_settings.argtypes = [
-            ctypes.POINTER(PairD3),  # PairD3* pair
-            ctypes.c_double,  # double rthr
-            ctypes.c_double,  # double cnthr
-            ctypes.c_char_p,  # const char* damp_name
-            ctypes.c_char_p,  # const char* func_name
-        ]
-        self._lib.pair_run_settings.restype = None
-
-        self._lib.pair_run_coeff.argtypes = [
-            ctypes.POINTER(PairD3),  # PairD3* pair
-            ctypes.POINTER(ctypes.c_int),  # int* atomic_numbers
-        ]
-        self._lib.pair_run_coeff.restype = None
-
-        self._lib.pair_run_compute.argtypes = [ctypes.POINTER(PairD3)]
-        self._lib.pair_run_compute.restype = None
-
-        self._lib.pair_get_energy.argtypes = [ctypes.POINTER(PairD3)]
-        self._lib.pair_get_energy.restype = ctypes.c_double
-
-        self._lib.pair_get_force.argtypes = [ctypes.POINTER(PairD3)]
-        self._lib.pair_get_force.restype = ctypes.POINTER(ctypes.c_double)
-
-        self._lib.pair_get_stress.argtypes = [ctypes.POINTER(PairD3)]
-        self._lib.pair_get_stress.restype = ctypes.POINTER(ctypes.c_double * 6)
-
-        self._lib.pair_fin.argtypes = [ctypes.POINTER(PairD3)]
-        self._lib.pair_fin.restype = None
-
-    def _idx_to_numbers(self, Z_of_atoms):
-        unique_numbers = list(dict.fromkeys(Z_of_atoms))
-        return unique_numbers
-
-    def _idx_to_types(self, Z_of_atoms):
-        unique_numbers = list(dict.fromkeys(Z_of_atoms))
-        mapping = {num: idx + 1 for idx, num in enumerate(unique_numbers)}
-        atom_types = [mapping[num] for num in Z_of_atoms]
-        return atom_types
-
-    def _convert_domain_ase2lammps(self, cell):
-        qtrans, ltrans = np.linalg.qr(cell.T, mode='complete')
-        lammps_cell = ltrans.T
-        signs = np.sign(np.diag(lammps_cell))
-        lammps_cell = lammps_cell * signs
-        qtrans = qtrans * signs
-        lammps_cell = lammps_cell[(0, 1, 2, 1, 2, 2), (0, 1, 2, 0, 0, 1)]
-        rotator = qtrans.T
-        return lammps_cell, rotator
-
-    def _stress2tensor(self, stress):
-        tensor = np.array(
-            [
-                [stress[0], stress[3], stress[4]],
-                [stress[3], stress[1], stress[5]],
-                [stress[4], stress[5], stress[2]],
-            ]
-        )
-        return tensor
-
-    def _tensor2stress(self, tensor):
-        stress = -np.array(
-            [
-                tensor[0, 0],
-                tensor[1, 1],
-                tensor[2, 2],
-                tensor[1, 2],
-                tensor[0, 2],
-                tensor[0, 1],
-            ]
-        )
-        return stress
-
-    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
-        Calculator.calculate(self, atoms, properties, system_changes)
-        if atoms is None:
-            raise ValueError('No atoms to evaluate')
-
-        if atoms.get_cell().sum() == 0:
-            print(
-                'Warning: D3Calculator requires a cell.\n'
-                'Warning: An orthogonal cell large enough is generated.'
-            )
-            positions = atoms.get_positions()
-            min_pos = positions.min(axis=0)
-            max_pos = positions.max(axis=0)
-            max_cutoff = np.sqrt(max(self.rthr, self.cnthr)) * 0.52917726
-
-            cell_lengths = max_pos - min_pos + max_cutoff + 1.0  # extra margin
-            cell = np.eye(3) * cell_lengths
-
-            atoms.set_cell(cell)
-            atoms.set_pbc([True, True, True])  # for minus positions
-
-        cell, rotator = self._convert_domain_ase2lammps(atoms.get_cell())
-
-        Z_of_atoms = atoms.get_atomic_numbers()
-        natoms = len(atoms)
-        ntypes = len(set(Z_of_atoms))
-        types = (ctypes.c_int * natoms)(*self._idx_to_types(Z_of_atoms))
-
-        positions = atoms.get_positions() @ rotator.T
-        x_flat = (ctypes.c_double * (natoms * 3))(*positions.flatten())
-
-        atomic_numbers = (ctypes.c_int * ntypes)(*self._idx_to_numbers(Z_of_atoms))
-
-        boxlo = (ctypes.c_double * 3)(0.0, 0.0, 0.0)
-        boxhi = (ctypes.c_double * 3)(cell[0], cell[1], cell[2])
-        xy = cell[3]
-        xz = cell[4]
-        yz = cell[5]
-        xperiodic, yperiodic, zperiodic = atoms.get_pbc()
-
-        lib = self._lib
-        assert lib is not None
-        lib.pair_set_atom(self.pair, natoms, ntypes, types, x_flat)
-
-        xperiodic = xperiodic.astype(int)
-        yperiodic = yperiodic.astype(int)
-        zperiodic = zperiodic.astype(int)
-        lib.pair_set_domain(
-            self.pair, xperiodic, yperiodic, zperiodic, boxlo, boxhi, xy, xz, yz
-        )
-
-        lib.pair_run_settings(
-            self.pair,
-            self.rthr,
-            self.cnthr,
-            self.damp_name.encode('utf-8'),
-            self.func_name.encode('utf-8'),
-        )
-
-        lib.pair_run_coeff(self.pair, atomic_numbers)
-        lib.pair_run_compute(self.pair)
-
-        result_E = lib.pair_get_energy(self.pair)
-
-        result_F_ptr = lib.pair_get_force(self.pair)
-        result_F_size = natoms * 3
-        result_F = np.ctypeslib.as_array(
-            result_F_ptr, shape=(result_F_size,)
-        ).reshape((natoms, 3))
-        result_F = np.array(result_F)
-        result_F = result_F @ rotator
-
-        result_S = lib.pair_get_stress(self.pair)
-        result_S = np.array(result_S.contents)
-        result_S = (
-            self._tensor2stress(rotator.T @ self._stress2tensor(result_S) @ rotator)
-            / atoms.get_volume()
-        )
-
-        self.results = {
-            'free_energy': result_E,
-            'energy': result_E,
-            'forces': result_F,
-            'stress': result_S,
-        }
-
-    def predict_one(self, atoms):
-        atoms = atoms.copy()
-        if atoms is None:
-            raise ValueError('No atoms to evaluate')
-
-        if atoms.get_cell().sum() == 0:
-            print(
-                'Warning: D3Calculator requires a cell.\n'
-                'Warning: An orthogonal cell large enough is generated.'
-            )
-            positions = atoms.get_positions()
-            min_pos = positions.min(axis=0)
-            max_pos = positions.max(axis=0)
-            max_cutoff = np.sqrt(max(self.rthr, self.cnthr)) * 0.52917726
-
-            cell_lengths = max_pos - min_pos + max_cutoff + 1.0  # extra margin
-            cell = np.eye(3) * cell_lengths
-
-            atoms.set_cell(cell)
-            atoms.set_pbc([True, True, True])  # for minus positions
-
-        cell, rotator = self._convert_domain_ase2lammps(atoms.get_cell())
-
-        Z_of_atoms = atoms.get_atomic_numbers()
-        natoms = len(atoms)
-        ntypes = len(set(Z_of_atoms))
-        types = (ctypes.c_int * natoms)(*self._idx_to_types(Z_of_atoms))
-
-        positions = atoms.get_positions() @ rotator.T
-        x_flat = (ctypes.c_double * (natoms * 3))(*positions.flatten())
-
-        atomic_numbers = (ctypes.c_int * ntypes)(*self._idx_to_numbers(Z_of_atoms))
-
-        boxlo = (ctypes.c_double * 3)(0.0, 0.0, 0.0)
-        boxhi = (ctypes.c_double * 3)(cell[0], cell[1], cell[2])
-        xy = cell[3]
-        xz = cell[4]
-        yz = cell[5]
-        xperiodic, yperiodic, zperiodic = atoms.get_pbc()
-
-        lib = self._lib
-        assert lib is not None
-        lib.pair_set_atom(self.pair, natoms, ntypes, types, x_flat)
-
-        xperiodic = xperiodic.astype(int)
-        yperiodic = yperiodic.astype(int)
-        zperiodic = zperiodic.astype(int)
-        lib.pair_set_domain(
-            self.pair, xperiodic, yperiodic, zperiodic, boxlo, boxhi, xy, xz, yz
-        )
-
-        lib.pair_run_settings(
-            self.pair,
-            self.rthr,
-            self.cnthr,
-            self.damp_name.encode('utf-8'),
-            self.func_name.encode('utf-8'),
-        )
-
-        lib.pair_run_coeff(self.pair, atomic_numbers)
-        lib.pair_run_compute(self.pair)
-
-        result_E = lib.pair_get_energy(self.pair)
-
-        result_F_ptr = lib.pair_get_force(self.pair)
-        result_F_size = natoms * 3
-        result_F = np.ctypeslib.as_array(
-            result_F_ptr, shape=(result_F_size,)
-        ).reshape((natoms, 3))
-        result_F = np.array(result_F)
-        result_F = result_F @ rotator
-
-        result_S = lib.pair_get_stress(self.pair)
-        result_S = np.array(result_S.contents)
-        result_S = (
-            self._tensor2stress(rotator.T @ self._stress2tensor(result_S) @ rotator)
-            / atoms.get_volume()
-        )
-
-        prediction = {
-            'free_energy': float(result_E),
-            'energy': float(result_E),
-            'forces': result_F.copy(),
-            'stress': result_S.copy(),
-        }
-
-        return prediction
-
-
-    def __del__(self):
-        if self._lib is not None:
-            self._lib.pair_fin(self.pair)
-            self._lib = None
-            self.pair = None
-
+import ctypes
+import os
+import pathlib
+import warnings
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+import torch.jit
+import torch.jit._script
+from ase.calculators.calculator import Calculator, all_changes
+from ase.calculators.mixing import SumCalculator
+from ase.data import chemical_symbols
+
+import sevenn._keys as KEY
+import sevenn.util as util
+from sevenn.atom_graph_data import AtomGraphData
+from sevenn.nn.sequential import AtomGraphSequential
+from sevenn.train.dataload import unlabeled_atoms_to_graph
+import logging
+
+torch_script_type = torch.jit._script.RecursiveScriptModule
+
+
+class SevenNetCalculator(Calculator):
+    """Supporting properties:
+    'free_energy', 'energy', 'forces', 'stress', 'energies'
+    free_energy equals energy. 'energies' stores atomic energy.
+
+    Multi-GPU acceleration is not supported with ASE calculator.
+    You should use LAMMPS for the acceleration.
+    """
+
+    def __init__(
+        self,
+        model: Union[str, pathlib.PurePath, AtomGraphSequential] = '7net-0',
+        file_type: str = 'checkpoint',
+        device: Union[torch.device, str] = 'auto',
+        modal: Optional[str] = None,
+        enable_cueq: bool = False,
+        sevennet_config: Optional[Dict] = None,  # Not used in logic, just meta info
+        **kwargs,
+    ):
+        """Initialize SevenNetCalculator.
+
+        Parameters
+        ----------
+        model: str | Path | AtomGraphSequential, default='7net-0'
+            Name of pretrained models (7net-mf-ompa, 7net-omat, 7net-l3i5, 7net-0) or
+            path to the checkpoint, deployed model or the model itself
+        file_type: str, default='checkpoint'
+            one of 'checkpoint' | 'torchscript' | 'model_instance'
+        device: str | torch.device, default='auto'
+            if not given, use CUDA if available
+        modal: str | None, default=None
+            modal (fidelity) if given model is multi-modal model. for 7net-mf-ompa,
+            it should be one of 'mpa' (MPtrj + sAlex) or 'omat24' (OMat24)
+            case insensitive
+        enable_cueq: bool, default=False
+            if True, use cuEquivariant to accelerate inference.
+        sevennet_config: dict | None, default=None
+            Not used, but can be used to carry meta information of this calculator
+        """
+        print("&&& Initializing SevenNetCalculator")
+        super().__init__(**kwargs)
+        self.sevennet_config = None
+
+        if isinstance(model, pathlib.PurePath):
+            model = str(model)
+
+        allowed_file_types = ['checkpoint', 'torchscript', 'model_instance']
+        file_type = file_type.lower()
+        if file_type not in allowed_file_types:
+            raise ValueError(f'file_type not in {allowed_file_types}')
+
+        if enable_cueq and file_type in ['model_instance', 'torchscript']:
+            warnings.warn(
+                'file_type should be checkpoint to enable cueq. cueq set to False'
+            )
+            enable_cueq = False
+
+        if isinstance(device, str):  # TODO: do we really need this?
+            if device == 'auto':
+                self.device = torch.device(
+                    'cuda' if torch.cuda.is_available() else 'cpu'
+                )
+            else:
+                self.device = torch.device(device)
+        else:
+            self.device = device
+
+        if file_type == 'checkpoint' and isinstance(model, str):
+            cp = util.load_checkpoint(model)
+
+            backend = 'e3nn' if not enable_cueq else 'cueq'
+            model_loaded = cp.build_model(backend)
+            model_loaded.set_is_batch_data(False)
+
+            self.type_map = cp.config[KEY.TYPE_MAP]
+            self.cutoff = cp.config[KEY.CUTOFF]
+            self.sevennet_config = cp.config
+
+        elif file_type == 'torchscript' and isinstance(model, str):
+            if modal:
+                raise NotImplementedError()
+            extra_dict = {
+                'chemical_symbols_to_index': b'',
+                'cutoff': b'',
+                'num_species': b'',
+                'model_type': b'',
+                'version': b'',
+                'dtype': b'',
+                'time': b'',
+            }
+            model_loaded = torch.jit.load(
+                model, _extra_files=extra_dict, map_location=self.device
+            )
+            chem_symbols = extra_dict['chemical_symbols_to_index'].decode('utf-8')
+            sym_to_num = {sym: n for n, sym in enumerate(chemical_symbols)}
+            self.type_map = {
+                sym_to_num[sym]: i for i, sym in enumerate(chem_symbols.split())
+            }
+            self.cutoff = float(extra_dict['cutoff'].decode('utf-8'))
+
+        elif isinstance(model, AtomGraphSequential):
+            if model.type_map is None:
+                raise ValueError(
+                    'Model must have the type_map to be used with calculator'
+                )
+            if model.cutoff == 0.0:
+                raise ValueError('Model cutoff seems not initialized')
+            model.eval_type_map = torch.tensor(True)  # ?
+            model.set_is_batch_data(False)
+            model_loaded = model
+            self.type_map = model.type_map
+            self.cutoff = model.cutoff
+        else:
+            raise ValueError('Unexpected input combinations')
+
+        if self.sevennet_config is None and sevennet_config is not None:
+            self.sevennet_config = sevennet_config
+
+        self.model = model_loaded
+
+        self.modal = None
+        if isinstance(self.model, AtomGraphSequential):
+            modal_map = self.model.modal_map
+            if modal_map:
+                modal_ava = list(modal_map.keys())
+                if not modal:
+                    raise ValueError(f'modal argument missing (avail: {modal_ava})')
+                elif modal not in modal_ava:
+                    raise ValueError(f'unknown modal {modal} (not in {modal_ava})')
+                self.modal = modal
+            elif not self.model.modal_map and modal:
+                warnings.warn(f'modal={modal} is ignored as model has no modal_map')
+
+        self.model.to(self.device)
+        self.model.eval()
+        self.implemented_properties = [
+            'free_energy',
+            'energy',
+            'forces',
+            'stress',
+            'energies',
+        ]
+
+    def set_atoms(self, atoms):
+        # called by ase, when atoms.calc = calc
+        zs = tuple(set(atoms.get_atomic_numbers()))
+        for z in zs:
+            if z not in self.type_map:
+                sp = list(self.type_map.keys())
+                raise ValueError(
+                    f'Model do not know atomic number: {z}, (knows: {sp})'
+                )
+
+    def output_to_results(self, output):
+        energy = output[KEY.PRED_TOTAL_ENERGY].detach().cpu().item()
+        num_atoms = output['num_atoms'].item()
+        atomic_energies = output[KEY.ATOMIC_ENERGY].detach().cpu().numpy().flatten()
+        forces = output[KEY.PRED_FORCE].detach().cpu().numpy()[:num_atoms, :]
+        stress = np.array(
+            (-output[KEY.PRED_STRESS])
+            .detach()
+            .cpu()
+            .numpy()[[0, 1, 2, 4, 5, 3]]  # as voigt notation
+        )
+        # Store results
+        return {
+            'free_energy': energy,
+            'energy': energy,
+            'energies': atomic_energies,
+            'forces': forces,
+            'stress': stress,
+            'num_edges': output[KEY.EDGE_IDX].shape[1],
+        }
+
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+        # call parent class to set necessary atom attributes
+        Calculator.calculate(self, atoms, properties, system_changes)
+        if atoms is None:
+            raise ValueError('No atoms to evaluate')
+        data = AtomGraphData.from_numpy_dict(
+            unlabeled_atoms_to_graph(atoms, self.cutoff)
+        )
+        if self.modal:
+            data[KEY.DATA_MODALITY] = self.modal
+
+        data.to(self.device)  # type: ignore
+
+        if isinstance(self.model, torch_script_type):
+            data[KEY.NODE_FEATURE] = torch.tensor(
+                [self.type_map[z.item()] for z in data[KEY.NODE_FEATURE]],
+                dtype=torch.int64,
+                device=self.device,
+            )
+            data[KEY.POS].requires_grad_(True)  # backward compatibility
+            data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
+            data = data.to_dict()
+            del data['data_info']
+
+        import logging
+        logging.debug(f"data: {data}")
+        # logging.debug(f"data[pos]: {data['pos']}")
+        # logging.debug(f"data[x]: {data['x']}")
+        logging.debug(f"data[cell_lattice_vectors]: {data['cell_lattice_vectors']}")
+        logging.debug(f"data[cell_volume]: {data['cell_volume']}")
+        output = self.model(data)
+        # logging.info(f"input: {data}")
+        # logging.info(f"output[{KEY.PRED_TOTAL_ENERGY}] = {output[KEY.PRED_TOTAL_ENERGY]}")
+        # logging.info(f"output[{KEY.PRED_FORCE}] = {output[KEY.PRED_FORCE]}")
+        # logging.info(f"output[{KEY.PRED_STRESS}] = {output[KEY.PRED_STRESS]}")
+        self.results = self.output_to_results(output)
+        # logging.debug(f"results['energy'] = {self.results['energy']}")
+        # logging.debug(f"results['forces'] = {self.results['forces']}")
+        # logging.debug(f"results['stress'] = {self.results['stress']}")
+
+    def predict_one(self, atoms):
+        if atoms is None:
+            raise ValueError('No atoms to evaluate')
+        data = AtomGraphData.from_numpy_dict(
+            unlabeled_atoms_to_graph(atoms, self.cutoff)
+        )
+        if self.modal:
+            data[KEY.DATA_MODALITY] = self.modal
+
+        data.to(self.device)  # type: ignore
+
+        if isinstance(self.model, torch_script_type):
+            data[KEY.NODE_FEATURE] = torch.tensor(
+                [self.type_map[z.item()] for z in data[KEY.NODE_FEATURE]],
+                dtype=torch.int64,
+                device=self.device,
+            )
+            data[KEY.POS].requires_grad_(True)  # backward compatibility
+            data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
+            data = data.to_dict()
+            del data['data_info']
+
+        return self.model(data)
+
+
+
+    def predict(self, atoms_list, properties=None):
+        
+        # if len(atoms_list) == 1:
+        #     output = self.predict_one(atoms_list[0])
+        #     predictions = {}
+        #     predictions['energy'] = output[KEY.PRED_TOTAL_ENERGY].to(torch.float64).unsqueeze(0)
+        #     predictions['forces'] = output[KEY.PRED_FORCE].to(torch.float64).unsqueeze(0)
+        #     voigt = (-output[KEY.PRED_STRESS])[[0, 1, 2, 4, 5, 3]].to(torch.float64).unsqueeze(0)
+        #     stress_list = []
+        #     for i in range(voigt.shape[0]):
+        #         stress_list.append(self._stress2tensor(voigt[i,:]))
+        #     predictions['stress'] = torch.stack(stress_list, dim=0).view(-1,3,3)
+        #     return predictions
+
+
+        if not atoms_list:
+            raise ValueError("Empty atoms_list provided")
+            
+        if not isinstance(atoms_list, list):
+            atoms_list = [atoms_list]
+            
+        # Convert atoms to graph data
+        graph_list = []
+        for atoms in atoms_list:
+            data = AtomGraphData.from_numpy_dict(
+                unlabeled_atoms_to_graph(atoms, self.cutoff)
+            )
+            if self.modal:
+                data[KEY.DATA_MODALITY] = self.modal
+                
+            if isinstance(self.model, torch_script_type):
+                data[KEY.NODE_FEATURE] = torch.tensor(
+                    [self.type_map[z.item()] for z in data[KEY.NODE_FEATURE]],
+                    dtype=torch.int64,
+                    device=self.device,
+                )
+                data[KEY.POS].requires_grad_(True)  # backward compatibility
+                data[KEY.EDGE_VEC].requires_grad_(True)  # backward compatibility
+                
+            graph_list.append(data)
+            
+        # Process graphs based on model type
+        # was_batch_mode = True
+        if isinstance(self.model, AtomGraphSequential):
+            # was_batch_mode = self.model.is_batch_data
+            self.model.set_is_batch_data(True)
+            self.model.eval()
+            
+        # Batch the data if there are multiple atoms
+        from torch_geometric.loader.dataloader import Collater
+        batched_data = Collater(graph_list)(graph_list)
+        batched_data = batched_data.to(self.device)
+        
+        import logging
+        logging.debug(f"batched_data: {batched_data}")
+        # logging.debug(f"batched_data[pos]: {batched_data['pos']}")
+        # logging.debug(f"batched_data[x]: {batched_data['x']}")
+        logging.debug(f"batched_data[cell_lattice_vectors]: {batched_data['cell_lattice_vectors']}")
+        logging.debug(f"batched_data[cell_volume]: {batched_data['cell_volume']}")
+        # Run model on batched data
+        if isinstance(self.model, torch_script_type):
+            batched_dict = batched_data.to_dict()
+            if 'data_info' in batched_dict:
+                del batched_dict['data_info']
+            output = self.model(batched_dict)
+        else:
+            output = self.model(batched_data)
+            
+        # Convert to list of individual outputs using util.to_atom_graph_list
+        # logging.info(f"input: {batched_data}")
+        # logging.info(f"output[{KEY.PRED_TOTAL_ENERGY}] = {output[KEY.PRED_TOTAL_ENERGY]}")
+        # logging.info(f"output[{KEY.PRED_FORCE}] = {output[KEY.PRED_FORCE]}")
+        # logging.info(f"output[{KEY.PRED_STRESS}] = {output[KEY.PRED_STRESS]}")
+
+        predictions = {}
+        predictions['energy'] = output[KEY.PRED_TOTAL_ENERGY].to(torch.float64).detach()
+        predictions['forces'] = output[KEY.PRED_FORCE].to(torch.float64).detach()
+        voigt = (-output[KEY.PRED_STRESS])[:, [0, 1, 2, 4, 5, 3]].to(torch.float64).detach()
+        stress_list = []
+        for i in range(voigt.shape[0]):
+            stress_list.append(self._stress2tensor(voigt[i,:]))
+        predictions['stress'] = torch.stack(stress_list, dim=0).view(-1,3,3).detach()
+
+        # logging.debug(f"predictions['energy'] = {predictions['energy']}")
+        # logging.debug(f"predictions['forces'] = {predictions['forces']}")
+        # logging.debug(f"predictions['stress'] = {predictions['stress']}")
+        return predictions
+
+    def _stress2tensor(self, stress):
+        tensor = torch.tensor(
+            [
+                [stress[0], stress[5], stress[4]],
+                [stress[5], stress[1], stress[3]],
+                [stress[4], stress[3], stress[2]],
+            ], 
+            device=self.device
+        )
+        return tensor
+
+
+class SevenNetD3Calculator(SumCalculator):
+    def __init__(
+        self,
+        model: Union[str, pathlib.PurePath, AtomGraphSequential] = '7net-0',
+        file_type: str = 'checkpoint',
+        device: Union[torch.device, str] = 'auto',
+        sevennet_config: Optional[Any] = None,  # hold meta information
+        damping_type: str = 'damp_bj',
+        functional_name: str = 'pbe',
+        vdw_cutoff: float = 9000,  # au^2, 0.52917726 angstrom = 1 au
+        cn_cutoff: float = 1600,  # au^2, 0.52917726 angstrom = 1 au
+        batch_size=10,
+        **kwargs,
+    ):
+        """Initialize SevenNetD3Calculator. CUDA required.
+
+        Parameters
+        ----------
+        model: str | Path | AtomGraphSequential
+            Name of pretrained models (7net-mf-ompa, 7net-omat, 7net-l3i5, 7net-0) or
+            path to the checkpoint, deployed model or the model itself
+        file_type: str, default='checkpoint'
+            one of 'checkpoint' | 'torchscript' | 'model_instance'
+        device: str | torch.device, default='auto'
+            if not given, use CUDA if available
+        modal: str | None, default=None
+            modal (fidelity) if given model is multi-modal model. for 7net-mf-ompa,
+            it should be one of 'mpa' (MPtrj + sAlex) or 'omat24' (OMat24)
+        enable_cueq: bool, default=False
+            if True, use cuEquivariant to accelerate inference.
+        damping_type: str, default='damp_bj'
+            Damping type of D3, one of 'damp_bj' | 'damp_zero'
+        functional_name: str, default='pbe'
+            Target functional name of D3 parameters.
+        vdw_cutoff: float, default=9000
+            vdw cutoff of D3 calculator in au
+        cn_cutoff: float, default=1600
+            cn cutoff of D3 calculator in au
+        """
+        self.d3_calc = D3Calculator(
+            damping_type=damping_type,
+            functional_name=functional_name,
+            vdw_cutoff=vdw_cutoff,
+            cn_cutoff=cn_cutoff,
+            **kwargs,
+        )
+
+        self.sevennet_calc = SevenNetCalculator(
+            model=model,
+            file_type=file_type,
+            device=device,
+            sevennet_config=sevennet_config,
+            **kwargs,
+        )
+
+        super().__init__([self.sevennet_calc, self.d3_calc])
+
+        self.device = device
+        self.d3_calcs = []
+        for _ in range(batch_size):
+            self.d3_calcs.append(
+                D3Calculator(
+                    damping_type=damping_type,
+                    functional_name=functional_name,
+                    vdw_cutoff=vdw_cutoff,
+                    cn_cutoff=cn_cutoff,
+                    **kwargs,
+                )
+            )
+
+
+    def predict(self, atoms_list):
+        """Predict the energy and forces for a list of atoms.
+        """
+        # Call the predict method of the first calculator (SevenNetCalculator)
+        predictions = self.sevennet_calc.predict(atoms_list)
+        
+        energy_list = []
+        forces_list = []
+        stress_list = []
+        predictions3d = {}
+        for i, atoms in enumerate(atoms_list):
+            prediction = self.d3_calcs[i].predict_one(atoms)
+            energy_list.append(torch.tensor(prediction['energy']))
+            forces_list.append(torch.from_numpy(prediction['forces']).to(self.device))
+            stress_list.append(self._stress2tensor(torch.from_numpy(prediction['stress'])))
+
+        # Convert lists to tensors
+        predictions3d['energy'] = torch.stack(energy_list, dim=0).to(self.device)
+        predictions3d['forces'] = torch.cat(forces_list, dim=0).view(-1, 3)
+        predictions3d['stress'] = torch.stack(stress_list, dim=0).view(-1, 3, 3)
+
+        predictions['energy'] += predictions3d['energy'].detach()
+        predictions['forces'] += predictions3d['forces'].detach()
+        predictions['stress'] += predictions3d['stress'].detach()
+        
+        return predictions
+
+    def _stress2tensor(self, stress):
+        tensor = torch.tensor(
+            [
+                # [stress[0], stress[3], stress[4]],
+                # [stress[3], stress[1], stress[5]],
+                # [stress[4], stress[5], stress[2]],
+                [stress[0], stress[5], stress[4]],
+                [stress[5], stress[1], stress[3]],
+                [stress[4], stress[3], stress[2]],
+            ], 
+            device=self.device
+        )
+        return tensor
+
+
+
+def _load(name: str) -> ctypes.CDLL:
+    from torch.utils.cpp_extension import LIB_EXT, _get_build_directory, load
+
+    # Load the library from the candidate locations
+
+    package_dir = os.path.dirname(os.path.abspath(__file__))
+    try:
+        return ctypes.CDLL(os.path.join(package_dir, f'{name}{LIB_EXT}'))
+    except OSError:
+        pass
+
+    cache_dir = _get_build_directory(name, verbose=False)
+    try:
+        return ctypes.CDLL(os.path.join(cache_dir, f'{name}{LIB_EXT}'))
+    except OSError:
+        pass
+
+    # Compile the library if it is not found
+
+    if os.access(package_dir, os.W_OK):
+        compile_dir = package_dir
+    else:
+        print('Warning: package directory is not writable. Using cache directory.')
+        compile_dir = cache_dir
+
+    if 'TORCH_CUDA_ARCH_LIST' not in os.environ:
+        print('Warning: TORCH_CUDA_ARCH_LIST is not set.')
+        print('Warning: Use default CUDA architectures: 61, 70, 75, 80, 86, 89, 90')
+        os.environ['TORCH_CUDA_ARCH_LIST'] = '6.1;7.0;7.5;8.0;8.6;8.9;9.0'
+
+    load(
+        name=name,
+        sources=[os.path.join(package_dir, 'pair_e3gnn', 'pair_d3_for_ase.cu')],
+        extra_cuda_cflags=['-O3', '--expt-relaxed-constexpr', '-fmad=false'],
+        build_directory=compile_dir,
+        verbose=True,
+        is_python_module=False,
+    )
+
+    return ctypes.CDLL(os.path.join(compile_dir, f'{name}{LIB_EXT}'))
+
+
+class PairD3(ctypes.Structure):
+    pass  # Opaque structure; only used as a pointer
+
+
+class D3Calculator(Calculator):
+    """ASE calculator for accelerated D3 van der Waals (vdW) correction.
+
+    Example:
+        from ase.calculators.mixing import SumCalculator
+        calc_1 = SevenNetCalculator()
+        calc_2 = D3Calculator()
+        return SumCalculator([calc_1, calc_2])
+
+    This calculator interfaces with the `libpaird3.so` library,
+    which is compiled by nvcc during the package installation.
+    If you encounter any errors, please verify
+    the installation process and the compilation options in `setup.py`.
+    Note: Multi-GPU parallel MD is not supported in this mode.
+    Note: Cffi could be used, but it was avoided to reduce dependencies.
+    """
+
+    # Here, free_energy = energy
+    implemented_properties = ['free_energy', 'energy', 'forces', 'stress']
+
+    def __init__(
+        self,
+        damping_type: str = 'damp_bj',  # damp_bj, damp_zero
+        functional_name: str = 'pbe',  # check the source code
+        vdw_cutoff: float = 9000,  # au^2, 0.52917726 angstrom = 1 au
+        cn_cutoff: float = 1600,  # au^2, 0.52917726 angstrom = 1 au
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        if not torch.cuda.is_available():
+            raise NotImplementedError('CPU + D3 is not implemented yet')
+
+        self.rthr = vdw_cutoff
+        self.cnthr = cn_cutoff
+        self.damp_name = damping_type.lower()
+        self.func_name = functional_name.lower()
+
+        if self.damp_name not in ['damp_bj', 'damp_zero']:
+            raise ValueError('Error: Invalid damping type.')
+
+        self._lib = _load('pair_d3')
+
+        self._lib.pair_init.restype = ctypes.POINTER(PairD3)
+        self.pair = self._lib.pair_init()
+
+        self._lib.pair_set_atom.argtypes = [
+            ctypes.POINTER(PairD3),  # PairD3* pair
+            ctypes.c_int,  # int natoms
+            ctypes.c_int,  # int ntypes
+            ctypes.POINTER(ctypes.c_int),  # int* types
+            ctypes.POINTER(ctypes.c_double),  # double* x
+        ]
+        self._lib.pair_set_atom.restype = None
+
+        self._lib.pair_set_domain.argtypes = [
+            ctypes.POINTER(PairD3),  # PairD3* pair
+            ctypes.c_int,  # int xperiodic
+            ctypes.c_int,  # int yperiodic
+            ctypes.c_int,  # int zperiodic
+            ctypes.POINTER(ctypes.c_double),  # double* boxlo
+            ctypes.POINTER(ctypes.c_double),  # double* boxhi
+            ctypes.c_double,  # double xy
+            ctypes.c_double,  # double xz
+            ctypes.c_double,  # double yz
+        ]
+        self._lib.pair_set_domain.restype = None
+
+        self._lib.pair_run_settings.argtypes = [
+            ctypes.POINTER(PairD3),  # PairD3* pair
+            ctypes.c_double,  # double rthr
+            ctypes.c_double,  # double cnthr
+            ctypes.c_char_p,  # const char* damp_name
+            ctypes.c_char_p,  # const char* func_name
+        ]
+        self._lib.pair_run_settings.restype = None
+
+        self._lib.pair_run_coeff.argtypes = [
+            ctypes.POINTER(PairD3),  # PairD3* pair
+            ctypes.POINTER(ctypes.c_int),  # int* atomic_numbers
+        ]
+        self._lib.pair_run_coeff.restype = None
+
+        self._lib.pair_run_compute.argtypes = [ctypes.POINTER(PairD3)]
+        self._lib.pair_run_compute.restype = None
+
+        self._lib.pair_get_energy.argtypes = [ctypes.POINTER(PairD3)]
+        self._lib.pair_get_energy.restype = ctypes.c_double
+
+        self._lib.pair_get_force.argtypes = [ctypes.POINTER(PairD3)]
+        self._lib.pair_get_force.restype = ctypes.POINTER(ctypes.c_double)
+
+        self._lib.pair_get_stress.argtypes = [ctypes.POINTER(PairD3)]
+        self._lib.pair_get_stress.restype = ctypes.POINTER(ctypes.c_double * 6)
+
+        self._lib.pair_fin.argtypes = [ctypes.POINTER(PairD3)]
+        self._lib.pair_fin.restype = None
+
+    def _idx_to_numbers(self, Z_of_atoms):
+        unique_numbers = list(dict.fromkeys(Z_of_atoms))
+        return unique_numbers
+
+    def _idx_to_types(self, Z_of_atoms):
+        unique_numbers = list(dict.fromkeys(Z_of_atoms))
+        mapping = {num: idx + 1 for idx, num in enumerate(unique_numbers)}
+        atom_types = [mapping[num] for num in Z_of_atoms]
+        return atom_types
+
+    def _convert_domain_ase2lammps(self, cell):
+        qtrans, ltrans = np.linalg.qr(cell.T, mode='complete')
+        lammps_cell = ltrans.T
+        signs = np.sign(np.diag(lammps_cell))
+        lammps_cell = lammps_cell * signs
+        qtrans = qtrans * signs
+        lammps_cell = lammps_cell[(0, 1, 2, 1, 2, 2), (0, 1, 2, 0, 0, 1)]
+        rotator = qtrans.T
+        return lammps_cell, rotator
+
+    def _stress2tensor(self, stress):
+        tensor = np.array(
+            [
+                [stress[0], stress[3], stress[4]],
+                [stress[3], stress[1], stress[5]],
+                [stress[4], stress[5], stress[2]],
+            ]
+        )
+        return tensor
+
+    def _tensor2stress(self, tensor):
+        stress = -np.array(
+            [
+                tensor[0, 0],
+                tensor[1, 1],
+                tensor[2, 2],
+                tensor[1, 2],
+                tensor[0, 2],
+                tensor[0, 1],
+            ]
+        )
+        return stress
+
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+        Calculator.calculate(self, atoms, properties, system_changes)
+        if atoms is None:
+            raise ValueError('No atoms to evaluate')
+
+        if atoms.get_cell().sum() == 0:
+            print(
+                'Warning: D3Calculator requires a cell.\n'
+                'Warning: An orthogonal cell large enough is generated.'
+            )
+            positions = atoms.get_positions()
+            min_pos = positions.min(axis=0)
+            max_pos = positions.max(axis=0)
+            max_cutoff = np.sqrt(max(self.rthr, self.cnthr)) * 0.52917726
+
+            cell_lengths = max_pos - min_pos + max_cutoff + 1.0  # extra margin
+            cell = np.eye(3) * cell_lengths
+
+            atoms.set_cell(cell)
+            atoms.set_pbc([True, True, True])  # for minus positions
+
+        cell, rotator = self._convert_domain_ase2lammps(atoms.get_cell())
+
+        Z_of_atoms = atoms.get_atomic_numbers()
+        natoms = len(atoms)
+        ntypes = len(set(Z_of_atoms))
+        types = (ctypes.c_int * natoms)(*self._idx_to_types(Z_of_atoms))
+
+        positions = atoms.get_positions() @ rotator.T
+        x_flat = (ctypes.c_double * (natoms * 3))(*positions.flatten())
+
+        atomic_numbers = (ctypes.c_int * ntypes)(*self._idx_to_numbers(Z_of_atoms))
+
+        boxlo = (ctypes.c_double * 3)(0.0, 0.0, 0.0)
+        boxhi = (ctypes.c_double * 3)(cell[0], cell[1], cell[2])
+        xy = cell[3]
+        xz = cell[4]
+        yz = cell[5]
+        xperiodic, yperiodic, zperiodic = atoms.get_pbc()
+
+        lib = self._lib
+        assert lib is not None
+        lib.pair_set_atom(self.pair, natoms, ntypes, types, x_flat)
+
+        xperiodic = xperiodic.astype(int)
+        yperiodic = yperiodic.astype(int)
+        zperiodic = zperiodic.astype(int)
+        lib.pair_set_domain(
+            self.pair, xperiodic, yperiodic, zperiodic, boxlo, boxhi, xy, xz, yz
+        )
+
+        lib.pair_run_settings(
+            self.pair,
+            self.rthr,
+            self.cnthr,
+            self.damp_name.encode('utf-8'),
+            self.func_name.encode('utf-8'),
+        )
+
+        lib.pair_run_coeff(self.pair, atomic_numbers)
+        lib.pair_run_compute(self.pair)
+
+        result_E = lib.pair_get_energy(self.pair)
+
+        result_F_ptr = lib.pair_get_force(self.pair)
+        result_F_size = natoms * 3
+        result_F = np.ctypeslib.as_array(
+            result_F_ptr, shape=(result_F_size,)
+        ).reshape((natoms, 3))
+        result_F = np.array(result_F)
+        result_F = result_F @ rotator
+
+        result_S = lib.pair_get_stress(self.pair)
+        result_S = np.array(result_S.contents)
+        result_S = (
+            self._tensor2stress(rotator.T @ self._stress2tensor(result_S) @ rotator)
+            / atoms.get_volume()
+        )
+
+        self.results = {
+            'free_energy': result_E,
+            'energy': result_E,
+            'forces': result_F,
+            'stress': result_S,
+        }
+
+    def predict_one(self, atoms):
+        atoms = atoms.copy()
+        if atoms is None:
+            raise ValueError('No atoms to evaluate')
+
+        if atoms.get_cell().sum() == 0:
+            print(
+                'Warning: D3Calculator requires a cell.\n'
+                'Warning: An orthogonal cell large enough is generated.'
+            )
+            positions = atoms.get_positions()
+            min_pos = positions.min(axis=0)
+            max_pos = positions.max(axis=0)
+            max_cutoff = np.sqrt(max(self.rthr, self.cnthr)) * 0.52917726
+
+            cell_lengths = max_pos - min_pos + max_cutoff + 1.0  # extra margin
+            cell = np.eye(3) * cell_lengths
+
+            atoms.set_cell(cell)
+            atoms.set_pbc([True, True, True])  # for minus positions
+
+        cell, rotator = self._convert_domain_ase2lammps(atoms.get_cell())
+
+        Z_of_atoms = atoms.get_atomic_numbers()
+        natoms = len(atoms)
+        ntypes = len(set(Z_of_atoms))
+        types = (ctypes.c_int * natoms)(*self._idx_to_types(Z_of_atoms))
+
+        positions = atoms.get_positions() @ rotator.T
+        x_flat = (ctypes.c_double * (natoms * 3))(*positions.flatten())
+
+        atomic_numbers = (ctypes.c_int * ntypes)(*self._idx_to_numbers(Z_of_atoms))
+
+        boxlo = (ctypes.c_double * 3)(0.0, 0.0, 0.0)
+        boxhi = (ctypes.c_double * 3)(cell[0], cell[1], cell[2])
+        xy = cell[3]
+        xz = cell[4]
+        yz = cell[5]
+        xperiodic, yperiodic, zperiodic = atoms.get_pbc()
+
+        lib = self._lib
+        assert lib is not None
+        lib.pair_set_atom(self.pair, natoms, ntypes, types, x_flat)
+
+        xperiodic = xperiodic.astype(int)
+        yperiodic = yperiodic.astype(int)
+        zperiodic = zperiodic.astype(int)
+        lib.pair_set_domain(
+            self.pair, xperiodic, yperiodic, zperiodic, boxlo, boxhi, xy, xz, yz
+        )
+
+        lib.pair_run_settings(
+            self.pair,
+            self.rthr,
+            self.cnthr,
+            self.damp_name.encode('utf-8'),
+            self.func_name.encode('utf-8'),
+        )
+
+        lib.pair_run_coeff(self.pair, atomic_numbers)
+        lib.pair_run_compute(self.pair)
+
+        result_E = lib.pair_get_energy(self.pair)
+
+        result_F_ptr = lib.pair_get_force(self.pair)
+        result_F_size = natoms * 3
+        result_F = np.ctypeslib.as_array(
+            result_F_ptr, shape=(result_F_size,)
+        ).reshape((natoms, 3))
+        result_F = np.array(result_F)
+        result_F = result_F @ rotator
+
+        result_S = lib.pair_get_stress(self.pair)
+        result_S = np.array(result_S.contents)
+        result_S = (
+            self._tensor2stress(rotator.T @ self._stress2tensor(result_S) @ rotator)
+            / atoms.get_volume()
+        )
+
+        prediction = {
+            'free_energy': float(result_E),
+            'energy': float(result_E),
+            'forces': result_F.copy(),
+            'stress': result_S.copy(),
+        }
+
+        return prediction
+
+
+    def __del__(self):
+        if self._lib is not None:
+            self._lib.pair_fin(self.pair)
+            self._lib = None
+            self.pair = None
+

mace-bench/3rdparty/SevenNet/sevenn/checkpoint.py

-import os
-import pathlib
-import uuid
-import warnings
-from copy import deepcopy
-from datetime import datetime
-from typing import Any, Dict, Optional, Union
-
-import pandas as pd
-from packaging.version import Version
-from torch import Tensor
-from torch import load as torch_load
-
-import sevenn
-import sevenn._const as consts
-import sevenn._keys as KEY
-import sevenn.scripts.backward_compatibility as compat
-from sevenn import model_build
-from sevenn.nn.scale import get_resolved_shift_scale
-from sevenn.nn.sequential import AtomGraphSequential
-
-
-def assert_atoms(atoms1, atoms2, rtol=1e-5, atol=1e-6):
-    import numpy as np
-
-    def acl(a, b, rtol=rtol, atol=atol):
-        return np.allclose(a, b, rtol=rtol, atol=atol)
-
-    assert len(atoms1) == len(atoms2)
-    assert acl(atoms1.get_cell(), atoms2.get_cell())
-    assert acl(atoms1.get_potential_energy(), atoms2.get_potential_energy())
-    assert acl(atoms1.get_forces(), atoms2.get_forces(), rtol * 10, atol * 10)
-    assert acl(
-        atoms1.get_stress(voigt=False),
-        atoms2.get_stress(voigt=False),
-        rtol * 10,
-        atol * 10,
-    )
-    # assert acl(atoms1.get_potential_energies(), atoms2.get_potential_energies())
-
-
-def copy_state_dict(state_dict) -> dict:
-    if isinstance(state_dict, dict):
-        return {key: copy_state_dict(value) for key, value in state_dict.items()}
-    elif isinstance(state_dict, list):
-        return [copy_state_dict(item) for item in state_dict]  # type: ignore
-    elif isinstance(state_dict, Tensor):
-        return state_dict.clone()  # type: ignore
-    else:
-        # For non-tensor values (e.g., scalars, None), return as-is
-        return state_dict
-
-
-def _config_cp_routine(config):
-    cp_ver = Version(config.get('version', None))
-    this_ver = Version(sevenn.__version__)
-    if cp_ver > this_ver:
-        warnings.warn(f'The checkpoint version ({cp_ver}) is newer than this source'
-                      f'({this_ver}). This may cause unexpected behaviors')
-
-    defaults = {**consts.model_defaults(config)}
-    config = compat.patch_old_config(config)  # type: ignore
-
-    scaler = model_build.init_shift_scale(config)
-    shift, scale = get_resolved_shift_scale(
-        scaler, config.get(KEY.TYPE_MAP), config.get(KEY.MODAL_MAP, None)
-    )
-    config['shift'] = shift
-    config['scale'] = scale
-
-    for k, v in defaults.items():
-        if k in config:
-            continue
-        if os.getenv('SEVENN_DEBUG', False):
-            warnings.warn(f'{k} not in config, use default value {v}', UserWarning)
-        config[k] = v
-
-    for k, v in config.items():
-        if isinstance(v, Tensor):
-            config[k] = v.cpu()
-    return config
-
-
-def _convert_e3nn_and_cueq(stct_src, stct_dst, src_config, from_cueq):
-    """
-    manually check keys and assert if something unexpected happens
-    """
-    n_layer = src_config['num_convolution_layer']
-
-    linear_module_names = [
-        'onehot_to_feature_x',
-        'reduce_input_to_hidden',
-        'reduce_hidden_to_energy',
-    ]
-    convolution_module_names = []
-    fc_tensor_product_module_names = []
-    for i in range(n_layer):
-        linear_module_names.append(f'{i}_self_interaction_1')
-        linear_module_names.append(f'{i}_self_interaction_2')
-        if src_config.get(KEY.SELF_CONNECTION_TYPE) == 'linear':
-            linear_module_names.append(f'{i}_self_connection_intro')
-        elif src_config.get(KEY.SELF_CONNECTION_TYPE) == 'nequip':
-            fc_tensor_product_module_names.append(f'{i}_self_connection_intro')
-        convolution_module_names.append(f'{i}_convolution')
-
-    # Rule: those keys can be safely ignored before state dict load,
-    #       except for linear.bias. This should be aborted in advance to
-    #       this function. Others are not parameters but constants.
-    cue_only_linear_followers = ['linear.f.tp.f_fx.module.c']
-    e3nn_only_linear_followers = ['linear.bias', 'linear.output_mask']
-    ignores_in_linear = cue_only_linear_followers + e3nn_only_linear_followers
-
-    cue_only_conv_followers = [
-        'convolution.f.tp.f_fx.module.c',
-        'convolution.f.tp.module.module.f.module.module._f.data',
-    ]
-    e3nn_only_conv_followers = [
-        'convolution._compiled_main_left_right._w3j',
-        'convolution.weight',
-        'convolution.output_mask',
-    ]
-    ignores_in_conv = cue_only_conv_followers + e3nn_only_conv_followers
-
-    cue_only_fc_followers = ['fc_tensor_product.f.tp.f_fx.module.c']
-    e3nn_only_fc_followers = [
-        'fc_tensor_product.output_mask',
-    ]
-    ignores_in_fc = cue_only_fc_followers + e3nn_only_fc_followers
-
-    updated_keys = []
-    for k, v in stct_src.items():
-        module_name = k.split('.')[0]
-        flag = False
-        if module_name in linear_module_names:
-            for ignore in ignores_in_linear:
-                if '.'.join([module_name, ignore]) in k:
-                    flag = True
-                    break
-            if not flag and k == '.'.join([module_name, 'linear.weight']):
-                updated_keys.append(k)
-                stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
-                flag = True
-            assert flag, f'Unexpected key from linear: {k}'
-        elif module_name in convolution_module_names:
-            for ignore in ignores_in_conv:
-                if '.'.join([module_name, ignore]) in k:
-                    flag = True
-                    break
-            if not flag and (
-                k.startswith(f'{module_name}.weight_nn')
-                or k == '.'.join([module_name, 'denominator'])
-            ):
-                updated_keys.append(k)
-                stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
-                flag = True
-            assert flag, f'Unexpected key from linear: {k}'
-        elif module_name in fc_tensor_product_module_names:
-            for ignore in ignores_in_fc:
-                if '.'.join([module_name, ignore]) in k:
-                    flag = True
-                    break
-            if not flag and k == '.'.join([module_name, 'fc_tensor_product.weight']):
-                updated_keys.append(k)
-                stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
-                flag = True
-            assert flag, f'Unexpected key from fc tensor product: {k}'
-        else:
-            # assert k in stct_dst
-            updated_keys.append(k)
-            stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
-
-    return stct_dst
-
-
-class SevenNetCheckpoint:
-    """
-    Tool box for checkpoint processed from SevenNet.
-    """
-
-    def __init__(self, checkpoint_path: Union[pathlib.Path, str]):
-        self._checkpoint_path = os.path.abspath(checkpoint_path)
-        self._config = None
-        self._epoch = None
-        self._model_state_dict = None
-        self._optimizer_state_dict = None
-        self._scheduler_state_dict = None
-        self._hash = None
-        self._time = None
-
-        self._loaded = False
-
-    def __repr__(self) -> str:
-        cfg = self.config  # just alias
-        if len(cfg) == 0:
-            return ''
-        dct = {
-            'Sevennet version': cfg.get('version', 'Not found'),
-            'When': self.time,
-            'Hash': self.hash,
-            'Cutoff': cfg.get('cutoff'),
-            'Channel': cfg.get('channel'),
-            'Lmax': cfg.get('lmax'),
-            'Group (parity)': 'O3' if cfg.get('is_parity') else 'SO3',
-            'Interaction layers': cfg.get('num_convolution_layer'),
-            'Self connection type': cfg.get('self_connection_type', 'nequip'),
-            'Last epoch': self.epoch,
-            'Elements': len(cfg.get('chemical_species', [])),
-        }
-        if cfg.get('use_modality', False):
-            dct['Modality'] = ', '.join(list(cfg.get('_modal_map', {}).keys()))
-
-        df = pd.DataFrame.from_dict([dct]).T  # type: ignore
-        df.columns = ['']
-        return df.to_string()
-
-    @property
-    def checkpoint_path(self) -> str:
-        return str(self._checkpoint_path)
-
-    @property
-    def config(self) -> Dict[str, Any]:
-        if not self._loaded:
-            self._load()
-        assert isinstance(self._config, dict)
-        return deepcopy(self._config)
-
-    @property
-    def model_state_dict(self) -> Dict[str, Any]:
-        if not self._loaded:
-            self._load()
-        assert isinstance(self._model_state_dict, dict)
-        return copy_state_dict(self._model_state_dict)
-
-    @property
-    def optimizer_state_dict(self) -> Dict[str, Any]:
-        if not self._loaded:
-            self._load()
-        assert isinstance(self._optimizer_state_dict, dict)
-        return copy_state_dict(self._optimizer_state_dict)
-
-    @property
-    def scheduler_state_dict(self) -> Dict[str, Any]:
-        if not self._loaded:
-            self._load()
-        assert isinstance(self._scheduler_state_dict, dict)
-        return copy_state_dict(self._scheduler_state_dict)
-
-    @property
-    def epoch(self) -> Optional[int]:
-        if not self._loaded:
-            self._load()
-        return self._epoch
-
-    @property
-    def time(self) -> str:
-        if not self._loaded:
-            self._load()
-        assert isinstance(self._time, str)
-        return self._time
-
-    @property
-    def hash(self) -> str:
-        if not self._loaded:
-            self._load()
-        assert isinstance(self._hash, str)
-        return self._hash
-
-    def _load(self) -> None:
-        assert not self._loaded
-        cp_path = self.checkpoint_path  # just alias
-
-        cp = torch_load(cp_path, weights_only=False, map_location='cpu')
-        self._config_original = cp.get('config', {})
-        self._model_state_dict = cp.get('model_state_dict', {})
-        self._optimizer_state_dict = cp.get('optimizer_state_dict', {})
-        self._scheduler_state_dict = cp.get('scheduler_state_dict', {})
-        self._epoch = cp.get('epoch', None)
-        self._time = cp.get('time', 'Not found')
-        self._hash = cp.get('hash', 'Not found')
-
-        if len(self._config_original) == 0:
-            warnings.warn(f'config is not found from {cp_path}')
-            self._config = {}
-        else:
-            self._config = _config_cp_routine(self._config_original)
-
-        if len(self._model_state_dict) == 0:
-            warnings.warn(f'model_state_dict is not found from {cp_path}')
-
-        self._loaded = True
-
-    def build_model(self, backend: Optional[str] = None) -> AtomGraphSequential:
-        from .model_build import build_E3_equivariant_model
-
-        use_cue = not backend or backend.lower() in ['cue', 'cueq']
-        try:
-            cp_using_cue = self.config[KEY.CUEQUIVARIANCE_CONFIG]['use']
-        except KeyError:
-            cp_using_cue = False
-
-        if (not backend) or (use_cue == cp_using_cue):
-            # backend not given, or checkpoint backend is same as requested
-            model = build_E3_equivariant_model(self.config)
-            state_dict = compat.patch_state_dict_if_old(
-                self.model_state_dict, self.config, model
-            )
-        else:
-            cfg_new = self.config
-            cfg_new[KEY.CUEQUIVARIANCE_CONFIG] = {'use': use_cue}
-            model = build_E3_equivariant_model(cfg_new)
-            stct_src = compat.patch_state_dict_if_old(
-                self.model_state_dict, self.config, model
-            )
-            state_dict = _convert_e3nn_and_cueq(
-                stct_src, model.state_dict(), self.config, from_cueq=cp_using_cue
-            )
-
-        missing, not_used = model.load_state_dict(state_dict, strict=False)
-        if len(not_used) > 0:
-            warnings.warn(f'Some keys are not used: {not_used}', UserWarning)
-
-        assert len(missing) == 0, f'Missing keys: {missing}'
-        return model
-
-    def yaml_dict(self, mode: str) -> dict:
-        """
-        Return dict for input.yaml from checkpoint config
-        Dataset paths and statistic values are removed intentionally
-        """
-        if mode not in ['reproduce', 'continue', 'continue_modal']:
-            raise ValueError(f'Unknown mode: {mode}')
-
-        ignore = [
-            'when',
-            KEY.DDP_BACKEND,
-            KEY.LOCAL_RANK,
-            KEY.IS_DDP,
-            KEY.DEVICE,
-            KEY.MODEL_TYPE,
-            KEY.SHIFT,
-            KEY.SCALE,
-            KEY.CONV_DENOMINATOR,
-            KEY.SAVE_DATASET,
-            KEY.SAVE_BY_LABEL,
-            KEY.SAVE_BY_TRAIN_VALID,
-            KEY.CONTINUE,
-            KEY.LOAD_DATASET,  # old
-        ]
-
-        cfg = self.config
-        len_atoms = len(cfg[KEY.TYPE_MAP])
-
-        world_size = cfg.pop(KEY.WORLD_SIZE, 1)
-        cfg[KEY.BATCH_SIZE] = cfg[KEY.BATCH_SIZE] * world_size
-        cfg[KEY.LOAD_TRAINSET] = '**path_to_training_set**'
-
-        major, minor, _ = cfg.pop('version', '0.0.0').split('.')[:3]
-        if int(major) == 0 and int(minor) <= 9:
-            warnings.warn('checkpoint version too old, yaml may wrong')
-
-        ret = {'model': {}, 'train': {}, 'data': {}}
-        for k, v in cfg.items():
-            if k.startswith('_') or k in ignore or k.endswith('set_path'):
-                continue
-            if k in consts.DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG:
-                ret['model'][k] = v
-            elif k in consts.DEFAULT_TRAINING_CONFIG:
-                ret['train'][k] = v
-            elif k in consts.DEFAULT_DATA_CONFIG:
-                ret['data'][k] = v
-
-        ret['model'][KEY.CHEMICAL_SPECIES] = (
-            'univ' if len_atoms == consts.NUM_UNIV_ELEMENT else 'auto'
-        )
-        ret['data'][KEY.LOAD_TRAINSET] = '**path_to_trainset**'
-        ret['data'][KEY.LOAD_VALIDSET] = '**path_to_validset**'
-
-        # TODO
-        ret['data'][KEY.SHIFT] = '**failed to infer shift, should be set**'
-        ret['data'][KEY.SCALE] = '**failed to infer scale, should be set**'
-
-        if mode.startswith('continue'):
-            ret['train'].update(
-                {KEY.CONTINUE: {KEY.CHECKPOINT: self.checkpoint_path}}
-            )
-        modal_names = None
-        if mode == 'continue_modal' and not cfg.get(KEY.USE_MODALITY, False):
-            ret['train'][KEY.USE_MODALITY] = True
-
-            # suggest defaults
-            ret['model'][KEY.USE_MODAL_NODE_EMBEDDING] = False
-            ret['model'][KEY.USE_MODAL_SELF_INTER_INTRO] = True
-            ret['model'][KEY.USE_MODAL_SELF_INTER_OUTRO] = True
-            ret['model'][KEY.USE_MODAL_OUTPUT_BLOCK] = True
-
-            ret['data'][KEY.USE_MODAL_WISE_SHIFT] = True
-            ret['data'][KEY.USE_MODAL_WISE_SCALE] = False
-
-            modal_names = ['my_modal1', 'my_modal2']
-        elif cfg.get(KEY.USE_MODALITY, False):
-            modal_names = list(cfg[KEY.MODAL_MAP].keys())
-
-        if modal_names:
-            ret['data'][KEY.LOAD_TRAINSET] = [
-                {'data_modality': mm, 'file_list': [{'file': f'**path_to_{mm}**'}]}
-                for mm in modal_names
-            ]
-
-        return ret
-
-    def append_modal(
-        self,
-        dst_config,
-        original_modal_name: str = 'origin',
-        working_dir: str = os.getcwd(),
-    ):
-        """ """
-        import sevenn.train.modal_dataset as modal_dataset
-        from sevenn.model_build import init_shift_scale
-        from sevenn.scripts.convert_model_modality import _append_modal_weight
-
-        src_config = self.config
-        src_has_no_modal = not src_config.get(KEY.USE_MODALITY, False)
-
-        # inherit element things first
-        chem_keys = [
-            KEY.TYPE_MAP,
-            KEY.NUM_SPECIES,
-            KEY.CHEMICAL_SPECIES,
-            KEY.CHEMICAL_SPECIES_BY_ATOMIC_NUMBER,
-        ]
-        dst_config.update({k: src_config[k] for k in chem_keys})
-
-        if dst_config[KEY.USE_MODAL_WISE_SHIFT] and (
-            KEY.SHIFT not in dst_config or not isinstance(dst_config[KEY.SHIFT], str)
-        ):
-            raise ValueError('To use modal wise shift, keyword shift is required')
-        if dst_config[KEY.USE_MODAL_WISE_SCALE] and (
-            KEY.SCALE not in dst_config or not isinstance(dst_config[KEY.SCALE], str)
-        ):
-            raise ValueError('To use modal wise scale, keyword scale is required')
-
-        if src_has_no_modal and not dst_config[KEY.USE_MODAL_WISE_SHIFT]:
-            dst_config[KEY.SHIFT] = src_config[KEY.SHIFT]
-        if src_has_no_modal and not dst_config[KEY.USE_MODAL_WISE_SCALE]:
-            dst_config[KEY.SCALE] = src_config[KEY.SCALE]
-
-        # get statistics of given datasets of yaml
-        # dst_config updated
-        _ = modal_dataset.from_config(dst_config, working_dir=working_dir)
-        dst_modal_map = dst_config[KEY.MODAL_MAP]
-
-        found_modal_names = list(dst_modal_map.keys())
-        if len(found_modal_names) == 0:
-            raise ValueError('No modality is found from config')
-
-        # Check difference btw given modals and new modal map
-        orig_modal_map = src_config.get(KEY.MODAL_MAP, {original_modal_name: 0})
-        assert isinstance(orig_modal_map, dict)
-        new_modal_map = orig_modal_map.copy()
-        for modal_name in found_modal_names:
-            if modal_name in orig_modal_map:  # duplicate, skipping
-                continue
-            new_modal_map[modal_name] = len(new_modal_map)  # assign new
-        print(f'New modals: {list(new_modal_map.keys())}')
-
-        if src_has_no_modal:
-            append_num = len(new_modal_map)
-        else:
-            append_num = len(new_modal_map) - len(orig_modal_map)
-        if append_num == 0:
-            raise ValueError('Nothing to append from checkpoint')
-
-        dst_config[KEY.NUM_MODALITIES] = len(new_modal_map)
-        dst_config[KEY.MODAL_MAP] = new_modal_map
-
-        # update dst_config's shift scales based on src_config
-        for ss_key, use_mw in (
-            (KEY.SHIFT, dst_config[KEY.USE_MODAL_WISE_SHIFT]),
-            (KEY.SCALE, dst_config[KEY.USE_MODAL_WISE_SCALE]),
-        ):
-            if not use_mw:  # not using mw ss, just assign
-                assert not isinstance(dst_config[ss_key], dict)
-                dst_config[ss_key] = src_config[ss_key]
-            elif src_has_no_modal:
-                assert isinstance(dst_config[ss_key], dict)
-                # mw ss, update by dict but use original_modal_name
-                dst_config[ss_key].update({original_modal_name: src_config[ss_key]})
-            else:
-                assert isinstance(dst_config[ss_key], dict)
-                # mw ss, update by dict
-                dst_config[ss_key].update(src_config[ss_key])
-        scaler = init_shift_scale(dst_config)
-
-        # finally, prepare updated continuable state dict using above
-        orig_model = self.build_model()
-        orig_state_dict = orig_model.state_dict()
-
-        new_state_dict = copy_state_dict(orig_state_dict)
-        for stct_key in orig_state_dict:
-            sp = stct_key.split('.')
-            k, follower = sp[0], '.'.join(sp[1:])
-            if k == 'rescale_atomic_energy' and follower == 'shift':
-                new_state_dict[stct_key] = scaler.shift.clone()
-            elif k == 'rescale_atomic_energy' and follower == 'scale':
-                new_state_dict[stct_key] = scaler.scale.clone()
-            elif follower == 'linear.weight' and (  # append linear layer
-                (
-                    dst_config[KEY.USE_MODAL_NODE_EMBEDDING]
-                    and k.endswith('onehot_to_feature_x')
-                )
-                or (
-                    dst_config[KEY.USE_MODAL_SELF_INTER_INTRO]
-                    and k.endswith('self_interaction_1')
-                )
-                or (
-                    dst_config[KEY.USE_MODAL_SELF_INTER_OUTRO]
-                    and k.endswith('self_interaction_2')
-                )
-                or (
-                    dst_config[KEY.USE_MODAL_OUTPUT_BLOCK]
-                    and k == 'reduce_input_to_hidden'
-                )
-            ):
-                orig_linear = getattr(orig_model._modules[k], 'linear')
-                # assert normalization element
-                new_state_dict[stct_key] = _append_modal_weight(
-                    orig_state_dict,
-                    k,
-                    orig_linear.irreps_in,
-                    orig_linear.irreps_out,
-                    append_num,
-                )
-
-        dst_config['version'] = sevenn.__version__
-
-        return new_state_dict
-
-    def get_checkpoint_dict(self) -> dict:
-        """
-        Return duplicate of this checkpoint with new hash and time.
-        Convenient for creating variant of the checkpoint
-        """
-        return {
-            'config': self.config,
-            'epoch': self.epoch,
-            'model_state_dict': self.model_state_dict,
-            'optimizer_state_dict': self.optimizer_state_dict,
-            'scheduler_state_dict': self.scheduler_state_dict,
-            'time': datetime.now().strftime('%Y-%m-%d %H:%M'),
-            'hash': uuid.uuid4().hex,
-        }
+import os
+import pathlib
+import uuid
+import warnings
+from copy import deepcopy
+from datetime import datetime
+from typing import Any, Dict, Optional, Union
+
+import pandas as pd
+from packaging.version import Version
+from torch import Tensor
+from torch import load as torch_load
+
+import sevenn
+import sevenn._const as consts
+import sevenn._keys as KEY
+import sevenn.scripts.backward_compatibility as compat
+from sevenn import model_build
+from sevenn.nn.scale import get_resolved_shift_scale
+from sevenn.nn.sequential import AtomGraphSequential
+
+
+def assert_atoms(atoms1, atoms2, rtol=1e-5, atol=1e-6):
+    import numpy as np
+
+    def acl(a, b, rtol=rtol, atol=atol):
+        return np.allclose(a, b, rtol=rtol, atol=atol)
+
+    assert len(atoms1) == len(atoms2)
+    assert acl(atoms1.get_cell(), atoms2.get_cell())
+    assert acl(atoms1.get_potential_energy(), atoms2.get_potential_energy())
+    assert acl(atoms1.get_forces(), atoms2.get_forces(), rtol * 10, atol * 10)
+    assert acl(
+        atoms1.get_stress(voigt=False),
+        atoms2.get_stress(voigt=False),
+        rtol * 10,
+        atol * 10,
+    )
+    # assert acl(atoms1.get_potential_energies(), atoms2.get_potential_energies())
+
+
+def copy_state_dict(state_dict) -> dict:
+    if isinstance(state_dict, dict):
+        return {key: copy_state_dict(value) for key, value in state_dict.items()}
+    elif isinstance(state_dict, list):
+        return [copy_state_dict(item) for item in state_dict]  # type: ignore
+    elif isinstance(state_dict, Tensor):
+        return state_dict.clone()  # type: ignore
+    else:
+        # For non-tensor values (e.g., scalars, None), return as-is
+        return state_dict
+
+
+def _config_cp_routine(config):
+    cp_ver = Version(config.get('version', None))
+    this_ver = Version(sevenn.__version__)
+    if cp_ver > this_ver:
+        warnings.warn(f'The checkpoint version ({cp_ver}) is newer than this source'
+                      f'({this_ver}). This may cause unexpected behaviors')
+
+    defaults = {**consts.model_defaults(config)}
+    config = compat.patch_old_config(config)  # type: ignore
+
+    scaler = model_build.init_shift_scale(config)
+    shift, scale = get_resolved_shift_scale(
+        scaler, config.get(KEY.TYPE_MAP), config.get(KEY.MODAL_MAP, None)
+    )
+    config['shift'] = shift
+    config['scale'] = scale
+
+    for k, v in defaults.items():
+        if k in config:
+            continue
+        if os.getenv('SEVENN_DEBUG', False):
+            warnings.warn(f'{k} not in config, use default value {v}', UserWarning)
+        config[k] = v
+
+    for k, v in config.items():
+        if isinstance(v, Tensor):
+            config[k] = v.cpu()
+    return config
+
+
+def _convert_e3nn_and_cueq(stct_src, stct_dst, src_config, from_cueq):
+    """
+    manually check keys and assert if something unexpected happens
+    """
+    n_layer = src_config['num_convolution_layer']
+
+    linear_module_names = [
+        'onehot_to_feature_x',
+        'reduce_input_to_hidden',
+        'reduce_hidden_to_energy',
+    ]
+    convolution_module_names = []
+    fc_tensor_product_module_names = []
+    for i in range(n_layer):
+        linear_module_names.append(f'{i}_self_interaction_1')
+        linear_module_names.append(f'{i}_self_interaction_2')
+        if src_config.get(KEY.SELF_CONNECTION_TYPE) == 'linear':
+            linear_module_names.append(f'{i}_self_connection_intro')
+        elif src_config.get(KEY.SELF_CONNECTION_TYPE) == 'nequip':
+            fc_tensor_product_module_names.append(f'{i}_self_connection_intro')
+        convolution_module_names.append(f'{i}_convolution')
+
+    # Rule: those keys can be safely ignored before state dict load,
+    #       except for linear.bias. This should be aborted in advance to
+    #       this function. Others are not parameters but constants.
+    cue_only_linear_followers = ['linear.f.tp.f_fx.module.c']
+    e3nn_only_linear_followers = ['linear.bias', 'linear.output_mask']
+    ignores_in_linear = cue_only_linear_followers + e3nn_only_linear_followers
+
+    cue_only_conv_followers = [
+        'convolution.f.tp.f_fx.module.c',
+        'convolution.f.tp.module.module.f.module.module._f.data',
+    ]
+    e3nn_only_conv_followers = [
+        'convolution._compiled_main_left_right._w3j',
+        'convolution.weight',
+        'convolution.output_mask',
+    ]
+    ignores_in_conv = cue_only_conv_followers + e3nn_only_conv_followers
+
+    cue_only_fc_followers = ['fc_tensor_product.f.tp.f_fx.module.c']
+    e3nn_only_fc_followers = [
+        'fc_tensor_product.output_mask',
+    ]
+    ignores_in_fc = cue_only_fc_followers + e3nn_only_fc_followers
+
+    updated_keys = []
+    for k, v in stct_src.items():
+        module_name = k.split('.')[0]
+        flag = False
+        if module_name in linear_module_names:
+            for ignore in ignores_in_linear:
+                if '.'.join([module_name, ignore]) in k:
+                    flag = True
+                    break
+            if not flag and k == '.'.join([module_name, 'linear.weight']):
+                updated_keys.append(k)
+                stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
+                flag = True
+            assert flag, f'Unexpected key from linear: {k}'
+        elif module_name in convolution_module_names:
+            for ignore in ignores_in_conv:
+                if '.'.join([module_name, ignore]) in k:
+                    flag = True
+                    break
+            if not flag and (
+                k.startswith(f'{module_name}.weight_nn')
+                or k == '.'.join([module_name, 'denominator'])
+            ):
+                updated_keys.append(k)
+                stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
+                flag = True
+            assert flag, f'Unexpected key from linear: {k}'
+        elif module_name in fc_tensor_product_module_names:
+            for ignore in ignores_in_fc:
+                if '.'.join([module_name, ignore]) in k:
+                    flag = True
+                    break
+            if not flag and k == '.'.join([module_name, 'fc_tensor_product.weight']):
+                updated_keys.append(k)
+                stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
+                flag = True
+            assert flag, f'Unexpected key from fc tensor product: {k}'
+        else:
+            # assert k in stct_dst
+            updated_keys.append(k)
+            stct_dst[k] = v.clone().reshape(stct_dst[k].shape)
+
+    return stct_dst
+
+
+class SevenNetCheckpoint:
+    """
+    Tool box for checkpoint processed from SevenNet.
+    """
+
+    def __init__(self, checkpoint_path: Union[pathlib.Path, str]):
+        self._checkpoint_path = os.path.abspath(checkpoint_path)
+        self._config = None
+        self._epoch = None
+        self._model_state_dict = None
+        self._optimizer_state_dict = None
+        self._scheduler_state_dict = None
+        self._hash = None
+        self._time = None
+
+        self._loaded = False
+
+    def __repr__(self) -> str:
+        cfg = self.config  # just alias
+        if len(cfg) == 0:
+            return ''
+        dct = {
+            'Sevennet version': cfg.get('version', 'Not found'),
+            'When': self.time,
+            'Hash': self.hash,
+            'Cutoff': cfg.get('cutoff'),
+            'Channel': cfg.get('channel'),
+            'Lmax': cfg.get('lmax'),
+            'Group (parity)': 'O3' if cfg.get('is_parity') else 'SO3',
+            'Interaction layers': cfg.get('num_convolution_layer'),
+            'Self connection type': cfg.get('self_connection_type', 'nequip'),
+            'Last epoch': self.epoch,
+            'Elements': len(cfg.get('chemical_species', [])),
+        }
+        if cfg.get('use_modality', False):
+            dct['Modality'] = ', '.join(list(cfg.get('_modal_map', {}).keys()))
+
+        df = pd.DataFrame.from_dict([dct]).T  # type: ignore
+        df.columns = ['']
+        return df.to_string()
+
+    @property
+    def checkpoint_path(self) -> str:
+        return str(self._checkpoint_path)
+
+    @property
+    def config(self) -> Dict[str, Any]:
+        if not self._loaded:
+            self._load()
+        assert isinstance(self._config, dict)
+        return deepcopy(self._config)
+
+    @property
+    def model_state_dict(self) -> Dict[str, Any]:
+        if not self._loaded:
+            self._load()
+        assert isinstance(self._model_state_dict, dict)
+        return copy_state_dict(self._model_state_dict)
+
+    @property
+    def optimizer_state_dict(self) -> Dict[str, Any]:
+        if not self._loaded:
+            self._load()
+        assert isinstance(self._optimizer_state_dict, dict)
+        return copy_state_dict(self._optimizer_state_dict)
+
+    @property
+    def scheduler_state_dict(self) -> Dict[str, Any]:
+        if not self._loaded:
+            self._load()
+        assert isinstance(self._scheduler_state_dict, dict)
+        return copy_state_dict(self._scheduler_state_dict)
+
+    @property
+    def epoch(self) -> Optional[int]:
+        if not self._loaded:
+            self._load()
+        return self._epoch
+
+    @property
+    def time(self) -> str:
+        if not self._loaded:
+            self._load()
+        assert isinstance(self._time, str)
+        return self._time
+
+    @property
+    def hash(self) -> str:
+        if not self._loaded:
+            self._load()
+        assert isinstance(self._hash, str)
+        return self._hash
+
+    def _load(self) -> None:
+        assert not self._loaded
+        cp_path = self.checkpoint_path  # just alias
+
+        cp = torch_load(cp_path, weights_only=False, map_location='cpu')
+        self._config_original = cp.get('config', {})
+        self._model_state_dict = cp.get('model_state_dict', {})
+        self._optimizer_state_dict = cp.get('optimizer_state_dict', {})
+        self._scheduler_state_dict = cp.get('scheduler_state_dict', {})
+        self._epoch = cp.get('epoch', None)
+        self._time = cp.get('time', 'Not found')
+        self._hash = cp.get('hash', 'Not found')
+
+        if len(self._config_original) == 0:
+            warnings.warn(f'config is not found from {cp_path}')
+            self._config = {}
+        else:
+            self._config = _config_cp_routine(self._config_original)
+
+        if len(self._model_state_dict) == 0:
+            warnings.warn(f'model_state_dict is not found from {cp_path}')
+
+        self._loaded = True
+
+    def build_model(self, backend: Optional[str] = None) -> AtomGraphSequential:
+        from .model_build import build_E3_equivariant_model
+
+        use_cue = not backend or backend.lower() in ['cue', 'cueq']
+        try:
+            cp_using_cue = self.config[KEY.CUEQUIVARIANCE_CONFIG]['use']
+        except KeyError:
+            cp_using_cue = False
+
+        if (not backend) or (use_cue == cp_using_cue):
+            # backend not given, or checkpoint backend is same as requested
+            model = build_E3_equivariant_model(self.config)
+            state_dict = compat.patch_state_dict_if_old(
+                self.model_state_dict, self.config, model
+            )
+        else:
+            cfg_new = self.config
+            cfg_new[KEY.CUEQUIVARIANCE_CONFIG] = {'use': use_cue}
+            model = build_E3_equivariant_model(cfg_new)
+            stct_src = compat.patch_state_dict_if_old(
+                self.model_state_dict, self.config, model
+            )
+            state_dict = _convert_e3nn_and_cueq(
+                stct_src, model.state_dict(), self.config, from_cueq=cp_using_cue
+            )
+
+        missing, not_used = model.load_state_dict(state_dict, strict=False)
+        if len(not_used) > 0:
+            warnings.warn(f'Some keys are not used: {not_used}', UserWarning)
+
+        assert len(missing) == 0, f'Missing keys: {missing}'
+        return model
+
+    def yaml_dict(self, mode: str) -> dict:
+        """
+        Return dict for input.yaml from checkpoint config
+        Dataset paths and statistic values are removed intentionally
+        """
+        if mode not in ['reproduce', 'continue', 'continue_modal']:
+            raise ValueError(f'Unknown mode: {mode}')
+
+        ignore = [
+            'when',
+            KEY.DDP_BACKEND,
+            KEY.LOCAL_RANK,
+            KEY.IS_DDP,
+            KEY.DEVICE,
+            KEY.MODEL_TYPE,
+            KEY.SHIFT,
+            KEY.SCALE,
+            KEY.CONV_DENOMINATOR,
+            KEY.SAVE_DATASET,
+            KEY.SAVE_BY_LABEL,
+            KEY.SAVE_BY_TRAIN_VALID,
+            KEY.CONTINUE,
+            KEY.LOAD_DATASET,  # old
+        ]
+
+        cfg = self.config
+        len_atoms = len(cfg[KEY.TYPE_MAP])
+
+        world_size = cfg.pop(KEY.WORLD_SIZE, 1)
+        cfg[KEY.BATCH_SIZE] = cfg[KEY.BATCH_SIZE] * world_size
+        cfg[KEY.LOAD_TRAINSET] = '**path_to_training_set**'
+
+        major, minor, _ = cfg.pop('version', '0.0.0').split('.')[:3]
+        if int(major) == 0 and int(minor) <= 9:
+            warnings.warn('checkpoint version too old, yaml may wrong')
+
+        ret = {'model': {}, 'train': {}, 'data': {}}
+        for k, v in cfg.items():
+            if k.startswith('_') or k in ignore or k.endswith('set_path'):
+                continue
+            if k in consts.DEFAULT_E3_EQUIVARIANT_MODEL_CONFIG:
+                ret['model'][k] = v
+            elif k in consts.DEFAULT_TRAINING_CONFIG:
+                ret['train'][k] = v
+            elif k in consts.DEFAULT_DATA_CONFIG:
+                ret['data'][k] = v
+
+        ret['model'][KEY.CHEMICAL_SPECIES] = (
+            'univ' if len_atoms == consts.NUM_UNIV_ELEMENT else 'auto'
+        )
+        ret['data'][KEY.LOAD_TRAINSET] = '**path_to_trainset**'
+        ret['data'][KEY.LOAD_VALIDSET] = '**path_to_validset**'
+
+        # TODO
+        ret['data'][KEY.SHIFT] = '**failed to infer shift, should be set**'
+        ret['data'][KEY.SCALE] = '**failed to infer scale, should be set**'
+
+        if mode.startswith('continue'):
+            ret['train'].update(
+                {KEY.CONTINUE: {KEY.CHECKPOINT: self.checkpoint_path}}
+            )
+        modal_names = None
+        if mode == 'continue_modal' and not cfg.get(KEY.USE_MODALITY, False):
+            ret['train'][KEY.USE_MODALITY] = True
+
+            # suggest defaults
+            ret['model'][KEY.USE_MODAL_NODE_EMBEDDING] = False
+            ret['model'][KEY.USE_MODAL_SELF_INTER_INTRO] = True
+            ret['model'][KEY.USE_MODAL_SELF_INTER_OUTRO] = True
+            ret['model'][KEY.USE_MODAL_OUTPUT_BLOCK] = True
+
+            ret['data'][KEY.USE_MODAL_WISE_SHIFT] = True
+            ret['data'][KEY.USE_MODAL_WISE_SCALE] = False
+
+            modal_names = ['my_modal1', 'my_modal2']
+        elif cfg.get(KEY.USE_MODALITY, False):
+            modal_names = list(cfg[KEY.MODAL_MAP].keys())
+
+        if modal_names:
+            ret['data'][KEY.LOAD_TRAINSET] = [
+                {'data_modality': mm, 'file_list': [{'file': f'**path_to_{mm}**'}]}
+                for mm in modal_names
+            ]
+
+        return ret
+
+    def append_modal(
+        self,
+        dst_config,
+        original_modal_name: str = 'origin',
+        working_dir: str = os.getcwd(),
+    ):
+        """ """
+        import sevenn.train.modal_dataset as modal_dataset
+        from sevenn.model_build import init_shift_scale
+        from sevenn.scripts.convert_model_modality import _append_modal_weight
+
+        src_config = self.config
+        src_has_no_modal = not src_config.get(KEY.USE_MODALITY, False)
+
+        # inherit element things first
+        chem_keys = [
+            KEY.TYPE_MAP,
+            KEY.NUM_SPECIES,
+            KEY.CHEMICAL_SPECIES,
+            KEY.CHEMICAL_SPECIES_BY_ATOMIC_NUMBER,
+        ]
+        dst_config.update({k: src_config[k] for k in chem_keys})
+
+        if dst_config[KEY.USE_MODAL_WISE_SHIFT] and (
+            KEY.SHIFT not in dst_config or not isinstance(dst_config[KEY.SHIFT], str)
+        ):
+            raise ValueError('To use modal wise shift, keyword shift is required')
+        if dst_config[KEY.USE_MODAL_WISE_SCALE] and (
+            KEY.SCALE not in dst_config or not isinstance(dst_config[KEY.SCALE], str)
+        ):
+            raise ValueError('To use modal wise scale, keyword scale is required')
+
+        if src_has_no_modal and not dst_config[KEY.USE_MODAL_WISE_SHIFT]:
+            dst_config[KEY.SHIFT] = src_config[KEY.SHIFT]
+        if src_has_no_modal and not dst_config[KEY.USE_MODAL_WISE_SCALE]:
+            dst_config[KEY.SCALE] = src_config[KEY.SCALE]
+
+        # get statistics of given datasets of yaml
+        # dst_config updated
+        _ = modal_dataset.from_config(dst_config, working_dir=working_dir)
+        dst_modal_map = dst_config[KEY.MODAL_MAP]
+
+        found_modal_names = list(dst_modal_map.keys())
+        if len(found_modal_names) == 0:
+            raise ValueError('No modality is found from config')
+
+        # Check difference btw given modals and new modal map
+        orig_modal_map = src_config.get(KEY.MODAL_MAP, {original_modal_name: 0})
+        assert isinstance(orig_modal_map, dict)
+        new_modal_map = orig_modal_map.copy()
+        for modal_name in found_modal_names:
+            if modal_name in orig_modal_map:  # duplicate, skipping
+                continue
+            new_modal_map[modal_name] = len(new_modal_map)  # assign new
+        print(f'New modals: {list(new_modal_map.keys())}')
+
+        if src_has_no_modal:
+            append_num = len(new_modal_map)
+        else:
+            append_num = len(new_modal_map) - len(orig_modal_map)
+        if append_num == 0:
+            raise ValueError('Nothing to append from checkpoint')
+
+        dst_config[KEY.NUM_MODALITIES] = len(new_modal_map)
+        dst_config[KEY.MODAL_MAP] = new_modal_map
+
+        # update dst_config's shift scales based on src_config
+        for ss_key, use_mw in (
+            (KEY.SHIFT, dst_config[KEY.USE_MODAL_WISE_SHIFT]),
+            (KEY.SCALE, dst_config[KEY.USE_MODAL_WISE_SCALE]),
+        ):
+            if not use_mw:  # not using mw ss, just assign
+                assert not isinstance(dst_config[ss_key], dict)
+                dst_config[ss_key] = src_config[ss_key]
+            elif src_has_no_modal:
+                assert isinstance(dst_config[ss_key], dict)
+                # mw ss, update by dict but use original_modal_name
+                dst_config[ss_key].update({original_modal_name: src_config[ss_key]})
+            else:
+                assert isinstance(dst_config[ss_key], dict)
+                # mw ss, update by dict
+                dst_config[ss_key].update(src_config[ss_key])
+        scaler = init_shift_scale(dst_config)
+
+        # finally, prepare updated continuable state dict using above
+        orig_model = self.build_model()
+        orig_state_dict = orig_model.state_dict()
+
+        new_state_dict = copy_state_dict(orig_state_dict)
+        for stct_key in orig_state_dict:
+            sp = stct_key.split('.')
+            k, follower = sp[0], '.'.join(sp[1:])
+            if k == 'rescale_atomic_energy' and follower == 'shift':
+                new_state_dict[stct_key] = scaler.shift.clone()
+            elif k == 'rescale_atomic_energy' and follower == 'scale':
+                new_state_dict[stct_key] = scaler.scale.clone()
+            elif follower == 'linear.weight' and (  # append linear layer
+                (
+                    dst_config[KEY.USE_MODAL_NODE_EMBEDDING]
+                    and k.endswith('onehot_to_feature_x')
+                )
+                or (
+                    dst_config[KEY.USE_MODAL_SELF_INTER_INTRO]
+                    and k.endswith('self_interaction_1')
+                )
+                or (
+                    dst_config[KEY.USE_MODAL_SELF_INTER_OUTRO]
+                    and k.endswith('self_interaction_2')
+                )
+                or (
+                    dst_config[KEY.USE_MODAL_OUTPUT_BLOCK]
+                    and k == 'reduce_input_to_hidden'
+                )
+            ):
+                orig_linear = getattr(orig_model._modules[k], 'linear')
+                # assert normalization element
+                new_state_dict[stct_key] = _append_modal_weight(
+                    orig_state_dict,
+                    k,
+                    orig_linear.irreps_in,
+                    orig_linear.irreps_out,
+                    append_num,
+                )
+
+        dst_config['version'] = sevenn.__version__
+
+        return new_state_dict
+
+    def get_checkpoint_dict(self) -> dict:
+        """
+        Return duplicate of this checkpoint with new hash and time.
+        Convenient for creating variant of the checkpoint
+        """
+        return {
+            'config': self.config,
+            'epoch': self.epoch,
+            'model_state_dict': self.model_state_dict,
+            'optimizer_state_dict': self.optimizer_state_dict,
+            'scheduler_state_dict': self.scheduler_state_dict,
+            'time': datetime.now().strftime('%Y-%m-%d %H:%M'),
+            'hash': uuid.uuid4().hex,
+        }

mace-bench/3rdparty/SevenNet/sevenn/error_recorder.py

-from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Tuple
-
-import torch
-import torch.distributed as dist
-
-import sevenn._keys as KEY
-from sevenn.train.loss import LossDefinition
-
-from .atom_graph_data import AtomGraphData
-from .train.optim import loss_dict
-
-_ERROR_TYPES = {
-    'TotalEnergy': {
-        'name': 'Energy',
-        'ref_key': KEY.ENERGY,
-        'pred_key': KEY.PRED_TOTAL_ENERGY,
-        'unit': 'eV',
-        'vdim': 1,
-    },
-    'Energy': {  # by default per-atom for energy
-        'name': 'Energy',
-        'ref_key': KEY.ENERGY,
-        'pred_key': KEY.PRED_TOTAL_ENERGY,
-        'unit': 'eV/atom',
-        'per_atom': True,
-        'vdim': 1,
-    },
-    'Force': {
-        'name': 'Force',
-        'ref_key': KEY.FORCE,
-        'pred_key': KEY.PRED_FORCE,
-        'unit': 'eV/Å',
-        'vdim': 3,
-    },
-    'Stress': {
-        'name': 'Stress',
-        'ref_key': KEY.STRESS,
-        'pred_key': KEY.PRED_STRESS,
-        'unit': 'kbar',
-        'coeff': 1602.1766208,
-        'vdim': 6,
-    },
-    'Stress_GPa': {
-        'name': 'Stress',
-        'ref_key': KEY.STRESS,
-        'pred_key': KEY.PRED_STRESS,
-        'unit': 'GPa',
-        'coeff': 160.21766208,
-        'vdim': 6,
-    },
-    'TotalLoss': {
-        'name': 'TotalLoss',
-        'unit': None,
-    },
-}
-
-
-def get_err_type(name: str) -> Dict[str, Any]:
-    return deepcopy(_ERROR_TYPES[name])
-
-
-def _get_loss_function_from_name(loss_functions, name):
-    for loss_def, w in loss_functions:
-        if loss_def.name.lower() == name.lower():
-            return loss_def, w
-    return None, None
-
-
-class AverageNumber:
-    def __init__(self):
-        self._sum = 0.0
-        self._count = 0
-
-    def update(self, values: torch.Tensor):
-        self._sum += values.sum().item()
-        self._count += values.numel()
-
-    def _ddp_reduce(self, device):
-        _sum = torch.tensor(self._sum, device=device)
-        _count = torch.tensor(self._count, device=device)
-        dist.all_reduce(_sum, op=dist.ReduceOp.SUM)
-        dist.all_reduce(_count, op=dist.ReduceOp.SUM)
-        self._sum = _sum.item()
-        self._count = _count.item()
-
-    def get(self):
-        if self._count == 0:
-            return torch.nan
-        return self._sum / self._count
-
-
-class ErrorMetric:
-    """
-    Base class for error metrics We always average error by # of structures,
-    and designed to collect errors in the middle of iteration (by AverageNumber)
-    """
-
-    def __init__(
-        self,
-        name: str,
-        ref_key: str,
-        pred_key: str,
-        coeff: float = 1.0,
-        unit: Optional[str] = None,
-        per_atom: bool = False,
-        ignore_unlabeled: bool = True,
-        **kwargs,
-    ):
-        self.name = name
-        self.unit = unit
-        self.coeff = coeff
-        self.ref_key = ref_key
-        self.pred_key = pred_key
-        self.per_atom = per_atom
-        self.ignore_unlabeled = ignore_unlabeled
-        self.value = AverageNumber()
-
-    def update(self, output: AtomGraphData):
-        raise NotImplementedError
-
-    def _retrieve(self, output: AtomGraphData):
-        y_ref = output[self.ref_key] * self.coeff
-        y_pred = output[self.pred_key] * self.coeff
-        if self.per_atom:
-            assert y_ref.dim() == 1 and y_pred.dim() == 1
-            natoms = output[KEY.NUM_ATOMS]
-            y_ref = y_ref / natoms
-            y_pred = y_pred / natoms
-        if self.ignore_unlabeled:
-            unlabelled_idx = torch.isnan(y_ref)
-            y_ref = y_ref[~unlabelled_idx]
-            y_pred = y_pred[~unlabelled_idx]
-        return y_ref, y_pred
-
-    def ddp_reduce(self, device):
-        self.value._ddp_reduce(device)
-
-    def reset(self):
-        self.value = AverageNumber()
-
-    def get(self):
-        return self.value.get()
-
-    def key_str(self, with_unit=True):
-        if self.unit is None or not with_unit:
-            return self.name
-        else:
-            return f'{self.name} ({self.unit})'
-
-    def __str__(self):
-        return f'{self.key_str()}: {self.value.get():.6f}'
-
-
-class RMSError(ErrorMetric):
-    """
-    Vector squared error
-    """
-
-    def __init__(self, vdim: int = 1, **kwargs):
-        super().__init__(**kwargs)
-        self.vdim = vdim
-        self._se = torch.nn.MSELoss(reduction='none')
-
-    def _square_error(self, y_ref, y_pred, vdim: int):
-        return self._se(y_ref.view(-1, vdim), y_pred.view(-1, vdim)).sum(dim=1)
-
-    def update(self, output: AtomGraphData):
-        y_ref, y_pred = self._retrieve(output)
-        se = self._square_error(y_ref, y_pred, self.vdim)
-        self.value.update(se)
-
-    def get(self):
-        return self.value.get() ** 0.5
-
-
-class ComponentRMSError(ErrorMetric):
-    """
-    Ignore vector dim and just average over components
-    Results smaller error
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-        self._se = torch.nn.MSELoss(reduction='none')
-
-    def _square_error(self, y_ref, y_pred):
-        return self._se(y_ref, y_pred)
-
-    def update(self, output: AtomGraphData):
-        y_ref, y_pred = self._retrieve(output)
-        y_ref = y_ref.view(-1)
-        y_pred = y_pred.view(-1)
-        se = self._square_error(y_ref, y_pred)
-        self.value.update(se)
-
-    def get(self):
-        return self.value.get() ** 0.5
-
-
-class MAError(ErrorMetric):
-    """
-    Average over all component
-    """
-
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def _square_error(self, y_ref, y_pred):
-        return torch.abs(y_ref - y_pred)
-
-    def update(self, output: AtomGraphData):
-        y_ref, y_pred = self._retrieve(output)
-        y_ref = y_ref.reshape((-1,))
-        y_pred = y_pred.reshape((-1,))
-        se = self._square_error(y_ref, y_pred)
-        self.value.update(se)
-
-
-class CustomError(ErrorMetric):
-    """
-    Custom error metric
-    Args:
-        func: a function that takes y_ref and y_pred
-              and returns a list of errors
-    """
-
-    def __init__(self, func: Callable, **kwargs):
-        super().__init__(**kwargs)
-        self.func = func
-
-    def update(self, output: AtomGraphData):
-        y_ref, y_pred = self._retrieve(output)
-        se = self.func(y_ref, y_pred) if len(y_ref) > 0 else torch.tensor([])
-        self.value.update(se)
-
-
-class LossError(ErrorMetric):
-    """
-    Error metric that record loss
-    """
-
-    def __init__(
-        self,
-        name: str,
-        loss_def: LossDefinition,
-        **kwargs,
-    ):
-        super().__init__(
-            name,
-            ignore_unlabeld=loss_def.ignore_unlabeled,
-            **kwargs,
-        )
-        self.loss_def = loss_def
-
-    def update(self, output: AtomGraphData):
-        loss = self.loss_def.get_loss(output)  # type: ignore
-        self.value.update(loss)  # type: ignore
-
-
-class CombinedError(ErrorMetric):
-    """
-    Combine multiple error metrics with weights
-    corresponds to a weighted sum of errors (normally used in loss)
-    """
-
-    def __init__(self, metrics: List[Tuple[ErrorMetric, float]], **kwargs):
-        super().__init__(**kwargs)
-        self.metrics = metrics
-        assert kwargs['unit'] is None
-
-    def update(self, output: AtomGraphData):
-        for metric, _ in self.metrics:
-            metric.update(output)
-
-    def reset(self):
-        for metric, _ in self.metrics:
-            metric.reset()
-
-    def ddp_reduce(self, device):  # override
-        for metric, _ in self.metrics:
-            metric.value._ddp_reduce(device)
-
-    def get(self):
-        val = 0.0
-        for metric, weight in self.metrics:
-            val += metric.get() * weight
-        return val
-
-
-class ErrorRecorder:
-    """
-    record errors of a model
-    """
-
-    METRIC_DICT = {
-        'RMSE': RMSError,
-        'ComponentRMSE': ComponentRMSError,
-        'MAE': MAError,
-        'Loss': LossError,
-    }
-
-    def __init__(self, metrics: List[ErrorMetric]):
-        self.history = []
-        self.metrics = metrics
-
-    def _update(self, output: AtomGraphData):
-        for metric in self.metrics:
-            metric.update(output)
-
-    def update(self, output: AtomGraphData, no_grad=True):
-        if no_grad:
-            with torch.no_grad():
-                self._update(output)
-        else:
-            self._update(output)
-
-    def get_metric_dict(self, with_unit=True):
-        return {metric.key_str(with_unit): metric.get() for metric in self.metrics}
-
-    def get_current(self):
-        dct = {}
-        for metric in self.metrics:
-            dct[metric.name] = {
-                'value': metric.get(),
-                'unit': metric.unit,
-                'ref_key': metric.ref_key,
-                'pred_key': metric.pred_key,
-            }
-        return dct
-
-    def get_dct(self, prefix=''):
-        dct = {}
-        if prefix.endswith('_') is False and prefix != '':
-            prefix = prefix + '_'
-        for metric in self.metrics:
-            dct[f'{prefix}{metric.name}'] = f'{metric.get():6f}'
-        return dct
-
-    def get_key_str(self, name: str):
-        for metric in self.metrics:
-            if name == metric.name:
-                return metric.key_str()
-        return None
-
-    def epoch_forward(self):
-        self.history.append(self.get_current())
-        pretty = self.get_metric_dict(with_unit=True)
-        for metric in self.metrics:
-            metric.reset()
-        return pretty  # for print
-
-    @staticmethod
-    def init_total_loss_metric(
-        config,
-        criteria: Optional[Callable] = None,
-        loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None,
-    ):
-        if criteria is None and loss_functions is None:
-            raise ValueError('both criteria and loss functions not given')
-
-        is_stress = config[KEY.IS_TRAIN_STRESS]
-        metrics = []
-        if criteria is not None:
-            energy_metric = CustomError(criteria, **get_err_type('Energy'))
-            metrics.append((energy_metric, 1))
-            force_metric = CustomError(criteria, **get_err_type('Force'))
-            metrics.append((force_metric, config[KEY.FORCE_WEIGHT]))
-            if is_stress:
-                stress_metric = CustomError(criteria, **get_err_type('Stress'))
-                metrics.append((stress_metric, config[KEY.STRESS_WEIGHT]))
-        else:  # TODO: this is hard-coded
-            for efs in ['Energy', 'Force', 'Stress']:
-                if efs == 'Stress' and not is_stress:
-                    continue
-                lf, w = _get_loss_function_from_name(loss_functions, efs)
-                if lf is None:
-                    raise ValueError(f'{efs} not found from loss_functions')
-                metric = LossError(loss_def=lf, **get_err_type(efs))
-                metrics.append((metric, w))
-
-        total_loss_metric = CombinedError(
-            metrics, name='TotalLoss', unit=None, ref_key=None, pred_key=None
-        )
-        return total_loss_metric
-
-    @staticmethod
-    def from_config(config: dict, loss_functions=None):
-        loss_cls = loss_dict[config.get(KEY.LOSS, 'mse').lower()]
-        loss_param = config.get(KEY.LOSS_PARAM, {})
-        criteria = loss_cls(**loss_param) if loss_functions is None else None
-
-        err_config = config.get(KEY.ERROR_RECORD, False)
-        if not err_config:
-            raise ValueError(
-                'No error_record config found. Consider util.get_error_recorder'
-            )
-        err_config_n = []
-        if not config.get(KEY.IS_TRAIN_STRESS, True):
-            for err_type, metric_name in err_config:
-                if 'Stress' in err_type:
-                    continue
-                err_config_n.append((err_type, metric_name))
-            err_config = err_config_n
-
-        err_metrics = []
-        for err_type, metric_name in err_config:
-            metric_kwargs = get_err_type(err_type)
-            if err_type == 'TotalLoss':  # special case
-                err_metrics.append(
-                    ErrorRecorder.init_total_loss_metric(
-                        config, criteria, loss_functions
-                    )
-                )
-                continue
-            metric_cls = ErrorRecorder.METRIC_DICT[metric_name]
-            assert isinstance(metric_kwargs['name'], str)
-            if metric_name == 'Loss':
-                if loss_functions is not None:
-                    metric_cls = LossError
-                    metric_kwargs['loss_def'], _ = _get_loss_function_from_name(
-                        loss_functions, metric_kwargs['name']
-                    )
-                else:
-                    metric_cls = CustomError
-                    metric_kwargs['func'] = criteria
-                metric_kwargs.pop('unit', None)
-            metric_kwargs['name'] += f'_{metric_name}'
-            err_metrics.append(metric_cls(**metric_kwargs))
-        return ErrorRecorder(err_metrics)
+from copy import deepcopy
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+
+import sevenn._keys as KEY
+from sevenn.train.loss import LossDefinition
+
+from .atom_graph_data import AtomGraphData
+from .train.optim import loss_dict
+
+_ERROR_TYPES = {
+    'TotalEnergy': {
+        'name': 'Energy',
+        'ref_key': KEY.ENERGY,
+        'pred_key': KEY.PRED_TOTAL_ENERGY,
+        'unit': 'eV',
+        'vdim': 1,
+    },
+    'Energy': {  # by default per-atom for energy
+        'name': 'Energy',
+        'ref_key': KEY.ENERGY,
+        'pred_key': KEY.PRED_TOTAL_ENERGY,
+        'unit': 'eV/atom',
+        'per_atom': True,
+        'vdim': 1,
+    },
+    'Force': {
+        'name': 'Force',
+        'ref_key': KEY.FORCE,
+        'pred_key': KEY.PRED_FORCE,
+        'unit': 'eV/Å',
+        'vdim': 3,
+    },
+    'Stress': {
+        'name': 'Stress',
+        'ref_key': KEY.STRESS,
+        'pred_key': KEY.PRED_STRESS,
+        'unit': 'kbar',
+        'coeff': 1602.1766208,
+        'vdim': 6,
+    },
+    'Stress_GPa': {
+        'name': 'Stress',
+        'ref_key': KEY.STRESS,
+        'pred_key': KEY.PRED_STRESS,
+        'unit': 'GPa',
+        'coeff': 160.21766208,
+        'vdim': 6,
+    },
+    'TotalLoss': {
+        'name': 'TotalLoss',
+        'unit': None,
+    },
+}
+
+
+def get_err_type(name: str) -> Dict[str, Any]:
+    return deepcopy(_ERROR_TYPES[name])
+
+
+def _get_loss_function_from_name(loss_functions, name):
+    for loss_def, w in loss_functions:
+        if loss_def.name.lower() == name.lower():
+            return loss_def, w
+    return None, None
+
+
+class AverageNumber:
+    def __init__(self):
+        self._sum = 0.0
+        self._count = 0
+
+    def update(self, values: torch.Tensor):
+        self._sum += values.sum().item()
+        self._count += values.numel()
+
+    def _ddp_reduce(self, device):
+        _sum = torch.tensor(self._sum, device=device)
+        _count = torch.tensor(self._count, device=device)
+        dist.all_reduce(_sum, op=dist.ReduceOp.SUM)
+        dist.all_reduce(_count, op=dist.ReduceOp.SUM)
+        self._sum = _sum.item()
+        self._count = _count.item()
+
+    def get(self):
+        if self._count == 0:
+            return torch.nan
+        return self._sum / self._count
+
+
+class ErrorMetric:
+    """
+    Base class for error metrics We always average error by # of structures,
+    and designed to collect errors in the middle of iteration (by AverageNumber)
+    """
+
+    def __init__(
+        self,
+        name: str,
+        ref_key: str,
+        pred_key: str,
+        coeff: float = 1.0,
+        unit: Optional[str] = None,
+        per_atom: bool = False,
+        ignore_unlabeled: bool = True,
+        **kwargs,
+    ):
+        self.name = name
+        self.unit = unit
+        self.coeff = coeff
+        self.ref_key = ref_key
+        self.pred_key = pred_key
+        self.per_atom = per_atom
+        self.ignore_unlabeled = ignore_unlabeled
+        self.value = AverageNumber()
+
+    def update(self, output: AtomGraphData):
+        raise NotImplementedError
+
+    def _retrieve(self, output: AtomGraphData):
+        y_ref = output[self.ref_key] * self.coeff
+        y_pred = output[self.pred_key] * self.coeff
+        if self.per_atom:
+            assert y_ref.dim() == 1 and y_pred.dim() == 1
+            natoms = output[KEY.NUM_ATOMS]
+            y_ref = y_ref / natoms
+            y_pred = y_pred / natoms
+        if self.ignore_unlabeled:
+            unlabelled_idx = torch.isnan(y_ref)
+            y_ref = y_ref[~unlabelled_idx]
+            y_pred = y_pred[~unlabelled_idx]
+        return y_ref, y_pred
+
+    def ddp_reduce(self, device):
+        self.value._ddp_reduce(device)
+
+    def reset(self):
+        self.value = AverageNumber()
+
+    def get(self):
+        return self.value.get()
+
+    def key_str(self, with_unit=True):
+        if self.unit is None or not with_unit:
+            return self.name
+        else:
+            return f'{self.name} ({self.unit})'
+
+    def __str__(self):
+        return f'{self.key_str()}: {self.value.get():.6f}'
+
+
+class RMSError(ErrorMetric):
+    """
+    Vector squared error
+    """
+
+    def __init__(self, vdim: int = 1, **kwargs):
+        super().__init__(**kwargs)
+        self.vdim = vdim
+        self._se = torch.nn.MSELoss(reduction='none')
+
+    def _square_error(self, y_ref, y_pred, vdim: int):
+        return self._se(y_ref.view(-1, vdim), y_pred.view(-1, vdim)).sum(dim=1)
+
+    def update(self, output: AtomGraphData):
+        y_ref, y_pred = self._retrieve(output)
+        se = self._square_error(y_ref, y_pred, self.vdim)
+        self.value.update(se)
+
+    def get(self):
+        return self.value.get() ** 0.5
+
+
+class ComponentRMSError(ErrorMetric):
+    """
+    Ignore vector dim and just average over components
+    Results smaller error
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self._se = torch.nn.MSELoss(reduction='none')
+
+    def _square_error(self, y_ref, y_pred):
+        return self._se(y_ref, y_pred)
+
+    def update(self, output: AtomGraphData):
+        y_ref, y_pred = self._retrieve(output)
+        y_ref = y_ref.view(-1)
+        y_pred = y_pred.view(-1)
+        se = self._square_error(y_ref, y_pred)
+        self.value.update(se)
+
+    def get(self):
+        return self.value.get() ** 0.5
+
+
+class MAError(ErrorMetric):
+    """
+    Average over all component
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def _square_error(self, y_ref, y_pred):
+        return torch.abs(y_ref - y_pred)
+
+    def update(self, output: AtomGraphData):
+        y_ref, y_pred = self._retrieve(output)
+        y_ref = y_ref.reshape((-1,))
+        y_pred = y_pred.reshape((-1,))
+        se = self._square_error(y_ref, y_pred)
+        self.value.update(se)
+
+
+class CustomError(ErrorMetric):
+    """
+    Custom error metric
+    Args:
+        func: a function that takes y_ref and y_pred
+              and returns a list of errors
+    """
+
+    def __init__(self, func: Callable, **kwargs):
+        super().__init__(**kwargs)
+        self.func = func
+
+    def update(self, output: AtomGraphData):
+        y_ref, y_pred = self._retrieve(output)
+        se = self.func(y_ref, y_pred) if len(y_ref) > 0 else torch.tensor([])
+        self.value.update(se)
+
+
+class LossError(ErrorMetric):
+    """
+    Error metric that record loss
+    """
+
+    def __init__(
+        self,
+        name: str,
+        loss_def: LossDefinition,
+        **kwargs,
+    ):
+        super().__init__(
+            name,
+            ignore_unlabeld=loss_def.ignore_unlabeled,
+            **kwargs,
+        )
+        self.loss_def = loss_def
+
+    def update(self, output: AtomGraphData):
+        loss = self.loss_def.get_loss(output)  # type: ignore
+        self.value.update(loss)  # type: ignore
+
+
+class CombinedError(ErrorMetric):
+    """
+    Combine multiple error metrics with weights
+    corresponds to a weighted sum of errors (normally used in loss)
+    """
+
+    def __init__(self, metrics: List[Tuple[ErrorMetric, float]], **kwargs):
+        super().__init__(**kwargs)
+        self.metrics = metrics
+        assert kwargs['unit'] is None
+
+    def update(self, output: AtomGraphData):
+        for metric, _ in self.metrics:
+            metric.update(output)
+
+    def reset(self):
+        for metric, _ in self.metrics:
+            metric.reset()
+
+    def ddp_reduce(self, device):  # override
+        for metric, _ in self.metrics:
+            metric.value._ddp_reduce(device)
+
+    def get(self):
+        val = 0.0
+        for metric, weight in self.metrics:
+            val += metric.get() * weight
+        return val
+
+
+class ErrorRecorder:
+    """
+    record errors of a model
+    """
+
+    METRIC_DICT = {
+        'RMSE': RMSError,
+        'ComponentRMSE': ComponentRMSError,
+        'MAE': MAError,
+        'Loss': LossError,
+    }
+
+    def __init__(self, metrics: List[ErrorMetric]):
+        self.history = []
+        self.metrics = metrics
+
+    def _update(self, output: AtomGraphData):
+        for metric in self.metrics:
+            metric.update(output)
+
+    def update(self, output: AtomGraphData, no_grad=True):
+        if no_grad:
+            with torch.no_grad():
+                self._update(output)
+        else:
+            self._update(output)
+
+    def get_metric_dict(self, with_unit=True):
+        return {metric.key_str(with_unit): metric.get() for metric in self.metrics}
+
+    def get_current(self):
+        dct = {}
+        for metric in self.metrics:
+            dct[metric.name] = {
+                'value': metric.get(),
+                'unit': metric.unit,
+                'ref_key': metric.ref_key,
+                'pred_key': metric.pred_key,
+            }
+        return dct
+
+    def get_dct(self, prefix=''):
+        dct = {}
+        if prefix.endswith('_') is False and prefix != '':
+            prefix = prefix + '_'
+        for metric in self.metrics:
+            dct[f'{prefix}{metric.name}'] = f'{metric.get():6f}'
+        return dct
+
+    def get_key_str(self, name: str):
+        for metric in self.metrics:
+            if name == metric.name:
+                return metric.key_str()
+        return None
+
+    def epoch_forward(self):
+        self.history.append(self.get_current())
+        pretty = self.get_metric_dict(with_unit=True)
+        for metric in self.metrics:
+            metric.reset()
+        return pretty  # for print
+
+    @staticmethod
+    def init_total_loss_metric(
+        config,
+        criteria: Optional[Callable] = None,
+        loss_functions: Optional[List[Tuple[LossDefinition, float]]] = None,
+    ):
+        if criteria is None and loss_functions is None:
+            raise ValueError('both criteria and loss functions not given')
+
+        is_stress = config[KEY.IS_TRAIN_STRESS]
+        metrics = []
+        if criteria is not None:
+            energy_metric = CustomError(criteria, **get_err_type('Energy'))
+            metrics.append((energy_metric, 1))
+            force_metric = CustomError(criteria, **get_err_type('Force'))
+            metrics.append((force_metric, config[KEY.FORCE_WEIGHT]))
+            if is_stress:
+                stress_metric = CustomError(criteria, **get_err_type('Stress'))
+                metrics.append((stress_metric, config[KEY.STRESS_WEIGHT]))
+        else:  # TODO: this is hard-coded
+            for efs in ['Energy', 'Force', 'Stress']:
+                if efs == 'Stress' and not is_stress:
+                    continue
+                lf, w = _get_loss_function_from_name(loss_functions, efs)
+                if lf is None:
+                    raise ValueError(f'{efs} not found from loss_functions')
+                metric = LossError(loss_def=lf, **get_err_type(efs))
+                metrics.append((metric, w))
+
+        total_loss_metric = CombinedError(
+            metrics, name='TotalLoss', unit=None, ref_key=None, pred_key=None
+        )
+        return total_loss_metric
+
+    @staticmethod
+    def from_config(config: dict, loss_functions=None):
+        loss_cls = loss_dict[config.get(KEY.LOSS, 'mse').lower()]
+        loss_param = config.get(KEY.LOSS_PARAM, {})
+        criteria = loss_cls(**loss_param) if loss_functions is None else None
+
+        err_config = config.get(KEY.ERROR_RECORD, False)
+        if not err_config:
+            raise ValueError(
+                'No error_record config found. Consider util.get_error_recorder'
+            )
+        err_config_n = []
+        if not config.get(KEY.IS_TRAIN_STRESS, True):
+            for err_type, metric_name in err_config:
+                if 'Stress' in err_type:
+                    continue
+                err_config_n.append((err_type, metric_name))
+            err_config = err_config_n
+
+        err_metrics = []
+        for err_type, metric_name in err_config:
+            metric_kwargs = get_err_type(err_type)
+            if err_type == 'TotalLoss':  # special case
+                err_metrics.append(
+                    ErrorRecorder.init_total_loss_metric(
+                        config, criteria, loss_functions
+                    )
+                )
+                continue
+            metric_cls = ErrorRecorder.METRIC_DICT[metric_name]
+            assert isinstance(metric_kwargs['name'], str)
+            if metric_name == 'Loss':
+                if loss_functions is not None:
+                    metric_cls = LossError
+                    metric_kwargs['loss_def'], _ = _get_loss_function_from_name(
+                        loss_functions, metric_kwargs['name']
+                    )
+                else:
+                    metric_cls = CustomError
+                    metric_kwargs['func'] = criteria
+                metric_kwargs.pop('unit', None)
+            metric_kwargs['name'] += f'_{metric_name}'
+            err_metrics.append(metric_cls(**metric_kwargs))
+        return ErrorRecorder(err_metrics)