Skip to content

fix(dpmodel): use graph forward for output statistics #5946

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Stardust0831 wants to merge 1 commit into
base: master
Choose a base branch
from
+209 −25
Open

fix(dpmodel): use graph forward for output statistics #5946

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
78 changes: 53 additions & 25 deletions deepmd/dpmodel/atomic_model/base_atomic_model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -808,6 +808,9 @@ def _get_forward_wrapper_func(self) -> Callable[..., dict[str, np.ndarray]]:
"""Get a forward wrapper of the atomic model for output bias calculation."""
import array_api_compat

from deepmd.dpmodel.utils.neighbor_graph import (
build_neighbor_graph,
)
from deepmd.dpmodel.utils.nlist import (
extend_input_and_build_neighbor_list,
)
Expand Down Expand Up @@ -841,31 +844,56 @@ def model_forward(
if charge_spin is not None:
charge_spin = xp.asarray(charge_spin, device=device)

(
extended_coord,
extended_atype,
mapping,
nlist,
) = extend_input_and_build_neighbor_list(
coord,
atype,
self.get_rcut(),
self.get_sel(),
mixed_types=self.mixed_types(),
box=box,
# exclusion is a nlist-BUILD transform (decision #18/A4);
# forward_common_atomic consumes a pre-excluded nlist.
pair_excl=self.pair_excl,
)
atomic_ret = self.forward_common_atomic(
extended_coord,
extended_atype,
nlist,
mapping=mapping,
fparam=fparam,
aparam=aparam,
charge_spin=charge_spin,
)
if self.uses_graph_lower() and atype.shape[1] > 0:
nf, nloc = atype.shape[:2]
graph = build_neighbor_graph(
coord,
atype,
box,
self.get_rcut(),
pair_excl=self.pair_excl,
)
atomic_ret = self.forward_common_atomic_graph(
graph,
xp.reshape(atype, (nf * nloc,)),
fparam=fparam,
aparam=(
xp.reshape(aparam, (nf * nloc, -1))
if aparam is not None
else None
),
charge_spin=charge_spin,
)
atomic_ret = {
kk: xp.reshape(vv, (nf, nloc, *vv.shape[1:]))
for kk, vv in atomic_ret.items()
}
else:
(
extended_coord,
extended_atype,
mapping,
nlist,
) = extend_input_and_build_neighbor_list(
coord,
atype,
self.get_rcut(),
self.get_sel(),
mixed_types=self.mixed_types(),
box=box,
# exclusion is a nlist-BUILD transform (decision #18/A4);
# forward_common_atomic consumes a pre-excluded nlist.
pair_excl=self.pair_excl,
)
atomic_ret = self.forward_common_atomic(
extended_coord,
extended_atype,
nlist,
mapping=mapping,
fparam=fparam,
aparam=aparam,
charge_spin=charge_spin,
)
# Convert outputs back to numpy arrays
return {kk: to_numpy_array(vv) for kk, vv in atomic_ret.items()}

Expand Down
156 changes: 156 additions & 0 deletions source/tests/pt_expt/model/test_dpa1_graph_lower.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@

from deepmd.dpmodel.utils.neighbor_graph import (
apply_pair_exclusion,
build_neighbor_graph,
from_dense_quartet,
)
from deepmd.dpmodel.utils.nlist import (
Expand Down Expand Up @@ -101,6 +102,8 @@ def _make_model(
smooth: bool = False,
pair_excl_types: list | None = None,
descr_excl_types: list | None = None,
numb_fparam: int = 0,
numb_aparam: int = 0,
) -> EnergyModel:
ds = DescrptDPA1(
self.rcut,
Expand Down Expand Up @@ -130,6 +133,8 @@ def _make_model(
ds.get_dim_out(),
1,
mixed_types=ds.mixed_types(),
numb_fparam=numb_fparam,
numb_aparam=numb_aparam,
precision="float64",
seed=GLOBAL_SEED,
).to(self.device)
Expand All @@ -140,6 +145,157 @@ def _make_model(
pair_exclude_types=pair_excl_types or [],
).to(self.device)

@pytest.mark.parametrize("nframes", [1, 2])
def test_output_stat_forward_uses_graph_lower(
self, monkeypatch, nframes: int
) -> None:
model = self._make_model(pair_excl_types=[(0, 1)], numb_fparam=2, numb_aparam=3)
model.eval()
atomic_model = model.atomic_model
coord = self.coord.repeat(nframes, 1, 1)
atype = self.atype.repeat(nframes, 1)
box = self.cell.reshape(1, 9).repeat(nframes, 1)
fparam = torch.arange(
nframes * 2, dtype=torch.float64, device=self.device
).reshape(nframes, 2)
aparam = torch.arange(
nframes * self.natoms * 3,
dtype=torch.float64,
device=self.device,
).reshape(nframes, self.natoms, 3)
charge_spin = torch.arange(
nframes * 2, dtype=torch.float64, device=self.device
).reshape(nframes, 2)

graph = build_neighbor_graph(
coord,
atype,
box,
atomic_model.get_rcut(),
pair_excl=atomic_model.pair_excl,
)
original_forward = atomic_model.forward_common_atomic_graph
expected = original_forward(
graph,
atype.reshape(-1),
fparam=fparam,
aparam=aparam.reshape(nframes * self.natoms, 3),
charge_spin=charge_spin,
)
forwarded = {}

def record_forward(graph, atype, fparam=None, aparam=None, charge_spin=None):
forwarded.update(
graph=graph,
atype=atype,
fparam=fparam,
aparam=aparam,
charge_spin=charge_spin,
)
return original_forward(
graph,
atype,
fparam=fparam,
aparam=aparam,
charge_spin=charge_spin,
)

def fail_dense(*args, **kwargs):
raise AssertionError("the graph statistics route must not use dense sel")

monkeypatch.setattr(atomic_model, "forward_common_atomic_graph", record_forward)
monkeypatch.setattr(atomic_model, "get_sel", fail_dense)
monkeypatch.setattr(
"deepmd.dpmodel.utils.nlist.extend_input_and_build_neighbor_list",
fail_dense,
)

actual = atomic_model._get_forward_wrapper_func()(
coord.detach().cpu().numpy(),
atype.detach().cpu().numpy(),
box.detach().cpu().numpy(),
fparam=fparam.detach().cpu().numpy(),
aparam=aparam.detach().cpu().numpy(),
charge_spin=charge_spin.detach().cpu().numpy(),
)

assert forwarded["atype"].shape == (nframes * self.natoms,)
assert forwarded["fparam"].shape == (nframes, 2)
assert forwarded["aparam"].shape == (nframes * self.natoms, 3)
assert forwarded["charge_spin"].shape == (nframes, 2)
for key in ("energy", "mask"):
expected_array = (
expected[key]
.detach()
.cpu()
.numpy()
.reshape(nframes, self.natoms, *expected[key].shape[1:])
)
np.testing.assert_allclose(
actual[key], expected_array, rtol=1e-12, atol=1e-12
)

def test_change_out_bias_uses_graph_lower(self, monkeypatch) -> None:
model = self._make_model(pair_excl_types=[(0, 1)])
model.eval()
atomic_model = model.atomic_model
nframes = 2
coord = self.coord.repeat(nframes, 1, 1)
atype = self.atype.repeat(nframes, 1)
box = self.cell.reshape(1, 9).repeat(nframes, 1)
graph = build_neighbor_graph(
coord,
atype,
box,
atomic_model.get_rcut(),
pair_excl=atomic_model.pair_excl,
)
atom_energy = atomic_model.forward_common_atomic_graph(
graph, atype.reshape(-1)
)["energy"].reshape(nframes, self.natoms, -1)
sampled = [
{
"coord": coord.detach().cpu().numpy(),
"atype": atype.detach().cpu().numpy(),
"box": box.detach().cpu().numpy(),
"natoms": np.tile(
np.array([[self.natoms, self.natoms, 3, 2]]), (nframes, 1)
),
"energy": atom_energy.sum(dim=1).detach().cpu().numpy(),
"find_energy": np.float32(1.0),
}
]
original_bias = atomic_model.get_out_bias().detach().clone()

def fail_dense(*args, **kwargs):
raise AssertionError("the graph statistics route must not use dense sel")

monkeypatch.setattr(atomic_model, "get_sel", fail_dense)
monkeypatch.setattr(
"deepmd.dpmodel.utils.nlist.extend_input_and_build_neighbor_list",
fail_dense,
)

atomic_model.change_out_bias(sampled, bias_adjust_mode="change-by-statistic")

torch.testing.assert_close(
atomic_model.get_out_bias(), original_bias, rtol=1e-12, atol=1e-12
)

def test_output_stat_forward_empty_system(self) -> None:
atomic_model = self._make_model(numb_aparam=3).atomic_model
nframes = 2

actual = atomic_model._get_forward_wrapper_func()(
np.empty((nframes, 0, 3)),
np.empty((nframes, 0), dtype=np.int64),
np.zeros((nframes, 9)),
aparam=np.empty((nframes, 0, 3)),
)

assert actual["energy"].shape == (nframes, 0, 1)
assert actual["mask"].shape == (nframes, 0)

def _prepare_lower_inputs(self, periodic: bool):
"""Build extended coords, atype, nlist, mapping as torch tensors."""
coord_np = self.coord.detach().cpu().numpy()
Expand Down