diff --git a/lib/scholar/covariance/ledoit_wolf.ex b/lib/scholar/covariance/ledoit_wolf.ex new file mode 100644 index 00000000..6aa567b3 --- /dev/null +++ b/lib/scholar/covariance/ledoit_wolf.ex @@ -0,0 +1,222 @@ +defmodule Scholar.Covariance.LedoitWolf do + @moduledoc """ + Ledoit-Wolf is a particular form of shrinkage covariance estimator, where the shrinkage coefficient is computed using O. Ledoit and M. Wolf’s formula. + + Ledoit and M. Wolf's formula as + described in "A Well-Conditioned Estimator for Large-Dimensional + Covariance Matrices", Ledoit and Wolf, Journal of Multivariate + Analysis, Volume 88, Issue 2, February 2004, pages 365-411. + """ + import Nx.Defn + + @derive {Nx.Container, containers: [:covariance, :shrinkage, :location]} + defstruct [:covariance, :shrinkage, :location] + + opts_schema = [ + assume_centered: [ + default: false, + type: :boolean, + doc: """ + If `true`, data will not be centered before computation. + Useful when working with data whose mean is almost, but not exactly + zero. + If `false`, data will be centered before computation. + """ + ] + ] + + @opts_schema NimbleOptions.new!(opts_schema) + @doc """ + Estimate the shrunk Ledoit-Wolf covariance matrix. + + ## Options + + #{NimbleOptions.docs(@opts_schema)} + + ## Return Values + + The function returns a struct with the following parameters: + + * `:covariance` - Tensor of shape `{num_features, num_features}`. Estimated covariance matrix. + + * `:shrinkage` - Coefficient in the convex combination used for the computation of the shrunken estimate. Range is `[0, 1]`. + + * `:location` - Tensor of shape `{num_features,}`. + Estimated location, i.e. the estimated mean. + + ## Examples + + iex> key = Nx.Random.key(0) + iex> {x, _new_key} = Nx.Random.multivariate_normal(key, Nx.tensor([0.0, 0.0]), Nx.tensor([[0.4, 0.2], [0.2, 0.8]]), shape: {50}, type: :f32) + iex> model = Scholar.Covariance.LedoitWolf.fit(x) + iex> model.covariance + #Nx.Tensor< + f32[2][2] + [ + [0.3557686507701874, 0.17340737581253052], + [0.17340737581253052, 1.0300586223602295] + ] + > + iex> model.shrinkage + #Nx.Tensor< + f32 + 0.15034137666225433 + > + iex> model.location + #Nx.Tensor< + f32[2] + [0.17184630036354065, 0.3276958167552948] + > + + iex> key = Nx.Random.key(0) + iex> {x, _new_key} = Nx.Random.multivariate_normal(key, Nx.tensor([0.0, 0.0, 0.0]), Nx.tensor([[3.0, 2.0, 1.0], [1.0, 2.0, 3.0], [1.3, 1.0, 2.2]]), shape: {10}, type: :f32) + iex> model = Scholar.Covariance.LedoitWolf.fit(x) + iex> model.covariance + #Nx.Tensor< + f32[3][3] + [ + [2.5945029258728027, 1.5078359842300415, 1.1623677015304565], + [1.5078359842300415, 2.106797456741333, 1.1812156438827515], + [1.1623677015304565, 1.1812156438827515, 1.4606266021728516] + ] + > + iex> model.shrinkage + #Nx.Tensor< + f32 + 0.1908363401889801 + > + iex> model.location + #Nx.Tensor< + f32[3] + [1.1228725910186768, 0.5419300198554993, 0.8678852319717407] + > + + iex> key = Nx.Random.key(0) + iex> {x, _new_key} = Nx.Random.multivariate_normal(key, Nx.tensor([0.0, 0.0, 0.0]), Nx.tensor([[3.0, 2.0, 1.0], [1.0, 2.0, 3.0], [1.3, 1.0, 2.2]]), shape: {10}, type: :f32) + iex> cov = Scholar.Covariance.LedoitWolf.fit(x, assume_centered: true) + iex> cov.covariance + #Nx.Tensor< + f32[3][3] + [ + [3.8574986457824707, 2.2048025131225586, 2.1504499912261963], + [2.2048025131225586, 2.4572863578796387, 1.7215262651443481], + [2.1504499912261963, 1.7215262651443481, 2.154898166656494] + ] + > + """ + + deftransform fit(x, opts \\ []) do + fit_n(x, NimbleOptions.validate!(opts, @opts_schema)) + end + + defnp fit_n(x, opts) do + {x, location} = center(x, opts) + + {covariance, shrinkage} = + ledoit_wolf(x) + + %__MODULE__{ + covariance: covariance, + shrinkage: shrinkage, + location: location + } + end + + defnp center(x, opts) do + x = + case Nx.shape(x) do + {_} -> Nx.new_axis(x, 1) + _ -> x + end + + location = + if opts[:assume_centered] do + 0 + else + Nx.mean(x, axes: [0]) + end + + {x - location, location} + end + + defnp ledoit_wolf(x) do + case Nx.shape(x) do + {_n, 1} -> + {Nx.mean(x ** 2) |> Nx.reshape({1, 1}), 0.0} + + _ -> + ledoit_wolf_shrinkage(x) + end + end + + defnp empirical_covariance(x) do + n = Nx.axis_size(x, 0) + + covariance = Nx.dot(x, [0], x, [0]) / n + + case Nx.shape(covariance) do + {} -> Nx.reshape(covariance, {1, 1}) + _ -> covariance + end + end + + defnp trace(x) do + x + |> Nx.take_diagonal() + |> Nx.sum() + end + + defnp ledoit_wolf_shrinkage(x) do + case Nx.shape(x) do + {_, 1} -> + 0 + + {n} -> + Nx.reshape(x, {1, n}) + |> ledoit_wolf_shrinkage_complex() + + _ -> + ledoit_wolf_shrinkage_complex(x) + end + end + + defnp ledoit_wolf_shrinkage_complex(x) do + {num_samples, num_features} = Nx.shape(x) + emp_cov = empirical_covariance(x) + + emp_cov_trace = trace(emp_cov) + mu = Nx.sum(emp_cov_trace) / num_features + + flatten_delta = Nx.flatten(emp_cov) + + indices = + Nx.shape(flatten_delta) + |> Nx.iota() + + subtract = Nx.select(Nx.remainder(indices, num_features + 1) == 0, mu, 0) + + delta = + (flatten_delta - subtract) + |> Nx.pow(2) + |> Nx.sum() + + delta = delta / num_features + + x2 = Nx.pow(x, 2) + + beta = + (Nx.dot(x2, [0], x2, [0]) / num_samples - emp_cov ** 2) + |> Nx.sum() + |> Nx.divide(num_features * num_samples) + + beta = Nx.min(beta, delta) + shrinkage = beta / delta + + shrunk_cov = (1.0 - shrinkage) * emp_cov + mask = Nx.iota(Nx.shape(shrunk_cov)) + selector = Nx.remainder(mask, num_features + 1) == 0 + shrunk_cov = Nx.select(selector, shrunk_cov + shrinkage * mu, shrunk_cov) + + {shrunk_cov, shrinkage} + end +end diff --git a/test/scholar/covariance/ledoit_wolf_test.exs b/test/scholar/covariance/ledoit_wolf_test.exs new file mode 100644 index 00000000..27e1893c --- /dev/null +++ b/test/scholar/covariance/ledoit_wolf_test.exs @@ -0,0 +1,127 @@ +defmodule Scholar.Covariance.LedoitWolfTest do + use Scholar.Case, async: true + alias Scholar.Covariance.LedoitWolf + doctest LedoitWolf + + defp key do + Nx.Random.key(1) + end + + test "fit test - all default options" do + key = key() + + {x, _new_key} = + Nx.Random.multivariate_normal( + key, + Nx.tensor([0.0, 0.0, 0.0]), + Nx.tensor([[3.0, 2.0, 1.0], [1.0, 2.0, 3.0], [1.3, 1.0, 2.2]]), + shape: {10}, + type: :f32 + ) + + model = LedoitWolf.fit(x) + + assert_all_close( + model.covariance, + Nx.tensor([ + [1.439786434173584, -0.0, 0.0], + [-0.0, 1.439786434173584, 0.0], + [0.0, 0.0, 1.439786434173584] + ]), + atol: 1.0e-3 + ) + + assert_all_close(model.shrinkage, Nx.tensor(1.0), atol: 1.0e-3) + + assert_all_close( + model.location, + Nx.tensor([-1.015519142150879, -0.4495307505130768, 0.06475571542978287]), + atol: 1.0e-3 + ) + end + + test "fit test - :assume_centered is true" do + key = key() + + {x, _new_key} = + Nx.Random.multivariate_normal( + key, + Nx.tensor([0.0, 0.0, 0.0]), + Nx.tensor([[3.0, 2.0, 1.0], [1.0, 2.0, 3.0], [1.3, 1.0, 2.2]]), + shape: {10}, + type: :f32 + ) + + model = LedoitWolf.fit(x, assume_centered: true) + + assert_all_close( + model.covariance, + Nx.tensor([ + [1.852303147315979, 0.0, 0.0], + [0.0, 1.852303147315979, 0.0], + [0.0, 0.0, 1.852303147315979] + ]), + atol: 1.0e-3 + ) + + assert_all_close(model.shrinkage, Nx.tensor(1.0), atol: 1.0e-3) + + assert_all_close(model.location, Nx.tensor([0, 0, 0]), atol: 1.0e-3) + end + + test "fit test 2" do + key = key() + + {x, _new_key} = + Nx.Random.multivariate_normal( + key, + Nx.tensor([0.0, 0.0]), + Nx.tensor([[2.2, 1.5], [0.7, 1.1]]), + shape: {50}, + type: :f32 + ) + + model = LedoitWolf.fit(x) + + assert_all_close( + model.covariance, + Nx.tensor([ + [1.8378269672393799, 0.27215731143951416], + [0.27215731143951416, 1.2268550395965576] + ]), + atol: 1.0e-3 + ) + + assert_all_close(model.shrinkage, Nx.tensor(0.38731059432029724), atol: 1.0e-3) + + assert_all_close(model.location, Nx.tensor([0.06882287561893463, 0.13750331103801727]), + atol: 1.0e-3 + ) + end + + test "fit test - 1 dim x" do + key = key() + + {x, _new_key} = + Nx.Random.multivariate_normal(key, Nx.tensor([0.0]), Nx.tensor([[0.4]]), + shape: {15}, + type: :f32 + ) + + x = Nx.flatten(x) + + model = LedoitWolf.fit(x) + + assert_all_close( + model.covariance, + Nx.tensor([ + [0.5322133302688599] + ]), + atol: 1.0e-3 + ) + + assert_all_close(model.shrinkage, Nx.tensor(0.0), atol: 1.0e-3) + + assert_all_close(model.location, Nx.tensor([0.060818854719400406]), atol: 1.0e-3) + end +end