diff --git a/lib/scholar/interpolation/monotonic_cubic_spline.ex b/lib/scholar/interpolation/monotonic_cubic_spline.ex new file mode 100644 index 00000000..49187be1 --- /dev/null +++ b/lib/scholar/interpolation/monotonic_cubic_spline.ex @@ -0,0 +1,228 @@ +defmodule Scholar.Interpolation.MonotonicCubicSpline do + @moduledoc """ + Monotonic cubic spline interpolation (also known as PCHIP, from + *Piecewise Cubic Hermite Interpolating Polynomial*). + + Like `Scholar.Interpolation.CubicSpline`, this fits a piecewise cubic + polynomial to the given points. The derivatives at each point, however, + are chosen with the Fritsch-Carlson method so that the resulting curve is + monotonic wherever the data is monotonic and has no overshoot around local + extrema. In exchange, only the first derivative of the curve is guaranteed + to be continuous (the second one is not). + + This is a good choice when the shape of the data matters more than its + smoothness, for example to avoid interpolated values that fall outside the + range of the surrounding samples. + + Monotonic cubic spline interpolation is $O(N)$ where $N$ is the number of points. + + References: + + * [1] - [Monotone cubic interpolation](https://en.wikipedia.org/wiki/Monotone_cubic_interpolation) + * [2] - [Fritsch, F. N.; Carlson, R. E. (1980). "Monotone Piecewise Cubic Interpolation"](https://doi.org/10.1137/0717021) + * [3] - [SciPy implementation (PchipInterpolator)](https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.PchipInterpolator.html) + """ + import Nx.Defn + import Scholar.Shared + + @derive {Nx.Container, containers: [:coefficients, :x]} + defstruct [:coefficients, :x] + + @doc """ + Fits a monotonic cubic spline interpolation of the given `(x, y)` points. + + Inputs are expected to be rank-1 tensors with the same shape + and at least 2 entries. + + ## Examples + + iex> x = Nx.iota({3}) + iex> y = Nx.tensor([0.0, 1.0, 0.0]) + iex> Scholar.Interpolation.MonotonicCubicSpline.fit(x, y) + %Scholar.Interpolation.MonotonicCubicSpline{ + coefficients: Nx.tensor( + [ + [0.0, -1.0, 2.0, 0.0], + [0.0, -1.0, 0.0, 1.0] + ] + ), + x: Nx.tensor( + [0, 1, 2] + ) + } + """ + deftransform fit(x, y) do + fit_n(x, y) + end + + defnp fit_n(x, y) do + # https://en.wikipedia.org/wiki/Monotone_cubic_interpolation + # Reference implementation in SciPy (PchipInterpolator) + + n = + case Nx.shape(x) do + {n} when n > 1 -> + n + + shape -> + raise ArgumentError, + "expected x to be a tensor with shape {n}, where n > 1, got: #{inspect(shape)}" + end + + case {Nx.shape(y), Nx.shape(x)} do + {shape, shape} -> + :ok + + {y_shape, x_shape} -> + raise ArgumentError, + "expected y to have shape #{inspect(x_shape)}, got: #{inspect(y_shape)}" + end + + sort_idx = Nx.argsort(x) + x = Nx.take(x, sort_idx) + y = Nx.take(y, sort_idx) + + dx = Nx.diff(x) + slope = Nx.diff(y) / dx + + s = find_derivatives(dx, slope, n) + + # convert to power-basis coefficients, as in `Scholar.Interpolation.CubicSpline` + t = (s[0..-2//1] + s[1..-1//1] - 2 * slope) / dx + + c_3 = t / dx + c_2 = (slope - s[0..-2//1]) / dx - t + c_1 = s[0..-2//1] + c_0 = y[0..-2//1] + + c = Nx.stack([c_3, c_2, c_1, c_0], axis: 1) + + %__MODULE__{coefficients: c, x: x} + end + + defnp find_derivatives(dx, slope, n) do + case n do + 2 -> + # a single interval is linear: both derivatives equal the secant slope + Nx.broadcast(slope[0], {2}) + + _ -> + sign = Nx.sign(slope) + slope_left = slope[0..-2//1] + slope_right = slope[1..-1//1] + + # zero the derivative where the data is not locally monotonic + # (secant slopes change sign or one is zero) to avoid overshoot + non_monotonic? = + sign[1..-1//1] != sign[0..-2//1] or slope_right == 0 or slope_left == 0 + + dx_left = dx[0..-2//1] + dx_right = dx[1..-1//1] + + w1 = 2 * dx_right + dx_left + w2 = dx_right + 2 * dx_left + + # avoid dividing by a zero slope: some backends raise on `x / 0` + # instead of returning infinity, so masking afterwards is not enough + safe_slope_left = Nx.select(slope_left == 0, 1.0, slope_left) + safe_slope_right = Nx.select(slope_right == 0, 1.0, slope_right) + + weighted_harmonic_mean = (w1 / safe_slope_left + w2 / safe_slope_right) / (w1 + w2) + safe_weighted_harmonic_mean = Nx.select(non_monotonic?, 1.0, weighted_harmonic_mean) + + interior = Nx.select(non_monotonic?, 0.0, 1.0 / safe_weighted_harmonic_mean) + + first = edge_derivative(dx[0], dx[1], slope[0], slope[1]) + last = edge_derivative(dx[-1], dx[-2], slope[-1], slope[-2]) + + Nx.concatenate([Nx.new_axis(first, 0), interior, Nx.new_axis(last, 0)]) + end + end + + # one-sided three-point estimate at the boundary, capped to avoid overshoot + defnp edge_derivative(h0, h1, m0, m1) do + d = ((2 * h0 + h1) * m0 - h0 * m1) / (h0 + h1) + + opposite_sign? = Nx.sign(d) != Nx.sign(m0) + overshoot? = Nx.sign(m0) != Nx.sign(m1) and Nx.abs(d) > 3 * Nx.abs(m0) + + d = Nx.select(opposite_sign?, 0.0, d) + Nx.select(not opposite_sign? and overshoot?, 3 * m0, d) + end + + predict_opts = [ + extrapolate: [ + required: false, + default: true, + type: :boolean, + doc: "if false, out-of-bounds x return NaN." + ] + ] + + @predict_opts_schema NimbleOptions.new!(predict_opts) + + @doc """ + Returns the value fit by `fit/2` corresponding to the `target_x` input. + + ### Options + + #{NimbleOptions.docs(@predict_opts_schema)} + + ## Examples + + iex> x = Nx.iota({3}) + iex> y = Nx.tensor([0.0, 1.0, 0.0]) + iex> model = Scholar.Interpolation.MonotonicCubicSpline.fit(x, y) + iex> Scholar.Interpolation.MonotonicCubicSpline.predict(model, Nx.tensor([0.5, 1.5])) + Nx.tensor( + [0.75, 0.75] + ) + """ + deftransform predict(%__MODULE__{} = model, target_x, opts \\ []) do + predict_n(model, target_x, NimbleOptions.validate!(opts, @predict_opts_schema)) + end + + defnp predict_n(%__MODULE__{x: x, coefficients: coefficients}, target_x, opts) do + original_shape = Nx.shape(target_x) + target_x = Nx.flatten(target_x) + + idx_selector = Nx.new_axis(target_x, 1) > Nx.new_axis(x, 0) + + idx_poly = + idx_selector + |> Nx.argmax(axis: 1, tie_break: :high) + |> Nx.min(Nx.size(x) - 2) + + # deal with the case where no valid index is found + # means that we're in the first interval + # _poly suffix because we're selecting a specific polynomial + # for each target_x value + idx_poly = + Nx.all(idx_selector == 0, axes: [1]) + |> Nx.select(0, idx_poly) + + coef_poly = Nx.take(coefficients, idx_poly) + + # each polynomial is calculated as if the origin was moved to the + # x value that represents the start of the interval + x_poly = target_x - Nx.take(x, idx_poly) + + result = + x_poly + |> Nx.new_axis(1) + |> Nx.pow(Nx.tensor([3, 2, 1, 0])) + |> Nx.dot([1], [0], coef_poly, [1], [0]) + + result = + if opts[:extrapolate] do + result + else + nan_selector = target_x < x[0] or target_x > x[-1] + + nan = Nx.tensor(:nan, type: to_float_type(target_x)) + Nx.select(nan_selector, nan, result) + end + + Nx.reshape(result, original_shape) + end +end diff --git a/mix.exs b/mix.exs index 38ef0cda..c57c855c 100644 --- a/mix.exs +++ b/mix.exs @@ -77,6 +77,7 @@ defmodule Scholar.MixProject do Scholar.Interpolation.BezierSpline, Scholar.Interpolation.CubicSpline, Scholar.Interpolation.Linear, + Scholar.Interpolation.MonotonicCubicSpline, Scholar.Linear.BayesianRidgeRegression, Scholar.Linear.IsotonicRegression, Scholar.Linear.LinearRegression, diff --git a/test/scholar/interpolation/monotonic_cubic_spline_test.exs b/test/scholar/interpolation/monotonic_cubic_spline_test.exs new file mode 100644 index 00000000..048fa569 --- /dev/null +++ b/test/scholar/interpolation/monotonic_cubic_spline_test.exs @@ -0,0 +1,202 @@ +defmodule Scholar.Interpolation.MonotonicCubicSplineTest do + use Scholar.Case, async: true + import Nx, only: :sigils + + alias Scholar.Interpolation.MonotonicCubicSpline + doctest MonotonicCubicSpline + + describe "monotonic cubic spline" do + test "fit/2" do + # Reference values taken from SciPy (scipy.interpolate.PchipInterpolator) + x = Nx.iota({5}) + y = Nx.tensor([1, 2, 3, -10, -1]) + + model = MonotonicCubicSpline.fit(x, y) + + assert_all_close( + model.coefficients, + Nx.tensor([ + [0.0, 0.0, 1.0, 1.0], + [-1.0, 1.0, 1.0, 2.0], + [26.0, -39.0, 0.0, 3.0], + [2.0, 7.0, 0.0, -10.0] + ]) + ) + end + + test "fit/2 with only two points" do + x = Nx.tensor([0, 1]) + y = Nx.tensor([0.0, 2.0]) + + model = MonotonicCubicSpline.fit(x, y) + + # a single interval collapses to a straight line + assert_all_close(model.coefficients, Nx.tensor([[0.0, 0.0, 2.0, 0.0]])) + assert_all_close(MonotonicCubicSpline.predict(model, Nx.tensor(0.5)), Nx.tensor(1.0)) + end + + test "input validation error cases" do + assert_raise ArgumentError, + "expected x to be a tensor with shape {n}, where n > 1, got: {1, 1, 1}", + fn -> + MonotonicCubicSpline.fit(Nx.iota({1, 1, 1}), Nx.iota({1, 1, 1})) + end + + assert_raise ArgumentError, + "expected x to be a tensor with shape {n}, where n > 1, got: {}", + fn -> + MonotonicCubicSpline.fit(Nx.iota({}), Nx.iota({})) + end + + assert_raise ArgumentError, + "expected x to be a tensor with shape {n}, where n > 1, got: {1}", + fn -> + MonotonicCubicSpline.fit(Nx.iota({1}), Nx.iota({1})) + end + + assert_raise ArgumentError, "expected y to have shape {4}, got: {3}", fn -> + MonotonicCubicSpline.fit(Nx.iota({4}), Nx.iota({3})) + end + end + + test "predict/2" do + # Reference values taken from SciPy (scipy.interpolate.PchipInterpolator) + x = Nx.iota({5}) + y = Nx.tensor([1, 2, 3, -10, -1]) + + model = MonotonicCubicSpline.fit(x, y) + + # ensure given values are predicted accurately + # also ensures that the code works for scalar tensors + assert_all_close(MonotonicCubicSpline.predict(model, 0), Nx.tensor(1.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 1), Nx.tensor(2.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 2), Nx.tensor(3.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 3), Nx.tensor(-10.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 4), Nx.tensor(-1.0)) + + # Test for continuity over the given point's boundaries + # (helps ensure no off-by-one's are happening when selecting polynomials) + assert_all_close( + MonotonicCubicSpline.predict( + model, + Nx.tensor([-0.001, 0.001, 0.999, 1.001, 1.999, 2.001, 2.999, 3.001, 3.999, 4.001]) + ), + Nx.tensor([ + 0.999, + 1.001, + 1.999, + 2.001000999, + 2.999998001, + 2.999961026, + -9.999961026, + -9.999992998, + -1.019987002, + -0.979986998 + ]) + ) + + # ensure reference values are calculated accordingly + x_predict = Nx.tensor([-1, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5]) + + assert_all_close( + MonotonicCubicSpline.predict(model, x_predict), + Nx.tensor([0.0, 0.5, 1.5, 2.625, -3.5, -8.0, 12.5, 34.0]) + ) + end + + test "predict/2 preserves monotonicity (no overshoot)" do + # For monotonic data the interpolant must stay monotonic, unlike a + # regular cubic spline which can overshoot around the flat region. + x = Nx.iota({5}) + y = Nx.tensor([0.0, 1.0, 1.0, 1.0, 5.0]) + + model = MonotonicCubicSpline.fit(x, y) + + dense = Nx.linspace(0.0, 4.0, n: 81) + values = MonotonicCubicSpline.predict(model, dense) + + # non-decreasing everywhere + assert Nx.reduce_min(Nx.diff(values)) |> Nx.to_number() >= -1.0e-6 + + # the flat region [1, 3] never rises above the sampled value of 1.0 + in_flat = Nx.logical_and(Nx.greater_equal(dense, 1.0), Nx.less_equal(dense, 3.0)) + max_in_flat = Nx.select(in_flat, values, 0.0) |> Nx.reduce_max() |> Nx.to_number() + assert max_in_flat <= 1.0 + 1.0e-6 + end + + test "predict/2 returns NaN if out-of-bounds and extrapolate: false" do + model = MonotonicCubicSpline.fit(Nx.tensor([0, 1, 2]), Nx.tensor([0, 1, 2])) + + assert MonotonicCubicSpline.predict(model, Nx.tensor([-1, 0, 1, 2, 3]), extrapolate: false) == + ~VEC[NaN 0 1 2 NaN] + end + + test "predict/2 returns the same shape as the input" do + # train a straight line + x = y = Nx.iota({3}) + + model = MonotonicCubicSpline.fit(x, y) + + # a straight line has zero second/third order coefficients + assert_all_close( + model.coefficients, + Nx.tensor([ + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 1.0] + ]) + ) + + for shape <- [{}, {2}, {3, 4}, {5, 6, 7}] do + target_x = Nx.iota(shape, type: :f32) + # a straight line will output the input + assert MonotonicCubicSpline.predict(model, target_x) == target_x + end + end + + test "fit/2 and predict/2 work with jit_apply" do + x = Nx.iota({5}) + y = Nx.tensor([1, 2, 3, -10, -1]) + + model = Nx.Defn.jit_apply(&MonotonicCubicSpline.fit/2, [x, y]) + + prediction = + Nx.Defn.jit_apply(&MonotonicCubicSpline.predict/3, [model, Nx.tensor(1.0), []]) + + assert prediction == Nx.tensor(2.0) + end + + test "fit/2 propagates input precision (f64)" do + # f64 inputs must not be downcast to f32 + x = Nx.tensor([0, 1, 2, 3, 4], type: :f64) + y = Nx.tensor([1, 2, 3, -10, -1], type: :f64) + + model = MonotonicCubicSpline.fit(x, y) + assert Nx.type(model.coefficients) == {:f, 64} + + prediction = MonotonicCubicSpline.predict(model, Nx.tensor([0.5, 2.5], type: :f64)) + assert Nx.type(prediction) == {:f, 64} + assert_all_close(prediction, Nx.tensor([1.5, -3.5], type: :f64)) + end + + test "not sorted x" do + x = Nx.tensor([3, 2, 4, 1, 0]) + y = Nx.tensor([-10, 3, -1, 2, 1]) + + model = MonotonicCubicSpline.fit(x, y) + + # sorting recovers the same points as the "predict/2" test + assert_all_close(MonotonicCubicSpline.predict(model, 0), Nx.tensor(1.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 1), Nx.tensor(2.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 2), Nx.tensor(3.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 3), Nx.tensor(-10.0)) + assert_all_close(MonotonicCubicSpline.predict(model, 4), Nx.tensor(-1.0)) + + x_predict = Nx.tensor([-1, -0.5, 0.5, 1.5, 2.5, 3.5, 4.5, 5]) + + assert_all_close( + MonotonicCubicSpline.predict(model, x_predict), + Nx.tensor([0.0, 0.5, 1.5, 2.625, -3.5, -8.0, 12.5, 34.0]) + ) + end + end +end