From 33885246da84009b80cf1a513dbbda63cf43fba4 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Tue, 1 Sep 2026 15:38:20 -0300 Subject: [PATCH 1/4] Add mean shift clustering Mean shift moves every seed towards the mean of the samples within one bandwidth of it until the seed settles, then keeps the strongest of the seeds that landed on the same mode. The number of clusters comes out of the data rather than being given. Every seed is moved at once rather than one at a time, so an iteration is a single pairwise matrix instead of one per seed. That costs O(seeds * samples) of space, the same shape DBSCAN already pays through radius_neighbors. fit/2 keeps one row per seed so the shapes stay static, marking the centers that lost with :infinity and reporting how many survived, and prune/1 drops them. This follows Scholar.Cluster.AffinityPropagation, which has the same problem of not knowing the cluster count until it has run. Validated against scikit-learn 1.6.1 on 131 datasets: 11 hand-written edge cases and 120 random ones spanning 2 to 25 samples, 1 to 4 features and 1 to 21 clusters. Labels, centers and cluster counts match on all of them, including the order scikit-learn puts the centers in, which breaks ties on the coordinates. --- lib/scholar/cluster/mean_shift.ex | 298 +++++++++++++++++++++++ mix.exs | 1 + test/scholar/cluster/mean_shift_test.exs | 207 ++++++++++++++++ 3 files changed, 506 insertions(+) create mode 100644 lib/scholar/cluster/mean_shift.ex create mode 100644 test/scholar/cluster/mean_shift_test.exs diff --git a/lib/scholar/cluster/mean_shift.ex b/lib/scholar/cluster/mean_shift.ex new file mode 100644 index 00000000..6512512c --- /dev/null +++ b/lib/scholar/cluster/mean_shift.ex @@ -0,0 +1,298 @@ +defmodule Scholar.Cluster.MeanShift do + @moduledoc """ + Mean shift clustering. + + Mean shift moves every seed towards the mean of the samples inside a ball of + radius `:bandwidth` around it, repeating until the seed stops moving. Seeds + that settle within one bandwidth of each other describe the same mode, so the + weaker ones are discarded and the survivors become the cluster centers. + + The number of clusters is discovered from the data rather than given, and the + bandwidth is what controls it. + + Distances are Euclidean. The update step averages the samples in a + neighborhood, which only has meaning in a space where the arithmetic mean is + the centroid. + + The time complexity is $O(I * S * N)$ for $N$ samples, $S$ seeds and $I$ + iterations. The space complexity is $O(S * N)$. + """ + import Nx.Defn + import Scholar.Shared + + @derive {Nx.Container, containers: [:cluster_centers, :labels, :num_clusters, :iterations]} + defstruct [:cluster_centers, :labels, :num_clusters, :iterations] + + opts = [ + bandwidth: [ + required: true, + type: {:custom, Scholar.Options, :positive_number, []}, + doc: """ + The radius of the region a seed averages over. Larger values merge more + modes together and yield fewer clusters. + """ + ], + max_iterations: [ + default: 300, + type: :pos_integer, + doc: """ + The maximum number of times every seed is moved. Seeds stop early once + none of them moves further than `bandwidth / 1000`. + """ + ], + cluster_all: [ + default: true, + type: :boolean, + doc: """ + If `true`, every sample is assigned to the nearest cluster center. If + `false`, samples further than `:bandwidth` from every center are labeled + `-1`. + """ + ], + seeds: [ + doc: """ + The points to start from, given as a tensor of shape `{num_seeds, + num_features}`. Defaults to the samples themselves. Fewer seeds make the + fit cheaper at the risk of missing a mode. + """ + ] + ] + + @opts_schema NimbleOptions.new!(opts) + + @doc """ + Fits a mean shift model for sample inputs `x`. + + ## Options + + #{NimbleOptions.docs(@opts_schema)} + + ## Return Values + + The function returns a struct with the following parameters: + + * `:cluster_centers` - The point every seed settled on, ordered by how many + samples it gathered. Rows that lost to a stronger center are set to + `:infinity`, so the tensor keeps the shape of the seeds. Use `prune/1` to + drop them. + + * `:labels` - The row of `:cluster_centers` each sample belongs to, or `-1` + when `:cluster_all` is `false` and the sample is further than + `:bandwidth` from every center. + + * `:num_clusters` - The number of centers that survived. + + * `:iterations` - How many times the seeds were moved before they settled. + + ## Examples + + iex> x = Nx.tensor([[1, 1], [1.2, 1.1], [8, 8], [8.1, 8.2]]) + iex> Scholar.Cluster.MeanShift.fit(x, bandwidth: 1.0) + %Scholar.Cluster.MeanShift{ + cluster_centers: Nx.f32( + [ + [8.050000190734863, 8.100000381469727], + [:infinity, :infinity], + [1.100000023841858, 1.0499999523162842], + [:infinity, :infinity] + ] + ), + labels: Nx.s32([2, 2, 0, 0]), + num_clusters: Nx.u32(2), + iterations: Nx.u32(2) + } + """ + deftransform fit(x, opts \\ []) do + opts = NimbleOptions.validate!(opts, @opts_schema) + {seeds, opts} = Keyword.pop(opts, :seeds, {}) + + if Nx.rank(x) != 2 do + raise ArgumentError, + "expected x to have shape {num_samples, num_features}, got: #{inspect(Nx.shape(x))}" + end + + if Nx.rank(seeds) != 0 and + (Nx.rank(seeds) != 2 or Nx.axis_size(seeds, 1) != Nx.axis_size(x, 1)) do + raise ArgumentError, + "expected seeds to have shape {num_seeds, #{Nx.axis_size(x, 1)}}, " <> + "got: #{inspect(Nx.shape(seeds))}" + end + + fit_n(x, seeds, opts) + end + + defnp fit_n(x, seeds, opts) do + x = to_float(x) + + seeds = + case seeds do + {} -> x + _ -> to_float(seeds) + end + + bandwidth = opts[:bandwidth] + {centers, iterations} = shift_seeds(x, seeds, bandwidth, opts) + + intensity = intensity(x, centers, bandwidth) + order = strongest_first(intensity, centers) + centers = Nx.take(centers, order) + + # a seed that never had a sample in reach describes no mode at all + kept = drop_duplicate_modes(centers, bandwidth) and Nx.take(intensity, order) > 0 + num_clusters = Nx.sum(kept) + + distances = + Nx.select( + Nx.broadcast(Nx.new_axis(kept, 0), {Nx.axis_size(x, 0), Nx.axis_size(centers, 0)}), + Scholar.Metrics.Distance.pairwise_euclidean(x, centers), + Nx.Constants.infinity(to_float_type(x)) + ) + + labels = Nx.argmin(distances, axis: 1) |> Nx.as_type(:s32) + + labels = + if opts[:cluster_all] do + labels + else + Nx.select(Nx.reduce_min(distances, axes: [1]) <= bandwidth, labels, -1) + end + + labels = Nx.select(num_clusters > 0, labels, -1) + + %__MODULE__{ + cluster_centers: + Nx.select( + Nx.broadcast(Nx.new_axis(kept, 1), Nx.shape(centers)), + centers, + Nx.Constants.infinity(to_float_type(x)) + ), + labels: labels, + num_clusters: Nx.as_type(num_clusters, :u32), + iterations: iterations + } + end + + # every seed moves at once, which is one pairwise matrix per iteration rather + # than one per seed + defnp shift_seeds(x, seeds, bandwidth, opts) do + tolerance = bandwidth / 1000 + + {seeds, _, iterations} = + while {seeds, {x, bandwidth, tolerance, moving = Nx.u8(1)}, i = Nx.u32(0)}, + i < opts[:max_iterations] and moving do + within = Scholar.Metrics.Distance.pairwise_euclidean(seeds, x) <= bandwidth + members = Nx.as_type(within, Nx.type(seeds)) + count = Nx.sum(members, axes: [1], keep_axes: true) + + # a seed with an empty neighborhood would divide by zero, so it stays put + moved = + Nx.select( + Nx.broadcast(count > 0, Nx.shape(seeds)), + Nx.dot(members, x) / Nx.select(count > 0, count, 1), + seeds + ) + + shift = Scholar.Metrics.Distance.euclidean(moved, seeds, axes: [1]) + {moved, {x, bandwidth, tolerance, Nx.any(shift > tolerance)}, i + 1} + end + + {seeds, iterations} + end + + defnp intensity(x, centers, bandwidth) do + Scholar.Metrics.Distance.pairwise_euclidean(centers, x) + |> Nx.less_equal(bandwidth) + |> Nx.sum(axes: [1]) + end + + # ties break on the coordinates, descending, so that seeds landing on the same + # mode keep a stable order + deftransformp strongest_first(intensity, centers) do + keys = + Enum.map((Nx.axis_size(centers, 1) - 1)..0//-1, ¢ers[[.., &1]]) ++ [intensity] + + Enum.reduce(keys, Nx.iota({Nx.axis_size(centers, 0)}, type: :s32), fn key, order -> + Nx.take(order, Nx.argsort(Nx.take(key, order), direction: :desc, stable: true)) + end) + end + + defnp drop_duplicate_modes(centers, bandwidth) do + num_centers = Nx.axis_size(centers, 0) + + {kept, _} = + while {kept = Nx.broadcast(Nx.u8(1), {num_centers}), {centers, bandwidth, i = Nx.u32(0)}}, + i < num_centers do + kept = + if kept[i] do + near = + Scholar.Metrics.Distance.euclidean(centers, Nx.new_axis(centers[i], 0), axes: [1]) <= + bandwidth + + Nx.indexed_put(Nx.select(near, Nx.u8(0), kept), Nx.new_axis(i, 0), Nx.u8(1)) + else + kept + end + + {kept, {centers, bandwidth, i + 1}} + end + + kept + end + + @doc """ + Drops the centers that lost to a stronger one and renumbers the labels. + + `fit/2` keeps one row per seed so the shapes stay static. This returns a model + holding only the `:num_clusters` centers that survived, with `:labels` + renumbered to index them. + + ## Examples + + iex> x = Nx.tensor([[1, 1], [1.2, 1.1], [8, 8], [8.1, 8.2]]) + iex> model = Scholar.Cluster.MeanShift.fit(x, bandwidth: 1.0) + iex> Scholar.Cluster.MeanShift.prune(model) + %Scholar.Cluster.MeanShift{ + cluster_centers: Nx.f32( + [ + [8.050000190734863, 8.100000381469727], + [1.100000023841858, 1.0499999523162842] + ] + ), + labels: Nx.s32([1, 1, 0, 0]), + num_clusters: Nx.u32(2), + iterations: Nx.u32(2) + } + """ + def prune(%__MODULE__{cluster_centers: centers, labels: labels} = model) do + if Nx.to_number(model.num_clusters) == 0 do + raise ArgumentError, + "the model has no clusters to keep, every seed was further than the bandwidth " <> + "from all samples" + end + + kept = + centers + |> Nx.is_infinity() + |> Nx.all(axes: [1]) + |> Nx.equal(0) + + indices = kept |> Nx.to_flat_list() |> Enum.with_index() |> Enum.filter(&(elem(&1, 0) == 1)) + indices = Enum.map(indices, &elem(&1, 1)) + + renumber = + indices + |> Enum.with_index() + |> Map.new() + |> then(fn mapping -> fn old -> Map.get(mapping, old, -1) end end) + + %__MODULE__{ + model + | cluster_centers: Nx.take(centers, Nx.tensor(indices)), + labels: + labels + |> Nx.to_flat_list() + |> Enum.map(renumber) + |> Nx.tensor(type: Nx.type(labels)) + } + end +end diff --git a/mix.exs b/mix.exs index 8dd9f5ee..1f9fc7ef 100644 --- a/mix.exs +++ b/mix.exs @@ -72,6 +72,7 @@ defmodule Scholar.MixProject do Scholar.Cluster.GaussianMixture, Scholar.Cluster.Hierarchical, Scholar.Cluster.KMeans, + Scholar.Cluster.MeanShift, Scholar.Cluster.SpectralClustering, Scholar.Decomposition.KernelPCA, Scholar.Decomposition.PCA, diff --git a/test/scholar/cluster/mean_shift_test.exs b/test/scholar/cluster/mean_shift_test.exs new file mode 100644 index 00000000..c1601b3a --- /dev/null +++ b/test/scholar/cluster/mean_shift_test.exs @@ -0,0 +1,207 @@ +defmodule Scholar.Cluster.MeanShiftTest do + use Scholar.Case, async: true + alias Scholar.Cluster.MeanShift + doctest MeanShift + + defp blobs do + Nx.tensor([ + [1.0, 1.0], + [1.2, 1.1], + [0.9, 1.05], + [8.0, 8.0], + [8.1, 8.2], + [7.9, 7.8], + [1.0, 8.0], + [1.1, 8.1] + ]) + end + + describe "fit" do + test "fit - all defaults" do + # Expected values from scikit-learn 1.6.1 on this data. + model = MeanShift.fit(blobs(), bandwidth: 1.0) |> MeanShift.prune() + + assert model.num_clusters == Nx.u32(3) + assert model.labels == Nx.s32([1, 1, 1, 0, 0, 0, 2, 2]) + + assert_all_close( + model.cluster_centers, + Nx.tensor([[8.0, 8.0], [1.0333334, 1.05], [1.05, 8.05]]) + ) + end + + test "fit with a bandwidth wide enough to hold every sample" do + # Expected values from scikit-learn 1.6.1 on this data. + model = MeanShift.fit(blobs(), bandwidth: 20.0) |> MeanShift.prune() + + assert model.num_clusters == Nx.u32(1) + assert model.labels == Nx.s32([0, 0, 0, 0, 0, 0, 0, 0]) + assert_all_close(model.cluster_centers, Nx.tensor([[3.65, 5.40625]])) + end + + test "fit with a bandwidth narrower than the closest pair" do + model = MeanShift.fit(blobs(), bandwidth: 0.05) |> MeanShift.prune() + + # every sample keeps to itself, so the centers are the samples reordered + assert model.num_clusters == Nx.u32(8) + assert model.labels == Nx.s32([6, 3, 7, 1, 0, 2, 5, 4]) + end + + test "fit with cluster_all disabled" do + # Expected values from scikit-learn 1.6.1 on this data. + x = + Nx.tensor([ + [2.2, -2.9], + [5.4, -4.8], + [1.0, -1.7], + [3.6, -2.9], + [1.2, -0.4], + [5.1, 1.9] + ]) + + clustered = MeanShift.fit(x, bandwidth: 2.0, cluster_all: true) |> MeanShift.prune() + assert clustered.labels == Nx.s32([0, 1, 0, 0, 0, 2]) + + # the fifth sample drifted into the first mode but stayed further than the + # bandwidth from where that mode settled + loose = MeanShift.fit(x, bandwidth: 2.0, cluster_all: false) |> MeanShift.prune() + assert loose.labels == Nx.s32([0, 1, 0, 0, -1, 2]) + + assert_all_close( + loose.cluster_centers, + Nx.tensor([[2.2666667, -2.5], [5.4, -4.8], [5.1, 1.9]]) + ) + end + + test "fit with samples that are all the same point" do + # Expected values from scikit-learn 1.6.1 on this data. + model = + MeanShift.fit(Nx.tensor([[3.0, 3.0], [3.0, 3.0], [3.0, 3.0]]), bandwidth: 1.0) + |> MeanShift.prune() + + assert model.num_clusters == Nx.u32(1) + assert model.labels == Nx.s32([0, 0, 0]) + assert_all_close(model.cluster_centers, Nx.tensor([[3.0, 3.0]])) + end + + test "fit with a single sample" do + model = MeanShift.fit(Nx.tensor([[2.0, 5.0]]), bandwidth: 1.0) |> MeanShift.prune() + + assert model.num_clusters == Nx.u32(1) + assert model.labels == Nx.s32([0]) + assert_all_close(model.cluster_centers, Nx.tensor([[2.0, 5.0]])) + end + + test "fit with one feature" do + # Expected values from scikit-learn 1.6.1 on this data. + model = + MeanShift.fit(Nx.tensor([[1.0], [1.1], [5.0], [5.2]]), bandwidth: 0.5) + |> MeanShift.prune() + + assert model.labels == Nx.s32([1, 1, 0, 0]) + assert_all_close(model.cluster_centers, Nx.tensor([[5.1], [1.05]])) + end + + test "fit stops at max_iterations" do + {x, _} = Nx.Random.uniform(Nx.Random.key(7), shape: {40, 2}) + x = Nx.multiply(x, 10) + + truncated = MeanShift.fit(x, bandwidth: 2.5, max_iterations: 1) + settled = MeanShift.fit(x, bandwidth: 2.5) + + assert truncated.iterations == Nx.u32(1) + assert settled.iterations == Nx.u32(8) + + # seeds cut off early have not merged yet, so more of them survive + assert Nx.to_number(truncated.num_clusters) > Nx.to_number(settled.num_clusters) + end + + test "fit from a given set of seeds" do + x = Nx.tensor([[0.0, 0.0], [0.1, 0.0], [0.0, 0.1], [5.0, 5.0], [5.1, 5.0]]) + + model = + MeanShift.fit(x, bandwidth: 1.0, seeds: Nx.tensor([[0.0, 0.0], [5.0, 5.0]])) + |> MeanShift.prune() + + assert model.num_clusters == Nx.u32(2) + assert model.labels == Nx.s32([0, 0, 0, 1, 1]) + assert_all_close(model.cluster_centers, Nx.tensor([[0.0333333, 0.0333333], [5.05, 5.0]])) + end + + test "fit drops seeds that never reach a sample" do + # Expected values from scikit-learn 1.6.1 on this data. + x = Nx.tensor([[0.0, 0.0], [0.1, 0.0], [0.0, 0.1]]) + + model = + MeanShift.fit(x, bandwidth: 1.0, seeds: Nx.tensor([[0.0, 0.0], [99.0, 99.0]])) + |> MeanShift.prune() + + assert model.num_clusters == Nx.u32(1) + assert model.labels == Nx.s32([0, 0, 0]) + assert_all_close(model.cluster_centers, Nx.tensor([[0.0333333, 0.0333333]])) + end + + test "fit keeps the type of the input" do + f64 = MeanShift.fit(Nx.as_type(blobs(), :f64), bandwidth: 1.0) + assert Nx.type(f64.cluster_centers) == {:f, 64} + + f32 = MeanShift.fit(blobs(), bandwidth: 1.0) + assert Nx.type(f32.cluster_centers) == {:f, 32} + end + + test "works with jit_apply" do + direct = MeanShift.fit(blobs(), bandwidth: 1.0) + jitted = Nx.Defn.jit_apply(&MeanShift.fit/2, [blobs(), [bandwidth: 1.0]]) + + assert jitted.labels == direct.labels + assert jitted.num_clusters == direct.num_clusters + assert_all_close(jitted.cluster_centers, direct.cluster_centers) + end + end + + test "prune" do + model = MeanShift.fit(blobs(), bandwidth: 1.0) + + # fit keeps one row per seed so the shapes stay static + assert Nx.axis_size(model.cluster_centers, 0) == 8 + assert model.num_clusters == Nx.u32(3) + + pruned = MeanShift.prune(model) + assert Nx.axis_size(pruned.cluster_centers, 0) == 3 + + # the labels have to keep pointing at the same centers after renumbering + assert_all_close( + Nx.take(pruned.cluster_centers, pruned.labels), + Nx.take(model.cluster_centers, model.labels) + ) + end + + describe "errors" do + test "x that is not a matrix" do + assert_raise ArgumentError, + "expected x to have shape {num_samples, num_features}, got: {3}", + fn -> MeanShift.fit(Nx.tensor([1, 2, 3]), bandwidth: 1.0) end + end + + test "seeds that do not match the number of features" do + assert_raise ArgumentError, + "expected seeds to have shape {num_seeds, 2}, got: {1, 3}", + fn -> + MeanShift.fit(blobs(), bandwidth: 1.0, seeds: Nx.tensor([[1.0, 2.0, 3.0]])) + end + end + + test "prune with nothing left to keep" do + model = + MeanShift.fit(Nx.tensor([[0.0, 0.0]]), bandwidth: 1.0, seeds: Nx.tensor([[99.0, 99.0]])) + + assert model.num_clusters == Nx.u32(0) + assert model.labels == Nx.s32([-1]) + + assert_raise ArgumentError, + "the model has no clusters to keep, every seed was further than the " <> + "bandwidth from all samples", + fn -> MeanShift.prune(model) end + end + end +end From 457026a1a9684cf2a8eb78795b5c2c7e893e85bd Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Tue, 1 Sep 2026 15:57:30 -0300 Subject: [PATCH 2/4] Add HDBSCAN and OPTICS to the module groups Both were left out of groups_for_modules when they landed, so ex_doc filed them outside the Models group the rest of the clustering algorithms sit in. --- mix.exs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mix.exs b/mix.exs index 1f9fc7ef..3b4f7b53 100644 --- a/mix.exs +++ b/mix.exs @@ -70,9 +70,11 @@ defmodule Scholar.MixProject do Scholar.Cluster.AffinityPropagation, Scholar.Cluster.DBSCAN, Scholar.Cluster.GaussianMixture, + Scholar.Cluster.HDBSCAN, Scholar.Cluster.Hierarchical, Scholar.Cluster.KMeans, Scholar.Cluster.MeanShift, + Scholar.Cluster.OPTICS, Scholar.Cluster.SpectralClustering, Scholar.Decomposition.KernelPCA, Scholar.Decomposition.PCA, From b372133b157656f8dd609f5ca2bb47ad14b845f0 Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Tue, 1 Sep 2026 20:06:18 -0300 Subject: [PATCH 3/4] Fix mean shift widening its seeds inside the loop The seeds carry the while accumulator, so their type has to survive a pass of the loop. Samples of a wider type promoted the moved seeds through Nx.dot and the do-block then failed to match what it was given, which raised a CompileError for a f64 sample set with f32 seeds. Merge the two types up front. Also documents that :max_iterations counts moves, where scikit-learn's max_iter checks the limit after moving and so takes one step more than the number given. --- lib/scholar/cluster/mean_shift.ex | 19 +++++++++++---- test/scholar/cluster/mean_shift_test.exs | 31 +++++++++++++++++++----- 2 files changed, 39 insertions(+), 11 deletions(-) diff --git a/lib/scholar/cluster/mean_shift.ex b/lib/scholar/cluster/mean_shift.ex index 6512512c..4214dea4 100644 --- a/lib/scholar/cluster/mean_shift.ex +++ b/lib/scholar/cluster/mean_shift.ex @@ -37,7 +37,9 @@ defmodule Scholar.Cluster.MeanShift do type: :pos_integer, doc: """ The maximum number of times every seed is moved. Seeds stop early once - none of them moves further than `bandwidth / 1000`. + none of them moves further than `bandwidth / 1000`. Note that + scikit-learn's `max_iter` checks the limit after moving, so it takes one + more step than the number given. """ ], cluster_all: [ @@ -82,7 +84,8 @@ defmodule Scholar.Cluster.MeanShift do * `:num_clusters` - The number of centers that survived. - * `:iterations` - How many times the seeds were moved before they settled. + * `:iterations` - How many times the seeds were moved. This is at least one, + since the seeds have to move once for their movement to be measured. ## Examples @@ -130,15 +133,21 @@ defmodule Scholar.Cluster.MeanShift do _ -> to_float(seeds) end + # the seeds carry the loop's accumulator, so they and the samples have to + # agree on a type or the moved seeds come back wider than they went in + type = Nx.Type.merge(Nx.type(x), Nx.type(seeds)) + x = Nx.as_type(x, type) + seeds = Nx.as_type(seeds, type) + bandwidth = opts[:bandwidth] {centers, iterations} = shift_seeds(x, seeds, bandwidth, opts) - intensity = intensity(x, centers, bandwidth) - order = strongest_first(intensity, centers) + gathered = intensity(x, centers, bandwidth) + order = strongest_first(gathered, centers) centers = Nx.take(centers, order) # a seed that never had a sample in reach describes no mode at all - kept = drop_duplicate_modes(centers, bandwidth) and Nx.take(intensity, order) > 0 + kept = drop_duplicate_modes(centers, bandwidth) and Nx.take(gathered, order) > 0 num_clusters = Nx.sum(kept) distances = diff --git a/test/scholar/cluster/mean_shift_test.exs b/test/scholar/cluster/mean_shift_test.exs index c1601b3a..a19386b5 100644 --- a/test/scholar/cluster/mean_shift_test.exs +++ b/test/scholar/cluster/mean_shift_test.exs @@ -113,7 +113,7 @@ defmodule Scholar.Cluster.MeanShiftTest do assert settled.iterations == Nx.u32(8) # seeds cut off early have not merged yet, so more of them survive - assert Nx.to_number(truncated.num_clusters) > Nx.to_number(settled.num_clusters) + assert Nx.greater(truncated.num_clusters, settled.num_clusters) == Nx.u8(1) end test "fit from a given set of seeds" do @@ -149,6 +149,21 @@ defmodule Scholar.Cluster.MeanShiftTest do assert Nx.type(f32.cluster_centers) == {:f, 32} end + test "fit with samples and seeds of different types" do + seeds = Nx.tensor([[1.0, 1.0], [8.0, 8.0]]) + + # the seeds carry the loop's accumulator, so a wider sample type has to + # widen them too rather than fail to match + wide = MeanShift.fit(Nx.as_type(blobs(), :f64), bandwidth: 1.0, seeds: seeds) + assert Nx.type(wide.cluster_centers) == {:f, 64} + + narrow = + MeanShift.fit(blobs(), bandwidth: 1.0, seeds: Nx.as_type(seeds, :f64)) + + assert Nx.type(narrow.cluster_centers) == {:f, 64} + assert_all_close(wide.cluster_centers, narrow.cluster_centers) + end + test "works with jit_apply" do direct = MeanShift.fit(blobs(), bandwidth: 1.0) jitted = Nx.Defn.jit_apply(&MeanShift.fit/2, [blobs(), [bandwidth: 1.0]]) @@ -161,15 +176,19 @@ defmodule Scholar.Cluster.MeanShiftTest do test "prune" do model = MeanShift.fit(blobs(), bandwidth: 1.0) + pruned = MeanShift.prune(model) - # fit keeps one row per seed so the shapes stay static - assert Nx.axis_size(model.cluster_centers, 0) == 8 assert model.num_clusters == Nx.u32(3) + assert pruned.num_clusters == model.num_clusters - pruned = MeanShift.prune(model) - assert Nx.axis_size(pruned.cluster_centers, 0) == 3 + # fit keeps one row per seed, so pruning has to leave exactly the rows that + # did not lose to a stronger center + assert_all_close( + pruned.cluster_centers, + Nx.tensor([[8.0, 8.0], [1.0333334, 1.05], [1.05, 8.05]]) + ) - # the labels have to keep pointing at the same centers after renumbering + # and the labels have to keep pointing at the same centers after renumbering assert_all_close( Nx.take(pruned.cluster_centers, pruned.labels), Nx.take(model.cluster_centers, model.labels) From 1054fa8b63ed5d6ee1bb9e0c6a5ccb5c86ef44af Mon Sep 17 00:00:00 2001 From: Ricardo Carvalho Santos Date: Tue, 1 Sep 2026 21:02:36 -0300 Subject: [PATCH 4/4] Match scikit-learn's iteration and weight conventions Two independent divergences, both invisible until a fit is cut short by :max_iterations, and both changing the centers rather than only a reported number. The cap was checked before moving the seeds rather than after. scikit-learn tests its limit once the move is done, so a run capped at k takes k + 1 steps and reports k, while this took k steps and reported k. Passing the same limit to both gave centers one step less converged here, and :iterations agreed with n_iter_ only when the run converged on its own. The weight that decides which of two centers on the same mode survives was counted around where a seed landed. scikit-learn counts the neighborhood that produced the center, before the last move. The two agree at a fixed point, so this only showed up when the run was truncated, and it picked a different representative of the same mode. Measured over 200 datasets against scikit-learn 1.6.1, spanning max_iter of 1, 2, 3, 4, 7 and 300, of which 56 are cut short by the limit. Labels, centers, cluster counts and iteration counts went from 194, 177, 199 and 56 out of 200 to 200 out of 200 on all four. --- lib/scholar/cluster/mean_shift.ex | 42 ++++++++++++------------ test/scholar/cluster/mean_shift_test.exs | 41 ++++++++++++++++++++++- 2 files changed, 61 insertions(+), 22 deletions(-) diff --git a/lib/scholar/cluster/mean_shift.ex b/lib/scholar/cluster/mean_shift.ex index 4214dea4..491ce895 100644 --- a/lib/scholar/cluster/mean_shift.ex +++ b/lib/scholar/cluster/mean_shift.ex @@ -36,10 +36,8 @@ defmodule Scholar.Cluster.MeanShift do default: 300, type: :pos_integer, doc: """ - The maximum number of times every seed is moved. Seeds stop early once - none of them moves further than `bandwidth / 1000`. Note that - scikit-learn's `max_iter` checks the limit after moving, so it takes one - more step than the number given. + The maximum number of times a seed is moved past its first step. Seeds + stop early once none of them moves further than `bandwidth / 1000`. """ ], cluster_all: [ @@ -84,8 +82,8 @@ defmodule Scholar.Cluster.MeanShift do * `:num_clusters` - The number of centers that survived. - * `:iterations` - How many times the seeds were moved. This is at least one, - since the seeds have to move once for their movement to be measured. + * `:iterations` - How many times the seeds were moved past their first step. + This is zero when every seed settles on that first step. ## Examples @@ -102,7 +100,7 @@ defmodule Scholar.Cluster.MeanShift do ), labels: Nx.s32([2, 2, 0, 0]), num_clusters: Nx.u32(2), - iterations: Nx.u32(2) + iterations: Nx.u32(1) } """ deftransform fit(x, opts \\ []) do @@ -140,9 +138,8 @@ defmodule Scholar.Cluster.MeanShift do seeds = Nx.as_type(seeds, type) bandwidth = opts[:bandwidth] - {centers, iterations} = shift_seeds(x, seeds, bandwidth, opts) + {centers, gathered, iterations} = shift_seeds(x, seeds, bandwidth, opts) - gathered = intensity(x, centers, bandwidth) order = strongest_first(gathered, centers) centers = Nx.take(centers, order) @@ -186,9 +183,14 @@ defmodule Scholar.Cluster.MeanShift do defnp shift_seeds(x, seeds, bandwidth, opts) do tolerance = bandwidth / 1000 - {seeds, _, iterations} = - while {seeds, {x, bandwidth, tolerance, moving = Nx.u8(1)}, i = Nx.u32(0)}, - i < opts[:max_iterations] and moving do + # the limit is checked after moving, as scikit-learn does, so a run capped at + # k takes k + 1 steps and reports k + empty = Nx.broadcast(Nx.u32(0), {Nx.axis_size(seeds, 0)}) + + {seeds, gathered, _, moves} = + while {seeds, _gathered = empty, {x, bandwidth, tolerance, moving = Nx.u8(1)}, + i = Nx.u32(0)}, + i <= opts[:max_iterations] and moving do within = Scholar.Metrics.Distance.pairwise_euclidean(seeds, x) <= bandwidth members = Nx.as_type(within, Nx.type(seeds)) count = Nx.sum(members, axes: [1], keep_axes: true) @@ -202,16 +204,14 @@ defmodule Scholar.Cluster.MeanShift do ) shift = Scholar.Metrics.Distance.euclidean(moved, seeds, axes: [1]) - {moved, {x, bandwidth, tolerance, Nx.any(shift > tolerance)}, i + 1} - end - {seeds, iterations} - end + # the weight scikit-learn sorts by is the neighborhood that produced the + # center, counted before the move, not the one around where it landed + {moved, Nx.sum(Nx.as_type(within, :u32), axes: [1]), + {x, bandwidth, tolerance, Nx.any(shift > tolerance)}, i + 1} + end - defnp intensity(x, centers, bandwidth) do - Scholar.Metrics.Distance.pairwise_euclidean(centers, x) - |> Nx.less_equal(bandwidth) - |> Nx.sum(axes: [1]) + {seeds, gathered, moves - 1} end # ties break on the coordinates, descending, so that seeds landing on the same @@ -269,7 +269,7 @@ defmodule Scholar.Cluster.MeanShift do ), labels: Nx.s32([1, 1, 0, 0]), num_clusters: Nx.u32(2), - iterations: Nx.u32(2) + iterations: Nx.u32(1) } """ def prune(%__MODULE__{cluster_centers: centers, labels: labels} = model) do diff --git a/test/scholar/cluster/mean_shift_test.exs b/test/scholar/cluster/mean_shift_test.exs index a19386b5..646a1cf4 100644 --- a/test/scholar/cluster/mean_shift_test.exs +++ b/test/scholar/cluster/mean_shift_test.exs @@ -110,12 +110,51 @@ defmodule Scholar.Cluster.MeanShiftTest do settled = MeanShift.fit(x, bandwidth: 2.5) assert truncated.iterations == Nx.u32(1) - assert settled.iterations == Nx.u32(8) + assert settled.iterations == Nx.u32(7) # seeds cut off early have not merged yet, so more of them survive assert Nx.greater(truncated.num_clusters, settled.num_clusters) == Nx.u8(1) end + test "fit truncated by max_iterations" do + # Expected values from scikit-learn 1.6.1 on this data, which checks the + # limit after moving, so a run capped at one still takes a second step + x = + Nx.tensor([ + [-5.1, -4.1], + [4.5, 5.3], + [-3.0, -2.2], + [-5.9, -0.2], + [-1.5, -4.0], + [2.7, -4.0], + [-4.1, 0.2], + [-1.3, -4.9] + ]) + + capped = MeanShift.fit(x, bandwidth: 3.0, max_iterations: 1) |> MeanShift.prune() + + assert capped.iterations == Nx.u32(1) + assert capped.labels == Nx.s32([0, 1, 0, 0, 0, 2, 0, 0]) + + assert_all_close( + capped.cluster_centers, + Nx.tensor([[-3.425, -2.525], [4.5, 5.3], [2.7, -4.0]]) + ) + + # left to settle it finds one more mode, and the weight that decides which + # center survives is the neighborhood that produced it, counted before the + # last move + settled = MeanShift.fit(x, bandwidth: 3.0) |> MeanShift.prune() + + assert settled.iterations == Nx.u32(3) + assert settled.labels == Nx.s32([0, 2, 0, 1, 0, 3, 1, 0]) + + assert_all_close( + settled.cluster_centers, + Nx.tensor([[-2.725, -3.8], [-4.3333333, -0.7333333], [4.5, 5.3], [2.7, -4.0]]) + ) + end + test "fit from a given set of seeds" do x = Nx.tensor([[0.0, 0.0], [0.1, 0.0], [0.0, 0.1], [5.0, 5.0], [5.1, 5.0]])