Fix bugs in Scholar.Manifold.Trimap

lib/scholar/manifold/trimap.ex

@@ -47,14 +47,17 @@ defmodule Scholar.Manifold.Trimap do
       type: :pos_integer,
       default: 5,
       doc: ~S"""
-      Number of outliers to sample.
+      Number of outliers to sample per inlier. `num_inliers * num_outliers`
+      cannot exceed `num_samples - num_inliers - 1`, the number of points
+      outside an anchor's own neighborhood.
       """
     ],
     num_random: [
-      type: :pos_integer,
+      type: :non_neg_integer,
       default: 3,
       doc: ~S"""
-      Number of random triplets to sample.
+      Number of random triplets to sample per point. Set to 0 to use only the
+      nearest neighbor triplets.
       """
     ],
     weight_temp: [
@@ -181,7 +184,8 @@ defmodule Scholar.Manifold.Trimap do
         end
       end
 
-    Nx.take(is_in, order1)
+    # order1 sorted tensor1; argsort inverts it to put the answers back in the input's order
+    Nx.take(is_in, Nx.argsort(order1))
   end
 
   defnp rejection_sample(key, shape, rejects, opts \\ []) do
@@ -277,7 +281,9 @@ defmodule Scholar.Manifold.Trimap do
     distances =
       (handle_dist(inputs[anchors], inputs[hits], opts) ** 2) |> Nx.reshape({num_points, :auto})
 
-    sigmas = Nx.max(Nx.mean(Nx.sqrt(distances[[.., 3..5]]), axes: [1]), 1.0e-10)
+    # the scale is the mean distance to the 4th-6th neighbor, or as many of those as exist
+    sigmas =
+      Nx.max(Nx.mean(Nx.sqrt(distances[[.., 3..min(5, num_neighbors - 1)]]), axes: [1]), 1.0e-10)
 
     scaled_distances = distances / (Nx.reshape(sigmas, {:auto, 1}) * sigmas[neighbors])
     sort_indices = Nx.argsort(scaled_distances, axis: 1)
@@ -483,6 +489,24 @@ defmodule Scholar.Manifold.Trimap do
       raise ArgumentError, "Number of points must be greater than 2"
     end
 
+    num_points = Nx.axis_size(inputs, 0)
+    available = num_points - opts[:num_inliers] - 1
+
+    # sampling outliers rejects the anchor and its inliers, and never repeats a draw,
+    # so a wider request than the remaining pool loops forever
+    if Nx.rank(triplets) == 0 and opts[:num_inliers] * opts[:num_outliers] > available do
+      raise ArgumentError,
+            "num_inliers * num_outliers must be at most #{available}, the number of points " <>
+              "outside an anchor's own neighborhood, got: #{opts[:num_inliers] * opts[:num_outliers]}"
+    end
+
+    if Nx.rank(init_embeddings) != 0 and
+         Nx.axis_size(init_embeddings, 1) != opts[:num_components] do
+      raise ArgumentError,
+            "init_embeddings must have num_components (#{opts[:num_components]}) columns, " <>
+              "got: #{Nx.axis_size(init_embeddings, 1)}"
+    end
+
     unless (Nx.rank(triplets) == Nx.rank(weights) and Nx.rank(triplets) == 0) or
              (Nx.rank(triplets) == 2 and Nx.rank(weights) == 1 and
                 Nx.axis_size(triplets, 0) == Nx.axis_size(weights, 0)) do
@@ -501,26 +525,26 @@ defmodule Scholar.Manifold.Trimap do
   end
 
   defnp transform_n(inputs, key, triplets, weights, init_embeddings, opts \\ []) do
-    {num_points, num_components} = Nx.shape(inputs)
+    {num_points, num_features} = Nx.shape(inputs)
 
-    {triplets, weights, key, applied_pca?} =
+    {inputs, triplets, weights, key, applied_pca?} =
       case triplets do
         {} ->
           {inputs, applied_pca} =
-            if num_components > @dim_pca do
+            if num_features > @dim_pca do
+              num_pca = min(@dim_pca, min(num_points, num_features))
               inputs = inputs - Nx.mean(inputs, axes: [0])
-              {u, s, vt} = Nx.LinAlg.SVD.svd(inputs, full_matrices: false)
-              inputs = Nx.dot(u[[.., 0..@dim_pca]] * s[0..@dim_pca], vt[[0..@dim_pca, ..]])
-              {inputs, Nx.u8(1)}
+              {u, s, _vt} = Nx.LinAlg.svd(inputs, full_matrices?: false)
+              {u[[.., 0..(num_pca - 1)]] * s[0..(num_pca - 1)], Nx.u8(1)}
             else
               {inputs, Nx.u8(0)}
             end
 
           {triplets, weights, key} = generate_triplets(key, inputs, opts)
-          {triplets, weights, key, applied_pca}
+          {inputs, triplets, weights, key, applied_pca}
 
         _ ->
-          {triplets, weights, key, Nx.u8(0)}
+          {inputs, triplets, weights, key, Nx.u8(0)}
       end
 
     embeddings =

test/scholar/manifold/trimap_test.exs

@@ -202,6 +202,140 @@ defmodule Scholar.Manifold.TrimapTest do
 
       assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3)
     end
