diff --git a/gnm/shape/gnm_tensorflow.py b/gnm/shape/gnm_tensorflow.py index ec7833f..0a55c11 100644 --- a/gnm/shape/gnm_tensorflow.py +++ b/gnm/shape/gnm_tensorflow.py @@ -125,14 +125,18 @@ def compute_vertex_normals( v2 = face_vertices[..., 2, :] face_normals_area = tf.experimental.numpy.cross(v1 - v0, v2 - v0, axis=-1) - data = tf.reshape(face_normals_area[:, :, None, :], (-1, 3)) + # Replicate each face normal for the three corner vertices it contributes + # to, matching the flattened (triangle, corner) order of the segment ids. + data = tf.reshape( + tf.tile(face_normals_area[:, :, None, :], [1, 1, 3, 1]), (-1, 3) + ) segment_ids = tf.reshape(self.triangles, (-1,)) segment_ids = tf.tile(segment_ids[None, :], [batch_size, 1]) batch_offset = ( tf.range(batch_size, dtype=segment_ids.dtype)[:, None] * num_vertices ) - flat_segment_ids = segment_ids + batch_offset + flat_segment_ids = tf.reshape(segment_ids + batch_offset, (-1,)) flat_normals = tf.math.unsorted_segment_sum( data, flat_segment_ids, num_segments=batch_size * num_vertices ) diff --git a/gnm/shape/gnm_tensorflow_test.py b/gnm/shape/gnm_tensorflow_test.py index 2244d92..17934eb 100644 --- a/gnm/shape/gnm_tensorflow_test.py +++ b/gnm/shape/gnm_tensorflow_test.py @@ -108,6 +108,32 @@ def test_parity_with_gnm_numpy( np.testing.assert_almost_equal(actual, desired, decimal=4) + @parameterized.product( + version=_MAINTAINED_MAJOR_GNM_VERSIONS, + variant=tuple(_SUPPORTED_VARIANTS), + batch_dims=BATCH_DIMS, + ) + def test_compute_vertex_normals_numpy_parity( + self, version: str, variant: str, batch_dims: tuple[int, ...] + ): + """Tests that TensorFlow vertex normals match the NumPy implementation.""" + if variant not in self.gnms_np[version]: + self.skipTest(f'variant {variant} not supported in {version}.') + gnm_np = self.gnms_np[version][variant] + gnm_tf = self.gnms_tf[version][variant] + + parameters_np = gnm_test_utils.random_gnm_parameters( + gnm_np, batch_shape=batch_dims, seed=self.rng + ) + vertices_np = gnm_np(**parameters_np) + desired = gnm_np.compute_vertex_normals(vertices_np) + + vertices_tf = tf.convert_to_tensor(vertices_np, dtype=tf.float32) + actual = gnm_tf.compute_vertex_normals(vertices_tf) + + self.assertEqual(tuple(actual.shape), vertices_np.shape) + np.testing.assert_allclose(actual.numpy(), desired, atol=1e-4) + @parameterized.product( version=_MAINTAINED_MAJOR_GNM_VERSIONS, variant=(gnm_tensorflow.GNMVariant.HEAD,),