diff --git a/lib/scholar/cluster/optics.ex b/lib/scholar/cluster/optics.ex index 4f62664c..7a43b8f8 100644 --- a/lib/scholar/cluster/optics.ex +++ b/lib/scholar/cluster/optics.ex @@ -10,8 +10,29 @@ defmodule Scholar.Cluster.OPTICS do """ import Nx.Defn - @derive {Nx.Container, containers: [:labels, :min_samples, :max_eps, :eps, :algorithm]} - defstruct [:labels, :min_samples, :max_eps, :eps, :algorithm] + @derive {Nx.Container, + containers: [ + :labels, + :min_samples, + :max_eps, + :eps, + :algorithm, + :reachability, + :core_distances, + :ordering, + :predecessor + ]} + defstruct [ + :labels, + :min_samples, + :max_eps, + :eps, + :algorithm, + :reachability, + :core_distances, + :ordering, + :predecessor + ] opts = [ min_samples: [ @@ -65,9 +86,26 @@ defmodule Scholar.Cluster.OPTICS do ## Return Values - The function returns a labels tensor of shape `{n_samples}`. - Cluster labels for each point in the dataset given to `fit`. - Noisy samples are labeled as -1. + The function returns a struct with the following fields: + + * `:labels` - tensor of shape `{n_samples}`. Cluster labels for each + point in the dataset given to `fit`. Noisy samples are labeled as -1. + + * `:reachability` - tensor of shape `{n_samples}`, indexed by sample. + Reachability distance of each sample. Use `reachability[ordering]` + to read it in cluster order. + + * `:core_distances` - tensor of shape `{n_samples}`, indexed by sample. + Distance at which each sample becomes a core point. A sample that + never becomes core has a distance of infinity. + + * `:ordering` - tensor of shape `{n_samples}`. The cluster-ordered + list of sample indices, which is what a reachability plot is drawn + against. + + * `:predecessor` - tensor of shape `{n_samples}`, indexed by sample. + Sample that each sample was reached from. Seed samples have a + predecessor of -1. ## Examples @@ -115,10 +153,15 @@ defmodule Scholar.Cluster.OPTICS do x = Scholar.Shared.to_float(x) module = validate_options(x, opts) + {labels, reachability, core_distances, ordering, predecessor} = fit_p(x, module) %__MODULE__{ module - | labels: fit_p(x, module) + | labels: labels, + reachability: reachability, + core_distances: core_distances, + ordering: ordering, + predecessor: predecessor } end @@ -172,19 +215,28 @@ defmodule Scholar.Cluster.OPTICS do """ end + n_samples = Nx.axis_size(x, 0) + %__MODULE__{ - labels: Nx.broadcast(-1, {Nx.axis_size(x, 0)}), + labels: Nx.broadcast(-1, {n_samples}), min_samples: min_samples, max_eps: max_eps, eps: eps, - algorithm: model + algorithm: model, + reachability: Nx.broadcast(Nx.Constants.infinity(Nx.type(x)), {n_samples}), + core_distances: Nx.broadcast(Nx.Constants.infinity(Nx.type(x)), {n_samples}), + ordering: Nx.broadcast(0, {n_samples}), + predecessor: Nx.broadcast(-1, {n_samples}) } end defnp fit_p(x, module) do - {core_distances, reachability, _predecessor, ordering} = compute_optics_graph(x, module) + {core_distances, reachability, predecessor, ordering} = compute_optics_graph(x, module) + labels = cluster_optics_dbscan(reachability, core_distances, ordering, module) - cluster_optics_dbscan(reachability, core_distances, ordering, module) + # core_distances carries a trailing size-1 axis from the neighbor search + # it is sliced out of; every other field here is already flat. + {labels, reachability, Nx.flatten(core_distances), ordering, predecessor} end defnp compute_optics_graph(x, %__MODULE__{max_eps: max_eps, min_samples: min_samples} = module) do diff --git a/test/scholar/cluster/optics_test.exs b/test/scholar/cluster/optics_test.exs new file mode 100644 index 00000000..89392e24 --- /dev/null +++ b/test/scholar/cluster/optics_test.exs @@ -0,0 +1,42 @@ +defmodule Scholar.Cluster.OPTICSTest do + use Scholar.Case, async: true + + alias Scholar.Cluster.OPTICS + doctest OPTICS + + # Expected values from scikit-learn 1.6.1 on this data. + defp x, do: Nx.tensor([[1, 2], [2, 5], [3, 6], [8, 7], [8, 8], [7, 3]], type: :f64) + + test "fit/2 returns the reachability graph" do + model = OPTICS.fit(x(), eps: 4.5, min_samples: 2) + + assert model.ordering == Nx.tensor([0, 1, 2, 5, 3, 4], type: :s32) + assert model.predecessor == Nx.tensor([-1, 0, 1, 5, 3, 2], type: :s32) + + assert_all_close( + model.core_distances, + Nx.tensor([3.16227766, 1.41421356, 1.41421356, 1.0, 1.0, 4.12310563], type: :f64) + ) + + assert_all_close( + model.reachability, + Nx.tensor([:infinity, 3.16227766, 1.41421356, 4.12310563, 1.0, 5.0], type: :f64) + ) + end + + test "fit/2 caps the reachability graph at max_eps" do + model = OPTICS.fit(x(), max_eps: 2, min_samples: 2) + + assert model.labels == Nx.tensor([-1, 0, 0, 1, 1, -1], type: :s32) + + assert_all_close( + model.core_distances, + Nx.tensor([:infinity, 1.41421356, 1.41421356, 1.0, 1.0, :infinity], type: :f64) + ) + + assert_all_close( + model.reachability, + Nx.tensor([:infinity, :infinity, 1.41421356, :infinity, 1.0, :infinity], type: :f64) + ) + end +end