+
+    test "inputs with more than 100 features are reduced to their principal components" do
+      key = Nx.Random.key(42)
+      {x, _} = Nx.Random.uniform(Nx.Random.key(11), shape: {30, 105}, type: :f64)
+
+      # 30 points span at most 30 components, so the reduction keeps all of u * s
+      {u, s, _vt} = Nx.LinAlg.svd(Nx.subtract(x, Nx.mean(x, axes: [0])), full_matrices?: false)
+      projection = Nx.multiply(u, s)
+
+      opts = [num_components: 2, key: key, num_inliers: 3, num_outliers: 2, num_iters: 10]
+
+      # everything downstream, the default initialization included, must see the
+      # projection rather than the original features
+      assert_all_close(
+        Trimap.transform(x, opts),
+        Trimap.transform(
+          projection,
+          opts ++ [init_embeddings: Nx.multiply(0.01, projection[[.., 0..1]])]
+        )
+      )
+    end
+
+    test "three points" do
+      key = Nx.Random.key(42)
+      {x, _} = Nx.Random.uniform(Nx.Random.key(1), shape: {3, 4})
+
+      res =
+        Trimap.transform(x,
+          num_components: 2,
+          key: key,
+          num_inliers: 1,
+          num_outliers: 1,
+          num_random: 1,
+          num_iters: 100
+        )
+
+      expected =
+        Nx.tensor([
+          [5.157116413116455, 5.203591823577881],
+          [20.418094635009766, 20.411230087280273],
+          [1.4575116634368896, 1.4591535329818726]
+        ])
+
+      assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3)
+    end
+
+    test "four points" do
+      key = Nx.Random.key(42)
+      {x, _} = Nx.Random.uniform(Nx.Random.key(2), shape: {4, 4})
+
+      res =
+        Trimap.transform(x,
+          num_components: 2,
+          key: key,
+          num_inliers: 1,
+          num_outliers: 2,
+          num_random: 1,
+          num_iters: 100
+        )
+
+      expected =
+        Nx.tensor([
+          [1.459843397140503, 6.352163791656494],
+          [20.399131774902344, 20.404085159301758],
+          [20.421730041503906, 20.41226577758789],
+          [1.4555573463439941, 1.45782470703125]
+        ])
+
+      assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3)
+    end
+
+    test "num_random set to zero" do
+      key = Nx.Random.key(42)
+      {x, _} = Nx.Random.uniform(Nx.Random.key(3), shape: {10, 6})
+
+      res =
+        Trimap.transform(x,
+          num_components: 2,
+          key: key,
+          num_inliers: 3,
+          num_outliers: 1,
+          num_random: 0,
+          num_iters: 100
+        )
+
+      expected =
+        Nx.tensor([
+          [14.64638900756836, 4.130362510681152],
+          [14.39341926574707, 11.968405723571777],
+          [10.877017974853516, 12.163773536682129],
+          [20.414138793945312, 20.422441482543945],
+          [13.235034942626953, 3.089409351348877],
+          [8.24937915802002, 10.1371431350708],
+          [12.3439302444458, 12.452646255493164],
+          [1.4550071954727173, 1.601865530014038],
+          [15.539610862731934, 5.656279563903809],
+          [1.4515089988708496, 1.460447072982788]
+        ])
+
+      assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3)
+    end
+
+    test "outlier sampling that has to retry" do
+      key = Nx.Random.key(42)
+      {x, _} = Nx.Random.uniform(Nx.Random.key(4), shape: {10, 3})
+
+      # 6 outliers per anchor out of the 6 points outside its neighborhood, so the
+      # first draw always collides and rejection sampling has to resample
+      res =
+        Trimap.transform(x,
+          num_components: 2,
+          key: key,
+          num_inliers: 3,
+          num_outliers: 2,
+          num_random: 2,
+          num_iters: 100
+        )
+
+      expected =
+        Nx.tensor([
+          [10.7711763381958, 4.751033782958984],
+          [20.177593231201172, 15.467020988464355],
+          [2.4794745445251465, 19.923784255981445],
+          [9.068700790405273, 4.932981491088867],
+          [10.389301300048828, 4.25684928894043],
+          [20.413013458251953, 13.008362770080566],
+          [4.308704853057861, 19.454608917236328],
+          [18.78362464904785, 18.64022445678711],
+          [1.4547451734542847, 20.409406661987305],
+          [8.286545753479004, 1.9438791275024414]
+        ])
+
+      assert_all_close(res, expected, atol: 1.0e-3, rtol: 1.0e-3)
+    end
   end
 
   describe "errors" do
@@ -241,5 +375,40 @@ defmodule Scholar.Manifold.TrimapTest do
                      )
                    end
     end
+
+    test "num_inliers times num_outliers larger than the pool of non-neighbors" do
+      x = Nx.iota({20, 6})
+      key = Nx.Random.key(42)
+
+      assert_raise ArgumentError,
+                   "num_inliers * num_outliers must be at most 16, the number of points " <>
+                     "outside an anchor's own neighborhood, got: 18",
+                   fn ->
+                     Trimap.transform(x,
+                       num_components: 2,
+                       key: key,
+                       num_inliers: 3,
+                       num_outliers: 6
+                     )
+                   end
+    end
+
+    test "init_embeddings with the wrong number of columns" do
+      x = Nx.iota({20, 6})
+      key = Nx.Random.key(42)
+      {init_embeddings, _} = Nx.Random.uniform(Nx.Random.key(3), shape: {20, 3})
+
+      assert_raise ArgumentError,
+                   "init_embeddings must have num_components (2) columns, got: 3",
+                   fn ->
+                     Trimap.transform(x,
+                       num_components: 2,
+                       key: key,
+                       num_inliers: 3,
+                       num_outliers: 1,
+                       init_embeddings: init_embeddings
+                     )
+                   end
+    end
   end
 end