diff --git a/lib/scholar/linear/isotonic_regression.ex b/lib/scholar/linear/isotonic_regression.ex index 87469cbd..604c0f1b 100644 --- a/lib/scholar/linear/isotonic_regression.ex +++ b/lib/scholar/linear/isotonic_regression.ex @@ -15,6 +15,7 @@ defmodule Scholar.Linear.IsotonicRegression do @derive { Nx.Container, + keep: [:out_of_bounds], containers: [ :increasing, :x_min, @@ -32,7 +33,8 @@ defmodule Scholar.Linear.IsotonicRegression do :y_thresholds, :increasing, :cutoff_index, - :preprocess + :preprocess, + :out_of_bounds ] @type t() :: %__MODULE__{ @@ -42,7 +44,8 @@ defmodule Scholar.Linear.IsotonicRegression do y_thresholds: Nx.Tensor.t(), increasing: Nx.Tensor.t(), cutoff_index: Nx.Tensor.t(), - preprocess: tuple() | Scholar.Interpolation.Linear.t() + preprocess: tuple() | Scholar.Interpolation.Linear.t(), + out_of_bounds: :clip | :nan } opts = [ @@ -140,11 +143,13 @@ defmodule Scholar.Linear.IsotonicRegression do cutoff_index: Nx.tensor( 5 ), - preprocess: {} + preprocess: {}, + out_of_bounds: :nan } """ deftransform fit(x, y, opts \\ []) do - {n_samples} = Nx.shape(x) + check_input_shape(x) + n_samples = Nx.axis_size(x, 0) y = LinearHelpers.validate_y_shape(y, n_samples, __MODULE__) opts = NimbleOptions.validate!(opts, @opts_schema) @@ -194,7 +199,8 @@ defmodule Scholar.Linear.IsotonicRegression do y_thresholds: y, increasing: increasing, cutoff_index: index_cut, - preprocess: {} + preprocess: {}, + out_of_bounds: opts[:out_of_bounds] } end @@ -222,12 +228,25 @@ defmodule Scholar.Linear.IsotonicRegression do check_preprocess(model) x = Nx.flatten(x) - x = Nx.clip(x, model.x_min, model.x_max) - Scholar.Interpolation.Linear.predict( - model.preprocess, - x - ) + predictions = + Scholar.Interpolation.Linear.predict( + model.preprocess, + Nx.clip(x, model.x_min, model.x_max) + ) + + handle_out_of_bounds(predictions, x, model) + end + + deftransformp handle_out_of_bounds(predictions, x, model) do + case model.out_of_bounds do + :clip -> predictions + :nan -> nan_out_of_bounds(predictions, x, model.x_min, model.x_max) + end + end + + defnp nan_out_of_bounds(predictions, x, x_min, x_max) do + Nx.select(x < x_min or x > x_max, Nx.Constants.nan(), predictions) end @doc """ @@ -271,7 +290,8 @@ defmodule Scholar.Linear.IsotonicRegression do x: Nx.tensor( [1.0, 4.0, 7.0, 9.0, 10.0, 11.0] ) - } + }, + out_of_bounds: :nan } """ def preprocess(%__MODULE__{} = model, trim_duplicates \\ true) do @@ -280,7 +300,7 @@ defmodule Scholar.Linear.IsotonicRegression do y = model.y_thresholds[0..cutoff] {x, y} = - if trim_duplicates do + if trim_duplicates and Nx.axis_size(y, 0) > 2 do keep_mask = Nx.logical_or( Nx.not_equal(y[1..-2//1], y[0..-3//1]), @@ -307,6 +327,7 @@ defmodule Scholar.Linear.IsotonicRegression do model | x_thresholds: x, y_thresholds: y, + cutoff_index: Nx.subtract(Nx.axis_size(x, 0), 1), preprocess: Scholar.Interpolation.Linear.fit(x, y) } end @@ -393,38 +414,47 @@ defmodule Scholar.Linear.IsotonicRegression do current_weight = Nx.tensor(0, type: Nx.type(sample_weights)) index = 0 + started = Nx.u8(0) - {{x_output, y_output, sample_weights_output, index, current_x, current_y, current_weight}, _} = + {{x_output, y_output, sample_weights_output, index, current_x, current_y, current_weight, + _started}, _} = while {{x_output, y_output, sample_weights_output, index, current_x, current_y, - current_weight}, {j = 0, eps = 1.0e-10, y, x, sample_weights}}, + current_weight, started}, {j = 0, eps = 1.0e-10, y, x, sample_weights}}, j < Nx.axis_size(x, 0) do x_j = x[j] - {x_output, y_output, sample_weights_output, index, current_x, current_weight, current_y} = - if x_j - current_x >= eps do - x_output = Nx.indexed_put(x_output, Nx.new_axis(index, 0), current_x) - y_output = Nx.indexed_put(y_output, Nx.new_axis(index, 0), current_y / current_weight) + {x_output, y_output, sample_weights_output, index, current_x, current_weight, current_y, + started} = + cond do + # sklearn drops these samples before fitting; keeping them makes a group + # whose total weight is zero, and its mean is then 0 / 0 + sample_weights[j] <= 0 -> + {x_output, y_output, sample_weights_output, index, current_x, current_weight, + current_y, started} + + not started -> + {x_output, y_output, sample_weights_output, index, x_j, sample_weights[j], + y[j] * sample_weights[j], Nx.u8(1)} - sample_weights_output = - Nx.indexed_put(sample_weights_output, Nx.new_axis(index, 0), current_weight) + x_j - current_x >= eps -> + x_output = Nx.indexed_put(x_output, Nx.new_axis(index, 0), current_x) - index = index + 1 - current_x = x_j - current_weight = sample_weights[j] - current_y = y[j] * sample_weights[j] + y_output = + Nx.indexed_put(y_output, Nx.new_axis(index, 0), current_y / current_weight) - {x_output, y_output, sample_weights_output, index, current_x, current_weight, - current_y} - else - current_weight = current_weight + sample_weights[j] - current_y = current_y + y[j] * sample_weights[j] + sample_weights_output = + Nx.indexed_put(sample_weights_output, Nx.new_axis(index, 0), current_weight) - {x_output, y_output, sample_weights_output, index, current_x, current_weight, - current_y} + {x_output, y_output, sample_weights_output, index + 1, x_j, sample_weights[j], + y[j] * sample_weights[j], started} + + true -> + {x_output, y_output, sample_weights_output, index, current_x, + current_weight + sample_weights[j], current_y + y[j] * sample_weights[j], started} end - {{x_output, y_output, sample_weights_output, index, current_x, current_y, current_weight}, - {j + 1, eps, y, x, sample_weights}} + {{x_output, y_output, sample_weights_output, index, current_x, current_y, current_weight, + started}, {j + 1, eps, y, x, sample_weights}} end x_output = Nx.indexed_put(x_output, Nx.new_axis(index, 0), current_x) @@ -505,23 +535,41 @@ defmodule Scholar.Linear.IsotonicRegression do i = if(increasing, do: 0, else: Nx.axis_size(y, 0) - 1 - max_size) |> Nx.as_type(:u32) + # the blocks run to y_size, which trails the placeholder slots when decreasing {y, _} = - while {y, {target, i, _k = Nx.u32(0), max_size}}, i < max_size + 1 do + while {y, {target, i, _k = Nx.u32(0), y_size}}, i < y_size + 1 do k = target[i] + 1 indices = Nx.iota({Nx.axis_size(y, 0)}, type: :u32) in_range? = Nx.logical_and(i + 1 <= indices, indices < k) y = Nx.select(in_range?, y[i], y) i = k - {y, {target, i, k, max_size}} + {y, {target, i, k, y_size}} end if increasing, do: y, else: Nx.reverse(y) end + # sign of Spearman's rho, matching sklearn; a constant input makes it undefined + # and the comparison then picks decreasing defnp check_increasing(x, y) do - x = Nx.new_axis(x, -1) - y = Nx.new_axis(y, -1) - model = Scholar.Linear.LinearRegression.fit(x, y) - model.coefficients[0] >= 0 + rank_x = average_rank(Nx.flatten(x)) + rank_y = average_rank(Nx.flatten(y)) + + centered_x = rank_x - Nx.mean(rank_x) + centered_y = rank_y - Nx.mean(rank_y) + + rho = + Nx.sum(centered_x * centered_y) / + Nx.sqrt(Nx.sum(centered_x ** 2) * Nx.sum(centered_y ** 2)) + + rho >= 0 + end + + defnp average_rank(t) do + column = Nx.new_axis(t, 1) + row = Nx.new_axis(t, 0) + smaller = Nx.sum(row < column, axes: [1]) + tied = Nx.sum(row == column, axes: [1]) + smaller + (tied + 1) / 2 end end diff --git a/test/scholar/linear/isotonic_regression_test.exs b/test/scholar/linear/isotonic_regression_test.exs index ec4f12ee..692ada3e 100644 --- a/test/scholar/linear/isotonic_regression_test.exs +++ b/test/scholar/linear/isotonic_regression_test.exs @@ -150,6 +150,76 @@ defmodule Scholar.Linear.IsotonicRegressionTest do model = Scholar.Linear.IsotonicRegression.fit(x, y, sample_weights: sample_weights) assert model.increasing == Nx.u8(0) end + + test "fit decreasing with tied x stays monotonic" do + # Expected values from scikit-learn 1.6.1 on this data. + model = + IsotonicRegression.fit(Nx.tensor([0, 0, 1, 1, 2]), Nx.tensor([4, 4, 5, 5, 3]), + increasing: false + ) + |> IsotonicRegression.preprocess() + + assert_all_close(model.y_thresholds, Nx.tensor([4.5, 4.5, 3.0])) + + assert_all_close( + IsotonicRegression.predict(model, Nx.tensor([0, 0, 1, 1, 2])), + Nx.tensor([4.5, 4.5, 4.5, 4.5, 3.0]) + ) + end + + test "fit drops samples with zero weight" do + # Expected values from scikit-learn 1.6.1 on this data. + weighted = + IsotonicRegression.fit(Nx.tensor([1, 2, 3, 4, 5]), Nx.tensor([1, 5, 3, 4, 5]), + sample_weights: [1, 0, 1, 1, 1] + ) + |> IsotonicRegression.preprocess() + + # dropping the zero-weight sample by hand has to give the same model + dropped = + IsotonicRegression.fit(Nx.tensor([1, 3, 4, 5]), Nx.tensor([1, 3, 4, 5])) + |> IsotonicRegression.preprocess() + + assert_all_close(weighted.x_thresholds, dropped.x_thresholds) + assert_all_close(weighted.y_thresholds, dropped.y_thresholds) + + assert_all_close( + IsotonicRegression.predict(weighted, Nx.tensor([1, 2, 3, 4, 5])), + Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0]) + ) + end + + test "fit accepts the column shape check_input_shape allows" do + y = Nx.tensor([1, 3, 6, 8, 9, 10]) + + column = + IsotonicRegression.fit(Nx.tensor([[1], [4], [7], [9], [10], [11]]), y) + |> IsotonicRegression.preprocess() + + flat = + IsotonicRegression.fit(Nx.tensor([1, 4, 7, 9, 10, 11]), y) + |> IsotonicRegression.preprocess() + + assert_all_close(column.x_thresholds, flat.x_thresholds) + assert_all_close(column.y_thresholds, flat.y_thresholds) + end + + test "fit with :increasing? as :auto follows the rank correlation" do + # Expected values from scikit-learn 1.6.1 on this data. + # the outlier drags an ordinary least squares slope to -5 while Spearman's rho + # stays at 0.45, and scikit-learn calls this increasing + x = Nx.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + y = Nx.tensor([1, 2, 3, 4, 5, 6, 7, 8, 9, -100]) + + assert IsotonicRegression.fit(x, y).increasing == Nx.u8(1) + + assert IsotonicRegression.fit(Nx.tensor([1, 2, 3, 4, 5]), Nx.tensor([2, 2, 2, 3, 1])).increasing == + Nx.u8(0) + + # a constant input leaves the correlation undefined, and scikit-learn picks decreasing + assert IsotonicRegression.fit(Nx.tensor([1, 2, 3, 4]), Nx.tensor([5, 5, 5, 5])).increasing == + Nx.u8(0) + end end test "preprocess" do @@ -277,4 +347,56 @@ defmodule Scholar.Linear.IsotonicRegressionTest do # Fitted with the default :increasing?, so predictions must not decrease with x. assert a <= b and b <= c end + + test "preprocess with fewer than three thresholds" do + # Expected values from scikit-learn 1.6.1 on this data. + constant = + IsotonicRegression.fit(Nx.tensor([1, 2, 3, 4]), Nx.tensor([5, 5, 5, 5])) + |> IsotonicRegression.preprocess() + + assert_all_close(constant.x_thresholds, Nx.tensor([1.0, 4.0])) + assert_all_close(constant.y_thresholds, Nx.tensor([5.0, 5.0])) + + two_points = + IsotonicRegression.fit(Nx.tensor([1, 2]), Nx.tensor([3, 4])) + |> IsotonicRegression.preprocess() + + assert_all_close(two_points.y_thresholds, Nx.tensor([3.0, 4.0])) + end + + test "cutoff_index follows the thresholds preprocess kept" do + model = + IsotonicRegression.fit(Nx.tensor([1, 2, 3, 4]), Nx.tensor([5, 5, 5, 5])) + |> IsotonicRegression.preprocess() + + assert model.cutoff_index == Nx.tensor(1) + + # a stale index makes a second pass read past the end of the tensor + assert_all_close( + IsotonicRegression.preprocess(model).y_thresholds, + model.y_thresholds + ) + end + + test "out_of_bounds decides what happens outside the fitted range" do + # Expected values from scikit-learn 1.6.1 on this data. + x = Nx.tensor([1, 2, 3]) + y = Nx.tensor([1, 2, 3]) + to_predict = Nx.tensor([0.0, 2.0, 9.0]) + + clipped = + IsotonicRegression.fit(x, y, out_of_bounds: :clip) |> IsotonicRegression.preprocess() + + assert_all_close( + IsotonicRegression.predict(clipped, to_predict), + Nx.tensor([1.0, 2.0, 3.0]) + ) + + # :nan is the documented default + nan = IsotonicRegression.fit(x, y) |> IsotonicRegression.preprocess() + predictions = IsotonicRegression.predict(nan, to_predict) + + assert Nx.is_nan(predictions) == Nx.tensor([1, 0, 1], type: :u8) + assert_all_close(predictions[1], Nx.tensor(2.0)) + end end