Normalize custom maps in SL parametrizations

Author: lezcano

geotorch/sl.py

@@ -30,18 +30,15 @@ def __init__(self, size, f="softplus", triv="expm"):
         super().__init__(size, SL.parse_f(f), triv)
 
     @staticmethod
-    def parse_f(f_name):
-        if f_name in FixedRank.fs.keys():
-            f, inv = FixedRank.parse_f(f_name)
-
-            def f_sl(x):
-                y = f(x)
-                log_y = y.log()
-                return (log_y - log_y.mean(dim=-1, keepdim=True)).exp()
-
-            return (f_sl, inv)
-        else:
-            return f_name
+    def parse_f(f):
+        f, inv = FixedRank.parse_f(f)
+
+        def f_sl(x):
+            y = f(x)
+            log_y = y.log()
+            return (log_y - log_y.mean(dim=-1, keepdim=True)).exp()
+
+        return f_sl, inv
 
     def in_manifold_singular_values(self, S, eps=5e-3):
         rank_eps = torch.finfo(S.dtype).eps * max(self.n, self.k)

test/test_sl.py

@@ -26,6 +26,20 @@ def test_singular_values_are_normalized_in_log_domain(self):
             )
         )
 
+    def test_custom_map_is_normalized(self):
+        manifold = SL(size=(2, 3, 3), f=(torch.exp, torch.log)).double()
+        X = torch.randn(2, 3, 3, dtype=torch.float64)
+        matrix = manifold(X)
+        sign, logabsdet = torch.linalg.slogdet(matrix)
+
+        self.assertTrue(torch.equal(sign, torch.ones_like(sign)))
+        torch.testing.assert_close(
+            logabsdet, torch.zeros_like(logabsdet), atol=1e-6, rtol=0
+        )
+
+        inverse = manifold.right_inverse(matrix)
+        torch.testing.assert_close(manifold(inverse), matrix)
+
     def test_negative_determinant_is_rejected(self):
         manifold = SL(size=(2, 2))
         self.assertFalse(manifold.in_manifold(torch.diag(torch.tensor([-1.0, 1.0]))))