Skip to content

Commit

Permalink
[Paddle Backend] Add zinc protein (#3135)
Browse files Browse the repository at this point in the history 
Add zinc protein example and fix typos in zinc_protein config.

---------

Signed-off-by: HydrogenSulfate <[email protected]>
Co-authored-by: xusuyong <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
@HydrogenSulfate @xusuyong @pre-commit-ci
3 people authored Jan 12, 2024
1 parent f132a39 commit 4c618d7
Show file tree
Hide file tree
Showing 9 changed files with 1,241 additions and 114 deletions.
573 changes: 469 additions & 104 deletions deepmd/descriptor/se_a_mask.py
View file

Large diffs are not rendered by default.

11 changes: 9 additions & 2 deletions deepmd/fit/ener.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -247,7 +247,13 @@ def __init__(
else:
type_i_layers.append(
OneLayer_deepmd(
self.dim_descrpt + self.numb_fparam + self.numb_aparam,
self.dim_descrpt
+ self.numb_fparam
+ (
self.numb_aparam
if self.use_aparam_as_mask is False
else 0
),
self.n_neuron[ii],
activation_fn=self.fitting_activation_fn,
precision=self.fitting_precision,
Expand Down Expand Up @@ -691,7 +697,8 @@ def forward(
self.atom_ener_before = outs * atype_filter
self.add_type = paddle.reshape(
paddle.nn.functional.embedding(
self.atype_nloc, self.t_bias_atom_e.reshape([2, -1])
self.atype_nloc,
self.t_bias_atom_e.reshape([self.t_bias_atom_e.shape[0], -1]),
),
[paddle.shape(inputs)[0], paddle.sum(natoms[2 : 2 + ntypes_atom]).item()],
)
Expand Down
11 changes: 6 additions & 5 deletions deepmd/train/trainer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -151,11 +151,12 @@ def _init_param(self, jdata):
descrpt_param["multi_task"] = True
if descrpt_param["type"] in ["se_e2_a", "se_a", "se_e2_r", "se_r", "hybrid"]:
descrpt_param["spin"] = self.spin
descrpt_param.pop("type")
descrpt_param["mixed_prec"] = self.mixed_prec
if descrpt_param["mixed_prec"] is not None:
descrpt_param["precision"]: str = self.mixed_prec["output_prec"]
self.descrpt = deepmd.descriptor.se_a.DescrptSeA(**descrpt_param)
elif descrpt_param["type"] == "se_a_mask":
descrpt_param.pop("type")
self.descrpt = deepmd.descriptor.se_a_mask.DescrptSeAMask(**descrpt_param)
else:
descrpt_param.pop("type")
self.descrpt = deepmd.descriptor.se_a.DescrptSeA(**descrpt_param)

# fitting net
if not self.multi_task_mode:
Expand Down
6 changes: 3 additions & 3 deletions examples/zinc_protein/zinc_se_a_mask.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -66,14 +66,14 @@
"training": {
"training_data": {
"systems": [
"example/zinc_protein/train_data_dp_mask/"
"examples/zinc_protein/train_data_dp_mask/"
],
"batch_size": 2,
"_comment": "that's all"
},
"validation_data": {
"systems": [
"example/zinc_protein/val_data_dp_mask/"
"examples/zinc_protein/val_data_dp_mask/"
],
"batch_size": 2,
"_comment": "that's all"
Expand All @@ -82,7 +82,7 @@
"seed": 10,
"disp_freq": 100,
"save_freq": 1000,
"tensorboard": true,
"tensorboard": false,
"tensorboard_log_dir": "log4tensorboard",
"tensorboard_freq": 100,
"_comment": "that's all"
Expand Down
3 changes: 3 additions & 0 deletions source/lib/paddle_src/custom_op_install.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,9 @@
"./paddle_prod_force_grad.cu",
"./paddle_prod_force_grad.cc",
"./paddle_neighbor_stat.cc",
"./paddle_descrpt_se_a_mask.cc",
"./paddle_prod_force_se_a_mask.cc",
"./paddle_prod_force_se_a_mask_grad.cc",
],
include_dirs=[
"../../lib/include/",
Expand Down
103 changes: 103 additions & 0 deletions source/lib/paddle_src/custom_op_test.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -256,6 +256,107 @@ def test_prod_virial_se_a(place="cpu"):
print(np.allclose(atom_virial.numpy(), atom_virial_load))


def test_prod_force_se_a_mask(place="cpu"):
print("=" * 10, f"test_prod_force_se_a_mask [place={place}]", "=" * 10)
import numpy as np

net_deriv_reshape = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "prod_force_se_a_mask/net_deriv_reshape.npy"))
)
descrpt_deriv = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "prod_force_se_a_mask/descrpt_deriv.npy"))
)
mask = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "prod_force_se_a_mask/mask.npy"))
)
nlist = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "prod_force_se_a_mask/nlist.npy"))
)
total_atom_num = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "prod_force_se_a_mask/total_atom_num.npy"))
)

