Standard Scaler fit-transform interface

lib/scholar/preprocessing.ex

@@ -133,14 +133,8 @@ defmodule Scholar.Preprocessing do
       >
   """
   deftransform standard_scale(tensor, opts \\ []) do
-    standard_scale_n(tensor, NimbleOptions.validate!(opts, @general_schema))
-  end
-
-  defnp standard_scale_n(tensor, opts) do
-    std = Nx.standard_deviation(tensor, axes: opts[:axes], keep_axes: true)
-    mean_reduced = Nx.mean(tensor, axes: opts[:axes], keep_axes: true)
-    mean_reduced = Nx.select(std == 0, 0.0, mean_reduced)
-    (tensor - mean_reduced) / Nx.select(std == 0, 1.0, std)
+    opts = NimbleOptions.validate!(opts, @general_schema)
+    Scholar.Preprocessing.StandardScaler.fit_transform(tensor, opts)
   end
 
   @doc """

lib/scholar/preprocessing/standard_scaler.ex

@@ -0,0 +1,46 @@
+defmodule Scholar.Preprocessing.StandardScaler do
+  import Nx.Defn
+
+  defstruct [:deviation, :mean]
+
+  opts_schema = [
+    axes: [
+      type: {:custom, Scholar.Options, :axes, []},
+      doc: """
+      Axes to calculate the distance over. By default the distance
+      is calculated between the whole tensors.
+      """
+    ]
+  ]
+
+  @opts_schema NimbleOptions.new!(opts_schema)
+
+  deftransform fit(tensor, opts \\ []) do
+    NimbleOptions.validate!(opts, @opts_schema)
+    {std, mean} = fit_n(tensor, opts)
+
+    %__MODULE__{deviation: std, mean: mean}
+  end
+
+  defnp fit_n(tensor, opts) do
+    std = Nx.standard_deviation(tensor, axes: opts[:axes], keep_axes: true)
+    mean_reduced = Nx.mean(tensor, axes: opts[:axes], keep_axes: true)
+    mean_reduced = Nx.select(Nx.equal(std, 0), 0.0, mean_reduced)
+
+    {std, mean_reduced}
+  end
+
+  deftransform transform(%__MODULE__{deviation: std, mean: mean}, tensor) do
+    scale(tensor, std, mean)
+  end
+
+  defn fit_transform(tensor, opts \\ []) do
+    tensor
+    |> fit(opts)
+    |> transform(tensor)
+  end
+
+  defnp scale(tensor, std, mean) do
+    (tensor - mean) / Nx.select(std == 0, 1.0, std)
+  end
+end

test/scholar/preprocessing/standard_scaler_test.exs

@@ -0,0 +1,25 @@
+defmodule StandardScalerTest do
+  use Scholar.Case, async: true
+  alias Scholar.Preprocessing.StandardScaler
+
+  describe "fit_transform/2" do
+    test "applies standard scaling to data" do
+      data = Nx.tensor([[1, -1, 2], [2, 0, 0], [0, 1, -1]])
+
+      expected =
+        Nx.tensor([
+          [0.5212860703468323, -1.3553436994552612, 1.4596009254455566],
+          [1.4596009254455566, -0.4170288145542145, -0.4170288145542145],
+          [-0.4170288145542145, 0.5212860703468323, -1.3553436994552612]
+        ])
+
+      assert_all_close(StandardScaler.fit_transform(data), expected)
+    end
+
+    test "leaves data as it is when variance is zero" do
+      data = 42.0
+      expected = Nx.tensor(data)
+      assert StandardScaler.fit_transform(data) == expected
+    end
+  end
+end