From 47cff4b951821b7f171caa5b5c768bc9e69dae23 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Sun, 28 Jan 2024 15:17:51 +0800 Subject: [PATCH 01/20] feat: add pair table model to pytorch --- deepmd/pt/model/model/pair_tab.py | 351 ++++++++++++++++++++++++++++++ source/tests/pt/test_pairtab.py | 74 +++++++ 2 files changed, 425 insertions(+) create mode 100644 deepmd/pt/model/model/pair_tab.py create mode 100644 source/tests/pt/test_pairtab.py diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py new file mode 100644 index 0000000000..4694ce618e --- /dev/null +++ b/deepmd/pt/model/model/pair_tab.py @@ -0,0 +1,351 @@ +from .atomic_model import AtomicModel +from deepmd.utils.pair_tab import ( + PairTab, +) +import logging +import torch +from torch import nn +import numpy as np +from typing import Dict, List, Optional, Union + +from deepmd.model_format import FittingOutputDef, OutputVariableDef + +class PairTabModel(nn.Module, AtomicModel): + """Pairwise tabulation energy model. + + This model can be used to tabulate the pairwise energy between atoms for either + short-range or long-range interactions, such as D3, LJ, ZBL, etc. It should not + be used alone, but rather as one submodel of a linear (sum) model, such as + DP+D3. + + Do not put the model on the first model of a linear model, since the linear + model fetches the type map from the first model. + + At this moment, the model does not smooth the energy at the cutoff radius, so + one needs to make sure the energy has been smoothed to zero. + + Parameters + ---------- + tab_file : str + The path to the tabulation file. + rcut : float + The cutoff radius. + sel : int or list[int] + The maxmum number of atoms in the cut-off radius. + """ + + def __init__( + self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs + ): + super().__init__() + self.tab_file = tab_file + self.rcut = rcut + + # check table data against rcut and update tab_file if needed. + self._check_table_upper_boundary() + + self.tab = PairTab(self.tab_file) + self.ntypes = self.tab.ntypes + + + tab_info, tab_data = self.tab.get() # this returns -> Tuple[np.array, np.array] + self.tab_info = torch.from_numpy(tab_info) + self.tab_data = torch.from_numpy(tab_data) + + # self.model_type = "ener" + # self.model_version = MODEL_VERSION ## this shoud be in the parent class + + if isinstance(sel, int): + self.sel = sel + elif isinstance(sel, list): + self.sel = sum(sel) + else: + raise TypeError("sel must be int or list[int]") + + def get_fitting_output_def(self)->FittingOutputDef: + return FittingOutputDef( + [OutputVariableDef( + name="energy", + shape=[1], + reduciable=True, + differentiable=True + )] + ) + + def get_rcut(self)->float: + return self.rcut + + def get_sel(self)->int: + return self.sel + + def distinguish_types(self)->bool: + # to match DPA1 and DPA2. + return False + + def forward_atomic( + self, + extended_coord, + extended_atype, + nlist, + mapping: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> Dict[str, torch.Tensor]: + + + nframes, nloc, nnei = nlist.shape + + #this will mask all -1 in the nlist + masked_nlist = torch.clamp(nlist,0) + + atype = extended_atype[:, :nloc] #(nframes, nloc) + pairwise_dr = self._get_pairwise_dist(extended_coord) # (nframes, nall, nall, 3) + pairwise_rr = pairwise_dr.pow(2).sum(-1).sqrt() # (nframes, nall, nall) + + self.tab_data = self.tab_data.reshape(self.tab.ntypes,self.tab.ntypes,self.tab.nspline,4) + + #to calculate the atomic_energy, we need 3 tensors, i_type, j_type, rr + #i_type : (nframes, nloc), this is atype. + #j_type : (nframes, nloc, nnei) + j_type = extended_atype[torch.arange(extended_atype.size(0))[:, None, None], masked_nlist] + + #slice rr to get (nframes, nloc, nnei) + rr = torch.gather(pairwise_rr[:, :nloc, :],2, masked_nlist) + + raw_atomic_energy = self._pair_tabulated_inter(atype, j_type, rr) + + atomic_energy = 0.5 * torch.sum(torch.where(nlist != -1, raw_atomic_energy, torch.zeros_like(raw_atomic_energy)) ,dim=-1) + + return {"energy": atomic_energy} + + def _check_table_upper_boundary(self): + """Update User Provided Table Based on `rcut` + + This function checks the upper boundary provided in the table against rcut. + If the table upper boundary values decay to zero before rcut, padding zeros will + be add to the table to cover rcut; if the table upper boundary values do not decay to zero + before ruct, linear extrapolation will be performed to rcut. In both cases, the table file + will be overwritten. + + Example + ------- + + table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0. 1. 1.5 ]] + + rcut = 0.022 + + new_table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0. 1. 1.5 ] + [0.02 0. 0.5 0.75 ] + [0.025 0. 0. 0. ]] + + ---------------------------------------------- + + table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0.5 1. 1.5 ] + [0.02 0.25 0.4 0.75 ] + [0.025 0. 0.1 0. ] + [0.03 0. 0. 0. ]] + + rcut = 0.031 + + new_table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0.5 1. 1.5 ] + [0.02 0.25 0.4 0.75 ] + [0.025 0. 0.1 0. ] + [0.03 0. 0. 0. ] + [0.035 0. 0. 0. ]] + """ + + raw_data = np.loadtxt(self.tab_file) + upper = raw_data[-1][0] + upper_val = raw_data[-1][1:] + upper_idx = raw_data.shape[0] - 1 + increment = raw_data[1][0] - raw_data[0][0] + + #the index of table for the grid point right after rcut + rcut_idx = int(self.rcut/increment) + + if np.all(upper_val == 0): + # if table values decay to `0` after rcut + if self.rcut < upper and np.any(raw_data[rcut_idx-1][1:]!=0): + logging.warning( + "The energy provided in the table does not decay to 0 at rcut." + ) + # if table values decay to `0` at rcut, do nothing + + # if table values decay to `0` before rcut, pad table with `0`s. + elif self.rcut > upper: + pad_zero = np.zeros((rcut_idx - upper_idx,4)) + pad_zero[:,0] = np.linspace(upper + increment, increment*(rcut_idx+1), rcut_idx-upper_idx) + raw_data = np.concatenate((raw_data,pad_zero),axis=0) + else: + # if table values do not decay to `0` at rcut + if self.rcut <= upper: + logging.warning( + "The energy provided in the table does not decay to 0 at rcut." + ) + # if rcut goes beyond table upper bond, need extrapolation. + else: + logging.warning( + "The rcut goes beyond table upper boundary, performing linear extrapolation." + ) + pad_linear = np.zeros((rcut_idx - upper_idx+1,4)) + pad_linear[:,0] = np.linspace(upper, increment*(rcut_idx+1), rcut_idx-upper_idx+1) + pad_linear[:,1:] = np.array([ + np.linspace(start, 0, rcut_idx - upper_idx+1) for start in upper_val + ]).T + raw_data = np.concatenate((raw_data[:-1,:],pad_zero),axis=0) + + #over writing file with padding if applicable. + with open(self.tab_file, 'wb') as f: + np.savetxt(f, raw_data) + + def _pair_tabulated_inter(self, nlist: torch.Tensor,i_type: torch.Tensor, j_type: torch.Tensor, rr: torch.Tensor) -> torch.Tensor: + """Pairwise tabulated energy. + + Parameters + ---------- + nlist : torch.Tensor + The unmasked neighbour list. (nframes, nloc) + + i_type : torch.Tensor + The integer representation of atom type for all local atoms for all frames. (nframes, nloc) + + j_type : torch.Tensor + The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) + + rr : torch.Tensor + The salar distance vector between two atoms. (nframes, nloc, nnei) + + Returns + ------- + torch.Tensor + The masked atomic energy for all local atoms for all frames. (nframes, nloc, nnei) + + Raises + ------ + Exception + If the distance is beyond the table. + + Notes + ----- + This function is used to calculate the pairwise energy between two atoms. + It uses a table containing cubic spline coefficients calculated in PairTab. + """ + + rmin = self.tab_info[0] + hh = self.tab_info[1] + hi = 1. / hh + + self.nspline = int(self.tab_info[2] + 0.1) + + uu = (rr - rmin) * hi # this is broadcasted to (nframes,nloc,nnei) + + + # if nnei of atom 0 has -1 in the nlist, uu would be 0. + # this is to handel the nlist where the mask is set to 0. + # by replacing the values wiht nspline + 1, the energy contribution will be 0 + uu = torch.where(nlist != -1, uu, self.nspline+1) + + if torch.any(uu < 0): + raise Exception("coord go beyond table lower boundary") + + idx = uu.to(torch.int) + + uu -= idx + + + final_coef = self._extract_spline_coefficient(i_type, j_type, idx) + + a3, a2, a1, a0 = torch.unbind(final_coef, dim=-1) # 4 * (nframes, nloc, nnei) + + etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. + ener = etmp * uu + a0 + return ener + + @staticmethod + def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: + """Get pairwise distance `dr`. + + Parameters + ---------- + coords : torch.Tensor + The coordinate of the atoms shape of (nframes * nall * 3). + + Returns + ------- + torch.Tensor + The pairwise distance between the atoms (nframes * nall * nall * 3). + + Examples + -------- + coords = torch.tensor([[ + [0,0,0], + [1,3,5], + [2,4,6] + ]]) + + dist = tensor([[ + [[ 0, 0, 0], + [-1, -3, -5], + [-2, -4, -6]], + + [[ 1, 3, 5], + [ 0, 0, 0], + [-1, -1, -1]], + + [[ 2, 4, 6], + [ 1, 1, 1], + [ 0, 0, 0]] + ]]) + """ + return coords.unsqueeze(2) - coords.unsqueeze(1) + + def _extract_spline_coefficient(self, i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: + """Extract the spline coefficient from the table. + + Parameters + ---------- + i_type : torch.Tensor + The integer representation of atom type for all local atoms for all frames. (nframes, nloc) + + j_type : torch.Tensor + The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) + + idx : torch.Tensor + The index of the spline coefficient. (nframes, nloc, nnei) + + Returns + ------- + torch.Tensor + The spline coefficient. (nframes, nloc, nnei, 4) + + Example + ------- + + """ + + # (nframes, nloc, nnei) + expanded_i_type = i_type.unsqueeze(-1).expand(-1, -1, j_type.shape[-1]) + + # (nframes, nloc, nnei, nspline, 4) + expanded_tab_data = self.tab_data[expanded_i_type, j_type] + + # (nframes, nloc, nnei, 1, 4) + expanded_idx = idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1,-1, -1, 4) + + #handle the case where idx is beyond the number of splines + clipped_indices = torch.clamp(expanded_idx, 0, self.nspline - 1).to(torch.int64) + + # (nframes, nloc, nnei, 4) + final_coef = torch.gather(expanded_tab_data, 3, clipped_indices).squeeze() + + # when the spline idx is beyond the table, all spline coefficients are set to `0`, and the resulting ener corresponding to the idx is also `0`. + final_coef[expanded_idx.squeeze() >= self.nspline] = 0 + + return final_coef \ No newline at end of file diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py new file mode 100644 index 0000000000..69feb38c32 --- /dev/null +++ b/source/tests/pt/test_pairtab.py @@ -0,0 +1,74 @@ +import unittest +import torch +import numpy as np +from unittest.mock import patch +from deepmd.pt.model.model.pair_tab import PairTabModel + +class TestPairTab(unittest.TestCase): + + @patch('numpy.loadtxt') + def setUp(self, mock_loadtxt) -> None: + + file_path = 'dummy_path' + mock_loadtxt.return_value = np.array([ + [0.005, 1. , 2. , 3. ], + [0.01 , 0.8 , 1.6 , 2.4 ], + [0.015, 0.5 , 1. , 1.5 ], + [0.02 , 0.25 , 0.4 , 0.75 ]]) + + self.model = PairTabModel( + tab_file = file_path, + rcut = 0.1, + sel = 2 + ) + + self.extended_coord = torch.tensor([ + [[0.01,0.01,0.01], + [0.01,0.02,0.01], + [0.01,0.01,0.02], + [0.02,0.01,0.01]], + + [[0.01,0.01,0.01], + [0.01,0.02,0.01], + [0.01,0.01,0.02], + [0.05,0.01,0.01]], + ]) + + # nframes=2, nall=4 + self.extended_atype = torch.tensor([ + [0,1,0,1], + [0,0,1,1] + ]) + + # nframes=2, nloc=2, nnei=2 + self.nlist = torch.tensor([ + [[1,2],[0,2]], + [[1,2],[0,3]] + ]) + + def test_without_mask(self): + + result = self.model.forward_atomic(self.extended_coord, self.extended_atype,self.nlist) + expected_result = torch.tensor([[2.4000, 2.7085], + [2.4000, 0.8000]]) + + torch.testing.assert_allclose(result,expected_result) + + def test_with_mask(self): + + self.nlist = torch.tensor([ + [[1,-1],[0,2]], + [[1,2],[0,3]] + ]) + + result = self.model.forward_atomic(self.extended_coord, self.extended_atype,self.nlist) + expected_result = torch.tensor([[1.6000, 2.7085], + [2.4000, 0.8000]]) + + torch.testing.assert_allclose(result,expected_result) + + def test_jit(self): + model = torch.jit.script(self.model) + + if __name__ == '__main__': + unittest.main() \ No newline at end of file From 04b6f57d67a87e2889bec21ef8f0e8aceee6fafc Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 28 Jan 2024 07:26:35 +0000 Subject: [PATCH 02/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/pair_tab.py | 224 +++++++++++++++++------------- source/tests/pt/test_pairtab.py | 107 +++++++------- 2 files changed, 178 insertions(+), 153 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 4694ce618e..9c40b6b009 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -1,14 +1,30 @@ -from .atomic_model import AtomicModel +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + Dict, + List, + Optional, + Union, +) + +import numpy as np +import torch +from torch import ( + nn, +) + +from deepmd.model_format import ( + FittingOutputDef, + OutputVariableDef, +) from deepmd.utils.pair_tab import ( PairTab, ) -import logging -import torch -from torch import nn -import numpy as np -from typing import Dict, List, Optional, Union -from deepmd.model_format import FittingOutputDef, OutputVariableDef +from .atomic_model import ( + AtomicModel, +) + class PairTabModel(nn.Module, AtomicModel): """Pairwise tabulation energy model. @@ -37,19 +53,18 @@ class PairTabModel(nn.Module, AtomicModel): def __init__( self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs ): - super().__init__() + super().__init__() self.tab_file = tab_file self.rcut = rcut # check table data against rcut and update tab_file if needed. self._check_table_upper_boundary() - + self.tab = PairTab(self.tab_file) self.ntypes = self.tab.ntypes - - tab_info, tab_data = self.tab.get() # this returns -> Tuple[np.array, np.array] - self.tab_info = torch.from_numpy(tab_info) + tab_info, tab_data = self.tab.get() # this returns -> Tuple[np.array, np.array] + self.tab_info = torch.from_numpy(tab_info) self.tab_data = torch.from_numpy(tab_data) # self.model_type = "ener" @@ -61,59 +76,67 @@ def __init__( self.sel = sum(sel) else: raise TypeError("sel must be int or list[int]") - - def get_fitting_output_def(self)->FittingOutputDef: + + def get_fitting_output_def(self) -> FittingOutputDef: return FittingOutputDef( - [OutputVariableDef( - name="energy", - shape=[1], - reduciable=True, - differentiable=True - )] + [ + OutputVariableDef( + name="energy", shape=[1], reduciable=True, differentiable=True + ) + ] ) - def get_rcut(self)->float: + def get_rcut(self) -> float: return self.rcut - - def get_sel(self)->int: + + def get_sel(self) -> int: return self.sel - - def distinguish_types(self)->bool: + + def distinguish_types(self) -> bool: # to match DPA1 and DPA2. return False - + def forward_atomic( self, - extended_coord, - extended_atype, + extended_coord, + extended_atype, nlist, mapping: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, - ) -> Dict[str, torch.Tensor]: - - + ) -> Dict[str, torch.Tensor]: nframes, nloc, nnei = nlist.shape - #this will mask all -1 in the nlist - masked_nlist = torch.clamp(nlist,0) + # this will mask all -1 in the nlist + masked_nlist = torch.clamp(nlist, 0) + + atype = extended_atype[:, :nloc] # (nframes, nloc) + pairwise_dr = self._get_pairwise_dist( + extended_coord + ) # (nframes, nall, nall, 3) + pairwise_rr = pairwise_dr.pow(2).sum(-1).sqrt() # (nframes, nall, nall) - atype = extended_atype[:, :nloc] #(nframes, nloc) - pairwise_dr = self._get_pairwise_dist(extended_coord) # (nframes, nall, nall, 3) - pairwise_rr = pairwise_dr.pow(2).sum(-1).sqrt() # (nframes, nall, nall) + self.tab_data = self.tab_data.reshape( + self.tab.ntypes, self.tab.ntypes, self.tab.nspline, 4 + ) - self.tab_data = self.tab_data.reshape(self.tab.ntypes,self.tab.ntypes,self.tab.nspline,4) + # to calculate the atomic_energy, we need 3 tensors, i_type, j_type, rr + # i_type : (nframes, nloc), this is atype. + # j_type : (nframes, nloc, nnei) + j_type = extended_atype[ + torch.arange(extended_atype.size(0))[:, None, None], masked_nlist + ] - #to calculate the atomic_energy, we need 3 tensors, i_type, j_type, rr - #i_type : (nframes, nloc), this is atype. - #j_type : (nframes, nloc, nnei) - j_type = extended_atype[torch.arange(extended_atype.size(0))[:, None, None], masked_nlist] + # slice rr to get (nframes, nloc, nnei) + rr = torch.gather(pairwise_rr[:, :nloc, :], 2, masked_nlist) - #slice rr to get (nframes, nloc, nnei) - rr = torch.gather(pairwise_rr[:, :nloc, :],2, masked_nlist) - raw_atomic_energy = self._pair_tabulated_inter(atype, j_type, rr) - atomic_energy = 0.5 * torch.sum(torch.where(nlist != -1, raw_atomic_energy, torch.zeros_like(raw_atomic_energy)) ,dim=-1) + atomic_energy = 0.5 * torch.sum( + torch.where( + nlist != -1, raw_atomic_energy, torch.zeros_like(raw_atomic_energy) + ), + dim=-1, + ) return {"energy": atomic_energy} @@ -126,9 +149,8 @@ def _check_table_upper_boundary(self): before ruct, linear extrapolation will be performed to rcut. In both cases, the table file will be overwritten. - Example - ------- - + Examples + -------- table = [[0.005 1. 2. 3. ] [0.01 0.8 1.6 2.4 ] [0.015 0. 1. 1.5 ]] @@ -148,9 +170,9 @@ def _check_table_upper_boundary(self): [0.015 0.5 1. 1.5 ] [0.02 0.25 0.4 0.75 ] [0.025 0. 0.1 0. ] - [0.03 0. 0. 0. ]] + [0.03 0. 0. 0. ]] - rcut = 0.031 + rcut = 0.031 new_table = [[0.005 1. 2. 3. ] [0.01 0.8 1.6 2.4 ] @@ -158,31 +180,32 @@ def _check_table_upper_boundary(self): [0.02 0.25 0.4 0.75 ] [0.025 0. 0.1 0. ] [0.03 0. 0. 0. ] - [0.035 0. 0. 0. ]] + [0.035 0. 0. 0. ]] """ - raw_data = np.loadtxt(self.tab_file) upper = raw_data[-1][0] upper_val = raw_data[-1][1:] - upper_idx = raw_data.shape[0] - 1 + upper_idx = raw_data.shape[0] - 1 increment = raw_data[1][0] - raw_data[0][0] - #the index of table for the grid point right after rcut - rcut_idx = int(self.rcut/increment) + # the index of table for the grid point right after rcut + rcut_idx = int(self.rcut / increment) if np.all(upper_val == 0): # if table values decay to `0` after rcut - if self.rcut < upper and np.any(raw_data[rcut_idx-1][1:]!=0): + if self.rcut < upper and np.any(raw_data[rcut_idx - 1][1:] != 0): logging.warning( "The energy provided in the table does not decay to 0 at rcut." ) # if table values decay to `0` at rcut, do nothing - + # if table values decay to `0` before rcut, pad table with `0`s. elif self.rcut > upper: - pad_zero = np.zeros((rcut_idx - upper_idx,4)) - pad_zero[:,0] = np.linspace(upper + increment, increment*(rcut_idx+1), rcut_idx-upper_idx) - raw_data = np.concatenate((raw_data,pad_zero),axis=0) + pad_zero = np.zeros((rcut_idx - upper_idx, 4)) + pad_zero[:, 0] = np.linspace( + upper + increment, increment * (rcut_idx + 1), rcut_idx - upper_idx + ) + raw_data = np.concatenate((raw_data, pad_zero), axis=0) else: # if table values do not decay to `0` at rcut if self.rcut <= upper: @@ -194,63 +217,69 @@ def _check_table_upper_boundary(self): logging.warning( "The rcut goes beyond table upper boundary, performing linear extrapolation." ) - pad_linear = np.zeros((rcut_idx - upper_idx+1,4)) - pad_linear[:,0] = np.linspace(upper, increment*(rcut_idx+1), rcut_idx-upper_idx+1) - pad_linear[:,1:] = np.array([ - np.linspace(start, 0, rcut_idx - upper_idx+1) for start in upper_val - ]).T - raw_data = np.concatenate((raw_data[:-1,:],pad_zero),axis=0) - - #over writing file with padding if applicable. - with open(self.tab_file, 'wb') as f: + pad_linear = np.zeros((rcut_idx - upper_idx + 1, 4)) + pad_linear[:, 0] = np.linspace( + upper, increment * (rcut_idx + 1), rcut_idx - upper_idx + 1 + ) + pad_linear[:, 1:] = np.array( + [ + np.linspace(start, 0, rcut_idx - upper_idx + 1) + for start in upper_val + ] + ).T + raw_data = np.concatenate((raw_data[:-1, :], pad_zero), axis=0) + + # over writing file with padding if applicable. + with open(self.tab_file, "wb") as f: np.savetxt(f, raw_data) - def _pair_tabulated_inter(self, nlist: torch.Tensor,i_type: torch.Tensor, j_type: torch.Tensor, rr: torch.Tensor) -> torch.Tensor: + def _pair_tabulated_inter( + self, + nlist: torch.Tensor, + i_type: torch.Tensor, + j_type: torch.Tensor, + rr: torch.Tensor, + ) -> torch.Tensor: """Pairwise tabulated energy. - + Parameters ---------- nlist : torch.Tensor The unmasked neighbour list. (nframes, nloc) - i_type : torch.Tensor The integer representation of atom type for all local atoms for all frames. (nframes, nloc) - j_type : torch.Tensor The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) - rr : torch.Tensor The salar distance vector between two atoms. (nframes, nloc, nnei) - + Returns ------- torch.Tensor The masked atomic energy for all local atoms for all frames. (nframes, nloc, nnei) - + Raises ------ Exception If the distance is beyond the table. - + Notes ----- This function is used to calculate the pairwise energy between two atoms. It uses a table containing cubic spline coefficients calculated in PairTab. """ - - rmin = self.tab_info[0] + rmin = self.tab_info[0] hh = self.tab_info[1] - hi = 1. / hh + hi = 1.0 / hh self.nspline = int(self.tab_info[2] + 0.1) - uu = (rr - rmin) * hi # this is broadcasted to (nframes,nloc,nnei) + uu = (rr - rmin) * hi # this is broadcasted to (nframes,nloc,nnei) - # if nnei of atom 0 has -1 in the nlist, uu would be 0. # this is to handel the nlist where the mask is set to 0. # by replacing the values wiht nspline + 1, the energy contribution will be 0 - uu = torch.where(nlist != -1, uu, self.nspline+1) + uu = torch.where(nlist != -1, uu, self.nspline + 1) if torch.any(uu < 0): raise Exception("coord go beyond table lower boundary") @@ -258,13 +287,12 @@ def _pair_tabulated_inter(self, nlist: torch.Tensor,i_type: torch.Tensor, j_type idx = uu.to(torch.int) uu -= idx - - + final_coef = self._extract_spline_coefficient(i_type, j_type, idx) - a3, a2, a1, a0 = torch.unbind(final_coef, dim=-1) # 4 * (nframes, nloc, nnei) + a3, a2, a1, a0 = torch.unbind(final_coef, dim=-1) # 4 * (nframes, nloc, nnei) - etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. + etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. ener = etmp * uu + a0 return ener @@ -281,7 +309,7 @@ def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: ------- torch.Tensor The pairwise distance between the atoms (nframes * nall * nall * 3). - + Examples -------- coords = torch.tensor([[ @@ -289,7 +317,7 @@ def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: [1,3,5], [2,4,6] ]]) - + dist = tensor([[ [[ 0, 0, 0], [-1, -3, -5], @@ -306,17 +334,17 @@ def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: """ return coords.unsqueeze(2) - coords.unsqueeze(1) - def _extract_spline_coefficient(self, i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: + def _extract_spline_coefficient( + self, i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor + ) -> torch.Tensor: """Extract the spline coefficient from the table. Parameters ---------- i_type : torch.Tensor The integer representation of atom type for all local atoms for all frames. (nframes, nloc) - j_type : torch.Tensor The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) - idx : torch.Tensor The index of the spline coefficient. (nframes, nloc, nnei) @@ -325,11 +353,7 @@ def _extract_spline_coefficient(self, i_type: torch.Tensor, j_type: torch.Tensor torch.Tensor The spline coefficient. (nframes, nloc, nnei, 4) - Example - ------- - """ - # (nframes, nloc, nnei) expanded_i_type = i_type.unsqueeze(-1).expand(-1, -1, j_type.shape[-1]) @@ -337,9 +361,9 @@ def _extract_spline_coefficient(self, i_type: torch.Tensor, j_type: torch.Tensor expanded_tab_data = self.tab_data[expanded_i_type, j_type] # (nframes, nloc, nnei, 1, 4) - expanded_idx = idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1,-1, -1, 4) - - #handle the case where idx is beyond the number of splines + expanded_idx = idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, -1, 4) + + # handle the case where idx is beyond the number of splines clipped_indices = torch.clamp(expanded_idx, 0, self.nspline - 1).to(torch.int64) # (nframes, nloc, nnei, 4) @@ -348,4 +372,4 @@ def _extract_spline_coefficient(self, i_type: torch.Tensor, j_type: torch.Tensor # when the spline idx is beyond the table, all spline coefficients are set to `0`, and the resulting ener corresponding to the idx is also `0`. final_coef[expanded_idx.squeeze() >= self.nspline] = 0 - return final_coef \ No newline at end of file + return final_coef diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index 69feb38c32..60c1bc01f3 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -1,74 +1,75 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later import unittest -import torch -import numpy as np -from unittest.mock import patch -from deepmd.pt.model.model.pair_tab import PairTabModel +from unittest.mock import ( + patch, +) -class TestPairTab(unittest.TestCase): +import numpy as np +import torch - @patch('numpy.loadtxt') - def setUp(self, mock_loadtxt) -> None: +from deepmd.pt.model.model.pair_tab import ( + PairTabModel, +) - file_path = 'dummy_path' - mock_loadtxt.return_value = np.array([ - [0.005, 1. , 2. , 3. ], - [0.01 , 0.8 , 1.6 , 2.4 ], - [0.015, 0.5 , 1. , 1.5 ], - [0.02 , 0.25 , 0.4 , 0.75 ]]) - self.model = PairTabModel( - tab_file = file_path, - rcut = 0.1, - sel = 2 +class TestPairTab(unittest.TestCase): + @patch("numpy.loadtxt") + def setUp(self, mock_loadtxt) -> None: + file_path = "dummy_path" + mock_loadtxt.return_value = np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + ] ) - self.extended_coord = torch.tensor([ - [[0.01,0.01,0.01], - [0.01,0.02,0.01], - [0.01,0.01,0.02], - [0.02,0.01,0.01]], + self.model = PairTabModel(tab_file=file_path, rcut=0.1, sel=2) - [[0.01,0.01,0.01], - [0.01,0.02,0.01], - [0.01,0.01,0.02], - [0.05,0.01,0.01]], - ]) + self.extended_coord = torch.tensor( + [ + [ + [0.01, 0.01, 0.01], + [0.01, 0.02, 0.01], + [0.01, 0.01, 0.02], + [0.02, 0.01, 0.01], + ], + [ + [0.01, 0.01, 0.01], + [0.01, 0.02, 0.01], + [0.01, 0.01, 0.02], + [0.05, 0.01, 0.01], + ], + ] + ) # nframes=2, nall=4 - self.extended_atype = torch.tensor([ - [0,1,0,1], - [0,0,1,1] - ]) + self.extended_atype = torch.tensor([[0, 1, 0, 1], [0, 0, 1, 1]]) # nframes=2, nloc=2, nnei=2 - self.nlist = torch.tensor([ - [[1,2],[0,2]], - [[1,2],[0,3]] - ]) + self.nlist = torch.tensor([[[1, 2], [0, 2]], [[1, 2], [0, 3]]]) def test_without_mask(self): - - result = self.model.forward_atomic(self.extended_coord, self.extended_atype,self.nlist) - expected_result = torch.tensor([[2.4000, 2.7085], - [2.4000, 0.8000]]) - - torch.testing.assert_allclose(result,expected_result) + result = self.model.forward_atomic( + self.extended_coord, self.extended_atype, self.nlist + ) + expected_result = torch.tensor([[2.4000, 2.7085], [2.4000, 0.8000]]) + + torch.testing.assert_allclose(result, expected_result) def test_with_mask(self): + self.nlist = torch.tensor([[[1, -1], [0, 2]], [[1, 2], [0, 3]]]) + + result = self.model.forward_atomic( + self.extended_coord, self.extended_atype, self.nlist + ) + expected_result = torch.tensor([[1.6000, 2.7085], [2.4000, 0.8000]]) - self.nlist = torch.tensor([ - [[1,-1],[0,2]], - [[1,2],[0,3]] - ]) - - result = self.model.forward_atomic(self.extended_coord, self.extended_atype,self.nlist) - expected_result = torch.tensor([[1.6000, 2.7085], - [2.4000, 0.8000]]) - - torch.testing.assert_allclose(result,expected_result) + torch.testing.assert_allclose(result, expected_result) def test_jit(self): model = torch.jit.script(self.model) - if __name__ == '__main__': - unittest.main() \ No newline at end of file + if __name__ == "__main__": + unittest.main() From eb59d872b63e7d4697ad31e5cdbb1193f93e266a Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Sun, 28 Jan 2024 15:54:08 +0800 Subject: [PATCH 03/20] fix: typo --- deepmd/pt/model/model/pair_tab.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 9c40b6b009..29fedb920c 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -129,7 +129,7 @@ def forward_atomic( # slice rr to get (nframes, nloc, nnei) rr = torch.gather(pairwise_rr[:, :nloc, :], 2, masked_nlist) - raw_atomic_energy = self._pair_tabulated_inter(atype, j_type, rr) + raw_atomic_energy = self._pair_tabulated_inter(nlist, atype, j_type, rr) atomic_energy = 0.5 * torch.sum( torch.where( @@ -227,7 +227,7 @@ def _check_table_upper_boundary(self): for start in upper_val ] ).T - raw_data = np.concatenate((raw_data[:-1, :], pad_zero), axis=0) + raw_data = np.concatenate((raw_data[:-1, :], pad_linear), axis=0) # over writing file with padding if applicable. with open(self.tab_file, "wb") as f: @@ -277,8 +277,7 @@ def _pair_tabulated_inter( uu = (rr - rmin) * hi # this is broadcasted to (nframes,nloc,nnei) # if nnei of atom 0 has -1 in the nlist, uu would be 0. - # this is to handel the nlist where the mask is set to 0. - # by replacing the values wiht nspline + 1, the energy contribution will be 0 + # this is to handel the nlist where the mask is set to 0, so that we don't raise exception for those atoms. uu = torch.where(nlist != -1, uu, self.nspline + 1) if torch.any(uu < 0): From b7cbbd50979313e62755abc8303f1da6c6dc9257 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Sun, 28 Jan 2024 16:13:29 +0800 Subject: [PATCH 04/20] fix: typo --- deepmd/pt/model/model/pair_tab.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 29fedb920c..1d83c7e818 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -141,7 +141,7 @@ def forward_atomic( return {"energy": atomic_energy} def _check_table_upper_boundary(self): - """Update User Provided Table Based on `rcut` + """Update User Provided Table Based on `rcut`. This function checks the upper boundary provided in the table against rcut. If the table upper boundary values decay to zero before rcut, padding zeros will From 84767f3c2fc1a52f978664d6d467fe0233bee59b Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Sun, 28 Jan 2024 18:02:32 +0800 Subject: [PATCH 05/20] fix: update ruct extrapolation --- deepmd/pt/model/model/pair_tab.py | 12 ++++++------ source/tests/pt/test_pairtab.py | 10 +++++----- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 1d83c7e818..6ae6a051e9 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -145,8 +145,8 @@ def _check_table_upper_boundary(self): This function checks the upper boundary provided in the table against rcut. If the table upper boundary values decay to zero before rcut, padding zeros will - be add to the table to cover rcut; if the table upper boundary values do not decay to zero - before ruct, linear extrapolation will be performed to rcut. In both cases, the table file + be added to the table to cover rcut; if the table upper boundary values do not decay to zero + before ruct, linear extrapolation will be performed till rcut. In both cases, the table file will be overwritten. Examples @@ -160,7 +160,7 @@ def _check_table_upper_boundary(self): new_table = [[0.005 1. 2. 3. ] [0.01 0.8 1.6 2.4 ] [0.015 0. 1. 1.5 ] - [0.02 0. 0.5 0.75 ] + [0.02 0. 0. 0. ] [0.025 0. 0. 0. ]] ---------------------------------------------- @@ -212,7 +212,7 @@ def _check_table_upper_boundary(self): logging.warning( "The energy provided in the table does not decay to 0 at rcut." ) - # if rcut goes beyond table upper bond, need extrapolation. + # if rcut goes beyond table upper bond, need extrapolation, ensure values decay to `0` before rcut. else: logging.warning( "The rcut goes beyond table upper boundary, performing linear extrapolation." @@ -221,9 +221,9 @@ def _check_table_upper_boundary(self): pad_linear[:, 0] = np.linspace( upper, increment * (rcut_idx + 1), rcut_idx - upper_idx + 1 ) - pad_linear[:, 1:] = np.array( + pad_linear[:-1, 1:] = np.array( [ - np.linspace(start, 0, rcut_idx - upper_idx + 1) + np.linspace(start, 0, rcut_idx - upper_idx) for start in upper_val ] ).T diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index 60c1bc01f3..f4f5a60153 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -12,7 +12,7 @@ ) -class TestPairTab(unittest.TestCase): +class TestPairTabCalculation(unittest.TestCase): @patch("numpy.loadtxt") def setUp(self, mock_loadtxt) -> None: file_path = "dummy_path" @@ -54,9 +54,9 @@ def test_without_mask(self): result = self.model.forward_atomic( self.extended_coord, self.extended_atype, self.nlist ) - expected_result = torch.tensor([[2.4000, 2.7085], [2.4000, 0.8000]]) + expected_result = torch.tensor([[1.2000, 1.3542], [1.2000, 0.4000]]) - torch.testing.assert_allclose(result, expected_result) + torch.testing.assert_allclose(result['energy'], expected_result) def test_with_mask(self): self.nlist = torch.tensor([[[1, -1], [0, 2]], [[1, 2], [0, 3]]]) @@ -64,9 +64,9 @@ def test_with_mask(self): result = self.model.forward_atomic( self.extended_coord, self.extended_atype, self.nlist ) - expected_result = torch.tensor([[1.6000, 2.7085], [2.4000, 0.8000]]) + expected_result = torch.tensor([[0.8000, 1.3542], [1.2000, 0.4000]]) - torch.testing.assert_allclose(result, expected_result) + torch.testing.assert_allclose(result['energy'], expected_result) def test_jit(self): model = torch.jit.script(self.model) From 8fee8fa77c0aabd586b1e1618682bd0259105709 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 28 Jan 2024 10:03:09 +0000 Subject: [PATCH 06/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/pair_tab.py | 5 +---- source/tests/pt/test_pairtab.py | 4 ++-- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 6ae6a051e9..73343b0478 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -222,10 +222,7 @@ def _check_table_upper_boundary(self): upper, increment * (rcut_idx + 1), rcut_idx - upper_idx + 1 ) pad_linear[:-1, 1:] = np.array( - [ - np.linspace(start, 0, rcut_idx - upper_idx) - for start in upper_val - ] + [np.linspace(start, 0, rcut_idx - upper_idx) for start in upper_val] ).T raw_data = np.concatenate((raw_data[:-1, :], pad_linear), axis=0) diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index f4f5a60153..db28e863e0 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -56,7 +56,7 @@ def test_without_mask(self): ) expected_result = torch.tensor([[1.2000, 1.3542], [1.2000, 0.4000]]) - torch.testing.assert_allclose(result['energy'], expected_result) + torch.testing.assert_allclose(result["energy"], expected_result) def test_with_mask(self): self.nlist = torch.tensor([[[1, -1], [0, 2]], [[1, 2], [0, 3]]]) @@ -66,7 +66,7 @@ def test_with_mask(self): ) expected_result = torch.tensor([[0.8000, 1.3542], [1.2000, 0.4000]]) - torch.testing.assert_allclose(result['energy'], expected_result) + torch.testing.assert_allclose(result["energy"], expected_result) def test_jit(self): model = torch.jit.script(self.model) From ff0851570fc2a796df74efb5f42f512f0e02c375 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Sun, 28 Jan 2024 18:35:18 +0800 Subject: [PATCH 07/20] fix: update allclose precision --- source/tests/pt/test_pairtab.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index db28e863e0..98b6ca9806 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -12,7 +12,7 @@ ) -class TestPairTabCalculation(unittest.TestCase): +class TestPairTab(unittest.TestCase): @patch("numpy.loadtxt") def setUp(self, mock_loadtxt) -> None: file_path = "dummy_path" @@ -56,7 +56,7 @@ def test_without_mask(self): ) expected_result = torch.tensor([[1.2000, 1.3542], [1.2000, 0.4000]]) - torch.testing.assert_allclose(result["energy"], expected_result) + torch.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) def test_with_mask(self): self.nlist = torch.tensor([[[1, -1], [0, 2]], [[1, 2], [0, 3]]]) @@ -66,7 +66,7 @@ def test_with_mask(self): ) expected_result = torch.tensor([[0.8000, 1.3542], [1.2000, 0.4000]]) - torch.testing.assert_allclose(result["energy"], expected_result) + torch.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) def test_jit(self): model = torch.jit.script(self.model) From 8cbb98c940f26e59cbcb9cc14c9b5735476055e6 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Mon, 29 Jan 2024 13:31:27 +0800 Subject: [PATCH 08/20] chore: refactor common method to PairTab --- deepmd/pt/model/model/pair_tab.py | 100 ++---------------- deepmd/utils/pair_tab.py | 109 ++++++++++++++++++-- source/tests/pt/test_pairtab.py | 2 +- source/tests/test_pairtab_preprocess.py | 131 ++++++++++++++++++++++++ 4 files changed, 240 insertions(+), 102 deletions(-) create mode 100644 source/tests/test_pairtab_preprocess.py diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 73343b0478..e33716dbcf 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -57,10 +57,7 @@ def __init__( self.tab_file = tab_file self.rcut = rcut - # check table data against rcut and update tab_file if needed. - self._check_table_upper_boundary() - - self.tab = PairTab(self.tab_file) + self.tab = PairTab(self.tab_file, rcut=rcut) self.ntypes = self.tab.ntypes tab_info, tab_data = self.tab.get() # this returns -> Tuple[np.array, np.array] @@ -140,96 +137,6 @@ def forward_atomic( return {"energy": atomic_energy} - def _check_table_upper_boundary(self): - """Update User Provided Table Based on `rcut`. - - This function checks the upper boundary provided in the table against rcut. - If the table upper boundary values decay to zero before rcut, padding zeros will - be added to the table to cover rcut; if the table upper boundary values do not decay to zero - before ruct, linear extrapolation will be performed till rcut. In both cases, the table file - will be overwritten. - - Examples - -------- - table = [[0.005 1. 2. 3. ] - [0.01 0.8 1.6 2.4 ] - [0.015 0. 1. 1.5 ]] - - rcut = 0.022 - - new_table = [[0.005 1. 2. 3. ] - [0.01 0.8 1.6 2.4 ] - [0.015 0. 1. 1.5 ] - [0.02 0. 0. 0. ] - [0.025 0. 0. 0. ]] - - ---------------------------------------------- - - table = [[0.005 1. 2. 3. ] - [0.01 0.8 1.6 2.4 ] - [0.015 0.5 1. 1.5 ] - [0.02 0.25 0.4 0.75 ] - [0.025 0. 0.1 0. ] - [0.03 0. 0. 0. ]] - - rcut = 0.031 - - new_table = [[0.005 1. 2. 3. ] - [0.01 0.8 1.6 2.4 ] - [0.015 0.5 1. 1.5 ] - [0.02 0.25 0.4 0.75 ] - [0.025 0. 0.1 0. ] - [0.03 0. 0. 0. ] - [0.035 0. 0. 0. ]] - """ - raw_data = np.loadtxt(self.tab_file) - upper = raw_data[-1][0] - upper_val = raw_data[-1][1:] - upper_idx = raw_data.shape[0] - 1 - increment = raw_data[1][0] - raw_data[0][0] - - # the index of table for the grid point right after rcut - rcut_idx = int(self.rcut / increment) - - if np.all(upper_val == 0): - # if table values decay to `0` after rcut - if self.rcut < upper and np.any(raw_data[rcut_idx - 1][1:] != 0): - logging.warning( - "The energy provided in the table does not decay to 0 at rcut." - ) - # if table values decay to `0` at rcut, do nothing - - # if table values decay to `0` before rcut, pad table with `0`s. - elif self.rcut > upper: - pad_zero = np.zeros((rcut_idx - upper_idx, 4)) - pad_zero[:, 0] = np.linspace( - upper + increment, increment * (rcut_idx + 1), rcut_idx - upper_idx - ) - raw_data = np.concatenate((raw_data, pad_zero), axis=0) - else: - # if table values do not decay to `0` at rcut - if self.rcut <= upper: - logging.warning( - "The energy provided in the table does not decay to 0 at rcut." - ) - # if rcut goes beyond table upper bond, need extrapolation, ensure values decay to `0` before rcut. - else: - logging.warning( - "The rcut goes beyond table upper boundary, performing linear extrapolation." - ) - pad_linear = np.zeros((rcut_idx - upper_idx + 1, 4)) - pad_linear[:, 0] = np.linspace( - upper, increment * (rcut_idx + 1), rcut_idx - upper_idx + 1 - ) - pad_linear[:-1, 1:] = np.array( - [np.linspace(start, 0, rcut_idx - upper_idx) for start in upper_val] - ).T - raw_data = np.concatenate((raw_data[:-1, :], pad_linear), axis=0) - - # over writing file with padding if applicable. - with open(self.tab_file, "wb") as f: - np.savetxt(f, raw_data) - def _pair_tabulated_inter( self, nlist: torch.Tensor, @@ -286,6 +193,11 @@ def _pair_tabulated_inter( final_coef = self._extract_spline_coefficient(i_type, j_type, idx) + # here we need to do postprocess to overwrite coefficients when table values are not zero at rcut in linear extrapolation. + if self.tab.rmax < self.rcut: + post_mask = rr >= self.rcut + final_coef[post_mask] = torch.zeros(4) + a3, a2, a1, a0 = torch.unbind(final_coef, dim=-1) # 4 * (nframes, nloc, nnei) etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 4451f53379..9433332c35 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -3,6 +3,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Tuple, + Optional ) import numpy as np @@ -10,6 +11,8 @@ CubicSpline, ) +import logging + class PairTab: """Pairwise tabulated potential. @@ -25,11 +28,11 @@ class PairTab: The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. """ - def __init__(self, filename: str) -> None: + def __init__(self, filename: str, rcut: Optional[float] = None) -> None: """Constructor.""" - self.reinit(filename) + self.reinit(filename, rcut) - def reinit(self, filename: str) -> None: + def reinit(self, filename: str, rcut: Optional[float] = None) -> None: """Initialize the tabulated interaction. Parameters @@ -44,8 +47,8 @@ def reinit(self, filename: str) -> None: """ self.vdata = np.loadtxt(filename) self.rmin = self.vdata[0][0] + self.rmax = self.vdata[-1][0] self.hh = self.vdata[1][0] - self.vdata[0][0] - self.nspline = self.vdata.shape[0] - 1 ncol = self.vdata.shape[1] - 1 n0 = (-1 + np.sqrt(1 + 8 * ncol)) * 0.5 self.ntypes = int(n0 + 0.1) @@ -53,14 +56,105 @@ def reinit(self, filename: str) -> None: "number of volumes provided in %s does not match guessed number of types %d" % (filename, self.ntypes) ) + + # check table data against rcut and update tab_file if needed, table upper boundary is used as rcut if not provided. + self.rcut = rcut if rcut is not None else self.rmax + self._check_table_upper_boundary() + self.nspline = self.vdata.shape[0] - 1 self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) self.tab_data = self._make_data() + def _check_table_upper_boundary(self) -> None: + """Update User Provided Table Based on `rcut`. + + This function checks the upper boundary provided in the table against rcut. + If the table upper boundary values decay to zero before rcut, padding zeros will + be added to the table to cover rcut; if the table upper boundary values do not decay to zero + before ruct, linear extrapolation will be performed till rcut. + + Examples + -------- + table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0. 1. 1.5 ]] + + rcut = 0.022 + + new_table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0. 1. 1.5 ] + [0.02 0. 0. 0. ] + [0.025 0. 0. 0. ]] + + ---------------------------------------------- + + table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0.5 1. 1.5 ] + [0.02 0.25 0.4 0.75 ] + [0.025 0. 0.1 0. ] + [0.03 0. 0. 0. ]] + + rcut = 0.031 + + new_table = [[0.005 1. 2. 3. ] + [0.01 0.8 1.6 2.4 ] + [0.015 0.5 1. 1.5 ] + [0.02 0.25 0.4 0.75 ] + [0.025 0. 0.1 0. ] + [0.03 0. 0. 0. ] + [0.035 0. 0. 0. ]] + """ + + upper_val = self.vdata[-1][1:] + upper_idx = self.vdata.shape[0] - 1 + + # the index of table for the grid point right after rcut + rcut_idx = int(self.rcut / self.hh) + + if np.all(upper_val == 0): + # if table values decay to `0` after rcut + if self.rcut < self.rmax and np.any(self.vdata[rcut_idx - 1][1:] != 0): + logging.warning( + "The energy provided in the table does not decay to 0 at rcut." + ) + # if table values decay to `0` at rcut, do nothing + + # if table values decay to `0` before rcut, pad table with `0`s. + elif self.rcut > self.rmax: + pad_zero = np.zeros((rcut_idx - upper_idx, 4)) + pad_zero[:, 0] = np.linspace( + self.rmax + self.hh, self.hh * (rcut_idx + 1), rcut_idx - upper_idx + ) + self.vdata = np.concatenate((self.vdata, pad_zero), axis=0) + else: + # if table values do not decay to `0` at rcut + if self.rcut <= self.rmax: + logging.warning( + "The energy provided in the table does not decay to 0 at rcut." + ) + # if rcut goes beyond table upper bond, need extrapolation, ensure values decay to `0` before rcut. + else: + logging.warning( + "The rcut goes beyond table upper boundary, performing linear extrapolation." + ) + pad_linear = np.zeros((rcut_idx - upper_idx + 1, 4)) + pad_linear[:, 0] = np.linspace( + self.rmax, self.hh * (rcut_idx + 1), rcut_idx - upper_idx + 1 + ) + pad_linear[:-1, 1:] = np.array( + [np.linspace(start, 0, rcut_idx - upper_idx) for start in upper_val] + ).T + self.vdata = np.concatenate((self.vdata[:-1, :], pad_linear), axis=0) + + def get(self) -> Tuple[np.array, np.array]: """Get the serialized table.""" return self.tab_info, self.tab_data def _make_data(self): + + # here we need to do postprocess, to overwrite coefficients when padding zeros resulting in negative energies. data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) stride = 4 * self.nspline idx_iter = 0 @@ -73,11 +167,12 @@ def _make_data(self): dd *= self.hh dtmp = np.zeros(stride) for ii in range(self.nspline): - dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] + # check if vv is zero, if so, that's case 1, set all coefficients to 0, + dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] if vv[ii] != 0 else 0 dtmp[ii * 4 + 1] = ( -3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1] - ) - dtmp[ii * 4 + 2] = dd[ii] + ) if vv[ii] != 0 else 0 + dtmp[ii * 4 + 2] = dd[ii] if vv[ii] != 0 else 0 dtmp[ii * 4 + 3] = vv[ii] data[ (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index 98b6ca9806..b27eb1dbaf 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -25,7 +25,7 @@ def setUp(self, mock_loadtxt) -> None: ] ) - self.model = PairTabModel(tab_file=file_path, rcut=0.1, sel=2) + self.model = PairTabModel(tab_file=file_path, rcut=0.02, sel=2) self.extended_coord = torch.tensor( [ diff --git a/source/tests/test_pairtab_preprocess.py b/source/tests/test_pairtab_preprocess.py new file mode 100644 index 0000000000..62aa123b14 --- /dev/null +++ b/source/tests/test_pairtab_preprocess.py @@ -0,0 +1,131 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from unittest.mock import ( + patch, +) + +import numpy as np +from deepmd.utils.pair_tab import ( + PairTab, +) + + +class TestPairTabPreprocessLinear(unittest.TestCase): + @patch("numpy.loadtxt") + def setUp(self, mock_loadtxt) -> None: + file_path = "dummy_path" + mock_loadtxt.return_value = np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + ] + ) + + self.tab1 = PairTab(filename=file_path, rcut=0.01) + self.tab2 = PairTab(filename=file_path, rcut=0.02) + self.tab3 = PairTab(filename=file_path, rcut=0.022) + self.tab4 = PairTab(filename=file_path, rcut=0.03) + + + def test_preprocess(self): + + np.testing.assert_allclose(self.tab1.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + ] + )) + np.testing.assert_allclose(self.tab2.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + ] + )) + + # for this test case, the table does not decay to zero at rcut = 0.22, + # in the cubic spline code, we use a fixed size grid, if will be a problem if we introduce variable gird size. + # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. + np.testing.assert_allclose(self.tab3.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0., 0.], + ] + )) + np.testing.assert_allclose(self.tab4.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.125, 0.2, 0.375], + [0.03, 0., 0., 0.], + [0.035, 0., 0., 0.], + ] + )) + +class TestPairTabPreprocessZero(unittest.TestCase): + @patch("numpy.loadtxt") + def setUp(self, mock_loadtxt) -> None: + file_path = "dummy_path" + mock_loadtxt.return_value = np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0., 0.], + ] + ) + + self.tab1 = PairTab(filename=file_path, rcut=0.023) + self.tab2 = PairTab(filename=file_path, rcut=0.025) + self.tab3 = PairTab(filename=file_path, rcut=0.028) + + + def test_preprocess(self): + + np.testing.assert_allclose(self.tab1.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0., 0.], + ] + )) + np.testing.assert_allclose(self.tab2.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0., 0.], + ] + )) + + np.testing.assert_allclose(self.tab3.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0., 0.], + [0.03, 0., 0., 0.], + ] + )) + np.testing.assert_equal(self.tab3.nspline,5) + + # for this test case, padding zeros between 0.025 and 0.03 will cause the cubic spline go below zero and result in negative energy values, + # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. + temp_data = self.tab3.tab_data.reshape(2,2,-1,4) + np.testing.assert_allclose(temp_data[:,:,-1,:], np.zeros((2,2,4))) + \ No newline at end of file From a08092ce0b7f9d51d22e847e198b1cb73b3b061f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 29 Jan 2024 05:32:15 +0000 Subject: [PATCH 09/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/pair_tab.py | 2 - deepmd/utils/pair_tab.py | 20 +-- source/tests/test_pairtab_preprocess.py | 166 +++++++++++++----------- 3 files changed, 103 insertions(+), 85 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index e33716dbcf..cc3843e334 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import logging from typing import ( Dict, List, @@ -7,7 +6,6 @@ Union, ) -import numpy as np import torch from torch import ( nn, diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 9433332c35..9f7f07671d 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -1,9 +1,10 @@ #!/usr/bin/env python3 # SPDX-License-Identifier: LGPL-3.0-or-later +import logging from typing import ( + Optional, Tuple, - Optional ) import numpy as np @@ -11,8 +12,6 @@ CubicSpline, ) -import logging - class PairTab: """Pairwise tabulated potential. @@ -105,7 +104,6 @@ def _check_table_upper_boundary(self) -> None: [0.03 0. 0. 0. ] [0.035 0. 0. 0. ]] """ - upper_val = self.vdata[-1][1:] upper_idx = self.vdata.shape[0] - 1 @@ -147,13 +145,11 @@ def _check_table_upper_boundary(self) -> None: ).T self.vdata = np.concatenate((self.vdata[:-1, :], pad_linear), axis=0) - def get(self) -> Tuple[np.array, np.array]: """Get the serialized table.""" return self.tab_info, self.tab_data def _make_data(self): - # here we need to do postprocess, to overwrite coefficients when padding zeros resulting in negative energies. data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) stride = 4 * self.nspline @@ -168,10 +164,16 @@ def _make_data(self): dtmp = np.zeros(stride) for ii in range(self.nspline): # check if vv is zero, if so, that's case 1, set all coefficients to 0, - dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] if vv[ii] != 0 else 0 + dtmp[ii * 4 + 0] = ( + 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] + if vv[ii] != 0 + else 0 + ) dtmp[ii * 4 + 1] = ( - -3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1] - ) if vv[ii] != 0 else 0 + (-3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1]) + if vv[ii] != 0 + else 0 + ) dtmp[ii * 4 + 2] = dd[ii] if vv[ii] != 0 else 0 dtmp[ii * 4 + 3] = vv[ii] data[ diff --git a/source/tests/test_pairtab_preprocess.py b/source/tests/test_pairtab_preprocess.py index 62aa123b14..e5ebd6d028 100644 --- a/source/tests/test_pairtab_preprocess.py +++ b/source/tests/test_pairtab_preprocess.py @@ -5,6 +5,7 @@ ) import numpy as np + from deepmd.utils.pair_tab import ( PairTab, ) @@ -27,50 +28,61 @@ def setUp(self, mock_loadtxt) -> None: self.tab2 = PairTab(filename=file_path, rcut=0.02) self.tab3 = PairTab(filename=file_path, rcut=0.022) self.tab4 = PairTab(filename=file_path, rcut=0.03) - def test_preprocess(self): - - np.testing.assert_allclose(self.tab1.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - ] - )) - np.testing.assert_allclose(self.tab2.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - ] - )) + np.testing.assert_allclose( + self.tab1.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + ] + ), + ) + np.testing.assert_allclose( + self.tab2.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + ] + ), + ) # for this test case, the table does not decay to zero at rcut = 0.22, # in the cubic spline code, we use a fixed size grid, if will be a problem if we introduce variable gird size. # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. - np.testing.assert_allclose(self.tab3.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0., 0.], - ] - )) - np.testing.assert_allclose(self.tab4.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0.125, 0.2, 0.375], - [0.03, 0., 0., 0.], - [0.035, 0., 0., 0.], - ] - )) + np.testing.assert_allclose( + self.tab3.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.0, 0.0], + ] + ), + ) + np.testing.assert_allclose( + self.tab4.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.125, 0.2, 0.375], + [0.03, 0.0, 0.0, 0.0], + [0.035, 0.0, 0.0, 0.0], + ] + ), + ) + class TestPairTabPreprocessZero(unittest.TestCase): @patch("numpy.loadtxt") @@ -82,50 +94,56 @@ def setUp(self, mock_loadtxt) -> None: [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0., 0.], + [0.025, 0.0, 0.0, 0.0], ] ) self.tab1 = PairTab(filename=file_path, rcut=0.023) self.tab2 = PairTab(filename=file_path, rcut=0.025) self.tab3 = PairTab(filename=file_path, rcut=0.028) - def test_preprocess(self): - - np.testing.assert_allclose(self.tab1.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0., 0.], - ] - )) - np.testing.assert_allclose(self.tab2.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0., 0.], - ] - )) - - np.testing.assert_allclose(self.tab3.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0., 0.], - [0.03, 0., 0., 0.], - ] - )) - np.testing.assert_equal(self.tab3.nspline,5) + np.testing.assert_allclose( + self.tab1.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.0, 0.0], + ] + ), + ) + np.testing.assert_allclose( + self.tab2.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.0, 0.0], + ] + ), + ) + + np.testing.assert_allclose( + self.tab3.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.0, 0.0], + [0.03, 0.0, 0.0, 0.0], + ] + ), + ) + np.testing.assert_equal(self.tab3.nspline, 5) # for this test case, padding zeros between 0.025 and 0.03 will cause the cubic spline go below zero and result in negative energy values, # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. - temp_data = self.tab3.tab_data.reshape(2,2,-1,4) - np.testing.assert_allclose(temp_data[:,:,-1,:], np.zeros((2,2,4))) - \ No newline at end of file + temp_data = self.tab3.tab_data.reshape(2, 2, -1, 4) + np.testing.assert_allclose(temp_data[:, :, -1, :], np.zeros((2, 2, 4))) From d3090b97b21e3be42b0900a9f9adb38229661ff4 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Mon, 29 Jan 2024 14:12:35 +0800 Subject: [PATCH 10/20] fix: update unit tests --- deepmd/utils/pair_tab.py | 2 +- source/tests/test_pairtab_preprocess.py | 71 +++++++++++++++++++++++-- 2 files changed, 69 insertions(+), 4 deletions(-) diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 9f7f07671d..719a9a9ed7 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -127,7 +127,7 @@ def _check_table_upper_boundary(self) -> None: self.vdata = np.concatenate((self.vdata, pad_zero), axis=0) else: # if table values do not decay to `0` at rcut - if self.rcut <= self.rmax: + if self.rcut < self.rmax: logging.warning( "The energy provided in the table does not decay to 0 at rcut." ) diff --git a/source/tests/test_pairtab_preprocess.py b/source/tests/test_pairtab_preprocess.py index e5ebd6d028..a500ae1c24 100644 --- a/source/tests/test_pairtab_preprocess.py +++ b/source/tests/test_pairtab_preprocess.py @@ -24,7 +24,7 @@ def setUp(self, mock_loadtxt) -> None: ] ) - self.tab1 = PairTab(filename=file_path, rcut=0.01) + self.tab1 = PairTab(filename=file_path, rcut=0.028) self.tab2 = PairTab(filename=file_path, rcut=0.02) self.tab3 = PairTab(filename=file_path, rcut=0.022) self.tab4 = PairTab(filename=file_path, rcut=0.03) @@ -38,6 +38,8 @@ def test_preprocess(self): [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.0, 0.0], + [0.03, 0., 0., 0.], ] ), ) @@ -49,6 +51,7 @@ def test_preprocess(self): [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0., 0.], ] ), ) @@ -145,5 +148,67 @@ def test_preprocess(self): # for this test case, padding zeros between 0.025 and 0.03 will cause the cubic spline go below zero and result in negative energy values, # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. - temp_data = self.tab3.tab_data.reshape(2, 2, -1, 4) - np.testing.assert_allclose(temp_data[:, :, -1, :], np.zeros((2, 2, 4))) + temp_data = self.tab3.tab_data.reshape(2,2,-1,4) + np.testing.assert_allclose(temp_data[:,:,-1,:], np.zeros((2,2,4))) + + +class TestPairTabPreprocessUneven(unittest.TestCase): + @patch("numpy.loadtxt") + def setUp(self, mock_loadtxt) -> None: + file_path = "dummy_path" + mock_loadtxt.return_value = np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0.1, 0.], + ] + ) + + self.tab1 = PairTab(filename=file_path, rcut=0.025) + self.tab2 = PairTab(filename=file_path, rcut=0.028) + self.tab3 = PairTab(filename=file_path, rcut=0.03) + + + def test_preprocess(self): + + np.testing.assert_allclose(self.tab1.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0.1, 0.], + [0.03, 0., 0., 0.], + ] + )) + np.testing.assert_allclose(self.tab2.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0.1, 0.], + [0.03, 0., 0., 0.], + ] + )) + np.testing.assert_allclose(self.tab3.vdata, np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0., 0.1, 0.], + [0.03, 0., 0., 0.], + [0.035, 0., 0., 0.], + ] + )) + + temp_data = self.tab3.tab_data.reshape(2,2,-1,4) + + np.testing.assert_allclose(temp_data[0,0,-2:], np.zeros((2,4))) + np.testing.assert_allclose(temp_data[1,1,-2:], np.zeros((2,4))) + np.testing.assert_allclose(temp_data[0,1,-1:], np.zeros((1,4))) + np.testing.assert_allclose(temp_data[1,0,-1:], np.zeros((1,4))) + \ No newline at end of file From daf2fc6ab3d1240af4abfdd4ccceade35c90712f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 29 Jan 2024 06:19:35 +0000 Subject: [PATCH 11/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/test_pairtab_preprocess.py | 98 +++++++++++++------------ 1 file changed, 52 insertions(+), 46 deletions(-) diff --git a/source/tests/test_pairtab_preprocess.py b/source/tests/test_pairtab_preprocess.py index a500ae1c24..d9ce624b2c 100644 --- a/source/tests/test_pairtab_preprocess.py +++ b/source/tests/test_pairtab_preprocess.py @@ -38,8 +38,8 @@ def test_preprocess(self): [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], - [0.025, 0.0, 0.0, 0.0], - [0.03, 0., 0., 0.], + [0.025, 0.0, 0.0, 0.0], + [0.03, 0.0, 0.0, 0.0], ] ), ) @@ -51,7 +51,7 @@ def test_preprocess(self): [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0., 0.], + [0.025, 0.0, 0.0, 0.0], ] ), ) @@ -148,9 +148,9 @@ def test_preprocess(self): # for this test case, padding zeros between 0.025 and 0.03 will cause the cubic spline go below zero and result in negative energy values, # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. - temp_data = self.tab3.tab_data.reshape(2,2,-1,4) - np.testing.assert_allclose(temp_data[:,:,-1,:], np.zeros((2,2,4))) - + temp_data = self.tab3.tab_data.reshape(2, 2, -1, 4) + np.testing.assert_allclose(temp_data[:, :, -1, :], np.zeros((2, 2, 4))) + class TestPairTabPreprocessUneven(unittest.TestCase): @patch("numpy.loadtxt") @@ -162,53 +162,59 @@ def setUp(self, mock_loadtxt) -> None: [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0.1, 0.], + [0.025, 0.0, 0.1, 0.0], ] ) self.tab1 = PairTab(filename=file_path, rcut=0.025) self.tab2 = PairTab(filename=file_path, rcut=0.028) self.tab3 = PairTab(filename=file_path, rcut=0.03) - def test_preprocess(self): - - np.testing.assert_allclose(self.tab1.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0.1, 0.], - [0.03, 0., 0., 0.], - ] - )) - np.testing.assert_allclose(self.tab2.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0.1, 0.], - [0.03, 0., 0., 0.], - ] - )) - np.testing.assert_allclose(self.tab3.vdata, np.array( - [ - [0.005, 1.0, 2.0, 3.0], - [0.01, 0.8, 1.6, 2.4], - [0.015, 0.5, 1.0, 1.5], - [0.02, 0.25, 0.4, 0.75], - [0.025, 0., 0.1, 0.], - [0.03, 0., 0., 0.], - [0.035, 0., 0., 0.], - ] - )) + np.testing.assert_allclose( + self.tab1.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.1, 0.0], + [0.03, 0.0, 0.0, 0.0], + ] + ), + ) + np.testing.assert_allclose( + self.tab2.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.1, 0.0], + [0.03, 0.0, 0.0, 0.0], + ] + ), + ) + np.testing.assert_allclose( + self.tab3.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.1, 0.0], + [0.03, 0.0, 0.0, 0.0], + [0.035, 0.0, 0.0, 0.0], + ] + ), + ) - temp_data = self.tab3.tab_data.reshape(2,2,-1,4) + temp_data = self.tab3.tab_data.reshape(2, 2, -1, 4) - np.testing.assert_allclose(temp_data[0,0,-2:], np.zeros((2,4))) - np.testing.assert_allclose(temp_data[1,1,-2:], np.zeros((2,4))) - np.testing.assert_allclose(temp_data[0,1,-1:], np.zeros((1,4))) - np.testing.assert_allclose(temp_data[1,0,-1:], np.zeros((1,4))) - \ No newline at end of file + np.testing.assert_allclose(temp_data[0, 0, -2:], np.zeros((2, 4))) + np.testing.assert_allclose(temp_data[1, 1, -2:], np.zeros((2, 4))) + np.testing.assert_allclose(temp_data[0, 1, -1:], np.zeros((1, 4))) + np.testing.assert_allclose(temp_data[1, 0, -1:], np.zeros((1, 4))) From 399c2780662a01abe4aa819dbc9b663984449b61 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Mon, 29 Jan 2024 14:54:00 +0800 Subject: [PATCH 12/20] fix: revert padding zero mask change --- deepmd/utils/pair_tab.py | 6 +----- source/tests/{ => tf}/test_pairtab_preprocess.py | 12 ------------ 2 files changed, 1 insertion(+), 17 deletions(-) rename source/tests/{ => tf}/test_pairtab_preprocess.py (88%) diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 719a9a9ed7..30fd74ddde 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -166,15 +166,11 @@ def _make_data(self): # check if vv is zero, if so, that's case 1, set all coefficients to 0, dtmp[ii * 4 + 0] = ( 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] - if vv[ii] != 0 - else 0 ) dtmp[ii * 4 + 1] = ( (-3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1]) - if vv[ii] != 0 - else 0 ) - dtmp[ii * 4 + 2] = dd[ii] if vv[ii] != 0 else 0 + dtmp[ii * 4 + 2] = dd[ii] dtmp[ii * 4 + 3] = vv[ii] data[ (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride diff --git a/source/tests/test_pairtab_preprocess.py b/source/tests/tf/test_pairtab_preprocess.py similarity index 88% rename from source/tests/test_pairtab_preprocess.py rename to source/tests/tf/test_pairtab_preprocess.py index d9ce624b2c..e2b2183f2e 100644 --- a/source/tests/test_pairtab_preprocess.py +++ b/source/tests/tf/test_pairtab_preprocess.py @@ -144,12 +144,6 @@ def test_preprocess(self): ] ), ) - np.testing.assert_equal(self.tab3.nspline, 5) - - # for this test case, padding zeros between 0.025 and 0.03 will cause the cubic spline go below zero and result in negative energy values, - # we will do post process to overwrite spline coefficient `a3`,`a2`,`a1`,`a0`, to ensure energy decays to `0`. - temp_data = self.tab3.tab_data.reshape(2, 2, -1, 4) - np.testing.assert_allclose(temp_data[:, :, -1, :], np.zeros((2, 2, 4))) class TestPairTabPreprocessUneven(unittest.TestCase): @@ -212,9 +206,3 @@ def test_preprocess(self): ), ) - temp_data = self.tab3.tab_data.reshape(2, 2, -1, 4) - - np.testing.assert_allclose(temp_data[0, 0, -2:], np.zeros((2, 4))) - np.testing.assert_allclose(temp_data[1, 1, -2:], np.zeros((2, 4))) - np.testing.assert_allclose(temp_data[0, 1, -1:], np.zeros((1, 4))) - np.testing.assert_allclose(temp_data[1, 0, -1:], np.zeros((1, 4))) From 59abe433d0dd3934ee308fdea2532fcdb8254573 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 29 Jan 2024 06:55:19 +0000 Subject: [PATCH 13/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/utils/pair_tab.py | 8 +++----- source/tests/tf/test_pairtab_preprocess.py | 1 - 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 30fd74ddde..832f2a84fd 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -164,13 +164,11 @@ def _make_data(self): dtmp = np.zeros(stride) for ii in range(self.nspline): # check if vv is zero, if so, that's case 1, set all coefficients to 0, - dtmp[ii * 4 + 0] = ( - 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] - ) + dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] dtmp[ii * 4 + 1] = ( - (-3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1]) + -3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1] ) - dtmp[ii * 4 + 2] = dd[ii] + dtmp[ii * 4 + 2] = dd[ii] dtmp[ii * 4 + 3] = vv[ii] data[ (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride diff --git a/source/tests/tf/test_pairtab_preprocess.py b/source/tests/tf/test_pairtab_preprocess.py index e2b2183f2e..aee6b3889e 100644 --- a/source/tests/tf/test_pairtab_preprocess.py +++ b/source/tests/tf/test_pairtab_preprocess.py @@ -205,4 +205,3 @@ def test_preprocess(self): ] ), ) - From 1c4ee0dd3c621c4b2a793fa34752158a0ec24215 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Tue, 30 Jan 2024 13:52:10 +0800 Subject: [PATCH 14/20] feat: redo extrapolation with cubic spline for smoothness --- deepmd/pt/model/model/pair_tab.py | 139 +++++++++++++++++++++++++----- deepmd/utils/pair_tab.py | 11 ++- source/tests/pt/test_pairtab.py | 63 ++++++++++++++ 3 files changed, 187 insertions(+), 26 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index cc3843e334..ea43b66864 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -6,6 +6,11 @@ Union, ) +from scipy.interpolate import ( + CubicSpline, +) + +import numpy as np import torch from torch import ( nn, @@ -99,12 +104,12 @@ def forward_atomic( mapping: Optional[torch.Tensor] = None, do_atomic_virial: bool = False, ) -> Dict[str, torch.Tensor]: - nframes, nloc, nnei = nlist.shape + self.nframes, self.nloc, self.nnei = nlist.shape # this will mask all -1 in the nlist masked_nlist = torch.clamp(nlist, 0) - atype = extended_atype[:, :nloc] # (nframes, nloc) + atype = extended_atype[:, :self.nloc] # (nframes, nloc) pairwise_dr = self._get_pairwise_dist( extended_coord ) # (nframes, nall, nall, 3) @@ -122,7 +127,7 @@ def forward_atomic( ] # slice rr to get (nframes, nloc, nnei) - rr = torch.gather(pairwise_rr[:, :nloc, :], 2, masked_nlist) + rr = torch.gather(pairwise_rr[:, :self.nloc, :], 2, masked_nlist) raw_atomic_energy = self._pair_tabulated_inter(nlist, atype, j_type, rr) @@ -189,19 +194,102 @@ def _pair_tabulated_inter( uu -= idx - final_coef = self._extract_spline_coefficient(i_type, j_type, idx) - - # here we need to do postprocess to overwrite coefficients when table values are not zero at rcut in linear extrapolation. - if self.tab.rmax < self.rcut: - post_mask = rr >= self.rcut - final_coef[post_mask] = torch.zeros(4) - - a3, a2, a1, a0 = torch.unbind(final_coef, dim=-1) # 4 * (nframes, nloc, nnei) - - etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. - ener = etmp * uu + a0 + table_coef = self._extract_spline_coefficient(i_type, j_type, idx, self.tab_data, self.nspline) + table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4) + ener = self._calcualte_ener(table_coef, uu) + + # here we need to do postprocess to overwrite energy to zero beyond rcut. + if self.tab.rmax <= self.rcut: + mask_beyond_rcut = rr > self.rcut + ener[mask_beyond_rcut] = 0 + + # here we use smooth extrapolation to replace linear extrapolation. + extrapolation = self._extrapolate_rmax_rcut() + if extrapolation is not None: + uu_extrapolate = (rr - self.tab.rmax) / (self.rcut - self.tab.rmax) + clipped_uu = torch.clamp(uu_extrapolate, 0, 1) # clip rr within rmax. + extrapolate_coef = self._extract_spline_coefficient(i_type, j_type, torch.zeros_like(idx), extrapolation, 1) + extrapolate_coef = extrapolate_coef.reshape(self.nframes, self.nloc, self.nnei, 4) + ener_extrpolate = self._calcualte_ener(extrapolate_coef, clipped_uu) + mask_rmax_to_rcut = (self.tab.rmax < rr) & (rr <= self.rcut) + ener[mask_rmax_to_rcut] = ener_extrpolate[mask_rmax_to_rcut] return ener + def _extrapolate_rmax_rcut(self) -> torch.Tensor: + """Soomth extrapolation between table upper boundary and rcut. + + This method should only be used when the table upper boundary `rmax` is smaller than `rcut`, and + the table upper boundary values are not zeros. To simplify the problem, we use a single + cubic spline between `rmax` and `rcut` for each pair of atom types. One can substitute this extrapolation + to higher order polynomials if needed. + + There are two scenarios: + 1. `ruct` - `rmax` >= hh: + Set values at the grid point right before `rcut` to 0, and perform exterapolation between + the grid point and `rmax`, this allows smooth decay to 0 at `rcut`. + 2. `rcut` - `rmax` < hh: + Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`, + the enery beyond `rcut` will be overwritten to `0` latter. + + Returns + ------- + torch.Tensor + The cubic spline coefficients for each pair of atom types. (ntype, ntype, 1, 4) + """ + #check if decays to `0` at rmax, if yes, no extrapolation is needed. + rmax_val = torch.from_numpy(self.tab.vdata[self.tab.vdata[:,0] == self.tab.rmax]) + pre_rmax_val = torch.from_numpy(self.tab.vdata[self.tab.vdata[:,0] == self.tab.rmax - self.tab.hh]) + + if torch.all(rmax_val[:,1:] == 0): + return + else: + if self.rcut - self.tab.rmax >= self.tab.hh: + rcut_idx = int(self.rcut/self.tab.hh - self.tab.rmin/self.tab.hh) + rcut_val = torch.tensor(self.tab.vdata[rcut_idx,:]).reshape(1,-1) + grid = torch.concatenate([rmax_val, rcut_val],axis=0) + else: + # the last two rows will be the rmax, and rmax+hh + grid = torch.from_numpy(self.tab.vdata[-2:,:]) + passin_slope = ((rmax_val - pre_rmax_val)/self.tab.hh)[:,1:].squeeze(0) if ~np.all(pre_rmax_val == None) else 0 # the slope at the end of table for each ntype pairs (ntypes,ntypes,1) + extrapolate_coef = torch.from_numpy(self._calculate_spline_coef(grid, passin_slope)).reshape(self.ntypes,self.ntypes,4) + return extrapolate_coef.unsqueeze(2) + + # might be able to refactor this, combine with PairTab + def _calculate_spline_coef(self, grid, passin_slope): + data = np.zeros([self.ntypes * self.ntypes * 4]) + stride = 4 + idx_iter = 0 + + xx = grid[:, 0] + for t0 in range(self.ntypes): + for t1 in range(t0, self.ntypes): + vv = grid[:, 1 + idx_iter] + slope_idx = [t0 * (2 * self.ntypes - t0 - 1)//2 + t1] + + print(f"slope: {passin_slope[slope_idx]}") + cs = CubicSpline(xx, vv, bc_type=((1,passin_slope[slope_idx][0]),(1,0))) + dd = cs(xx, 1) + dd *= self.tab.hh + dtmp = np.zeros(stride) + dtmp[0] = ( + 2 * vv[0] - 2 * vv[1] + dd[0] + dd[1] + ) + dtmp[1] = ( + (-3 * vv[0] + 3 * vv[1] - 2 * dd[0] - dd[1]) + ) + dtmp[2] = dd[0] + dtmp[3] = vv[0] + data[ + (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride + + stride + ] = dtmp + data[ + (t1 * self.ntypes + t0) * stride : (t1 * self.ntypes + t0) * stride + + stride + ] = dtmp + idx_iter += 1 + return data + @staticmethod def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: """Get pairwise distance `dr`. @@ -240,8 +328,8 @@ def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: """ return coords.unsqueeze(2) - coords.unsqueeze(1) - def _extract_spline_coefficient( - self, i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor + @staticmethod + def _extract_spline_coefficient(i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor, tab_data: torch.Tensor, nspline: int ) -> torch.Tensor: """Extract the spline coefficient from the table. @@ -253,29 +341,40 @@ def _extract_spline_coefficient( The integer representation of atom type for all neighbour atoms of all local atoms for all frames. (nframes, nloc, nnei) idx : torch.Tensor The index of the spline coefficient. (nframes, nloc, nnei) + tab_data : torch.Tensor + The table storing all the spline coefficient. (ntype, ntype, nspline, 4) + nspline : int + The number of splines in the table. Returns ------- torch.Tensor - The spline coefficient. (nframes, nloc, nnei, 4) + The spline coefficient. (nframes, nloc, nnei, 4), shape maybe squeezed. """ # (nframes, nloc, nnei) expanded_i_type = i_type.unsqueeze(-1).expand(-1, -1, j_type.shape[-1]) # (nframes, nloc, nnei, nspline, 4) - expanded_tab_data = self.tab_data[expanded_i_type, j_type] + expanded_tab_data = tab_data[expanded_i_type, j_type] # (nframes, nloc, nnei, 1, 4) expanded_idx = idx.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, -1, -1, 4) # handle the case where idx is beyond the number of splines - clipped_indices = torch.clamp(expanded_idx, 0, self.nspline - 1).to(torch.int64) + clipped_indices = torch.clamp(expanded_idx, 0, nspline - 1).to(torch.int64) # (nframes, nloc, nnei, 4) final_coef = torch.gather(expanded_tab_data, 3, clipped_indices).squeeze() # when the spline idx is beyond the table, all spline coefficients are set to `0`, and the resulting ener corresponding to the idx is also `0`. - final_coef[expanded_idx.squeeze() >= self.nspline] = 0 + final_coef[expanded_idx.squeeze() >= nspline] = 0 return final_coef + + @staticmethod + def _calcualte_ener(coef, uu): + a3, a2, a1, a0 = torch.unbind(coef, dim=-1) # 4 * (nframes, nloc, nnei) + etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. + ener = etmp * uu + a0 # this energy has the linear extrapolated value when rcut > rmax + return ener \ No newline at end of file diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 832f2a84fd..b87ed04fd3 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -59,7 +59,7 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: # check table data against rcut and update tab_file if needed, table upper boundary is used as rcut if not provided. self.rcut = rcut if rcut is not None else self.rmax self._check_table_upper_boundary() - self.nspline = self.vdata.shape[0] - 1 + self.nspline = self.vdata.shape[0] - 1 # this nspline is updated based on the expanded table. self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) self.tab_data = self._make_data() @@ -106,7 +106,7 @@ def _check_table_upper_boundary(self) -> None: """ upper_val = self.vdata[-1][1:] upper_idx = self.vdata.shape[0] - 1 - + ncol = self.vdata.shape[1] # the index of table for the grid point right after rcut rcut_idx = int(self.rcut / self.hh) @@ -120,7 +120,7 @@ def _check_table_upper_boundary(self) -> None: # if table values decay to `0` before rcut, pad table with `0`s. elif self.rcut > self.rmax: - pad_zero = np.zeros((rcut_idx - upper_idx, 4)) + pad_zero = np.zeros((rcut_idx - upper_idx, ncol)) pad_zero[:, 0] = np.linspace( self.rmax + self.hh, self.hh * (rcut_idx + 1), rcut_idx - upper_idx ) @@ -134,9 +134,9 @@ def _check_table_upper_boundary(self) -> None: # if rcut goes beyond table upper bond, need extrapolation, ensure values decay to `0` before rcut. else: logging.warning( - "The rcut goes beyond table upper boundary, performing linear extrapolation." + "The rcut goes beyond table upper boundary, performing extrapolation." ) - pad_linear = np.zeros((rcut_idx - upper_idx + 1, 4)) + pad_linear = np.zeros((rcut_idx - upper_idx + 1, ncol)) pad_linear[:, 0] = np.linspace( self.rmax, self.hh * (rcut_idx + 1), rcut_idx - upper_idx + 1 ) @@ -150,7 +150,6 @@ def get(self) -> Tuple[np.array, np.array]: return self.tab_info, self.tab_data def _make_data(self): - # here we need to do postprocess, to overwrite coefficients when padding zeros resulting in negative energies. data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) stride = 4 * self.nspline idx_iter = 0 diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index b27eb1dbaf..b55e8f44b5 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -71,5 +71,68 @@ def test_with_mask(self): def test_jit(self): model = torch.jit.script(self.model) + +class TestPairTabTwoAtoms(unittest.TestCase): + @patch("numpy.loadtxt") + def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: + """Scenarios to test. + + rcut < rmax: + rr < rcut: use table values, or interpolate. + rr == rcut: use table values, or interpolate. + rr > rcut: should be 0 + rcut == rmax: + rr < rcut: use table values, or interpolate. + rr == rcut: use table values, or interpolate. + rr > rcut: should be 0 + rcut > rmax: + rr < rmax: use table values, or interpolate. + rr == rmax: use table values, or interpolate. + rmax < rr < rcut: extrapolate + rr >= rcut: should be 0 + + """ + file_path = "dummy_path" + mock_loadtxt.return_value = np.array( + [ + [0.005, 1.0], + [0.01, 0.8], + [0.015, 0.5], + [0.02, 0.25], + ] + ) + + # nframes=1, nall=2 + extended_atype = torch.tensor([[0, 0]]) + + # nframes=1, nloc=2, nnei=1 + nlist = torch.tensor([[[1], [-1]]]) + + results = [] + + for dist, rcut in zip( + [0.010, 0.015, 0.020, 0.015, 0.020, 0.021, 0.015, 0.020, 0.021, 0.025, 0.026, 0.025, 0.02999], + [0.015, 0.015, 0.015, 0.02, 0.020, 0.020, 0.022, 0.022, 0.022, 0.025, 0.025, 0.030, 0.030]): + extended_coord = torch.tensor( + [ + [ + [0.0, 0.0, 0.0], + [0.0, dist, 0.0], + ], + ] + ) + + model = PairTabModel(tab_file=file_path, rcut=rcut, sel=2) + results.append(model.forward_atomic( + extended_coord, extended_atype, nlist + )["energy"]) + + + expected_result = torch.stack([torch.tensor([[[0.4, 0], [0.25,0.], [0.,0], [0.25,0], [0.125,0], [0.,0], [0.25,0] ,[0.125,0], [0.0469,0], [0.,0], [0.,0],[0.0469,0],[0,0]]])]).reshape(13,2) + results = torch.stack(results).reshape(13,2) + + torch.testing.assert_allclose(results, expected_result, 0.0001, 0.0001) + + if __name__ == "__main__": unittest.main() From 5793828256907298d6c074e4aba2c8b25101527c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Jan 2024 05:56:38 +0000 Subject: [PATCH 15/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/pair_tab.py | 102 ++++++++++++++++++------------ deepmd/utils/pair_tab.py | 4 +- source/tests/pt/test_pairtab.py | 95 +++++++++++++++++++++------- 3 files changed, 136 insertions(+), 65 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index ea43b66864..d9d6d74282 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -6,12 +6,11 @@ Union, ) +import numpy as np +import torch from scipy.interpolate import ( CubicSpline, ) - -import numpy as np -import torch from torch import ( nn, ) @@ -109,7 +108,7 @@ def forward_atomic( # this will mask all -1 in the nlist masked_nlist = torch.clamp(nlist, 0) - atype = extended_atype[:, :self.nloc] # (nframes, nloc) + atype = extended_atype[:, : self.nloc] # (nframes, nloc) pairwise_dr = self._get_pairwise_dist( extended_coord ) # (nframes, nall, nall, 3) @@ -127,7 +126,7 @@ def forward_atomic( ] # slice rr to get (nframes, nloc, nnei) - rr = torch.gather(pairwise_rr[:, :self.nloc, :], 2, masked_nlist) + rr = torch.gather(pairwise_rr[:, : self.nloc, :], 2, masked_nlist) raw_atomic_energy = self._pair_tabulated_inter(nlist, atype, j_type, rr) @@ -194,12 +193,14 @@ def _pair_tabulated_inter( uu -= idx - table_coef = self._extract_spline_coefficient(i_type, j_type, idx, self.tab_data, self.nspline) + table_coef = self._extract_spline_coefficient( + i_type, j_type, idx, self.tab_data, self.nspline + ) table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4) ener = self._calcualte_ener(table_coef, uu) # here we need to do postprocess to overwrite energy to zero beyond rcut. - if self.tab.rmax <= self.rcut: + if self.tab.rmax <= self.rcut: mask_beyond_rcut = rr > self.rcut ener[mask_beyond_rcut] = 0 @@ -207,12 +208,16 @@ def _pair_tabulated_inter( extrapolation = self._extrapolate_rmax_rcut() if extrapolation is not None: uu_extrapolate = (rr - self.tab.rmax) / (self.rcut - self.tab.rmax) - clipped_uu = torch.clamp(uu_extrapolate, 0, 1) # clip rr within rmax. - extrapolate_coef = self._extract_spline_coefficient(i_type, j_type, torch.zeros_like(idx), extrapolation, 1) - extrapolate_coef = extrapolate_coef.reshape(self.nframes, self.nloc, self.nnei, 4) + clipped_uu = torch.clamp(uu_extrapolate, 0, 1) # clip rr within rmax. + extrapolate_coef = self._extract_spline_coefficient( + i_type, j_type, torch.zeros_like(idx), extrapolation, 1 + ) + extrapolate_coef = extrapolate_coef.reshape( + self.nframes, self.nloc, self.nnei, 4 + ) ener_extrpolate = self._calcualte_ener(extrapolate_coef, clipped_uu) mask_rmax_to_rcut = (self.tab.rmax < rr) & (rr <= self.rcut) - ener[mask_rmax_to_rcut] = ener_extrpolate[mask_rmax_to_rcut] + ener[mask_rmax_to_rcut] = ener_extrpolate[mask_rmax_to_rcut] return ener def _extrapolate_rmax_rcut(self) -> torch.Tensor: @@ -222,38 +227,48 @@ def _extrapolate_rmax_rcut(self) -> torch.Tensor: the table upper boundary values are not zeros. To simplify the problem, we use a single cubic spline between `rmax` and `rcut` for each pair of atom types. One can substitute this extrapolation to higher order polynomials if needed. - + There are two scenarios: - 1. `ruct` - `rmax` >= hh: + 1. `ruct` - `rmax` >= hh: Set values at the grid point right before `rcut` to 0, and perform exterapolation between the grid point and `rmax`, this allows smooth decay to 0 at `rcut`. - 2. `rcut` - `rmax` < hh: - Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`, + 2. `rcut` - `rmax` < hh: + Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`, the enery beyond `rcut` will be overwritten to `0` latter. - + Returns ------- torch.Tensor The cubic spline coefficients for each pair of atom types. (ntype, ntype, 1, 4) """ - #check if decays to `0` at rmax, if yes, no extrapolation is needed. - rmax_val = torch.from_numpy(self.tab.vdata[self.tab.vdata[:,0] == self.tab.rmax]) - pre_rmax_val = torch.from_numpy(self.tab.vdata[self.tab.vdata[:,0] == self.tab.rmax - self.tab.hh]) - - if torch.all(rmax_val[:,1:] == 0): - return + # check if decays to `0` at rmax, if yes, no extrapolation is needed. + rmax_val = torch.from_numpy( + self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax] + ) + pre_rmax_val = torch.from_numpy( + self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax - self.tab.hh] + ) + + if torch.all(rmax_val[:, 1:] == 0): + return else: if self.rcut - self.tab.rmax >= self.tab.hh: - rcut_idx = int(self.rcut/self.tab.hh - self.tab.rmin/self.tab.hh) - rcut_val = torch.tensor(self.tab.vdata[rcut_idx,:]).reshape(1,-1) - grid = torch.concatenate([rmax_val, rcut_val],axis=0) + rcut_idx = int(self.rcut / self.tab.hh - self.tab.rmin / self.tab.hh) + rcut_val = torch.tensor(self.tab.vdata[rcut_idx, :]).reshape(1, -1) + grid = torch.concatenate([rmax_val, rcut_val], axis=0) else: # the last two rows will be the rmax, and rmax+hh - grid = torch.from_numpy(self.tab.vdata[-2:,:]) - passin_slope = ((rmax_val - pre_rmax_val)/self.tab.hh)[:,1:].squeeze(0) if ~np.all(pre_rmax_val == None) else 0 # the slope at the end of table for each ntype pairs (ntypes,ntypes,1) - extrapolate_coef = torch.from_numpy(self._calculate_spline_coef(grid, passin_slope)).reshape(self.ntypes,self.ntypes,4) + grid = torch.from_numpy(self.tab.vdata[-2:, :]) + passin_slope = ( + ((rmax_val - pre_rmax_val) / self.tab.hh)[:, 1:].squeeze(0) + if ~np.all(pre_rmax_val == None) + else 0 + ) # the slope at the end of table for each ntype pairs (ntypes,ntypes,1) + extrapolate_coef = torch.from_numpy( + self._calculate_spline_coef(grid, passin_slope) + ).reshape(self.ntypes, self.ntypes, 4) return extrapolate_coef.unsqueeze(2) - + # might be able to refactor this, combine with PairTab def _calculate_spline_coef(self, grid, passin_slope): data = np.zeros([self.ntypes * self.ntypes * 4]) @@ -264,20 +279,18 @@ def _calculate_spline_coef(self, grid, passin_slope): for t0 in range(self.ntypes): for t1 in range(t0, self.ntypes): vv = grid[:, 1 + idx_iter] - slope_idx = [t0 * (2 * self.ntypes - t0 - 1)//2 + t1] + slope_idx = [t0 * (2 * self.ntypes - t0 - 1) // 2 + t1] print(f"slope: {passin_slope[slope_idx]}") - cs = CubicSpline(xx, vv, bc_type=((1,passin_slope[slope_idx][0]),(1,0))) + cs = CubicSpline( + xx, vv, bc_type=((1, passin_slope[slope_idx][0]), (1, 0)) + ) dd = cs(xx, 1) dd *= self.tab.hh dtmp = np.zeros(stride) - dtmp[0] = ( - 2 * vv[0] - 2 * vv[1] + dd[0] + dd[1] - ) - dtmp[1] = ( - (-3 * vv[0] + 3 * vv[1] - 2 * dd[0] - dd[1]) - ) - dtmp[2] = dd[0] + dtmp[0] = 2 * vv[0] - 2 * vv[1] + dd[0] + dd[1] + dtmp[1] = -3 * vv[0] + 3 * vv[1] - 2 * dd[0] - dd[1] + dtmp[2] = dd[0] dtmp[3] = vv[0] data[ (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride @@ -329,7 +342,12 @@ def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: return coords.unsqueeze(2) - coords.unsqueeze(1) @staticmethod - def _extract_spline_coefficient(i_type: torch.Tensor, j_type: torch.Tensor, idx: torch.Tensor, tab_data: torch.Tensor, nspline: int + def _extract_spline_coefficient( + i_type: torch.Tensor, + j_type: torch.Tensor, + idx: torch.Tensor, + tab_data: torch.Tensor, + nspline: int, ) -> torch.Tensor: """Extract the spline coefficient from the table. @@ -376,5 +394,7 @@ def _extract_spline_coefficient(i_type: torch.Tensor, j_type: torch.Tensor, idx: def _calcualte_ener(coef, uu): a3, a2, a1, a0 = torch.unbind(coef, dim=-1) # 4 * (nframes, nloc, nnei) etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. - ener = etmp * uu + a0 # this energy has the linear extrapolated value when rcut > rmax - return ener \ No newline at end of file + ener = ( + etmp * uu + a0 + ) # this energy has the linear extrapolated value when rcut > rmax + return ener diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index b87ed04fd3..c159df2004 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -59,7 +59,9 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: # check table data against rcut and update tab_file if needed, table upper boundary is used as rcut if not provided. self.rcut = rcut if rcut is not None else self.rmax self._check_table_upper_boundary() - self.nspline = self.vdata.shape[0] - 1 # this nspline is updated based on the expanded table. + self.nspline = ( + self.vdata.shape[0] - 1 + ) # this nspline is updated based on the expanded table. self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) self.tab_data = self._make_data() diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index b55e8f44b5..e9ce99f92e 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -77,20 +77,20 @@ class TestPairTabTwoAtoms(unittest.TestCase): def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: """Scenarios to test. - rcut < rmax: - rr < rcut: use table values, or interpolate. - rr == rcut: use table values, or interpolate. - rr > rcut: should be 0 - rcut == rmax: - rr < rcut: use table values, or interpolate. - rr == rcut: use table values, or interpolate. - rr > rcut: should be 0 - rcut > rmax: - rr < rmax: use table values, or interpolate. - rr == rmax: use table values, or interpolate. - rmax < rr < rcut: extrapolate - rr >= rcut: should be 0 - + rcut < rmax: + rr < rcut: use table values, or interpolate. + rr == rcut: use table values, or interpolate. + rr > rcut: should be 0 + rcut == rmax: + rr < rcut: use table values, or interpolate. + rr == rcut: use table values, or interpolate. + rr > rcut: should be 0 + rcut > rmax: + rr < rmax: use table values, or interpolate. + rr == rmax: use table values, or interpolate. + rmax < rr < rcut: extrapolate + rr >= rcut: should be 0 + """ file_path = "dummy_path" mock_loadtxt.return_value = np.array( @@ -111,8 +111,37 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: results = [] for dist, rcut in zip( - [0.010, 0.015, 0.020, 0.015, 0.020, 0.021, 0.015, 0.020, 0.021, 0.025, 0.026, 0.025, 0.02999], - [0.015, 0.015, 0.015, 0.02, 0.020, 0.020, 0.022, 0.022, 0.022, 0.025, 0.025, 0.030, 0.030]): + [ + 0.010, + 0.015, + 0.020, + 0.015, + 0.020, + 0.021, + 0.015, + 0.020, + 0.021, + 0.025, + 0.026, + 0.025, + 0.02999, + ], + [ + 0.015, + 0.015, + 0.015, + 0.02, + 0.020, + 0.020, + 0.022, + 0.022, + 0.022, + 0.025, + 0.025, + 0.030, + 0.030, + ], + ): extended_coord = torch.tensor( [ [ @@ -123,16 +152,36 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: ) model = PairTabModel(tab_file=file_path, rcut=rcut, sel=2) - results.append(model.forward_atomic( - extended_coord, extended_atype, nlist - )["energy"]) - + results.append( + model.forward_atomic(extended_coord, extended_atype, nlist)["energy"] + ) - expected_result = torch.stack([torch.tensor([[[0.4, 0], [0.25,0.], [0.,0], [0.25,0], [0.125,0], [0.,0], [0.25,0] ,[0.125,0], [0.0469,0], [0.,0], [0.,0],[0.0469,0],[0,0]]])]).reshape(13,2) - results = torch.stack(results).reshape(13,2) + expected_result = torch.stack( + [ + torch.tensor( + [ + [ + [0.4, 0], + [0.25, 0.0], + [0.0, 0], + [0.25, 0], + [0.125, 0], + [0.0, 0], + [0.25, 0], + [0.125, 0], + [0.0469, 0], + [0.0, 0], + [0.0, 0], + [0.0469, 0], + [0, 0], + ] + ] + ) + ] + ).reshape(13, 2) + results = torch.stack(results).reshape(13, 2) torch.testing.assert_allclose(results, expected_result, 0.0001, 0.0001) - if __name__ == "__main__": unittest.main() From 92dec182b56d0ac03039d80cf909e61d7c10f6da Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Tue, 30 Jan 2024 14:40:54 +0800 Subject: [PATCH 16/20] chore: refactor _make_data in PairTab --- deepmd/pt/model/model/pair_tab.py | 73 ++++++++++++------------------- deepmd/utils/pair_tab.py | 34 ++++++++------ 2 files changed, 48 insertions(+), 59 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index d9d6d74282..52cbbab4d2 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -183,7 +183,7 @@ def _pair_tabulated_inter( uu = (rr - rmin) * hi # this is broadcasted to (nframes,nloc,nnei) # if nnei of atom 0 has -1 in the nlist, uu would be 0. - # this is to handel the nlist where the mask is set to 0, so that we don't raise exception for those atoms. + # this is to handle the nlist where the mask is set to 0, so that we don't raise exception for those atoms. uu = torch.where(nlist != -1, uu, self.nspline + 1) if torch.any(uu < 0): @@ -199,8 +199,9 @@ def _pair_tabulated_inter( table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4) ener = self._calcualte_ener(table_coef, uu) - # here we need to do postprocess to overwrite energy to zero beyond rcut. + if self.tab.rmax <= self.rcut: + # here we need to overwrite energy to zero beyond rcut. mask_beyond_rcut = rr > self.rcut ener[mask_beyond_rcut] = 0 @@ -233,15 +234,14 @@ def _extrapolate_rmax_rcut(self) -> torch.Tensor: Set values at the grid point right before `rcut` to 0, and perform exterapolation between the grid point and `rmax`, this allows smooth decay to 0 at `rcut`. 2. `rcut` - `rmax` < hh: - Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`, - the enery beyond `rcut` will be overwritten to `0` latter. + Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`. Returns ------- torch.Tensor The cubic spline coefficients for each pair of atom types. (ntype, ntype, 1, 4) """ - # check if decays to `0` at rmax, if yes, no extrapolation is needed. + rmax_val = torch.from_numpy( self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax] ) @@ -249,6 +249,7 @@ def _extrapolate_rmax_rcut(self) -> torch.Tensor: self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax - self.tab.hh] ) + # check if decays to `0` at rmax, if yes, no extrapolation is needed. if torch.all(rmax_val[:, 1:] == 0): return else: @@ -261,47 +262,13 @@ def _extrapolate_rmax_rcut(self) -> torch.Tensor: grid = torch.from_numpy(self.tab.vdata[-2:, :]) passin_slope = ( ((rmax_val - pre_rmax_val) / self.tab.hh)[:, 1:].squeeze(0) - if ~np.all(pre_rmax_val == None) - else 0 + if self.tab.rmax > self.tab.hh + else torch.zeros_like(rmax_val[:, 1:]).squeeze(0) ) # the slope at the end of table for each ntype pairs (ntypes,ntypes,1) extrapolate_coef = torch.from_numpy( - self._calculate_spline_coef(grid, passin_slope) + self.tab._make_data(self.ntypes, 1, grid, self.tab.hh, passin_slope) ).reshape(self.ntypes, self.ntypes, 4) - return extrapolate_coef.unsqueeze(2) - - # might be able to refactor this, combine with PairTab - def _calculate_spline_coef(self, grid, passin_slope): - data = np.zeros([self.ntypes * self.ntypes * 4]) - stride = 4 - idx_iter = 0 - - xx = grid[:, 0] - for t0 in range(self.ntypes): - for t1 in range(t0, self.ntypes): - vv = grid[:, 1 + idx_iter] - slope_idx = [t0 * (2 * self.ntypes - t0 - 1) // 2 + t1] - - print(f"slope: {passin_slope[slope_idx]}") - cs = CubicSpline( - xx, vv, bc_type=((1, passin_slope[slope_idx][0]), (1, 0)) - ) - dd = cs(xx, 1) - dd *= self.tab.hh - dtmp = np.zeros(stride) - dtmp[0] = 2 * vv[0] - 2 * vv[1] + dd[0] + dd[1] - dtmp[1] = -3 * vv[0] + 3 * vv[1] - 2 * dd[0] - dd[1] - dtmp[2] = dd[0] - dtmp[3] = vv[0] - data[ - (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride - + stride - ] = dtmp - data[ - (t1 * self.ntypes + t0) * stride : (t1 * self.ntypes + t0) * stride - + stride - ] = dtmp - idx_iter += 1 - return data + return extrapolate_coef.unsqueeze(2) @staticmethod def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: @@ -367,7 +334,7 @@ def _extract_spline_coefficient( Returns ------- torch.Tensor - The spline coefficient. (nframes, nloc, nnei, 4), shape maybe squeezed. + The spline coefficient. (nframes, nloc, nnei, 4), shape may be squeezed. """ # (nframes, nloc, nnei) @@ -391,8 +358,22 @@ def _extract_spline_coefficient( return final_coef @staticmethod - def _calcualte_ener(coef, uu): - a3, a2, a1, a0 = torch.unbind(coef, dim=-1) # 4 * (nframes, nloc, nnei) + def _calcualte_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor: + """Calculate energy using spline coeeficients. + + Paramerters + ----------- + coef : torch.Tensor + The spline coefficients. (nframes, nloc, nnei, 4) + uu : torch.Tensor + The atom displancemnt used in interpolation and extrapolation (nframes, nloc, nnei) + + Returns + ------- + torch.Tensor + The atomic energy for all local atoms for all frames. (nframes, nloc, nnei) + """ + a3, a2, a1, a0 = torch.unbind(coef, dim=-1) etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. ener = ( etmp * uu + a0 diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index c159df2004..10c8eec84d 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -63,7 +63,7 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: self.vdata.shape[0] - 1 ) # this nspline is updated based on the expanded table. self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) - self.tab_data = self._make_data() + self.tab_data = self._make_data(self.ntypes, self.nspline, self.vdata, self.hh) def _check_table_upper_boundary(self) -> None: """Update User Provided Table Based on `rcut`. @@ -151,19 +151,27 @@ def get(self) -> Tuple[np.array, np.array]: """Get the serialized table.""" return self.tab_info, self.tab_data - def _make_data(self): - data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) - stride = 4 * self.nspline + @staticmethod + def _make_data(ntypes: int, nspline: int, vdata: np.array, hh: float, passin_slope: Optional[np.array] = None) -> np.array: + data = np.zeros([ntypes * ntypes * 4 * nspline]) + stride = 4 * nspline idx_iter = 0 - xx = self.vdata[:, 0] - for t0 in range(self.ntypes): - for t1 in range(t0, self.ntypes): - vv = self.vdata[:, 1 + idx_iter] - cs = CubicSpline(xx, vv) + xx = vdata[:, 0] + for t0 in range(ntypes): + for t1 in range(t0, ntypes): + vv = vdata[:, 1 + idx_iter] + if passin_slope is not None: + slope_idx = [t0 * (2 * ntypes - t0 - 1) // 2 + t1] + cs = CubicSpline( + # setting first order derivation and both end for extrapolation. + xx, vv, bc_type=((1, passin_slope[slope_idx][0]), (1, 0)) + ) + else: + cs = CubicSpline(xx, vv) dd = cs(xx, 1) - dd *= self.hh + dd *= hh dtmp = np.zeros(stride) - for ii in range(self.nspline): + for ii in range(nspline): # check if vv is zero, if so, that's case 1, set all coefficients to 0, dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] dtmp[ii * 4 + 1] = ( @@ -172,11 +180,11 @@ def _make_data(self): dtmp[ii * 4 + 2] = dd[ii] dtmp[ii * 4 + 3] = vv[ii] data[ - (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride + (t0 * ntypes + t1) * stride : (t0 * ntypes + t1) * stride + stride ] = dtmp data[ - (t1 * self.ntypes + t0) * stride : (t1 * self.ntypes + t0) * stride + (t1 * ntypes + t0) * stride : (t1 * ntypes + t0) * stride + stride ] = dtmp idx_iter += 1 From bc04359110b7935007f7c26022adeac81ea71de7 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Jan 2024 06:42:25 +0000 Subject: [PATCH 17/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/pair_tab.py | 14 ++++---------- deepmd/utils/pair_tab.py | 18 ++++++++++++------ 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 52cbbab4d2..4cd1be2f57 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -6,11 +6,7 @@ Union, ) -import numpy as np import torch -from scipy.interpolate import ( - CubicSpline, -) from torch import ( nn, ) @@ -199,7 +195,6 @@ def _pair_tabulated_inter( table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4) ener = self._calcualte_ener(table_coef, uu) - if self.tab.rmax <= self.rcut: # here we need to overwrite energy to zero beyond rcut. mask_beyond_rcut = rr > self.rcut @@ -241,7 +236,6 @@ def _extrapolate_rmax_rcut(self) -> torch.Tensor: torch.Tensor The cubic spline coefficients for each pair of atom types. (ntype, ntype, 1, 4) """ - rmax_val = torch.from_numpy( self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax] ) @@ -268,7 +262,7 @@ def _extrapolate_rmax_rcut(self) -> torch.Tensor: extrapolate_coef = torch.from_numpy( self.tab._make_data(self.ntypes, 1, grid, self.tab.hh, passin_slope) ).reshape(self.ntypes, self.ntypes, 4) - return extrapolate_coef.unsqueeze(2) + return extrapolate_coef.unsqueeze(2) @staticmethod def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: @@ -361,8 +355,8 @@ def _extract_spline_coefficient( def _calcualte_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor: """Calculate energy using spline coeeficients. - Paramerters - ----------- + Parameters + ---------- coef : torch.Tensor The spline coefficients. (nframes, nloc, nnei, 4) uu : torch.Tensor @@ -373,7 +367,7 @@ def _calcualte_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor: torch.Tensor The atomic energy for all local atoms for all frames. (nframes, nloc, nnei) """ - a3, a2, a1, a0 = torch.unbind(coef, dim=-1) + a3, a2, a1, a0 = torch.unbind(coef, dim=-1) etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. ener = ( etmp * uu + a0 diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 10c8eec84d..fda0e2597b 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -152,7 +152,13 @@ def get(self) -> Tuple[np.array, np.array]: return self.tab_info, self.tab_data @staticmethod - def _make_data(ntypes: int, nspline: int, vdata: np.array, hh: float, passin_slope: Optional[np.array] = None) -> np.array: + def _make_data( + ntypes: int, + nspline: int, + vdata: np.array, + hh: float, + passin_slope: Optional[np.array] = None, + ) -> np.array: data = np.zeros([ntypes * ntypes * 4 * nspline]) stride = 4 * nspline idx_iter = 0 @@ -164,7 +170,9 @@ def _make_data(ntypes: int, nspline: int, vdata: np.array, hh: float, passin_slo slope_idx = [t0 * (2 * ntypes - t0 - 1) // 2 + t1] cs = CubicSpline( # setting first order derivation and both end for extrapolation. - xx, vv, bc_type=((1, passin_slope[slope_idx][0]), (1, 0)) + xx, + vv, + bc_type=((1, passin_slope[slope_idx][0]), (1, 0)), ) else: cs = CubicSpline(xx, vv) @@ -180,12 +188,10 @@ def _make_data(ntypes: int, nspline: int, vdata: np.array, hh: float, passin_slo dtmp[ii * 4 + 2] = dd[ii] dtmp[ii * 4 + 3] = vv[ii] data[ - (t0 * ntypes + t1) * stride : (t0 * ntypes + t1) * stride - + stride + (t0 * ntypes + t1) * stride : (t0 * ntypes + t1) * stride + stride ] = dtmp data[ - (t1 * ntypes + t0) * stride : (t1 * ntypes + t0) * stride - + stride + (t1 * ntypes + t0) * stride : (t1 * ntypes + t0) * stride + stride ] = dtmp idx_iter += 1 return data From 4433035c3c6e79a0d6bc42d72e69a9db11b2cc4b Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Tue, 30 Jan 2024 16:21:03 +0800 Subject: [PATCH 18/20] chore: move file --- source/tests/{tf => common}/test_pairtab_preprocess.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename source/tests/{tf => common}/test_pairtab_preprocess.py (100%) diff --git a/source/tests/tf/test_pairtab_preprocess.py b/source/tests/common/test_pairtab_preprocess.py similarity index 100% rename from source/tests/tf/test_pairtab_preprocess.py rename to source/tests/common/test_pairtab_preprocess.py From 4851a0a67a79e34d4a9518cb89598286e95f40f8 Mon Sep 17 00:00:00 2001 From: Anyang Peng Date: Wed, 31 Jan 2024 16:51:37 +0800 Subject: [PATCH 19/20] chore: refactor extrapolation code --- deepmd/pt/model/model/pair_tab.py | 77 ++--------- deepmd/utils/pair_tab.py | 121 +++++++++++------- .../tests/common/test_pairtab_preprocess.py | 80 +++++++++--- source/tests/pt/test_pairtab.py | 45 ++++--- 4 files changed, 171 insertions(+), 152 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index 4cd1be2f57..bf93c4ff97 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -195,74 +195,16 @@ def _pair_tabulated_inter( table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4) ener = self._calcualte_ener(table_coef, uu) - if self.tab.rmax <= self.rcut: - # here we need to overwrite energy to zero beyond rcut. - mask_beyond_rcut = rr > self.rcut - ener[mask_beyond_rcut] = 0 - - # here we use smooth extrapolation to replace linear extrapolation. - extrapolation = self._extrapolate_rmax_rcut() - if extrapolation is not None: - uu_extrapolate = (rr - self.tab.rmax) / (self.rcut - self.tab.rmax) - clipped_uu = torch.clamp(uu_extrapolate, 0, 1) # clip rr within rmax. - extrapolate_coef = self._extract_spline_coefficient( - i_type, j_type, torch.zeros_like(idx), extrapolation, 1 - ) - extrapolate_coef = extrapolate_coef.reshape( - self.nframes, self.nloc, self.nnei, 4 - ) - ener_extrpolate = self._calcualte_ener(extrapolate_coef, clipped_uu) - mask_rmax_to_rcut = (self.tab.rmax < rr) & (rr <= self.rcut) - ener[mask_rmax_to_rcut] = ener_extrpolate[mask_rmax_to_rcut] - return ener - - def _extrapolate_rmax_rcut(self) -> torch.Tensor: - """Soomth extrapolation between table upper boundary and rcut. - This method should only be used when the table upper boundary `rmax` is smaller than `rcut`, and - the table upper boundary values are not zeros. To simplify the problem, we use a single - cubic spline between `rmax` and `rcut` for each pair of atom types. One can substitute this extrapolation - to higher order polynomials if needed. + # here we need to overwrite energy to zero at rcut and beyond. + mask_beyond_rcut = rr >= self.rcut + # also overwrite values beyond extrapolation to zero + extrapolation_mask = rr >= self.tab.rmin + self.nspline * self.tab.hh + ener[mask_beyond_rcut] = 0 + ener[extrapolation_mask] = 0 - There are two scenarios: - 1. `ruct` - `rmax` >= hh: - Set values at the grid point right before `rcut` to 0, and perform exterapolation between - the grid point and `rmax`, this allows smooth decay to 0 at `rcut`. - 2. `rcut` - `rmax` < hh: - Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`. - - Returns - ------- - torch.Tensor - The cubic spline coefficients for each pair of atom types. (ntype, ntype, 1, 4) - """ - rmax_val = torch.from_numpy( - self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax] - ) - pre_rmax_val = torch.from_numpy( - self.tab.vdata[self.tab.vdata[:, 0] == self.tab.rmax - self.tab.hh] - ) + return ener - # check if decays to `0` at rmax, if yes, no extrapolation is needed. - if torch.all(rmax_val[:, 1:] == 0): - return - else: - if self.rcut - self.tab.rmax >= self.tab.hh: - rcut_idx = int(self.rcut / self.tab.hh - self.tab.rmin / self.tab.hh) - rcut_val = torch.tensor(self.tab.vdata[rcut_idx, :]).reshape(1, -1) - grid = torch.concatenate([rmax_val, rcut_val], axis=0) - else: - # the last two rows will be the rmax, and rmax+hh - grid = torch.from_numpy(self.tab.vdata[-2:, :]) - passin_slope = ( - ((rmax_val - pre_rmax_val) / self.tab.hh)[:, 1:].squeeze(0) - if self.tab.rmax > self.tab.hh - else torch.zeros_like(rmax_val[:, 1:]).squeeze(0) - ) # the slope at the end of table for each ntype pairs (ntypes,ntypes,1) - extrapolate_coef = torch.from_numpy( - self.tab._make_data(self.ntypes, 1, grid, self.tab.hh, passin_slope) - ).reshape(self.ntypes, self.ntypes, 4) - return extrapolate_coef.unsqueeze(2) @staticmethod def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: @@ -347,8 +289,7 @@ def _extract_spline_coefficient( final_coef = torch.gather(expanded_tab_data, 3, clipped_indices).squeeze() # when the spline idx is beyond the table, all spline coefficients are set to `0`, and the resulting ener corresponding to the idx is also `0`. - final_coef[expanded_idx.squeeze() >= nspline] = 0 - + final_coef[expanded_idx.squeeze() > nspline] = 0 return final_coef @staticmethod @@ -371,5 +312,5 @@ def _calcualte_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor: etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. ener = ( etmp * uu + a0 - ) # this energy has the linear extrapolated value when rcut > rmax + ) # this energy has the extrapolated value when rcut > rmax return ener diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index fda0e2597b..9ced80fe95 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -63,15 +63,14 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: self.vdata.shape[0] - 1 ) # this nspline is updated based on the expanded table. self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) - self.tab_data = self._make_data(self.ntypes, self.nspline, self.vdata, self.hh) - + self.tab_data = self._make_data() def _check_table_upper_boundary(self) -> None: """Update User Provided Table Based on `rcut`. This function checks the upper boundary provided in the table against rcut. If the table upper boundary values decay to zero before rcut, padding zeros will be added to the table to cover rcut; if the table upper boundary values do not decay to zero - before ruct, linear extrapolation will be performed till rcut. + before ruct, extrapolation will be performed till rcut. Examples -------- @@ -85,7 +84,6 @@ def _check_table_upper_boundary(self) -> None: [0.01 0.8 1.6 2.4 ] [0.015 0. 1. 1.5 ] [0.02 0. 0. 0. ] - [0.025 0. 0. 0. ]] ---------------------------------------------- @@ -108,10 +106,10 @@ def _check_table_upper_boundary(self) -> None: """ upper_val = self.vdata[-1][1:] upper_idx = self.vdata.shape[0] - 1 - ncol = self.vdata.shape[1] - # the index of table for the grid point right after rcut - rcut_idx = int(self.rcut / self.hh) + self.ncol = self.vdata.shape[1] + # the index in table for the grid point of rcut, always give the point after rcut. + rcut_idx = int(np.ceil(self.rcut / self.hh - self.rmin / self.hh)) if np.all(upper_val == 0): # if table values decay to `0` after rcut if self.rcut < self.rmax and np.any(self.vdata[rcut_idx - 1][1:] != 0): @@ -122,14 +120,14 @@ def _check_table_upper_boundary(self) -> None: # if table values decay to `0` before rcut, pad table with `0`s. elif self.rcut > self.rmax: - pad_zero = np.zeros((rcut_idx - upper_idx, ncol)) + pad_zero = np.zeros((rcut_idx - upper_idx, self.ncol)) pad_zero[:, 0] = np.linspace( - self.rmax + self.hh, self.hh * (rcut_idx + 1), rcut_idx - upper_idx + self.rmax + self.hh, self.rmax + self.hh * (rcut_idx - upper_idx), rcut_idx - upper_idx ) self.vdata = np.concatenate((self.vdata, pad_zero), axis=0) else: # if table values do not decay to `0` at rcut - if self.rcut < self.rmax: + if self.rcut <= self.rmax: logging.warning( "The energy provided in the table does not decay to 0 at rcut." ) @@ -138,49 +136,78 @@ def _check_table_upper_boundary(self) -> None: logging.warning( "The rcut goes beyond table upper boundary, performing extrapolation." ) - pad_linear = np.zeros((rcut_idx - upper_idx + 1, ncol)) - pad_linear[:, 0] = np.linspace( - self.rmax, self.hh * (rcut_idx + 1), rcut_idx - upper_idx + 1 + pad_extrapolation = np.zeros((rcut_idx - upper_idx, self.ncol)) + + pad_extrapolation[:, 0] = np.linspace( + self.rmax + self.hh, self.rmax + self.hh * (rcut_idx - upper_idx), rcut_idx - upper_idx ) - pad_linear[:-1, 1:] = np.array( - [np.linspace(start, 0, rcut_idx - upper_idx) for start in upper_val] - ).T - self.vdata = np.concatenate((self.vdata[:-1, :], pad_linear), axis=0) + # need to calculate table values to fill in with cubic spline + pad_extrapolation = self._extrapolate_table(pad_extrapolation) + + self.vdata = np.concatenate((self.vdata, pad_extrapolation), axis=0) def get(self) -> Tuple[np.array, np.array]: """Get the serialized table.""" return self.tab_info, self.tab_data - @staticmethod - def _make_data( - ntypes: int, - nspline: int, - vdata: np.array, - hh: float, - passin_slope: Optional[np.array] = None, - ) -> np.array: - data = np.zeros([ntypes * ntypes * 4 * nspline]) - stride = 4 * nspline + def _extrapolate_table(self, pad_extrapolation: np.array) -> np.array: + """Soomth extrapolation between table upper boundary and rcut. + + This method should only be used when the table upper boundary `rmax` is smaller than `rcut`, and + the table upper boundary values are not zeros. To simplify the problem, we use a single + cubic spline between `rmax` and `rcut` for each pair of atom types. One can substitute this extrapolation + to higher order polynomials if needed. + + There are two scenarios: + 1. `ruct` - `rmax` >= hh: + Set values at the grid point right before `rcut` to 0, and perform exterapolation between + the grid point and `rmax`, this allows smooth decay to 0 at `rcut`. + 2. `rcut` - `rmax` < hh: + Set values at `rmax + hh` to 0, and perform extrapolation between `rmax` and `rmax + hh`. + + Parameters + ---------- + pad_extrapolation: np.array + The emepty grid that holds the extrapolation values. + + Returns + ------- + np.array + The cubic spline extrapolation. + """ + # in theory we should check if the table has at least two rows. + slope = (self.vdata[-1,1:] - self.vdata[-2,1:]) # shape of (ncol-1, ) + + # for extrapolation, we want values decay to `0` prior to `ruct` if possible + # here we try to find the grid point prior to `rcut` + grid_point = -2 if pad_extrapolation[-1,0]/self.hh - self.rmax/self.hh >= 2 else -1 + temp_grid = np.stack((self.vdata[-1,:], pad_extrapolation[grid_point,:])) + vv = temp_grid[:,1:] + xx = temp_grid[:,0] + cs = CubicSpline(xx,vv, bc_type=((1,slope),(1,np.zeros_like(slope)))) + xx_grid = pad_extrapolation[:,0] + res = cs(xx_grid) + + pad_extrapolation[:,1:] = res + + # Note: when doing cubic spline, if we want to ensure values decay to zero prior to `rcut` + # this may cause values be positive post `rcut`, we need to overwrite those values to zero + pad_extrapolation = pad_extrapolation if grid_point == -1 else pad_extrapolation[:-1,:] + return pad_extrapolation + + def _make_data(self): + data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) + stride = 4 * self.nspline idx_iter = 0 - xx = vdata[:, 0] - for t0 in range(ntypes): - for t1 in range(t0, ntypes): - vv = vdata[:, 1 + idx_iter] - if passin_slope is not None: - slope_idx = [t0 * (2 * ntypes - t0 - 1) // 2 + t1] - cs = CubicSpline( - # setting first order derivation and both end for extrapolation. - xx, - vv, - bc_type=((1, passin_slope[slope_idx][0]), (1, 0)), - ) - else: - cs = CubicSpline(xx, vv) + xx = self.vdata[:, 0] + for t0 in range(self.ntypes): + for t1 in range(t0, self.ntypes): + vv = self.vdata[:, 1 + idx_iter] + cs = CubicSpline(xx, vv, bc_type='clamped') dd = cs(xx, 1) - dd *= hh + dd *= self.hh dtmp = np.zeros(stride) - for ii in range(nspline): - # check if vv is zero, if so, that's case 1, set all coefficients to 0, + for ii in range(self.nspline): dtmp[ii * 4 + 0] = 2 * vv[ii] - 2 * vv[ii + 1] + dd[ii] + dd[ii + 1] dtmp[ii * 4 + 1] = ( -3 * vv[ii] + 3 * vv[ii + 1] - 2 * dd[ii] - dd[ii + 1] @@ -188,10 +215,12 @@ def _make_data( dtmp[ii * 4 + 2] = dd[ii] dtmp[ii * 4 + 3] = vv[ii] data[ - (t0 * ntypes + t1) * stride : (t0 * ntypes + t1) * stride + stride + (t0 * self.ntypes + t1) * stride : (t0 * self.ntypes + t1) * stride + + stride ] = dtmp data[ - (t1 * ntypes + t0) * stride : (t1 * ntypes + t0) * stride + stride + (t1 * self.ntypes + t0) * stride : (t1 * self.ntypes + t0) * stride + + stride ] = dtmp idx_iter += 1 return data diff --git a/source/tests/common/test_pairtab_preprocess.py b/source/tests/common/test_pairtab_preprocess.py index aee6b3889e..cc776383e1 100644 --- a/source/tests/common/test_pairtab_preprocess.py +++ b/source/tests/common/test_pairtab_preprocess.py @@ -11,7 +11,7 @@ ) -class TestPairTabPreprocessLinear(unittest.TestCase): +class TestPairTabPreprocessExtrapolate(unittest.TestCase): @patch("numpy.loadtxt") def setUp(self, mock_loadtxt) -> None: file_path = "dummy_path" @@ -28,6 +28,7 @@ def setUp(self, mock_loadtxt) -> None: self.tab2 = PairTab(filename=file_path, rcut=0.02) self.tab3 = PairTab(filename=file_path, rcut=0.022) self.tab4 = PairTab(filename=file_path, rcut=0.03) + self.tab5 = PairTab(filename=file_path, rcut=0.032) def test_preprocess(self): np.testing.assert_allclose( @@ -39,9 +40,9 @@ def test_preprocess(self): [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.0, 0.0], - [0.03, 0.0, 0.0, 0.0], + ] - ), + ), rtol=1e-04, atol = 1e-04 ) np.testing.assert_allclose( self.tab2.vdata, @@ -51,9 +52,8 @@ def test_preprocess(self): [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], - [0.025, 0.0, 0.0, 0.0], ] - ), + ), rtol=1e-04, atol = 1e-04 ) # for this test case, the table does not decay to zero at rcut = 0.22, @@ -69,8 +69,9 @@ def test_preprocess(self): [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.0, 0.0], ] - ), + ), rtol=1e-04, atol = 1e-04 ) + np.testing.assert_allclose( self.tab4.vdata, np.array( @@ -79,11 +80,24 @@ def test_preprocess(self): [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], - [0.025, 0.125, 0.2, 0.375], + [0.025, 0.0, 0.0, 0.0], + + ] + ), rtol = 1e-04, atol = 1e-04 + ) + + np.testing.assert_allclose( + self.tab5.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.12468, 0.1992, 0.3741], [0.03, 0.0, 0.0, 0.0], - [0.035, 0.0, 0.0, 0.0], ] - ), + ), rtol = 1e-04, atol = 1e-04 ) @@ -93,7 +107,6 @@ def setUp(self, mock_loadtxt) -> None: file_path = "dummy_path" mock_loadtxt.return_value = np.array( [ - [0.005, 1.0, 2.0, 3.0], [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -104,13 +117,14 @@ def setUp(self, mock_loadtxt) -> None: self.tab1 = PairTab(filename=file_path, rcut=0.023) self.tab2 = PairTab(filename=file_path, rcut=0.025) self.tab3 = PairTab(filename=file_path, rcut=0.028) + self.tab4 = PairTab(filename=file_path, rcut=0.033) def test_preprocess(self): np.testing.assert_allclose( self.tab1.vdata, np.array( [ - [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -122,7 +136,7 @@ def test_preprocess(self): self.tab2.vdata, np.array( [ - [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -135,7 +149,7 @@ def test_preprocess(self): self.tab3.vdata, np.array( [ - [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -145,6 +159,20 @@ def test_preprocess(self): ), ) + np.testing.assert_allclose( + self.tab4.vdata, + np.array( + [ + + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.0, 0.0], + [0.03, 0.0, 0.0, 0.0], + [0.035, 0.0, 0.0, 0.0], + ] + ), + ) class TestPairTabPreprocessUneven(unittest.TestCase): @patch("numpy.loadtxt") @@ -163,6 +191,7 @@ def setUp(self, mock_loadtxt) -> None: self.tab1 = PairTab(filename=file_path, rcut=0.025) self.tab2 = PairTab(filename=file_path, rcut=0.028) self.tab3 = PairTab(filename=file_path, rcut=0.03) + self.tab4 = PairTab(filename=file_path, rcut=0.037) def test_preprocess(self): np.testing.assert_allclose( @@ -174,9 +203,8 @@ def test_preprocess(self): [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.1, 0.0], - [0.03, 0.0, 0.0, 0.0], ] - ), + ), rtol=1e-04, atol = 1e-04 ) np.testing.assert_allclose( self.tab2.vdata, @@ -189,8 +217,9 @@ def test_preprocess(self): [0.025, 0.0, 0.1, 0.0], [0.03, 0.0, 0.0, 0.0], ] - ), + ), rtol=1e-04, atol = 1e-04 ) + np.testing.assert_allclose( self.tab3.vdata, np.array( @@ -201,7 +230,22 @@ def test_preprocess(self): [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.1, 0.0], [0.03, 0.0, 0.0, 0.0], - [0.035, 0.0, 0.0, 0.0], ] - ), + ), rtol=1e-04, atol = 1e-04 ) + + np.testing.assert_allclose( + self.tab4.vdata, + np.array( + [ + [0.005, 1.0, 2.0, 3.0], + [0.01, 0.8, 1.6, 2.4], + [0.015, 0.5, 1.0, 1.5], + [0.02, 0.25, 0.4, 0.75], + [0.025, 0.0, 0.1, 0.0], + [0.03, 0.0, 0.04963, 0.0], + [0.035, 0.0, 0.0, 0.0], + + ] + ), rtol=1e-03, atol = 1e-03 + ) \ No newline at end of file diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index e9ce99f92e..e169a6827c 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -54,7 +54,7 @@ def test_without_mask(self): result = self.model.forward_atomic( self.extended_coord, self.extended_atype, self.nlist ) - expected_result = torch.tensor([[1.2000, 1.3542], [1.2000, 0.4000]]) + expected_result = torch.tensor([[1.2000, 1.3614], [1.2000, 0.4000]]) torch.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) @@ -64,7 +64,7 @@ def test_with_mask(self): result = self.model.forward_atomic( self.extended_coord, self.extended_atype, self.nlist ) - expected_result = torch.tensor([[0.8000, 1.3542], [1.2000, 0.4000]]) + expected_result = torch.tensor([[0.8000, 1.3614], [1.2000, 0.4000]]) torch.testing.assert_allclose(result["energy"], expected_result, 0.0001, 0.0001) @@ -112,34 +112,38 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: for dist, rcut in zip( [ - 0.010, + 0.01, 0.015, 0.020, 0.015, - 0.020, + 0.02, 0.021, 0.015, - 0.020, + 0.02, 0.021, 0.025, 0.026, 0.025, - 0.02999, + 0.025, + 0.0216161 + ], [ 0.015, 0.015, 0.015, 0.02, - 0.020, - 0.020, + 0.02, + 0.02, 0.022, 0.022, 0.022, 0.025, 0.025, - 0.030, - 0.030, + 0.03, + 0.035, + 0.025 + ], ): extended_coord = torch.tensor( @@ -162,24 +166,25 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: [ [ [0.4, 0], - [0.25, 0.0], [0.0, 0], - [0.25, 0], - [0.125, 0], [0.0, 0], [0.25, 0], + [0, 0], + [0, 0], + [0.25, 0], [0.125, 0], - [0.0469, 0], - [0.0, 0], - [0.0, 0], - [0.0469, 0], - [0, 0], + [0.0922, 0], + [0, 0], + [0, 0], + [0, 0], + [0.0923, 0], + [0.0713, 0], ] ] ) ] - ).reshape(13, 2) - results = torch.stack(results).reshape(13, 2) + ).reshape(14, 2) + results = torch.stack(results).reshape(14, 2) torch.testing.assert_allclose(results, expected_result, 0.0001, 0.0001) From ddbe7dbc48e4a1705df4b40778920656a23b7546 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 31 Jan 2024 08:54:20 +0000 Subject: [PATCH 20/20] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/model/pair_tab.py | 6 +-- deepmd/utils/pair_tab.py | 45 ++++++++++-------- .../tests/common/test_pairtab_preprocess.py | 46 ++++++++++++------- source/tests/pt/test_pairtab.py | 20 ++++---- 4 files changed, 66 insertions(+), 51 deletions(-) diff --git a/deepmd/pt/model/model/pair_tab.py b/deepmd/pt/model/model/pair_tab.py index bf93c4ff97..6f0782289a 100644 --- a/deepmd/pt/model/model/pair_tab.py +++ b/deepmd/pt/model/model/pair_tab.py @@ -195,7 +195,6 @@ def _pair_tabulated_inter( table_coef = table_coef.reshape(self.nframes, self.nloc, self.nnei, 4) ener = self._calcualte_ener(table_coef, uu) - # here we need to overwrite energy to zero at rcut and beyond. mask_beyond_rcut = rr >= self.rcut # also overwrite values beyond extrapolation to zero @@ -205,7 +204,6 @@ def _pair_tabulated_inter( return ener - @staticmethod def _get_pairwise_dist(coords: torch.Tensor) -> torch.Tensor: """Get pairwise distance `dr`. @@ -310,7 +308,5 @@ def _calcualte_ener(coef: torch.Tensor, uu: torch.Tensor) -> torch.Tensor: """ a3, a2, a1, a0 = torch.unbind(coef, dim=-1) etmp = (a3 * uu + a2) * uu + a1 # this should be elementwise operations. - ener = ( - etmp * uu + a0 - ) # this energy has the extrapolated value when rcut > rmax + ener = etmp * uu + a0 # this energy has the extrapolated value when rcut > rmax return ener diff --git a/deepmd/utils/pair_tab.py b/deepmd/utils/pair_tab.py index 9ced80fe95..56f8e618df 100644 --- a/deepmd/utils/pair_tab.py +++ b/deepmd/utils/pair_tab.py @@ -64,6 +64,7 @@ def reinit(self, filename: str, rcut: Optional[float] = None) -> None: ) # this nspline is updated based on the expanded table. self.tab_info = np.array([self.rmin, self.hh, self.nspline, self.ntypes]) self.tab_data = self._make_data() + def _check_table_upper_boundary(self) -> None: """Update User Provided Table Based on `rcut`. @@ -122,7 +123,9 @@ def _check_table_upper_boundary(self) -> None: elif self.rcut > self.rmax: pad_zero = np.zeros((rcut_idx - upper_idx, self.ncol)) pad_zero[:, 0] = np.linspace( - self.rmax + self.hh, self.rmax + self.hh * (rcut_idx - upper_idx), rcut_idx - upper_idx + self.rmax + self.hh, + self.rmax + self.hh * (rcut_idx - upper_idx), + rcut_idx - upper_idx, ) self.vdata = np.concatenate((self.vdata, pad_zero), axis=0) else: @@ -139,11 +142,13 @@ def _check_table_upper_boundary(self) -> None: pad_extrapolation = np.zeros((rcut_idx - upper_idx, self.ncol)) pad_extrapolation[:, 0] = np.linspace( - self.rmax + self.hh, self.rmax + self.hh * (rcut_idx - upper_idx), rcut_idx - upper_idx + self.rmax + self.hh, + self.rmax + self.hh * (rcut_idx - upper_idx), + rcut_idx - upper_idx, ) - # need to calculate table values to fill in with cubic spline + # need to calculate table values to fill in with cubic spline pad_extrapolation = self._extrapolate_table(pad_extrapolation) - + self.vdata = np.concatenate((self.vdata, pad_extrapolation), axis=0) def get(self) -> Tuple[np.array, np.array]: @@ -167,34 +172,38 @@ def _extrapolate_table(self, pad_extrapolation: np.array) -> np.array: Parameters ---------- - pad_extrapolation: np.array + pad_extrapolation : np.array The emepty grid that holds the extrapolation values. Returns ------- np.array - The cubic spline extrapolation. + The cubic spline extrapolation. """ # in theory we should check if the table has at least two rows. - slope = (self.vdata[-1,1:] - self.vdata[-2,1:]) # shape of (ncol-1, ) + slope = self.vdata[-1, 1:] - self.vdata[-2, 1:] # shape of (ncol-1, ) # for extrapolation, we want values decay to `0` prior to `ruct` if possible # here we try to find the grid point prior to `rcut` - grid_point = -2 if pad_extrapolation[-1,0]/self.hh - self.rmax/self.hh >= 2 else -1 - temp_grid = np.stack((self.vdata[-1,:], pad_extrapolation[grid_point,:])) - vv = temp_grid[:,1:] - xx = temp_grid[:,0] - cs = CubicSpline(xx,vv, bc_type=((1,slope),(1,np.zeros_like(slope)))) - xx_grid = pad_extrapolation[:,0] + grid_point = ( + -2 if pad_extrapolation[-1, 0] / self.hh - self.rmax / self.hh >= 2 else -1 + ) + temp_grid = np.stack((self.vdata[-1, :], pad_extrapolation[grid_point, :])) + vv = temp_grid[:, 1:] + xx = temp_grid[:, 0] + cs = CubicSpline(xx, vv, bc_type=((1, slope), (1, np.zeros_like(slope)))) + xx_grid = pad_extrapolation[:, 0] res = cs(xx_grid) - - pad_extrapolation[:,1:] = res + + pad_extrapolation[:, 1:] = res # Note: when doing cubic spline, if we want to ensure values decay to zero prior to `rcut` # this may cause values be positive post `rcut`, we need to overwrite those values to zero - pad_extrapolation = pad_extrapolation if grid_point == -1 else pad_extrapolation[:-1,:] + pad_extrapolation = ( + pad_extrapolation if grid_point == -1 else pad_extrapolation[:-1, :] + ) return pad_extrapolation - + def _make_data(self): data = np.zeros([self.ntypes * self.ntypes * 4 * self.nspline]) stride = 4 * self.nspline @@ -203,7 +212,7 @@ def _make_data(self): for t0 in range(self.ntypes): for t1 in range(t0, self.ntypes): vv = self.vdata[:, 1 + idx_iter] - cs = CubicSpline(xx, vv, bc_type='clamped') + cs = CubicSpline(xx, vv, bc_type="clamped") dd = cs(xx, 1) dd *= self.hh dtmp = np.zeros(stride) diff --git a/source/tests/common/test_pairtab_preprocess.py b/source/tests/common/test_pairtab_preprocess.py index cc776383e1..a866c42236 100644 --- a/source/tests/common/test_pairtab_preprocess.py +++ b/source/tests/common/test_pairtab_preprocess.py @@ -40,9 +40,10 @@ def test_preprocess(self): [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.0, 0.0], - ] - ), rtol=1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) np.testing.assert_allclose( self.tab2.vdata, @@ -53,7 +54,9 @@ def test_preprocess(self): [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], ] - ), rtol=1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) # for this test case, the table does not decay to zero at rcut = 0.22, @@ -69,7 +72,9 @@ def test_preprocess(self): [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.0, 0.0], ] - ), rtol=1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) np.testing.assert_allclose( @@ -81,9 +86,10 @@ def test_preprocess(self): [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.0, 0.0], - ] - ), rtol = 1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) np.testing.assert_allclose( @@ -97,7 +103,9 @@ def test_preprocess(self): [0.025, 0.12468, 0.1992, 0.3741], [0.03, 0.0, 0.0, 0.0], ] - ), rtol = 1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) @@ -124,7 +132,6 @@ def test_preprocess(self): self.tab1.vdata, np.array( [ - [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -136,7 +143,6 @@ def test_preprocess(self): self.tab2.vdata, np.array( [ - [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -149,7 +155,6 @@ def test_preprocess(self): self.tab3.vdata, np.array( [ - [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -163,7 +168,6 @@ def test_preprocess(self): self.tab4.vdata, np.array( [ - [0.01, 0.8, 1.6, 2.4], [0.015, 0.5, 1.0, 1.5], [0.02, 0.25, 0.4, 0.75], @@ -174,6 +178,7 @@ def test_preprocess(self): ), ) + class TestPairTabPreprocessUneven(unittest.TestCase): @patch("numpy.loadtxt") def setUp(self, mock_loadtxt) -> None: @@ -204,7 +209,9 @@ def test_preprocess(self): [0.02, 0.25, 0.4, 0.75], [0.025, 0.0, 0.1, 0.0], ] - ), rtol=1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) np.testing.assert_allclose( self.tab2.vdata, @@ -217,7 +224,9 @@ def test_preprocess(self): [0.025, 0.0, 0.1, 0.0], [0.03, 0.0, 0.0, 0.0], ] - ), rtol=1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) np.testing.assert_allclose( @@ -231,7 +240,9 @@ def test_preprocess(self): [0.025, 0.0, 0.1, 0.0], [0.03, 0.0, 0.0, 0.0], ] - ), rtol=1e-04, atol = 1e-04 + ), + rtol=1e-04, + atol=1e-04, ) np.testing.assert_allclose( @@ -245,7 +256,8 @@ def test_preprocess(self): [0.025, 0.0, 0.1, 0.0], [0.03, 0.0, 0.04963, 0.0], [0.035, 0.0, 0.0, 0.0], - ] - ), rtol=1e-03, atol = 1e-03 - ) \ No newline at end of file + ), + rtol=1e-03, + atol=1e-03, + ) diff --git a/source/tests/pt/test_pairtab.py b/source/tests/pt/test_pairtab.py index e169a6827c..b4dbda6702 100644 --- a/source/tests/pt/test_pairtab.py +++ b/source/tests/pt/test_pairtab.py @@ -125,8 +125,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.026, 0.025, 0.025, - 0.0216161 - + 0.0216161, ], [ 0.015, @@ -142,8 +141,7 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: 0.025, 0.03, 0.035, - 0.025 - + 0.025, ], ): extended_coord = torch.tensor( @@ -169,16 +167,16 @@ def test_extrapolation_nonzero_rmax(self, mock_loadtxt) -> None: [0.0, 0], [0.0, 0], [0.25, 0], - [0, 0], - [0, 0], + [0, 0], + [0, 0], [0.25, 0], [0.125, 0], [0.0922, 0], - [0, 0], - [0, 0], - [0, 0], - [0.0923, 0], - [0.0713, 0], + [0, 0], + [0, 0], + [0, 0], + [0.0923, 0], + [0.0713, 0], ] ] )