net_deriv_reshape = paddle.to_tensor(
net_deriv_reshape, stop_gradient=False, place=place
)
descrpt_deriv = paddle.to_tensor(descrpt_deriv, "float64", place=place)
mask = paddle.to_tensor(mask, "int32", place=place)
nlist = paddle.to_tensor(nlist, "int32", place=place)

force = paddle_deepmd_lib.prod_force_se_a_mask(
net_deriv_reshape,
descrpt_deriv,
mask,
nlist,
total_atom_num,
)
force.sum().backward()
# print(f"net_deriv_reshape.grad.shape = {net_deriv_reshape.grad.shape}")

# print(mn.shape, mn.min().item()); print(mn.max().item()); print(mn.mean().item()); print(mn.var().item())
# print(mn_load.shape); print(mn_load.min().item()); print(mn_load.max().item()); print(mn_load.mean().item()); print(mn_load.var().item())
# print(dt.shape, dt.min().item(), dt.max().item(), dt.mean().item(), dt.var().item())
# print(dt_load.shape, dt_load.min().item(), dt_load.max().item(), dt_load.mean().item(), dt_load.var().item())
force_load = np.load(osp.join(unitest_dir, "prod_force_se_a_mask/force.npy"))
grad_load = np.load(
osp.join(unitest_dir, "prod_force_se_a_mask/net_deriv_reshape_grad.npy")
)

print(np.allclose(force.numpy(), force_load))
print(np.allclose(net_deriv_reshape.grad.numpy(), grad_load))


def test_descrpt_se_a_mask(place="cpu"):
print("=" * 10, f"test_descrpt_se_a_mask [place={place}]", "=" * 10)
import numpy as np

coord = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "descrpt_se_a_mask/coord.npy"))
) # float64 (2,441)
atype = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "descrpt_se_a_mask/atype.npy"))
) # int32 (2,147)
mask = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "descrpt_se_a_mask/mask.npy"))
) # int32 (2, 147)
box = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "descrpt_se_a_mask/box.npy"))
) # float64 (2, 9)
natoms = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "descrpt_se_a_mask/natoms.npy"))
) # int32 (8,)
mesh = np.ascontiguousarray(
np.load(osp.join(unitest_dir, "descrpt_se_a_mask/mesh.npy"))
) # int32 (0,)

coord = paddle.to_tensor(coord, "float64", place=place)
atype = paddle.to_tensor(atype, "int32", place=place)
mask = paddle.to_tensor(mask, "int32", place=place)
box = paddle.to_tensor(box, "float64", place=place)
natoms = paddle.to_tensor(natoms, "int32", place=place)
mesh = paddle.to_tensor(mesh, "int32", place=place)

descrpt, descrpt_deriv, rij, nlist = paddle_deepmd_lib.descrpt_se_a_mask(
coord, atype, mask, box, natoms, mesh
)

# print(mn.shape, mn.min().item()); print(mn.max().item()); print(mn.mean().item()); print(mn.var().item())
# print(mn_load.shape); print(mn_load.min().item()); print(mn_load.max().item()); print(mn_load.mean().item()); print(mn_load.var().item())
# print(dt.shape, dt.min().item(), dt.max().item(), dt.mean().item(), dt.var().item())
# print(dt_load.shape, dt_load.min().item(), dt_load.max().item(), dt_load.mean().item(), dt_load.var().item())
descrpt_load = np.load(osp.join(unitest_dir, "descrpt_se_a_mask/descrpt.npy"))
descrpt_deriv_load = np.load(
osp.join(unitest_dir, "descrpt_se_a_mask/descrpt_deriv.npy")
)
rij_load = np.load(osp.join(unitest_dir, "descrpt_se_a_mask/rij.npy"))
nlist_load = np.load(osp.join(unitest_dir, "descrpt_se_a_mask/nlist.npy"))

print(np.allclose(descrpt.numpy(), descrpt_load))
print(np.allclose(descrpt_deriv.numpy(), descrpt_deriv_load))
print(np.allclose(rij.numpy(), rij_load))
print(np.allclose(nlist.numpy(), nlist_load))


if __name__ == "__main__":
test_neighbor_stat()

Expand All @@ -266,3 +367,5 @@ def test_prod_virial_se_a(place="cpu"):
test_prod_env_mat_a("cpu")
test_prod_force_se_a("cpu")
test_prod_virial_se_a("cpu")
test_prod_force_se_a_mask()
test_descrpt_se_a_mask()
Loading

0 comments on commit 4c618d7

Please sign in to comment.