diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 2da752eaa7..0b4fe001b3 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -8,9 +8,13 @@ from .make_base_fitting import ( make_base_fitting, ) +from .polarizability_fitting import ( + PolarFitting, +) __all__ = [ "InvarFitting", "make_base_fitting", "DipoleFitting", + "PolarFitting", ] diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index 5a4952edb8..c210945e76 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -24,7 +24,7 @@ @fitting_check_output class DipoleFitting(GeneralFitting): - r"""Fitting rotationally invariant diploe of the system. + r"""Fitting rotationally equivariant diploe of the system. Parameters ---------- @@ -34,7 +34,7 @@ class DipoleFitting(GeneralFitting): The number of atom types. dim_descrpt The dimension of the input descriptor. - dim_rot_mat : int + embedding_width : int The dimension of rotation matrix, m1. neuron Number of neurons :math:`N` in each hidden layer of the fitting net @@ -78,7 +78,7 @@ def __init__( var_name: str, ntypes: int, dim_descrpt: int, - dim_rot_mat: int, + embedding_width: int, neuron: List[int] = [120, 120, 120], resnet_dt: bool = True, numb_fparam: int = 0, @@ -108,7 +108,7 @@ def __init__( if atom_ener is not None and atom_ener != []: raise NotImplementedError("atom_ener is not implemented") - self.dim_rot_mat = dim_rot_mat + self.embedding_width = embedding_width super().__init__( var_name=var_name, ntypes=ntypes, @@ -133,11 +133,11 @@ def __init__( def _net_out_dim(self): """Set the FittingNet output dim.""" - return self.dim_rot_mat + return self.embedding_width def serialize(self) -> dict: data = super().serialize() - data["dim_rot_mat"] = self.dim_rot_mat + data["embedding_width"] = self.embedding_width data["old_impl"] = self.old_impl return data @@ -194,7 +194,7 @@ def call( self.var_name ] # (nframes * nloc, 1, m1) - out = out.reshape(-1, 1, self.dim_rot_mat) + out = out.reshape(-1, 1, self.embedding_width) # (nframes * nloc, m1, 3) gr = gr.reshape(nframes * nloc, -1, 3) # (nframes, nloc, 3) diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 03f8f25237..890a065f15 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -189,6 +189,8 @@ def __setitem__(self, key, value): self.aparam_avg = value elif key in ["aparam_inv_std"]: self.aparam_inv_std = value + elif key in ["scale"]: + self.scale = value else: raise KeyError(key) @@ -203,6 +205,8 @@ def __getitem__(self, key): return self.aparam_avg elif key in ["aparam_inv_std"]: return self.aparam_inv_std + elif key in ["scale"]: + return self.scale else: raise KeyError(key) @@ -327,10 +331,10 @@ def _call_common( mask = np.tile( (atype == type_i).reshape([nf, nloc, 1]), [1, 1, net_dim_out] ) - atom_energy = self.nets[(type_i,)](xx) - atom_energy = atom_energy + self.bias_atom_e[type_i] - atom_energy = atom_energy * mask - outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + atom_property = self.nets[(type_i,)](xx) + atom_property = atom_property + self.bias_atom_e[type_i] + atom_property = atom_property * mask + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: outs = self.nets[()](xx) + self.bias_atom_e[atype] # nf x nloc diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py new file mode 100644 index 0000000000..d828693fe0 --- /dev/null +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -0,0 +1,241 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + Dict, + List, + Optional, +) + +import numpy as np + +from deepmd.common import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.dpmodel import ( + DEFAULT_PRECISION, +) +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) + +from .general_fitting import ( + GeneralFitting, +) + + +@fitting_check_output +class PolarFitting(GeneralFitting): + r"""Fitting rotationally equivariant polarizability of the system. + + Parameters + ---------- + var_name + The name of the output variable. + ntypes + The number of atom types. + dim_descrpt + The dimension of the input descriptor. + embedding_width : int + The dimension of rotation matrix, m1. + neuron + Number of neurons :math:`N` in each hidden layer of the fitting net + resnet_dt + Time-step `dt` in the resnet construction: + :math:`y = x + dt * \phi (Wx + b)` + numb_fparam + Number of frame parameter + numb_aparam + Number of atomic parameter + rcond + The condition number for the regression of atomic energy. + tot_ener_zero + Force the total energy to zero. Useful for the charge fitting. + trainable + If the weights of fitting net are trainable. + Suppose that we have :math:`N_l` hidden layers in the fitting net, + this list is of length :math:`N_l + 1`, specifying if the hidden layers and the output layer are trainable. + atom_ener + Specifying atomic energy contribution in vacuum. The `set_davg_zero` key in the descrptor should be set. + activation_function + The activation function :math:`\boldsymbol{\phi}` in the embedding net. Supported options are |ACTIVATION_FN| + precision + The precision of the embedding net parameters. Supported options are |PRECISION| + layer_name : list[Optional[str]], optional + The name of the each layer. If two layers, either in the same fitting or different fittings, + have the same name, they will share the same neural network parameters. + use_aparam_as_mask: bool, optional + If True, the atomic parameters will be used as a mask that determines the atom is real/virtual. + And the aparam will not be used as the atomic parameters for embedding. + mixed_types + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + fit_diag : bool + Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to + normal polarizability matrix by contracting with the rotation matrix. + scale : List[float] + The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i] + shift_diag : bool + Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale. + """ + + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + embedding_width: int, + neuron: List[int] = [120, 120, 120], + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + rcond: Optional[float] = None, + tot_ener_zero: bool = False, + trainable: Optional[List[bool]] = None, + atom_ener: Optional[List[Optional[float]]] = None, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + layer_name: Optional[List[Optional[str]]] = None, + use_aparam_as_mask: bool = False, + spin: Any = None, + mixed_types: bool = False, + exclude_types: List[int] = [], + old_impl: bool = False, + fit_diag: bool = True, + scale: Optional[List[float]] = None, + shift_diag: bool = True, + ): + # seed, uniform_seed are not included + if tot_ener_zero: + raise NotImplementedError("tot_ener_zero is not implemented") + if spin is not None: + raise NotImplementedError("spin is not implemented") + if use_aparam_as_mask: + raise NotImplementedError("use_aparam_as_mask is not implemented") + if layer_name is not None: + raise NotImplementedError("layer_name is not implemented") + if atom_ener is not None and atom_ener != []: + raise NotImplementedError("atom_ener is not implemented") + + self.embedding_width = embedding_width + self.fit_diag = fit_diag + self.scale = scale + if self.scale is None: + self.scale = [1.0 for _ in range(ntypes)] + else: + assert ( + isinstance(self.scale, list) and len(self.scale) == ntypes + ), "Scale should be a list of length ntypes." + self.scale = np.array(self.scale, dtype=GLOBAL_NP_FLOAT_PRECISION).reshape( + ntypes, 1 + ) + self.shift_diag = shift_diag + super().__init__( + var_name=var_name, + ntypes=ntypes, + dim_descrpt=dim_descrpt, + neuron=neuron, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + rcond=rcond, + tot_ener_zero=tot_ener_zero, + trainable=trainable, + atom_ener=atom_ener, + activation_function=activation_function, + precision=precision, + layer_name=layer_name, + use_aparam_as_mask=use_aparam_as_mask, + spin=spin, + mixed_types=mixed_types, + exclude_types=exclude_types, + ) + self.old_impl = False + + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return ( + self.embedding_width + if self.fit_diag + else self.embedding_width * self.embedding_width + ) + + def serialize(self) -> dict: + data = super().serialize() + data["embedding_width"] = self.embedding_width + data["old_impl"] = self.old_impl + data["fit_diag"] = self.fit_diag + data["@variables"]["scale"] = self.scale + return data + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [3, 3], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def call( + self, + descriptor: np.ndarray, + atype: np.ndarray, + gr: Optional[np.ndarray] = None, + g2: Optional[np.ndarray] = None, + h2: Optional[np.ndarray] = None, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + ) -> Dict[str, np.ndarray]: + """Calculate the fitting. + + Parameters + ---------- + descriptor + input descriptor. shape: nf x nloc x nd + atype + the atom type. shape: nf x nloc + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + g2 + The rotationally invariant pair-partical representation. + shape: nf x nloc x nnei x ng + h2 + The rotationally equivariant pair-partical representation. + shape: nf x nloc x nnei x 3 + fparam + The frame parameter. shape: nf x nfp. nfp being `numb_fparam` + aparam + The atomic parameter. shape: nf x nloc x nap. nap being `numb_aparam` + + """ + nframes, nloc, _ = descriptor.shape + assert ( + gr is not None + ), "Must provide the rotation matrix for polarizability fitting." + # (nframes, nloc, _net_out_dim) + out = self._call_common(descriptor, atype, gr, g2, h2, fparam, aparam)[ + self.var_name + ] + out = out * self.scale[atype] + # (nframes * nloc, m1, 3) + gr = gr.reshape(nframes * nloc, -1, 3) + + if self.fit_diag: + out = out.reshape(-1, self.embedding_width) + out = np.einsum("ij,ijk->ijk", out, gr) + else: + out = out.reshape(-1, self.embedding_width, self.embedding_width) + out = (out + np.transpose(out, axes=(0, 2, 1))) / 2 + out = np.einsum("bim,bmj->bij", out, gr) # (nframes * nloc, m1, 3) + out = np.einsum( + "bim,bmj->bij", np.transpose(gr, axes=(0, 2, 1)), out + ) # (nframes * nloc, 3, 3) + out = out.reshape(nframes, nloc, 3, 3) + return {self.var_name: out} diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index fedf4386c0..4ea66e2636 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -25,7 +25,7 @@ class DipoleFittingNet(GeneralFitting): - """Construct a general fitting net. + """Construct a dipole fitting net. Parameters ---------- @@ -35,9 +35,7 @@ class DipoleFittingNet(GeneralFitting): Element count. dim_descrpt : int Embedding width per atom. - dim_out : int - The output dimension of the fitting net. - dim_rot_mat : int + embedding_width : int The dimension of rotation matrix, m1. neuron : List[int] Number of neurons in each hidden layers of the fitting net. @@ -51,8 +49,9 @@ class DipoleFittingNet(GeneralFitting): Activation function. precision : str Numerical precision. - distinguish_types : bool - Neighbor list that distinguish different atomic types or not. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. rcond : float, optional The condition number for the regression of atomic energy. seed : int, optional @@ -64,20 +63,20 @@ def __init__( var_name: str, ntypes: int, dim_descrpt: int, - dim_rot_mat: int, + embedding_width: int, neuron: List[int] = [128, 128, 128], resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, - distinguish_types: bool = False, + mixed_types: bool = True, rcond: Optional[float] = None, seed: Optional[int] = None, exclude_types: List[int] = [], **kwargs, ): - self.dim_rot_mat = dim_rot_mat + self.embedding_width = embedding_width super().__init__( var_name=var_name, ntypes=ntypes, @@ -88,7 +87,7 @@ def __init__( numb_aparam=numb_aparam, activation_function=activation_function, precision=precision, - distinguish_types=distinguish_types, + mixed_types=mixed_types, rcond=rcond, seed=seed, exclude_types=exclude_types, @@ -98,11 +97,11 @@ def __init__( def _net_out_dim(self): """Set the FittingNet output dim.""" - return self.dim_rot_mat + return self.embedding_width def serialize(self) -> dict: data = super().serialize() - data["dim_rot_mat"] = self.dim_rot_mat + data["embedding_width"] = self.embedding_width data["old_impl"] = self.old_impl return data @@ -144,7 +143,7 @@ def forward( self.var_name ] # (nframes * nloc, 1, m1) - out = out.view(-1, 1, self.dim_rot_mat) + out = out.view(-1, 1, self.embedding_width) # (nframes * nloc, m1, 3) gr = gr.view(nframes * nloc, -1, 3) # (nframes, nloc, 3) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 3a904a0696..cade533f1a 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -476,6 +476,8 @@ def __setitem__(self, key, value): self.aparam_avg = value elif key in ["aparam_inv_std"]: self.aparam_inv_std = value + elif key in ["scale"]: + self.scale = value else: raise KeyError(key) @@ -490,6 +492,8 @@ def __getitem__(self, key): return self.aparam_avg elif key in ["aparam_inv_std"]: return self.aparam_inv_std + elif key in ["scale"]: + return self.scale else: raise KeyError(key) @@ -585,7 +589,9 @@ def _forward_common( atom_property = ( self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] ) - outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + outs = ( + outs + atom_property + ) # Shape is [nframes, natoms[0], net_dim_out] else: for type_i, ll in enumerate(self.filter_layers.networks): mask = (atype == type_i).unsqueeze(-1) @@ -593,7 +599,9 @@ def _forward_common( atom_property = ll(xx) atom_property = atom_property + self.bias_atom_e[type_i] atom_property = atom_property * mask - outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + outs = ( + outs + atom_property + ) # Shape is [nframes, natoms[0], net_dim_out] # nf x nloc mask = self.emask(atype) # nf x nloc x nod diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py new file mode 100644 index 0000000000..dc8d13ee84 --- /dev/null +++ b/deepmd/pt/model/task/polarizability.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + List, + Optional, +) + +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.pt.model.task.fitting import ( + GeneralFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) + +log = logging.getLogger(__name__) + + +class PolarFittingNet(GeneralFitting): + """Construct a polar fitting net. + + Parameters + ---------- + var_name : str + The atomic property to fit, 'polar'. + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + embedding_width : int + The dimension of rotation matrix, m1. + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + resnet_dt : bool + Using time-step in the ResNet construction. + numb_fparam : int + Number of frame parameters. + numb_aparam : int + Number of atomic parameters. + activation_function : str + Activation function. + precision : str + Numerical precision. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. + fit_diag : bool + Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to + normal polarizability matrix by contracting with the rotation matrix. + scale : List[float] + The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i] + shift_diag : bool + Whether to shift the diagonal part of the polarizability matrix. The shift operation is carried out after scale. + """ + + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + embedding_width: int, + neuron: List[int] = [128, 128, 128], + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + mixed_types: bool = True, + rcond: Optional[float] = None, + seed: Optional[int] = None, + exclude_types: List[int] = [], + fit_diag: bool = True, + scale: Optional[List[float]] = None, + shift_diag: bool = True, + **kwargs, + ): + self.embedding_width = embedding_width + self.fit_diag = fit_diag + self.scale = scale + if self.scale is None: + self.scale = [1.0 for _ in range(ntypes)] + else: + assert ( + isinstance(self.scale, list) and len(self.scale) == ntypes + ), "Scale should be a list of length ntypes." + self.scale = torch.tensor( + self.scale, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ).view(ntypes, 1) + self.shift_diag = shift_diag + super().__init__( + var_name=var_name, + ntypes=ntypes, + dim_descrpt=dim_descrpt, + neuron=neuron, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + rcond=rcond, + seed=seed, + exclude_types=exclude_types, + **kwargs, + ) + self.old_impl = False # this only supports the new implementation. + + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return ( + self.embedding_width + if self.fit_diag + else self.embedding_width * self.embedding_width + ) + + def serialize(self) -> dict: + data = super().serialize() + data["embedding_width"] = self.embedding_width + data["old_impl"] = self.old_impl + data["fit_diag"] = self.fit_diag + data["fit_diag"] = self.fit_diag + data["@variables"]["scale"] = to_numpy_array(self.scale) + return data + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [3, 3], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + @property + def data_stat_key(self): + """ + Get the keys for the data statistic of the fitting. + Return a list of statistic names needed, such as "bias_atom_e". + """ + return [] + + def forward( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + nframes, nloc, _ = descriptor.shape + assert ( + gr is not None + ), "Must provide the rotation matrix for polarizability fitting." + # (nframes, nloc, _net_out_dim) + out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[ + self.var_name + ] + out = out * self.scale[atype] + gr = gr.view(nframes * nloc, -1, 3) # (nframes * nloc, m1, 3) + + if self.fit_diag: + out = out.reshape(-1, self.embedding_width) + out = torch.einsum("ij,ijk->ijk", out, gr) + else: + out = out.reshape(-1, self.embedding_width, self.embedding_width) + out = (out + out.transpose(1, 2)) / 2 + out = torch.einsum("bim,bmj->bij", out, gr) # (nframes * nloc, m1, 3) + out = torch.einsum( + "bim,bmj->bij", gr.transpose(1, 2), out + ) # (nframes * nloc, 3, 3) + out = out.view(nframes, nloc, 3, 3) + + return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py index fffed123e0..3f67043767 100644 --- a/source/tests/pt/model/test_dipole_fitting.py +++ b/source/tests/pt/model/test_dipole_fitting.py @@ -36,7 +36,7 @@ class TestDipoleFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self): TestCaseSingleFrameWithNlist.setUp(self) self.rng = np.random.default_rng() - self.nf, self.nloc, nnei = self.nlist.shape + self.nf, self.nloc, _ = self.nlist.shape self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) def test_consistency( @@ -51,7 +51,7 @@ def test_consistency( self.atype_ext[:, : self.nloc], dtype=int, device=env.DEVICE ) - for distinguish_types, nfp, nap in itertools.product( + for mixed_types, nfp, nap in itertools.product( [True, False], [0, 3], [0, 4], @@ -60,10 +60,10 @@ def test_consistency( "foo", self.nt, self.dd0.dim_out, - dim_rot_mat=self.dd0.get_dim_emb(), + embedding_width=self.dd0.get_dim_emb(), numb_fparam=nfp, numb_aparam=nap, - use_tebd=(not distinguish_types), + mixed_types=mixed_types, ).to(env.DEVICE) ft1 = DPDipoleFitting.deserialize(ft0.serialize()) ft2 = DipoleFittingNet.deserialize(ft1.serialize()) @@ -104,7 +104,7 @@ def test_consistency( def test_jit( self, ): - for distinguish_types, nfp, nap in itertools.product( + for mixed_types, nfp, nap in itertools.product( [True, False], [0, 3], [0, 4], @@ -113,10 +113,10 @@ def test_jit( "foo", self.nt, self.dd0.dim_out, - dim_rot_mat=self.dd0.get_dim_emb(), + embedding_width=self.dd0.get_dim_emb(), numb_fparam=nfp, numb_aparam=nap, - use_tebd=(not distinguish_types), + mixed_types=mixed_types, ).to(env.DEVICE) torch.jit.script(ft0) @@ -140,7 +140,7 @@ def test_rot(self): rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype).to(env.DEVICE) coord_rot = torch.matmul(self.coord, rmat) rng = np.random.default_rng() - for distinguish_types, nfp, nap in itertools.product( + for mixed_types, nfp, nap in itertools.product( [True, False], [0, 3], [0, 4], @@ -149,10 +149,10 @@ def test_rot(self): "foo", 3, # ntype self.dd0.dim_out, # dim_descrpt - dim_rot_mat=self.dd0.get_dim_emb(), + embedding_width=self.dd0.get_dim_emb(), numb_fparam=nfp, numb_aparam=nap, - use_tebd=False, + mixed_types=mixed_types, ).to(env.DEVICE) if nfp > 0: ifp = torch.tensor( @@ -174,10 +174,10 @@ def test_rot(self): ( extended_coord, extended_atype, - mapping, + _, nlist, ) = extend_input_and_build_neighbor_list( - xyz + self.shift, atype, self.rcut, self.sel, distinguish_types + xyz + self.shift, atype, self.rcut, self.sel, not mixed_types ) rd0, gr0, _, _, _ = self.dd0( @@ -199,10 +199,10 @@ def test_permu(self): "foo", 3, # ntype self.dd0.dim_out, - dim_rot_mat=self.dd0.get_dim_emb(), + embedding_width=self.dd0.get_dim_emb(), numb_fparam=0, numb_aparam=0, - use_tebd=False, + mixed_types=False, ).to(env.DEVICE) res = [] for idx_perm in [[0, 1, 2, 3, 4], [1, 0, 4, 3, 2]]: @@ -210,10 +210,10 @@ def test_permu(self): ( extended_coord, extended_atype, - mapping, + _, nlist, ) = extend_input_and_build_neighbor_list( - coord[idx_perm], atype, self.rcut, self.sel, False + coord[idx_perm], atype, self.rcut, self.sel, True ) rd0, gr0, _, _, _ = self.dd0( @@ -241,17 +241,17 @@ def test_trans(self): "foo", 3, # ntype self.dd0.dim_out, - dim_rot_mat=self.dd0.get_dim_emb(), + embedding_width=self.dd0.get_dim_emb(), numb_fparam=0, numb_aparam=0, - use_tebd=False, + mixed_types=True, ).to(env.DEVICE) res = [] for xyz in [self.coord, coord_s]: ( extended_coord, extended_atype, - mapping, + _, nlist, ) = extend_input_and_build_neighbor_list( xyz, atype, self.rcut, self.sel, False diff --git a/source/tests/pt/model/test_polarizability_fitting.py b/source/tests/pt/model/test_polarizability_fitting.py new file mode 100644 index 0000000000..de43c57b8b --- /dev/null +++ b/source/tests/pt/model/test_polarizability_fitting.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import torch +from scipy.stats import ( + special_ortho_group, +) + +from deepmd.dpmodel.fitting import PolarFitting as DPPolarFitting +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.task.polarizability import ( + PolarFittingNet, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class TestDipoleFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + self.rng = np.random.default_rng() + self.nf, self.nloc, _ = self.nlist.shape + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() + + def test_consistency( + self, + ): + rd0, gr, _, _, _ = self.dd0( + torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), + torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), + torch.tensor(self.nlist, dtype=int, device=env.DEVICE), + ) + atype = torch.tensor( + self.atype_ext[:, : self.nloc], dtype=int, device=env.DEVICE + ) + + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( + [True, False], + [0, 3], + [0, 4], + [True, False], + [None, self.scale], + ): + ft0 = PolarFittingNet( + "foo", + self.nt, + self.dd0.dim_out, + embedding_width=self.dd0.get_dim_emb(), + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=mixed_types, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + ft1 = DPPolarFitting.deserialize(ft0.serialize()) + ft2 = PolarFittingNet.deserialize(ft0.serialize()) + ft3 = DPPolarFitting.deserialize(ft1.serialize()) + + if nfp > 0: + ifp = torch.tensor( + self.rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + self.rng.normal(size=(self.nf, self.nloc, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + ret0 = ft0(rd0, atype, gr, fparam=ifp, aparam=iap) + ret1 = ft1( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + gr.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + ret2 = ft2(rd0, atype, gr, fparam=ifp, aparam=iap) + ret3 = ft3( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + gr.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + ret1["foo"], + ) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + to_numpy_array(ret2["foo"]), + ) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + ret3["foo"], + ) + + def test_jit( + self, + ): + for mixed_types, nfp, nap, fit_diag in itertools.product( + [True, False], + [0, 3], + [0, 4], + [True, False], + ): + ft0 = PolarFittingNet( + "foo", + self.nt, + self.dd0.dim_out, + embedding_width=self.dd0.get_dim_emb(), + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=mixed_types, + fit_diag=fit_diag, + ).to(env.DEVICE) + torch.jit.script(ft0) + + +class TestEquivalence(unittest.TestCase): + def setUp(self) -> None: + self.natoms = 5 + self.rcut = 4 + self.rcut_smth = 0.5 + self.sel = [46, 92, 4] + self.nf = 1 + self.nt = 3 + self.rng = np.random.default_rng() + self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype).to(env.DEVICE) + self.shift = torch.tensor([4, 4, 4], dtype=dtype).to(env.DEVICE) + self.atype = torch.IntTensor([0, 0, 0, 1, 1]).to(env.DEVICE) + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.cell = torch.rand([3, 3], dtype=dtype).to(env.DEVICE) + self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3).to(env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() + + def test_rot(self): + atype = self.atype.reshape(1, 5) + rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype).to(env.DEVICE) + coord_rot = torch.matmul(self.coord, rmat) + + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( + [True, False], + [0, 3], + [0, 4], + [True, False], + [None, self.scale], + ): + ft0 = PolarFittingNet( + "foo", + self.nt, + self.dd0.dim_out, # dim_descrpt + embedding_width=self.dd0.get_dim_emb(), + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + if nfp > 0: + ifp = torch.tensor( + self.rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + self.rng.normal(size=(self.nf, self.natoms, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + res = [] + for xyz in [self.coord, coord_rot]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz + self.shift, atype, self.rcut, self.sel, mixed_types + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=ifp, aparam=iap) + res.append(ret0["foo"]) + print(res[1].shape) + np.testing.assert_allclose( + to_numpy_array(res[1]), + to_numpy_array( + torch.matmul( + rmat.T, + torch.matmul(res[0], rmat), + ) + ), + ) + + def test_permu(self): + coord = torch.matmul(self.coord, self.cell) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PolarFittingNet( + "foo", + self.nt, + self.dd0.dim_out, + embedding_width=self.dd0.get_dim_emb(), + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for idx_perm in [[0, 1, 2, 3, 4], [1, 0, 4, 3, 2]]: + atype = self.atype[idx_perm].reshape(1, 5) + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + coord[idx_perm], atype, self.rcut, self.sel, False + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=None, aparam=None) + res.append(ret0["foo"]) + + np.testing.assert_allclose( + to_numpy_array(res[0][:, idx_perm]), + to_numpy_array(res[1]), + ) + + def test_trans(self): + atype = self.atype.reshape(1, 5) + coord_s = torch.matmul( + torch.remainder( + torch.matmul(self.coord + self.shift, torch.linalg.inv(self.cell)), 1.0 + ), + self.cell, + ) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PolarFittingNet( + "foo", + self.nt, + self.dd0.dim_out, + embedding_width=self.dd0.get_dim_emb(), + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for xyz in [self.coord, coord_s]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz, atype, self.rcut, self.sel, False + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=0, aparam=0) + res.append(ret0["foo"]) + + np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) + + +if __name__ == "__main__": + unittest.main()