Skip to content

Commit

Permalink
add spin_energy python train/test code
Browse files Browse the repository at this point in the history
  • Loading branch information
@HydrogenSulfate
HydrogenSulfate committed Dec 4, 2023
1 parent b3c4a6f commit 691d85e
Show file tree
Hide file tree
Showing 13 changed files with 660 additions and 608 deletions.
79 changes: 30 additions & 49 deletions deepmd/descriptor/se_a.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ def __init__(
raise RuntimeError(
f"rcut_smth ({rcut_smth:f}) should be no more than rcut ({rcut:f})!"
)
self.sel_a = sel # [46(O), 92(H)]
self.sel_a = sel # [46(O), 92(H)] / [60, 60]
self.rcut_r = rcut # 6.0
# NOTE: register 'rcut' in buffer to be accessed in inference
self.register_buffer("buffer_rcut", paddle.to_tensor(rcut, dtype="float64"))
Expand Down Expand Up @@ -177,10 +177,6 @@ def __init__(
self.sel_a.extend(self.sel_a_spin)
else:
self.ntypes_spin = 0
# NOTE: register 'ntypes_spin' in buffer to be accessed in inference
self.register_buffer(
"buffer_ntypes_spin", paddle.to_tensor(self.ntypes_spin, dtype="int32")
)

# descrpt config
self.sel_r = [0 for ii in range(len(self.sel_a))] # [0, 0]
Expand All @@ -206,20 +202,11 @@ def __init__(
# self.mixed_prec = None
# self.place_holders = {}
# self.nei_type = np.repeat(np.arange(self.ntypes), self.sel_a)
"""
array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
1, 1, 1, 1, 1, 1])
"""
self.avg_zero = paddle.zeros(
[self.ntypes, self.ndescrpt], dtype="float32"
[self.ntypes, self.ndescrpt], dtype=GLOBAL_PD_FLOAT_PRECISION
) # [2, 552]
self.std_ones = paddle.ones(
[self.ntypes, self.ndescrpt], dtype="float32"
[self.ntypes, self.ndescrpt], dtype=GLOBAL_PD_FLOAT_PRECISION
) # [2, 552]
nets = []
# self._pass_filter => self._filter => self._filter_lower
Expand Down Expand Up @@ -296,11 +283,22 @@ def __init__(
}

self.t_rcut = paddle.to_tensor(
np.max([self.rcut_r, self.rcut_a]), dtype="float32"
np.max([self.rcut_r, self.rcut_a]), dtype=GLOBAL_PD_FLOAT_PRECISION
)
# self.t_ntypes = paddle.to_tensor(self.ntypes, dtype="int32")
# self.t_ndescrpt = paddle.to_tensor(self.ndescrpt, dtype="int32")
# self.t_sel = paddle.to_tensor(self.sel_a, dtype="int32")
self.register_buffer("buffer_sel", paddle.to_tensor(self.sel_a, dtype="int32"))
self.register_buffer(
"buffer_ndescrpt", paddle.to_tensor(self.ndescrpt, dtype="int32")
)
self.register_buffer(
"buffer_original_sel",
paddle.to_tensor(
self.original_sel if self.original_sel is not None else self.sel_a,
dtype="int32",
),
)
self.t_ntypes = paddle.to_tensor(self.ntypes, dtype="int32")
self.t_ndescrpt = paddle.to_tensor(self.ndescrpt, dtype="int32")
self.t_sel = paddle.to_tensor(self.sel_a, dtype="int32")

t_avg = paddle.to_tensor(
np.zeros([self.ntypes, self.ndescrpt]), dtype="float64"
Expand Down Expand Up @@ -747,28 +745,24 @@ def _pass_filter(
else:
type_embedding = None
start_index = 0
# print(inputs.shape) # [192, 552(nnei*4)]，每个原子和它周围nnei个原子的R矩阵(展平后)
inputs = paddle.reshape(inputs, [-1, int(natoms[0].item()), int(self.ndescrpt)])
output = []
output_qmat = []
if not self.type_one_side and type_embedding is None:
for type_i in range(self.ntypes):
# 按不同原子类型进行处理
inputs_i = paddle.slice(
inputs,
[0, 1, 2],
[0, start_index, 0],
[
inputs.shape[0],
start_index + natoms[2 + type_i],
start_index + natoms[2 + type_i].item(),
inputs.shape[2],
],
) # [1, 某种类型原子个数64/128, 552]
inputs_i = paddle.reshape(
inputs_i, [-1, self.ndescrpt]
) # [某种类型原子个数64/128, 552]
)
inputs_i = paddle.reshape(inputs_i, [-1, self.ndescrpt])
filter_name = "filter_type_" + str(type_i) + suffix
layer, qmat = self._filter( # 计算某个类型的原子的 result 和 qmat
layer, qmat = self._filter(
inputs_i,
type_i,
name=filter_name,
Expand All @@ -779,18 +773,18 @@ def _pass_filter(
) # [natom, M1*M2], qmat: [natom, M1, 3]
layer = paddle.reshape(
layer, [inputs.shape[0], natoms[2 + type_i], self.get_dim_out()]
) # [1, 某种类型原子个数64/128, M1*M2]
)
qmat = paddle.reshape(
qmat,
[
inputs.shape[0],
natoms[2 + type_i],
self.get_dim_rot_mat_1() * 3,
],
) # [1, 某种类型原子个数64/128, 100*3]
)
output.append(layer)
output_qmat.append(qmat)
start_index += natoms[2 + type_i]
start_index += natoms[2 + type_i].item()
else:
...
# This branch will not be excecuted at current
Expand Down Expand Up @@ -842,24 +836,11 @@ def _compute_dstats_sys_smth(
self, data_coord, data_box, data_atype, natoms_vec, mesh
):
input_dict = {}
# dd_all = run_sess(
# self.sub_sess,
# self.stat_descrpt,
# feed_dict={
# self.place_holders["coord"]: data_coord,
# self.place_holders["type"]: data_atype,
# self.place_holders["natoms_vec"]: natoms_vec,
# self.place_holders["box"]: data_box,
# self.place_holders["default_mesh"]: mesh,
# },
# )
input_dict["coord"] = paddle.to_tensor(data_coord, dtype="float32")
input_dict["box"] = paddle.to_tensor(data_box, dtype="float32")
input_dict["type"] = paddle.to_tensor(data_atype, dtype="int32")
input_dict["natoms_vec"] = paddle.to_tensor(
natoms_vec, dtype="int32", place="cpu"
)
input_dict["default_mesh"] = paddle.to_tensor(mesh, dtype="int32")
input_dict["coord"] = paddle.to_tensor(data_coord, GLOBAL_PD_FLOAT_PRECISION)
input_dict["box"] = paddle.to_tensor(data_box, GLOBAL_PD_FLOAT_PRECISION)
input_dict["type"] = paddle.to_tensor(data_atype, "int32")
input_dict["natoms_vec"] = paddle.to_tensor(natoms_vec, "int32", place="cpu")
input_dict["default_mesh"] = paddle.to_tensor(mesh, "int32")

self.stat_descrpt, descrpt_deriv, rij, nlist = op_module.prod_env_mat_a(
input_dict["coord"], # fp32
Expand Down
28 changes: 1 addition & 27 deletions deepmd/entrypoints/freeze.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -355,27 +355,11 @@ def freeze_graph(
input_spec=[
InputSpec(shape=[None], dtype="float64"), # coord_
InputSpec(shape=[None], dtype="int32"), # atype_
InputSpec(shape=[4], dtype="int32"), # natoms
InputSpec(shape=[2 + dp.model.descrpt.ntypes], dtype="int32"), # natoms
InputSpec(shape=[None], dtype="float64"), # box
InputSpec(shape=[6], dtype="int32"), # mesh
{
# "coord": InputSpec(
# shape=[2880],
# dtype="float64"
# ),
# "type": InputSpec(
# shape=[960],
# dtype="int32"
# ),
# "natoms_vec": InputSpec(
# shape=[4],
# dtype="int32"
# ),
"box": InputSpec(shape=[None], dtype="float64"),
# "default_mesh": InputSpec(
# shape=[6],
# dtype="int32"
# ),
},
"",
False,
Expand All @@ -385,16 +369,6 @@ def freeze_graph(
print(
f"[{name}, {param.shape}] generated name in static_model is: {param.name}"
)
# print(f"st_model.descrpt.buffer_rcut.name = {st_model.descrpt.buffer_rcut.name}")
# print(
# f"st_model.descrpt.buffer_ntypes.name = {st_model.descrpt.buffer_ntypes.name}"
# )
# print(
# f"st_model.fitting.buffer_dfparam.name = {st_model.fitting.buffer_dfparam.name}"
# )
# print(
# f"st_model.fitting.buffer_daparam.name = {st_model.fitting.buffer_daparam.name}"
# )
# 跳过对program的裁剪，从而保留rcut、ntypes等不参与前向的参数，从而在C++端可以获取这些参数
skip_prune_program = True
print(f"==>> Set skip_prune_program = {skip_prune_program}")
Expand Down
45 changes: 24 additions & 21 deletions deepmd/fit/ener.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,8 @@ def __init__(
self.atom_ener.append(None)
self.useBN = False
self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64)
ntypes_atom = self.ntypes - self.ntypes_spin
self.bias_atom_e = self.bias_atom_e[:ntypes_atom]
self.register_buffer(
"t_bias_atom_e",
paddle.to_tensor(self.bias_atom_e),
Expand Down Expand Up @@ -248,7 +250,6 @@ def __init__(
1,
activation_fn=None,
precision=self.fitting_precision,
bavg=self.bias_atom_e,
name=layer_suffix,
seed=self.seed,
trainable=self.trainable[-1],
Expand Down Expand Up @@ -310,6 +311,26 @@ def compute_output_stats(self, all_stat: dict, mixed_type: bool = False) -> None
self.bias_atom_e = self._compute_output_stats(
all_stat, rcond=self.rcond, mixed_type=mixed_type
)
ntypes_atom = self.ntypes - self.ntypes_spin
if self.spin is not None:
for type_i in range(ntypes_atom):
if self.bias_atom_e.shape[0] != self.ntypes:
self.bias_atom_e = np.pad(
self.bias_atom_e,
(0, self.ntypes_spin),
"constant",
constant_values=(0, 0),
)
bias_atom_e = self.bias_atom_e
if self.spin.use_spin[type_i]:
self.bias_atom_e[type_i] = (
self.bias_atom_e[type_i]
+ self.bias_atom_e[type_i + ntypes_atom]
)
else:
self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
self.bias_atom_e = self.bias_atom_e[:ntypes_atom]

paddle.assign(self.bias_atom_e, self.t_bias_atom_e)

def _compute_output_stats(self, all_stat, rcond=1e-3, mixed_type=False):
Expand Down Expand Up @@ -548,24 +569,6 @@ def forward(
self.aparam_inv_std = 1.0

ntypes_atom = self.ntypes - self.ntypes_spin
if self.spin is not None:
for type_i in range(ntypes_atom):
if self.bias_atom_e.shape[0] != self.ntypes:
self.bias_atom_e = np.pad(
self.bias_atom_e,
(0, self.ntypes_spin),
"constant",
constant_values=(0, 0),
)
bias_atom_e = self.bias_atom_e
if self.spin.use_spin[type_i]:
self.bias_atom_e[type_i] = (
self.bias_atom_e[type_i]
+ self.bias_atom_e[type_i + ntypes_atom]
)
else:
self.bias_atom_e[type_i] = self.bias_atom_e[type_i]
self.bias_atom_e = self.bias_atom_e[:ntypes_atom]

inputs = paddle.reshape(
inputs, [-1, natoms[0], self.dim_descrpt]
Expand All @@ -583,7 +586,7 @@ def forward(
inputs_zero = paddle.zeros_like(inputs, dtype=GLOBAL_PD_FLOAT_PRECISION)

if bias_atom_e is not None:
assert len(bias_atom_e) == self.ntypes
assert len(bias_atom_e) == self.ntypes - self.ntypes_spin

fparam = None
if self.numb_fparam > 0:
Expand Down Expand Up @@ -615,7 +618,7 @@ def forward(
atype_nall,
[0, 1],
[0, 0],
[-1, paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
[atype_nall.shape[0], paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
)
atype_filter = paddle.cast(self.atype_nloc >= 0, GLOBAL_PD_FLOAT_PRECISION)
self.atype_nloc = paddle.reshape(self.atype_nloc, [-1])
Expand Down
Loading

0 comments on commit 691d85e

Please sign in to comment.