diff --git a/lib/scholar/cluster/dbscan.ex b/lib/scholar/cluster/dbscan.ex index 39af5be5..87e2bf20 100644 --- a/lib/scholar/cluster/dbscan.ex +++ b/lib/scholar/cluster/dbscan.ex @@ -88,13 +88,15 @@ defmodule Scholar.Cluster.DBSCAN do } """ deftransform fit(x, opts \\ []) do - fit_n(x, NimbleOptions.validate!(opts, @opts_schema)) + opts = NimbleOptions.validate!(opts, @opts_schema) + {weights, opts} = Keyword.pop(opts, :weights, nil) + weights = validate_weights(weights, Nx.axis_size(x, 0), type: to_float_type(x)) + fit_n(x, weights, opts) end - defnp fit_n(x, opts) do + defnp fit_n(x, weights, opts) do num_samples = Nx.axis_size(x, 0) - weights = validate_weights(opts[:weights], num_samples, type: to_float_type(x)) - weights = if Nx.rank(weights) == 0, do: weights, else: Nx.new_axis(weights, -1) + weights = if Nx.rank(weights) == 0, do: weights, else: Nx.new_axis(weights, 0) y_dummy = Nx.broadcast(Nx.tensor(0), {num_samples}) neighbor_model = diff --git a/test/scholar/cluster/dbscan_test.exs b/test/scholar/cluster/dbscan_test.exs index 667a5389..f9add4ff 100644 --- a/test/scholar/cluster/dbscan_test.exs +++ b/test/scholar/cluster/dbscan_test.exs @@ -40,6 +40,50 @@ defmodule Scholar.Cluster.DBSCANTest do assert model.core_sample_indices == Nx.tensor([0, 1, 0, 1, 0, 1, 0, 1, 0, 0], type: :u8) end + test "fit with weights" do + # Expected values from scikit-learn 1.6.1 on this data. + x = Nx.tensor([[1, 2], [2, 2], [2, 3], [8, 7], [8, 8], [25, 80]]) + weights = Nx.tensor([5.0, 0.1, 0.1, 1.0, 1.0, 1.0]) + model = DBSCAN.fit(x, eps: 3, min_samples: 2, weights: weights) + + assert model.labels == Nx.tensor([0, 0, 0, 1, 1, -1]) + assert model.core_sample_indices == Nx.tensor([1, 1, 1, 1, 1, 0], type: :u8) + end + + test "fit with weights joining two clusters through a light sample" do + # Expected values from scikit-learn 1.6.1 on this data. + x = Nx.tensor([[0, 0], [1, 0], [2, 0]]) + weights = Nx.tensor([5.0, 0.1, 5.0]) + model = DBSCAN.fit(x, eps: 1.5, min_samples: 3, weights: weights) + + assert model.labels == Nx.tensor([0, 0, 0]) + assert model.core_sample_indices == Nx.tensor([1, 1, 1], type: :u8) + end + + test "fit with negative weights" do + # Expected values from scikit-learn 1.6.1 on this data. + x = Nx.tensor([[0, 0], [1, 0], [2, 0]]) + weights = Nx.tensor([3.0, -1.0, 3.0]) + model = DBSCAN.fit(x, eps: 1.5, min_samples: 3, weights: weights) + + assert model.labels == Nx.tensor([0, 0, 0]) + assert model.core_sample_indices == Nx.tensor([0, 1, 0], type: :u8) + end + + test "works with jit_apply" do + opts = [ + eps: 3.3, + min_samples: 2, + weights: Nx.tensor([5.0, 0.1, 0.1, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) + ] + + model = DBSCAN.fit(x(), opts) + jitted = Nx.Defn.jit_apply(&DBSCAN.fit(&1, opts), [x()]) + + assert jitted.labels == model.labels + assert jitted.core_sample_indices == model.core_sample_indices + end + test "test with artificially created 3 blobs" do x = Nx.tensor([