Dedupe test model lookups.

gnm/shape/fitting_utils/project_on_pca_test.py

@@ -19,6 +19,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 from gnm.shape import gnm_numpy
+from gnm.shape import gnm_test_utils
 from gnm.shape import gnm_utils
 from gnm.shape.data.versions import gnm_test_catalog
 from gnm.shape.fitting_utils import project_on_pca
@@ -37,16 +38,20 @@ def setUpClass(cls):
     super().setUpClass()
     cls.rng = np.random.default_rng(0)
 
-    cls.gnms = {}
-    for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS:
-      if (
-          gnm_numpy.GNMVariant.HEAD
-          in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version]
-      ):
-        cls.gnms[version] = gnm_numpy.GNM.from_remote(
-            gnm_numpy.GNMMajorVersion(version.removeprefix('v')),
-            gnm_numpy.GNMVariant.HEAD,
+    cls.gnms = {
+        version: gnm_test_utils.load_gnm(
+            gnm_numpy.GNM, version, gnm_numpy.GNMVariant.HEAD
         )
+        for version in gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS
+        if gnm_numpy.GNMVariant.HEAD
+        in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version]
+    }
+
+  def _get_gnm(self, version: str) -> gnm_numpy.GNM:
+    """Returns the head model of `version`, or skips if it is unavailable."""
+    if version not in self.gnms:
+      self.skipTest(f'Head variant not available in version {version}')
+    return self.gnms[version]
 
   @parameterized.product(
       version=gnm_test_catalog.MAINTAINED_MAJOR_VERSIONS,
@@ -55,7 +60,7 @@ def setUpClass(cls):
   )
   def test_can_recover_identity(self, version: str, index: int, value: float):
     """Tests that the estimated identity is close to the ground truth."""
-    gnm_np = self.gnms[version]
+    gnm_np = self._get_gnm(version)
     identity = np.zeros(gnm_np.identity_dim)
     expression = np.zeros(gnm_np.expression_dim)
     identity[index] = value
@@ -78,7 +83,7 @@ def test_can_recover_identity(self, version: str, index: int, value: float):
   )
   def test_can_recover_expression(self, version: str, index: int, value: float):
     """Tests that the estimated expression is close to the ground truth."""
-    gnm_np = self.gnms[version]
+    gnm_np = self._get_gnm(version)
     identity = np.zeros(gnm_np.identity_dim)
     expression = np.zeros(gnm_np.expression_dim)
     expression[index] = value
@@ -102,17 +107,10 @@ def test_can_fit_identity_using_subset_of_basis(
       self, version, num_components_and_threshold
   ):
     """Tests that we can estimate identity using a subset of the basis."""
-    if (
-        gnm_numpy.GNMVariant.HEAD
-        not in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version]
-    ):
-      self.skipTest(f'Head variant not available in version {version}')
+    gnm = self._get_gnm(version)
 
     num_components, threshold = num_components_and_threshold
 
-    major_version = gnm_numpy.GNMMajorVersion(version.removeprefix('v'))
-    gnm = gnm_numpy.GNM.from_remote(major_version, gnm_numpy.GNMVariant.HEAD)
-
     identity = self.rng.uniform(low=-1.5, high=1.5, size=gnm.identity_dim)
     expression = np.zeros(gnm.expression_dim)
 
@@ -149,17 +147,10 @@ def test_can_recover_identity_with_object(
       self, version, batch_size_num_components_threshold
   ):
     """Tests that the projection object can estimate identity ."""
-    if (
-        gnm_numpy.GNMVariant.HEAD
-        not in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version]
-    ):
-      self.skipTest(f'Head variant not available in version {version}')
+    gnm = self._get_gnm(version)
 
     batch_size, num_components, threshold = batch_size_num_components_threshold
 
-    major_version = gnm_numpy.GNMMajorVersion(version.removeprefix('v'))
-    gnm = gnm_numpy.GNM.from_remote(major_version, gnm_numpy.GNMVariant.HEAD)
-
     identity = self.rng.uniform(
         low=-1.5, high=1.5, size=(batch_size, gnm.identity_dim)
     )
@@ -204,15 +195,9 @@ def test_can_recover_expression_with_object(
       self, version, batch_size_num_components_threshold
   ):
     """Tests that the projection object can estimate expression."""
-    if (
-        gnm_numpy.GNMVariant.HEAD
-        not in gnm_test_catalog.MAJOR_VERSION_TO_VARIANTS_MAP[version]
-    ):
-      self.skipTest(f'Head variant not available in version {version}')
+    gnm = self._get_gnm(version)
 
     batch_size, num_components, threshold = batch_size_num_components_threshold
-    major_version = gnm_numpy.GNMMajorVersion(version.removeprefix('v'))
-    gnm = gnm_numpy.GNM.from_remote(major_version, gnm_numpy.GNMVariant.HEAD)
 
     expected_regions = [
         'left_eye_region',

gnm/shape/gnm_jax_test.py

@@ -81,6 +81,15 @@ def setUpClass(cls):
     )
     cls.addClassCleanup(cls.gnms_jax.clear)
 
