Add unit tests for core modules: latent explorer, attribute directions, and custom loss

tests/test_attribute_directions.py

@@ -0,0 +1,23 @@
+import unittest
+import numpy as np
+from faceforge_core.attribute_directions import LatentDirectionFinder
+
+class TestLatentDirectionFinder(unittest.TestCase):
+    def setUp(self):
+        # 100 samples, 5D latent
+        self.latents = np.random.randn(100, 5)
+        self.labels = [0]*50 + [1]*50
+        self.finder = LatentDirectionFinder(self.latents)
+
+    def test_pca_direction(self):
+        components, explained = self.finder.pca_direction(n_components=2)
+        self.assertEqual(components.shape, (2, 5))
+        self.assertEqual(explained.shape, (2,))
+
+    def test_classifier_direction(self):
+        direction = self.finder.classifier_direction(self.labels)
+        self.assertEqual(direction.shape, (5,))
+        self.assertAlmostEqual(np.linalg.norm(direction), 1.0, places=5)
+
+if __name__ == "__main__":
+    unittest.main() 

tests/test_custom_loss.py

@@ -0,0 +1,28 @@
+import unittest
+import torch
+from faceforge_core.custom_loss import attribute_preserving_loss
+
+class TestAttributePreservingLoss(unittest.TestCase):
+    def setUp(self):
+        self.generated = torch.ones((2, 3, 4, 4))
+        self.original = torch.zeros((2, 3, 4, 4))
+        self.y_target = torch.ones((2, 1))
+        self.attr_predictor = lambda x: torch.ones((2, 1))
+
+    def test_loss_value(self):
+        loss = attribute_preserving_loss(
+            self.generated, self.original, self.attr_predictor, self.y_target, lambda_pred=2.0, lambda_recon=3.0
+        )
+        # pred_loss = 0, recon_loss = mean((1-0)^2) = 1
+        self.assertAlmostEqual(loss.item(), 3.0)
+
+    def test_loss_with_nonzero_pred(self):
+        attr_predictor = lambda x: torch.zeros((2, 1))
+        loss = attribute_preserving_loss(
+            self.generated, self.original, attr_predictor, self.y_target, lambda_pred=2.0, lambda_recon=3.0
+        )
+        # pred_loss = mean((0-1)^2) = 1, recon_loss = 1
+        self.assertAlmostEqual(loss.item(), 5.0)
+
+if __name__ == "__main__":
+    unittest.main() 

tests/test_latent_explorer.py

@@ -0,0 +1,44 @@
+import unittest
+import numpy as np
+from faceforge_core.latent_explorer import LatentSpaceExplorer, LatentPoint
+
+class TestLatentSpaceExplorer(unittest.TestCase):
+    def setUp(self):
+        self.explorer = LatentSpaceExplorer()
+        self.dummy_encoding = np.array([1.0, 2.0])
+
+    def test_add_point(self):
+        self.explorer.add_point("test", self.dummy_encoding, (0.5, 0.5))
+        self.assertEqual(len(self.explorer.points), 1)
+        self.assertEqual(self.explorer.points[0].text, "test")
+        np.testing.assert_array_equal(self.explorer.points[0].encoding, self.dummy_encoding)
+        self.assertEqual(self.explorer.points[0].xy_pos, (0.5, 0.5))
+
+    def test_delete_point(self):
+        self.explorer.add_point("test", self.dummy_encoding)
+        self.explorer.delete_point(0)
+        self.assertEqual(len(self.explorer.points), 0)
+
+    def test_modify_point(self):
+        self.explorer.add_point("test", self.dummy_encoding)
+        new_encoding = np.array([3.0, 4.0])
+        self.explorer.modify_point(0, "new", new_encoding)
+        self.assertEqual(self.explorer.points[0].text, "new")
+        np.testing.assert_array_equal(self.explorer.points[0].encoding, new_encoding)
+
+    def test_sample_encoding_distance(self):
+        self.explorer.add_point("a", np.array([1.0, 0.0]), (0.0, 0.0))
+        self.explorer.add_point("b", np.array([0.0, 1.0]), (1.0, 0.0))
+        sampled = self.explorer.sample_encoding((0.5, 0.0), mode="distance")
+        self.assertIsNotNone(sampled)
+        self.assertEqual(sampled.shape, (2,))
+
+    def test_sample_encoding_circle(self):
+        self.explorer.add_point("a", np.array([1.0, 0.0]), (1.0, 0.0))
+        self.explorer.add_point("b", np.array([0.0, 1.0]), (0.0, 1.0))
+        sampled = self.explorer.sample_encoding((1.0, 1.0), mode="circle")
+        self.assertIsNotNone(sampled)
+        self.assertEqual(sampled.shape, (2,))
+
+if __name__ == "__main__":
+    unittest.main()