add compress

Author: iProzd

deepmd/dpmodel/descriptor/dpa3.py

@@ -14,6 +14,7 @@ def __init__(
         a_rcut: float = 4.0,
         a_rcut_smth: float = 3.5,
         a_sel: int = 20,
+        a_compress_rate: int = 0,
         axis_neuron: int = 4,
         update_angle: bool = True,
         update_style: str = "res_residual",
@@ -45,6 +46,10 @@ def __init__(
             Where to start smoothing for angle. For example the 1/r term is smoothed from rcut to rcut_smth.
         a_sel : int, optional
             Maximally possible number of selected angle neighbors.
+        a_compress_rate : int, optional
+            The compression rate for angular messages. The default value is 0, indicating no compression.
+            If a non-zero integer c is provided, the node and edge dimensions will be compressed
+            to n_dim/c and e_dim/2c, respectively, within the angular message.
         axis_neuron : int, optional
             The number of dimension of submatrix in the symmetrization ops.
         update_angle : bool, optional
@@ -72,6 +77,7 @@ def __init__(
         self.a_rcut = a_rcut
         self.a_rcut_smth = a_rcut_smth
         self.a_sel = a_sel
+        self.a_compress_rate = a_compress_rate
         self.axis_neuron = axis_neuron
         self.update_angle = update_angle
         self.update_style = update_style
@@ -97,6 +103,7 @@ def serialize(self) -> dict:
             "a_rcut": self.a_rcut,
             "a_rcut_smth": self.a_rcut_smth,
             "a_sel": self.a_sel,
+            "a_compress_rate": self.a_compress_rate,
             "axis_neuron": self.axis_neuron,
             "update_angle": self.update_angle,
             "update_style": self.update_style,

deepmd/pt/model/descriptor/dpa3.py

@@ -151,6 +151,7 @@ def init_subclass_params(sub_data, sub_class):
             n_dim=self.repflow_args.n_dim,
             e_dim=self.repflow_args.e_dim,
             a_dim=self.repflow_args.a_dim,
+            a_compress_rate=self.repflow_args.a_compress_rate,
             axis_neuron=self.repflow_args.axis_neuron,
             update_angle=self.repflow_args.update_angle,
             activation_function=self.activation_function,

deepmd/pt/model/descriptor/repflow_layer.py

@@ -45,6 +45,7 @@ def __init__(
         n_dim: int = 128,
         e_dim: int = 16,
         a_dim: int = 64,
+        a_compress_rate: int = 0,
         axis_neuron: int = 4,
         update_angle: bool = True,  # angle
         activation_function: str = "silu",
@@ -70,6 +71,12 @@ def __init__(
         self.n_dim = n_dim
         self.e_dim = e_dim
         self.a_dim = a_dim
+        self.a_compress_rate = a_compress_rate
+        if a_compress_rate != 0:
+            assert a_dim % (2 * a_compress_rate) == 0, (
+                f"For a_compress_rate of {a_compress_rate}, a_dim must be divisible by {2 * a_compress_rate}. "
+                f"Currently, a_dim={a_dim} is not valid."
+            )
         self.axis_neuron = axis_neuron
         self.update_angle = update_angle
         self.activation_function = activation_function
@@ -167,20 +174,42 @@ def __init__(
             )
 
         if self.update_angle:
-            self.angle_dim = self.a_dim + self.n_dim + 2 * self.e_dim
+            self.angle_dim = self.a_dim
+            if self.a_compress_rate == 0:
+                # angle + node + edge * 2
+                self.angle_dim += self.n_dim + 2 * self.e_dim
+                self.a_compress_n_linear = None
+                self.a_compress_e_linear = None
+            else:
+                # angle + node/c + edge/2c * 2
+                self.angle_dim += 2 * (self.a_dim // self.a_compress_rate)
+                self.a_compress_n_linear = MLPLayer(
+                    self.n_dim,
+                    self.a_dim // self.a_compress_rate,
+                    precision=precision,
+                    bias=False,
+                    seed=child_seed(seed, 8),
+                )
+                self.a_compress_e_linear = MLPLayer(
+                    self.e_dim,
+                    self.a_dim // (2 * self.a_compress_rate),
+                    precision=precision,
+                    bias=False,
+                    seed=child_seed(seed, 9),
+                )
 
             # edge angle message
             self.edge_angle_linear1 = MLPLayer(
                 self.angle_dim,
                 self.e_dim,
                 precision=precision,
-                seed=child_seed(seed, 8),
+                seed=child_seed(seed, 10),
             )
             self.edge_angle_linear2 = MLPLayer(
                 self.e_dim,
                 self.e_dim,
                 precision=precision,
-                seed=child_seed(seed, 9),
+                seed=child_seed(seed, 11),
             )
             if self.update_style == "res_residual":
                 self.e_residual.append(
@@ -189,7 +218,7 @@ def __init__(
                         self.update_residual,
                         self.update_residual_init,
                         precision=precision,
-                        seed=child_seed(seed, 10),
+                        seed=child_seed(seed, 12),
                     )
                 )
 
@@ -198,7 +227,7 @@ def __init__(
                 self.angle_dim,
                 self.a_dim,
                 precision=precision,
-                seed=child_seed(seed, 11),
+                seed=child_seed(seed, 13),
             )
             if self.update_style == "res_residual":
                 self.a_residual.append(
@@ -207,13 +236,15 @@ def __init__(
                         self.update_residual,
                         self.update_residual_init,
                         precision=precision,
-                        seed=child_seed(seed, 12),
+                        seed=child_seed(seed, 14),
                     )
                 )
         else:
             self.angle_self_linear = None
             self.edge_angle_linear1 = None
             self.edge_angle_linear2 = None
+            self.a_compress_n_linear = None
+            self.a_compress_e_linear = None
             self.angle_dim = 0
 
         self.n_residual = nn.ParameterList(self.n_residual)
@@ -448,12 +479,22 @@ def forward(
             assert self.edge_angle_linear1 is not None
             assert self.edge_angle_linear2 is not None
             # get angle info
+            if self.a_compress_rate != 0:
+                assert self.a_compress_n_linear is not None
+                assert self.a_compress_e_linear is not None
+                node_ebd_for_angle = self.a_compress_n_linear(node_ebd)
+                edge_ebd_for_angle = self.a_compress_e_linear(edge_ebd)
+            else:
+                node_ebd_for_angle = node_ebd
+                edge_ebd_for_angle = edge_ebd
+
             # nb x nloc x a_nnei x a_nnei x n_dim
             node_for_angle_info = torch.tile(
-                node_ebd.unsqueeze(2).unsqueeze(2), (1, 1, self.a_sel, self.a_sel, 1)
+                node_ebd_for_angle.unsqueeze(2).unsqueeze(2),
+                (1, 1, self.a_sel, self.a_sel, 1),
             )
             # nb x nloc x a_nnei x e_dim
-            edge_for_angle = edge_ebd[:, :, : self.a_sel, :]
+            edge_for_angle = edge_ebd_for_angle[:, :, : self.a_sel, :]
             # nb x nloc x a_nnei x e_dim
             edge_for_angle = torch.where(
                 a_nlist_mask.unsqueeze(-1), edge_for_angle, 0.0
@@ -471,7 +512,7 @@ def forward(
                 [edge_for_angle_i, edge_for_angle_j], dim=-1
             )
             angle_info_list = [angle_ebd, node_for_angle_info, edge_for_angle_info]
-            # nb x nloc x a_nnei x a_nnei x (a + n_dim + e_dim*2)
+            # nb x nloc x a_nnei x a_nnei x (a + n_dim + e_dim*2) or (a + a/c + a/c)
             angle_info = torch.cat(angle_info_list, dim=-1)
 
             # edge angle message
@@ -605,6 +646,7 @@ def serialize(self) -> dict:
             "n_dim": self.n_dim,
             "e_dim": self.e_dim,
             "a_dim": self.a_dim,
+            "a_compress_rate": self.a_compress_rate,
             "axis_neuron": self.axis_neuron,
             "activation_function": self.activation_function,
             "update_angle": self.update_angle,
@@ -625,6 +667,13 @@ def serialize(self) -> dict:
                     "angle_self_linear": self.angle_self_linear.serialize(),
                 }
             )
+            if self.a_compress_rate != 0:
+                data.update(
+                    {
+                        "a_compress_n_linear": self.a_compress_n_linear.serialize(),
+                        "a_compress_e_linear": self.a_compress_e_linear.serialize(),
+                    }
+                )
         if self.update_style == "res_residual":
             data.update(
                 {
@@ -650,13 +699,16 @@ def deserialize(cls, data: dict) -> "RepFlowLayer":
         check_version_compatibility(data.pop("@version"), 1, 1)
         data.pop("@class")
         update_angle = data["update_angle"]
+        a_compress_rate = data["a_compress_rate"]
         node_self_mlp = data.pop("node_self_mlp")
         node_sym_linear = data.pop("node_sym_linear")
         node_edge_linear = data.pop("node_edge_linear")
         edge_self_linear = data.pop("edge_self_linear")
         edge_angle_linear1 = data.pop("edge_angle_linear1", None)
         edge_angle_linear2 = data.pop("edge_angle_linear2", None)
         angle_self_linear = data.pop("angle_self_linear", None)
+        a_compress_n_linear = data.pop("a_compress_n_linear", None)
+        a_compress_e_linear = data.pop("a_compress_e_linear", None)
         update_style = data["update_style"]
         variables = data.pop("@variables", {})
         n_residual = variables.get("n_residual", data.pop("n_residual", []))
@@ -676,6 +728,11 @@ def deserialize(cls, data: dict) -> "RepFlowLayer":
             obj.edge_angle_linear1 = MLPLayer.deserialize(edge_angle_linear1)
             obj.edge_angle_linear2 = MLPLayer.deserialize(edge_angle_linear2)
             obj.angle_self_linear = MLPLayer.deserialize(angle_self_linear)
+            if a_compress_rate != 0:
+                assert isinstance(a_compress_n_linear, dict)
+                assert isinstance(a_compress_e_linear, dict)
+                obj.a_compress_n_linear = MLPLayer.deserialize(a_compress_n_linear)
+                obj.a_compress_e_linear = MLPLayer.deserialize(a_compress_e_linear)
 
         if update_style == "res_residual":
             for ii, t in enumerate(obj.n_residual):

deepmd/pt/model/descriptor/repflows.py

@@ -85,6 +85,7 @@ def __init__(
         n_dim: int = 128,
         e_dim: int = 64,
         a_dim: int = 64,
+        a_compress_rate: int = 0,
         axis_neuron: int = 4,
         update_angle: bool = True,
         activation_function: str = "silu",
@@ -123,6 +124,10 @@ def __init__(
             Where to start smoothing for angle. For example the 1/r term is smoothed from rcut to rcut_smth.
         a_sel : int, optional
             Maximally possible number of selected angle neighbors.
+        a_compress_rate : int, optional
+            The compression rate for angular messages. The default value is 0, indicating no compression.
+            If a non-zero integer c is provided, the node and edge dimensions will be compressed
+            to n_dim/c and e_dim/2c, respectively, within the angular message.
         axis_neuron : int, optional
             The number of dimension of submatrix in the symmetrization ops.
         update_angle : bool, optional
@@ -175,6 +180,7 @@ def __init__(
         self.rcut_smth = e_rcut_smth
         self.sec = self.sel
         self.split_sel = self.sel
+        self.a_compress_rate = a_compress_rate
         self.axis_neuron = axis_neuron
         self.set_davg_zero = set_davg_zero
         self.skip_stat = skip_stat
@@ -218,6 +224,7 @@ def __init__(
                     n_dim=self.n_dim,
                     e_dim=self.e_dim,
                     a_dim=self.a_dim,
+                    a_compress_rate=self.a_compress_rate,
                     axis_neuron=self.axis_neuron,
                     update_angle=self.update_angle,
                     activation_function=self.activation_function,

deepmd/utils/argcheck.py

@@ -1440,6 +1440,11 @@ def dpa3_repflow_args():
     doc_a_sel = 'Maximally possible number of selected angle neighbors. It can be:\n\n\
     - `int`. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200. \n\n\
     - `str`. Can be "auto:factor" or "auto". "factor" is a float number larger than 1. This option will automatically determine the `sel`. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the "factor". Finally the number is wrapped up to 4 divisible. The option "auto" is equivalent to "auto:1.1".'
+    doc_a_compress_rate = (
+        "The compression rate for angular messages. The default value is 0, indicating no compression. "
+        " If a non-zero integer c is provided, the node and edge dimensions will be compressed "
+        "to n_dim/c and e_dim/2c, respectively, within the angular message."
+    )
     doc_axis_neuron = "The number of dimension of submatrix in the symmetrization ops."
     doc_update_angle = (
         "Where to update the angle rep. If not, only node and edge rep will be used."
@@ -1475,6 +1480,9 @@ def dpa3_repflow_args():
         Argument("a_rcut", float, doc=doc_a_rcut),
         Argument("a_rcut_smth", float, doc=doc_a_rcut_smth),
         Argument("a_sel", [int, str], doc=doc_a_sel),
+        Argument(
+            "a_compress_rate", int, optional=True, default=0, doc=doc_a_compress_rate
+        ),
         Argument(
             "axis_neuron",
             int,

source/tests/pt/model/test_dpa3.py

@@ -48,12 +48,14 @@ def test_consistency(
             ua,
             rus,
             ruri,
+            acr,
             prec,
             ect,
         ) in itertools.product(
             [True, False],  # update_angle
             ["res_residual"],  # update_style
             ["norm", "const"],  # update_residual_init
+            [0, 1],  # a_compress_rate
             ["float64"],  # precision
             [False],  # use_econf_tebd
         ):
@@ -73,6 +75,7 @@ def test_consistency(
                 a_rcut=self.rcut - 0.1,
                 a_rcut_smth=self.rcut_smth,
                 a_sel=nnei - 1,
+                a_compress_rate=acr,
                 axis_neuron=4,
                 update_angle=ua,
                 update_style=rus,
@@ -127,12 +130,14 @@ def test_jit(
             ua,
             rus,
             ruri,
+            acr,
             prec,
             ect,
         ) in itertools.product(
             [True, False],  # update_angle
             ["res_residual"],  # update_style
             ["norm", "const"],  # update_residual_init
+            [0, 1],  # a_compress_rate
             ["float64"],  # precision
             [False],  # use_econf_tebd
         ):
@@ -150,6 +155,7 @@ def test_jit(
                 a_rcut=self.rcut - 0.1,
                 a_rcut_smth=self.rcut_smth,
                 a_sel=nnei - 1,
+                a_compress_rate=acr,
                 axis_neuron=4,
                 update_angle=ua,
                 update_style=rus,

source/tests/universal/dpmodel/descriptor/test_descriptor.py

@@ -475,6 +475,7 @@ def DescriptorParamDPA3(
     update_residual=0.1,
     update_residual_init="const",
     update_angle=True,
+    a_compress_rate=0,
     precision="float64",
 ):
     input_dict = {
@@ -491,6 +492,7 @@ def DescriptorParamDPA3(
                 "a_rcut": rcut / 2,
                 "a_rcut_smth": rcut_smth / 2,
                 "a_sel": sum(sel) // 4,
+                "a_compress_rate": a_compress_rate,
                 "axis_neuron": 4,
                 "update_angle": update_angle,
                 "update_style": update_style,
@@ -520,6 +522,7 @@ def DescriptorParamDPA3(
             "update_residual_init": ("const",),
             "exclude_types": ([], [[0, 1]]),
             "update_angle": (True, False),
+            "a_compress_rate": (0, 1),
             "env_protection": (0.0, 1e-8),
             "precision": ("float64",),
         }