Skip to content

GitLab

Commit 2409a22f authored by fanding2000's avatar fanding2000
Browse files

Format fix. More options in readme

parent ce29afea
Hide whitespace changes
Inline Side-by-side
Showing with 3619 additions and 3618 deletions
basic_function/format_parser.py
View file @ 2409a22f
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
View file @ 2409a22f
# -*- 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
View file @ 2409a22f
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
View file @ 2409a22f
# -*- 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
View file @ 2409a22f
#!/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
View file @ 2409a22f
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
View file @ 2409a22f
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
View file @ 2409a22f
"""
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
View file @ 2409a22f
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
View file @ 2409a22f
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
View file @ 2409a22f
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
View file @ 2409a22f
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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment