Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
5 changes: 4 additions & 1 deletion deepmd/dpmodel/descriptor/dpa4c_nn/spin.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,7 +327,10 @@ def call(self, spin_moments: Array, degree_two: Array) -> Array:
# anisotropy sum over neighbours.
xp.reshape(
xp.matmul(xp.permute_dims(quadrupole, (0, 2, 1)), degree_two),
(spin_moments.shape[0], -1),
(
spin_moments.shape[0],
self.quadrupole_width * self.degree_channels[2],
),
),
magnitude,
coordination,
Expand Down
26 changes: 26 additions & 0 deletions source/tests/common/dpmodel/test_descriptor_dpa4c.py
Original file line number Diff line number Diff line change
Expand Up @@ -804,6 +804,32 @@ def make_spin_descriptor(**overrides: Any) -> DescrptDPA4C:
return descriptor


@pytest.mark.parametrize("precision", ["float32", "float64"])
@pytest.mark.parametrize("channels", [8, 32])
@pytest.mark.parametrize("node_count", [0, 2])
def test_native_spin_graph_without_edges(
precision: str, channels: int, node_count: int
) -> None:
"""Both zero-node graphs and isolated atoms retain the feature width."""
descriptor = make_spin_descriptor(precision=precision, channels=channels)
graph = neighbor_graph.NeighborGraph(
n_node=np.array([node_count], dtype=np.int64),
n_local=np.array([node_count], dtype=np.int64),
edge_index=np.empty((2, 0), dtype=np.int64),
edge_vec=np.empty((0, 3), dtype=precision),
edge_mask=np.empty((0,), dtype=bool),
)
output, rot_mat = descriptor.call_graph(
graph,
np.arange(node_count, dtype=np.int64) % 2,
spin=np.ones((node_count, 3), dtype=precision),
)
assert output.shape == (node_count, descriptor.get_dim_out())
assert output.dtype == graph.edge_vec.dtype
assert np.isfinite(output).all()
assert rot_mat is None


def spin_reference_terms(
descriptor: DescrptDPA4C,
graph: Any,
Expand Down
39 changes: 39 additions & 0 deletions source/tests/pt_expt/descriptor/test_dpa4c.py
Original file line number Diff line number Diff line change
Expand Up @@ -458,6 +458,45 @@ def test_coincident_edge_has_finite_third_derivative(self) -> None:
)


@pytest.mark.parametrize("precision", ["float32", "float64"])
@pytest.mark.parametrize("channels", [8, 32])
@pytest.mark.parametrize("node_count", [0, 2])
def test_native_spin_graph_without_edges(
precision: str, channels: int, node_count: int
) -> None:
"""The reference spin path supports zero nodes as well as isolated atoms."""
descriptor = TestDPA4C.build(
precision=precision, channels=channels, use_spin=[True, False]
).eval()
with torch.no_grad():
descriptor.spin.spin_gate.fill_(1.0)
dtype = getattr(torch, precision)
graph = NeighborGraph(
n_node=torch.tensor([node_count], dtype=torch.int64, device=env.DEVICE),
n_local=torch.tensor([node_count], dtype=torch.int64, device=env.DEVICE),
edge_index=torch.empty((2, 0), dtype=torch.int64, device=env.DEVICE),
edge_vec=torch.empty((0, 3), dtype=dtype, device=env.DEVICE),
edge_mask=torch.empty((0,), dtype=torch.bool, device=env.DEVICE),
)
spin = torch.ones(
(node_count, 3), dtype=dtype, device=env.DEVICE, requires_grad=True
)
output, rot_mat = descriptor.call_graph(
graph,
torch.arange(node_count, dtype=torch.int64, device=env.DEVICE) % 2,
spin=spin,
)
assert output.shape == (node_count, descriptor.get_dim_out())
assert output.dtype == dtype
assert output.device == spin.device
assert torch.isfinite(output).all()
assert rot_mat is None
output.sum().backward()
assert spin.grad is not None
assert spin.grad.shape == spin.shape
assert torch.isfinite(spin.grad).all()


class TestDPA4CSpinGate:
"""Torch-side contracts of the spin branch gate."""

Expand Down