+  def _get_gnms(
+      self, version: str, variant: Any
+  ) -> tuple[gnm_jax.GNM, gnm_numpy.GNM]:
+    """Returns the JAX and NumPy GNMs, or skips the test if unavailable."""
+    return (
+        self.gnms_jax.get_or_skip(version, variant),
+        self.gnms_np.get_or_skip(version, variant),
+    )
+
   def _get_default_kwargs(
       self, gnm_np: gnm_numpy.GNM, n_batch: int = 1
   ) -> dict[str, jnp.ndarray]:
@@ -118,11 +127,7 @@ def _get_random_kwargs(
   )
   def test_properties_match(self, version: str, variant: Any):
     """Tests that important properties match between implementations."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     self.assertEqual(gnm.version, gnm_np.version)
     self.assertEqual(gnm.identity_dim, gnm_np.identity_dim)
@@ -148,11 +153,7 @@ def test_parity_with_gnm_numpy(
       dtype: jnp.dtype,
   ):
     """Tests that JAX GNM poses vertices the same as NumPy GNM."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     # Build a batch of random parameters.
     gnm_jax_parameters = self._get_random_kwargs(gnm_np, batch_dim)
@@ -177,11 +178,7 @@ def test_vertices_and_landmarks(
       batch_size: tuple[int, ...],
   ):
     """Tests extracting vertices and landmarks in JAX."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     gnm_jax_parameters = self._get_random_kwargs(gnm_np, batch_size)
     verts, landmarks = gnm.vertices_and_landmarks(
@@ -197,11 +194,7 @@ def test_vertices_and_landmarks(
       dtype=DTYPES,
   )
   def test_jax_jit_behavior(self, version: str, variant: Any, dtype: jnp.dtype):
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
     # Build a batch of random parameters.
     n_batch = 10
 
@@ -234,11 +227,7 @@ def jax_jit_wrapper(identity, expression, rotations, translation):
   )
   def test_bad_shape(self, version: str, variant: Any):
     """Badly shaped parameter should throw an error."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     bad_dimension = gnm.expression_dim + gnm.identity_dim + gnm.num_joints
     n_batch = 5
@@ -254,11 +243,7 @@ def test_bad_shape(self, version: str, variant: Any):
   )
   def test_bad_shape_joint_transforms(self, version: str, variant: Any):
     """Badly shaped parameter should throw an error."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     bad_dimension = gnm.identity_dim + gnm.num_joints
     bad_inputs = [
@@ -284,11 +269,7 @@ def test_joint_transforms_numpy_parity(
       self, version: str, variant: Any, dtype: jnp.dtype
   ):
     """Test that the joint transformations function matches Numpy."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     # Build a batch of random parameters.
     n_batch = 10
@@ -330,11 +311,7 @@ def test_gradients(
       # pyrefly: ignore[bad-function-definition]
       dtype: jnp.dtype = jnp.float32,
   ):
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
-    gnm_np = self.gnms_np[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     # Build a batch of random parameters.
     grad_func = jax.grad(
@@ -388,11 +365,7 @@ def test_gradients(
   )
   def test_transformed_local_offsets(self, version: str, variant: Any):
     """Tests that vertex offset is applied to the posed positions."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm_np = self.gnms_np[version][variant_str]
-    gnm = self.gnms_jax[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     gnm_jax_parameters = self._get_default_kwargs(n_batch=10, gnm_np=gnm_np)
     local_vertices = gnm.vertex_positions_bind_pose(
@@ -433,10 +406,7 @@ def test_transformed_local_offsets(self, version: str, variant: Any):
   )
   def test_to_numpy_data_dict(self, version: str, variant: Any):
     """Tests to_numpy_data_dict method."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
+    gnm = self.gnms_jax.get_or_skip(version, variant)
 
     data_dict = gnm.to_numpy_data_dict()
 
@@ -459,19 +429,14 @@ def test_to_numpy_data_dict(self, version: str, variant: Any):
   )
   def test_from_gnm(self, version: str, variant: Any):
     """Tests from_gnm factory method."""
-    variant_str = variant.value if hasattr(variant, 'value') else variant
-    if variant_str not in self.gnms_np[version]:
-      self.skipTest(f'variant {variant_str} not supported in {version}.')
-    gnm = self.gnms_jax[version][variant_str]
+    gnm, gnm_np = self._get_gnms(version, variant)
 
     new_gnm = gnm_jax.GNM.from_gnm(gnm)
 
     self.assertEqual(new_gnm.version, gnm.version)
     self.assertEqual(new_gnm.variant, gnm.variant)
 
-    parameters = self._get_random_kwargs(
-        self.gnms_np[version][variant_str], n_batch=(10,)
-    )
+    parameters = self._get_random_kwargs(gnm_np, n_batch=(10,))
     parameters_jax = jax.tree.map(jnp.asarray, parameters)
 
     vertices_orig = gnm(**parameters_jax)