refactor: merge TestDipoleSelType with TestDipole using parameterized sel_type

Co-authored-by: njzjz <9496702+njzjz@users.noreply.github.com>

Author: Copilot

source/tests/consistent/fitting/test_dipole.py

@@ -61,176 +61,9 @@
     (True, False),  # resnet_dt
     ("float64", "float32"),  # precision
     (True, False),  # mixed_types
+    ([], [0, 1]),  # sel_type
 )
 class TestDipole(CommonTest, DipoleFittingTest, unittest.TestCase):
-    @property
-    def data(self) -> dict:
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        return {
-            "neuron": [5, 5, 5],
-            "resnet_dt": resnet_dt,
-            "precision": precision,
-            "seed": 20240217,
-        }
-
-    @property
-    def skip_pt(self) -> bool:
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        return CommonTest.skip_pt
-
-    tf_class = DipoleFittingTF
-    dp_class = DipoleFittingDP
-    pt_class = DipoleFittingPT
-    jax_class = DipoleFittingJAX
-    array_api_strict_class = DipoleFittingArrayAPIStrict
-    args = fitting_dipole()
-    skip_jax = not INSTALLED_JAX
-    skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT
-
-    def setUp(self) -> None:
-        CommonTest.setUp(self)
-
-        self.ntypes = 2
-        self.natoms = np.array([6, 6, 2, 4], dtype=np.int32)
-        self.inputs = np.ones((1, 6, 20), dtype=GLOBAL_NP_FLOAT_PRECISION)
-        self.gr = np.ones((1, 6, 30, 3), dtype=GLOBAL_NP_FLOAT_PRECISION)
-        self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32)
-        # inconsistent if not sorted
-        self.atype.sort()
-
-    @property
-    def additional_data(self) -> dict:
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        return {
-            "ntypes": self.ntypes,
-            "dim_descrpt": self.inputs.shape[-1],
-            "mixed_types": mixed_types,
-            "embedding_width": 30,
-        }
-
-    def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]:
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        return self.build_tf_fitting(
-            obj,
-            self.inputs.ravel(),
-            self.gr,
-            self.natoms,
-            self.atype,
-            None,
-            suffix,
-        )
-
-    def eval_pt(self, pt_obj: Any) -> Any:
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        return (
-            pt_obj(
-                torch.from_numpy(self.inputs).to(device=PT_DEVICE),
-                torch.from_numpy(self.atype.reshape(1, -1)).to(device=PT_DEVICE),
-                torch.from_numpy(self.gr).to(device=PT_DEVICE),
-                None,
-            )["dipole"]
-            .detach()
-            .cpu()
-            .numpy()
-        )
-
-    def eval_dp(self, dp_obj: Any) -> Any:
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        return dp_obj(
-            self.inputs,
-            self.atype.reshape(1, -1),
-            self.gr,
-            None,
-        )["dipole"]
-
-    def eval_jax(self, jax_obj: Any) -> Any:
-        return np.asarray(
-            jax_obj(
-                jnp.asarray(self.inputs),
-                jnp.asarray(self.atype.reshape(1, -1)),
-                jnp.asarray(self.gr),
-                None,
-            )["dipole"]
-        )
-
-    def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any:
-        return to_numpy_array(
-            array_api_strict_obj(
-                array_api_strict.asarray(self.inputs),
-                array_api_strict.asarray(self.atype.reshape(1, -1)),
-                array_api_strict.asarray(self.gr),
-                None,
-            )["dipole"]
-        )
-
-    def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]:
-        if backend == self.RefBackend.TF:
-            # shape is not same
-            ret = ret[0].reshape(-1, self.natoms[0], 1)
-        return (ret,)
-
-    @property
-    def rtol(self) -> float:
-        """Relative tolerance for comparing the return value."""
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        if precision == "float64":
-            return 1e-10
-        elif precision == "float32":
-            return 1e-4
-        else:
-            raise ValueError(f"Unknown precision: {precision}")
-
-    @property
-    def atol(self) -> float:
-        """Absolute tolerance for comparing the return value."""
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-        ) = self.param
-        if precision == "float64":
-            return 1e-10
-        elif precision == "float32":
-            return 1e-4
-        else:
-            raise ValueError(f"Unknown precision: {precision}")
-
-
-@parameterized(
-    (True, False),  # resnet_dt
-    ("float64", "float32"),  # precision
-    (True, False),  # mixed_types
-    ([0, 1],),  # sel_type - only test with all types selected for consistency
-)
-class TestDipoleSelType(CommonTest, DipoleFittingTest, unittest.TestCase):
     @property
     def data(self) -> dict:
         (
@@ -239,13 +72,16 @@ def data(self) -> dict:
             mixed_types,
             sel_type,
         ) = self.param
-        return {
+        data = {
             "neuron": [5, 5, 5],
             "resnet_dt": resnet_dt,
             "precision": precision,
             "seed": 20240217,
-            "sel_type": sel_type,  # For TF backend
         }
+        # Only add sel_type if it's not empty (for TF backend compatibility)
+        if sel_type:
+            data["sel_type"] = sel_type
+        return data
 
     @property
     def skip_pt(self) -> bool:
@@ -285,16 +121,18 @@ def additional_data(self) -> dict:
             mixed_types,
             sel_type,
         ) = self.param
-        # For DP/PT backends, use exclude_types instead of sel_type
-        all_types = list(range(self.ntypes))
-        exclude_types = [t for t in all_types if t not in sel_type]
-        return {
+        additional = {
             "ntypes": self.ntypes,
             "dim_descrpt": self.inputs.shape[-1],
             "mixed_types": mixed_types,
             "embedding_width": 30,
-            "exclude_types": exclude_types,  # For DP/PT backends
         }
+        # For DP/PT backends, use exclude_types instead of sel_type
+        if sel_type:
+            all_types = list(range(self.ntypes))
+            exclude_types = [t for t in all_types if t not in sel_type]
+            additional["exclude_types"] = exclude_types
+        return additional
 
     def build_tf(self, obj: Any, suffix: str) -> tuple[list, dict]:
         (
@@ -404,37 +242,6 @@ def atol(self) -> float:
         else:
             raise ValueError(f"Unknown precision: {precision}")
 
-    def test_sel_type_behavior(self):
-        """Test that sel_type parameter works correctly across backends."""
-        (
-            resnet_dt,
-            precision,
-            mixed_types,
-            sel_type,
-        ) = self.param
-
-        # Test TF backend if available
-        if INSTALLED_TF:
-            tf_obj = self.tf_class(**{**self.data, **self.additional_data})
-
-            # Verify that only selected types have fitting nets
-            if hasattr(tf_obj, "sel_type"):
-                self.assertEqual(set(tf_obj.sel_type), set(sel_type))
-            if hasattr(tf_obj, "sel_mask"):
-                expected_mask = np.array([i in sel_type for i in range(self.ntypes)])
-                np.testing.assert_array_equal(tf_obj.sel_mask, expected_mask)
-
-        # Test DP backend
-        all_types = list(range(self.ntypes))
-        exclude_types = [t for t in all_types if t not in sel_type]
-        dp_data = {**self.data}
-        dp_data.pop("sel_type", None)  # Remove sel_type for DP backend
-        dp_obj = self.dp_class(**{**dp_data, **self.additional_data})
-
-        # Verify that exclude_types is set correctly
-        if hasattr(dp_obj, "exclude_types"):
-            self.assertEqual(set(dp_obj.exclude_types), set(exclude_types))
-
 
 class TestDipoleSelTypeBehavior(unittest.TestCase):
     """Test sel_type behavior specifically, without cross-backend consistency."""