diff --git a/lib/scholar/covariance/ledoit_wolf.ex b/lib/scholar/covariance/ledoit_wolf.ex index 3eb19994..086c79bc 100644 --- a/lib/scholar/covariance/ledoit_wolf.ex +++ b/lib/scholar/covariance/ledoit_wolf.ex @@ -13,11 +13,6 @@ defmodule Scholar.Covariance.LedoitWolf do defstruct [:covariance, :shrinkage, :location] opts_schema = [ - block_size: [ - default: 1000, - type: {:custom, Scholar.Options, :positive_number, []}, - doc: "Size of blocks into which the covariance matrix will be split." - ], assume_centered: [ default: false, type: :boolean, @@ -42,11 +37,11 @@ defmodule Scholar.Covariance.LedoitWolf do The function returns a struct with the following parameters: - * `:covariance` - Tensor of shape `{n_features, n_features}`. Estimated covariance matrix. + * `: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 `{n_features,}`. + * `:location` - Tensor of shape `{num_features,}`. Estimated location, i.e. the estimated mean. ## Examples @@ -58,14 +53,14 @@ defmodule Scholar.Covariance.LedoitWolf do #Nx.Tensor< f32[2][2] [ - [0.355768620967865, 0.17340737581253052], + [0.3557686507701874, 0.17340737581253052], [0.17340737581253052, 1.0300586223602295] ] > iex> model.shrinkage #Nx.Tensor< f32 - 0.15034136176109314 + 0.15034137666225433 > iex> model.location #Nx.Tensor< @@ -80,15 +75,15 @@ defmodule Scholar.Covariance.LedoitWolf do #Nx.Tensor< f32[3][3] [ - [2.5945029258728027, 1.507835865020752, 1.1623677015304565], - [1.507835865020752, 2.106797218322754, 1.181215524673462], - [1.1623677015304565, 1.181215524673462, 1.460626482963562] + [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.1908363550901413 + 0.1908363401889801 > iex> model.location #Nx.Tensor< @@ -118,7 +113,7 @@ defmodule Scholar.Covariance.LedoitWolf do {x, location} = center(x, opts) {covariance, shrinkage} = - ledoit_wolf(x, opts) + ledoit_wolf(x) %__MODULE__{ covariance: covariance, @@ -140,32 +135,21 @@ defmodule Scholar.Covariance.LedoitWolf do else Nx.mean(x, axes: [0]) end + {x - location, location} end - defnp ledoit_wolf(x, opts) do + defnp ledoit_wolf(x) do case Nx.shape(x) do {_n, 1} -> {Nx.mean(x ** 2) |> Nx.reshape({1, 1}), 0.0} _ -> - ledoit_wolf_complex(x, opts) + ledoit_wolf_shrinkage(x) end end - defnp ledoit_wolf_complex(x, opts) do - n_features = Nx.axis_size(x, 1) - shrinkage = ledoit_wolf_shrinkage(x, opts) - emp_cov = empirical_covariance(x, opts) - mu = Nx.sum(trace(emp_cov)) / n_features - shrunk_cov = (1.0 - shrinkage) * emp_cov - mask = Nx.iota(Nx.shape(shrunk_cov)) - selector = Nx.remainder(mask, n_features + 1) == 0 - shrunk_cov = Nx.select(selector, shrunk_cov + shrinkage * mu, shrunk_cov) - {shrunk_cov, shrinkage} - end - - defnp empirical_covariance(x, _opts) do + defnp empirical_covariance(x) do n = Nx.axis_size(x, 0) covariance = Nx.dot(x, [0], x, [0]) / n @@ -182,129 +166,57 @@ defmodule Scholar.Covariance.LedoitWolf do |> Nx.sum() end - defnp ledoit_wolf_shrinkage(x, opts) do + defnp ledoit_wolf_shrinkage(x) do case Nx.shape(x) do {_, 1} -> 0 {n} -> Nx.reshape(x, {1, n}) - |> ledoit_wolf_shrinkage_complex(opts) + |> ledoit_wolf_shrinkage_complex() _ -> - ledoit_wolf_shrinkage_complex(x, opts) + ledoit_wolf_shrinkage_complex(x) end end - defnp ledoit_wolf_shrinkage_complex(x, opts) do - {n_samples, n_features} = Nx.shape(x) - {_, size} = Nx.type(x) - block_size = opts[:block_size] - n_splits = (n_features / block_size) |> Nx.floor() |> Nx.as_type({:s, size}) + defnp ledoit_wolf_shrinkage_complex(x) do + {num_samples, num_features} = Nx.shape(x) + emp_cov = empirical_covariance(x) - x2 = Nx.pow(x, 2) - emp_cov_trace = Nx.sum(x2, axes: [0]) / n_samples - mu = Nx.sum(emp_cov_trace) / n_features - beta = Nx.tensor(0.0, type: {:f, size}) - i = Nx.tensor(0) - block = Nx.iota({block_size}) - - beta = if n_splits > 0 do - {beta, _} = - while {beta, {x, x2, block, i, n_splits}}, i < n_splits do - {block_size} = Nx.shape(block) - j = Nx.tensor(0) - - {beta, _} = - while {beta, {x, x2, block, i, j, n_splits}}, j < n_splits do - {block_size} = Nx.shape(block) - rows_from = block_size * i - cols_from = block_size * j - - to_beta = - Nx.slice_along_axis(x2, rows_from, block_size, axis: 1) - |> Nx.dot([0], Nx.slice_along_axis(x2, cols_from, block_size, axis: 1), [0]) - |> Nx.sum() - - beta = beta + to_beta - - {beta, {x, x2, block, i, j + 1, n_splits}} - end - - rows_from = block_size * i - x2_t = Nx.transpose(x2) - {m, n} = Nx.shape(x2) - mask = Nx.iota({1, n}) |> Nx.tile([m]) |> Nx.reshape({m, n}) - mask = Nx.select(mask >= block_size * n_splits, 1, 0) - - to_beta = - Nx.slice_along_axis(x2_t, rows_from, block_size, axis: 0) - |> Nx.dot(Nx.multiply(x2, mask)) - |> Nx.sum() - - beta = beta + to_beta - - {beta, {x, x2, block, i + 1, n_splits}} - end - - j = Nx.tensor(0) - - {beta, _} = - while {beta, {x, x2, block, j, n_splits}}, j < n_splits do - {block_size} = Nx.shape(block) - cols_from = block_size * j - x2_t = Nx.transpose(x2) - {rows, cols} = Nx.shape(x) - - mask = - Nx.iota({1, cols}) |> Nx.tile([rows]) |> Nx.reshape({rows, cols}) |> Nx.transpose() - - mask = Nx.select(mask >= block_size * n_splits, 1, 0) - - to_beta = - Nx.multiply(x2_t, mask) - |> Nx.dot(Nx.slice_along_axis(x2, cols_from, block_size, axis: 1)) - |> Nx.sum() - - beta = beta + to_beta - - {beta, {x, x2, block, j + 1, n_splits}} - end - beta - else - beta - end - delta = beta - x2_t = Nx.transpose(x2) - x_t = Nx.transpose(x) - {rows, cols} = Nx.shape(x) - mask = Nx.iota({1, cols}) |> Nx.tile([rows]) |> Nx.reshape({rows, cols}) - mask = Nx.select(mask >= block_size * n_splits, 1, 0) - mask_t = Nx.transpose(mask) - - to_delta = - Nx.multiply(x_t, mask_t) - |> Nx.dot(Nx.multiply(x, mask)) + 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 + to_delta - delta = delta / Nx.pow(n_samples, 2) + delta = delta / num_features - to_beta = - Nx.multiply(x2_t, mask_t) - |> Nx.dot(Nx.multiply(x2, mask)) + 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 = beta + to_beta - beta = 1.0 / (n_features * n_samples) * (beta / n_samples - delta) - delta = delta - 2.0 * mu * Nx.sum(emp_cov_trace) + n_features * Nx.pow(mu, 2) - delta = delta / n_features beta = Nx.min(beta, delta) + shrinkage = beta / delta - case beta do - 0 -> 0 - _ -> beta / delta - end + 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..bcaa187d --- /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 - set :block_size" 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, block_size: 20) + + 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