diff --git a/lib/scholar/naive_bayes/bernoulli.ex b/lib/scholar/naive_bayes/bernoulli.ex index 61602334..943e6b8d 100644 --- a/lib/scholar/naive_bayes/bernoulli.ex +++ b/lib/scholar/naive_bayes/bernoulli.ex @@ -10,13 +10,14 @@ defmodule Scholar.NaiveBayes.Bernoulli do import Scholar.Shared @derive {Nx.Container, + keep: [:binarize], containers: [ :feature_count, :class_count, :class_log_priors, :feature_log_probability ]} - defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability] + defstruct [:feature_count, :class_count, :class_log_priors, :feature_log_probability, :binarize] opts_schema = [ num_classes: [ @@ -27,7 +28,12 @@ defmodule Scholar.NaiveBayes.Bernoulli do """ ], alpha: [ - type: {:or, [:float, {:list, :float}]}, + type: + {:or, + [ + {:custom, Scholar.Options, :non_negative_number, []}, + {:list, {:custom, Scholar.Options, :non_negative_number, []}} + ]}, default: 1.0, doc: ~S""" Additive (Laplace/Lidstone) smoothing parameter @@ -127,7 +133,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-1.0986123085021973, -1.0986123085021973, -0.40546512603759766], [-0.28768205642700195, -0.28768205642700195, -0.28768205642700195] ] - ) + ), + binarize: 1.0 } iex> x = Nx.iota({4, 3}) @@ -153,7 +160,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do [-23.025850296020508, 0.0, 0.0], [0.0, 0.0, 0.0] ] - ) + ), + binarize: 0.0 } """ @@ -203,7 +211,7 @@ defmodule Scholar.NaiveBayes.Bernoulli do priors_flag = opts[:class_priors] != nil {class_priors, opts} = Keyword.pop(opts, :class_priors, :nan) - class_priors = Nx.tensor(class_priors) + class_priors = Nx.tensor(class_priors, type: type) if priors_flag and Nx.size(class_priors) != num_classes do raise ArgumentError, @@ -281,14 +289,16 @@ defmodule Scholar.NaiveBayes.Bernoulli do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ class_count: class_count, class_log_priors: class_log_priors, feature_count: feature_count, - feature_log_probability: feature_log_probability + feature_log_probability: feature_log_probability, + binarize: opts[:binarize] } end @@ -343,8 +353,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [-4.7047806, -12.329399, -0.009097099], - [-8.750494, -19.147701, -1.5830994e-4] + [-1.4696369, -2.162784, -0.42314053], + [-1.4696369, -2.162784, -0.42314053] ] > """ @@ -373,8 +383,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [0.0090519, 4.419875e-6, 0.99094415], - [1.5838306e-4, 4.8334696e-9, 0.9998417] + [0.23000899, 0.11500449, 0.65498656], + [0.23000899, 0.11500449, 0.65498656] ] > """ @@ -395,8 +405,8 @@ defmodule Scholar.NaiveBayes.Bernoulli do #Nx.Tensor< f32[2][3] [ - [3.6356335, -3.988985, 8.331317], - [10.567104, 0.16989732, 19.31744] + [-2.6026897, -3.295837, -1.5561934], + [-2.6026897, -3.295837, -1.5561934] ] > """ @@ -421,10 +431,17 @@ defmodule Scholar.NaiveBayes.Bernoulli do defnp joint_log_likelihood( %__MODULE__{ feature_log_probability: feature_log_probability, - class_log_priors: class_log_priors + class_log_priors: class_log_priors, + binarize: binarize }, x ) do + x = + case binarize do + nil -> x + threshold -> Scholar.Preprocessing.Binarizer.fit_transform(x, threshold: threshold) + end + neg_prob = Nx.log(1 - Nx.exp(feature_log_probability)) jll = Nx.dot(x, [1], feature_log_probability - neg_prob, [1]) jll + class_log_priors + Nx.sum(neg_prob, axes: [1]) diff --git a/lib/scholar/naive_bayes/categorical.ex b/lib/scholar/naive_bayes/categorical.ex index 3fdd9d66..85d152e2 100644 --- a/lib/scholar/naive_bayes/categorical.ex +++ b/lib/scholar/naive_bayes/categorical.ex @@ -29,7 +29,12 @@ defmodule Scholar.NaiveBayes.Categorical do """ ], alpha: [ - type: {:or, [:float, {:list, :float}]}, + type: + {:or, + [ + {:custom, Scholar.Options, :non_negative_number, []}, + {:list, {:custom, Scholar.Options, :non_negative_number, []}} + ]}, default: 1.0, doc: ~S""" Additive (Laplace/Lidstone) smoothing parameter @@ -235,7 +240,7 @@ defmodule Scholar.NaiveBayes.Categorical do priors_flag = opts[:class_priors] != nil {class_priors, opts} = Keyword.pop(opts, :class_priors, :nan) - class_priors = Nx.tensor(class_priors) + class_priors = Nx.tensor(class_priors, type: type) if priors_flag and Nx.size(class_priors) != num_classes do raise ArgumentError, @@ -272,10 +277,11 @@ defmodule Scholar.NaiveBayes.Categorical do end num_categories = - (opts[:min_categories] || x) - |> Nx.reduce_max() - |> Nx.add(1) - |> Nx.to_number() + if min_categories_flag do + min_categories |> Nx.reduce_max() |> Nx.to_number() |> trunc() + else + x |> Nx.reduce_max() |> Nx.add(1) |> Nx.to_number() |> trunc() + end opts = opts ++ @@ -332,7 +338,8 @@ defmodule Scholar.NaiveBayes.Categorical do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ @@ -360,7 +367,7 @@ defmodule Scholar.NaiveBayes.Categorical do """ defn predict(%__MODULE__{} = model, x, classes) do - check_dim(x, Nx.axis_size(model.feature_count, 1)) + check_dim(x, Nx.axis_size(model.feature_count, 0)) if Nx.rank(classes) != 1 do raise ArgumentError, @@ -401,7 +408,7 @@ defmodule Scholar.NaiveBayes.Categorical do """ defn predict_log_probability(%__MODULE__{} = model, x) do - check_dim(x, Nx.axis_size(model.feature_count, 1)) + check_dim(x, Nx.axis_size(model.feature_count, 0)) jll = joint_log_likelihood(model, x) log_proba_x = @@ -453,7 +460,7 @@ defmodule Scholar.NaiveBayes.Categorical do """ defn predict_joint_log_probability(%__MODULE__{} = model, x) do - check_dim(x, Nx.axis_size(model.feature_count, 1)) + check_dim(x, Nx.axis_size(model.feature_count, 0)) joint_log_likelihood(model, x) end @@ -476,14 +483,22 @@ defmodule Scholar.NaiveBayes.Categorical do }, x ) do + # jll accumulates one {num_samples, num_classes} term per feature. Note + # that this is not the shape of x: feature_log_probability is + # {num_features, num_classes, num_categories}, so num_classes comes from + # its axis 1, not from x's second axis. + num_samples = Nx.axis_size(x, 0) + num_classes = Nx.axis_size(feature_log_probability, 1) + {_, jll} = - while {{i = 0, feature_log_probability, x}, jll = Nx.broadcast(0.0, Nx.shape(x))}, + while {{i = 0, feature_log_probability, x}, + jll = Nx.broadcast(0.0, {num_samples, num_classes})}, i < Nx.axis_size(x, 1) do indices = Nx.slice_along_axis(x, i, 1, axis: 1) |> Nx.squeeze(axes: [1]) jll = Nx.slice_along_axis(feature_log_probability, i, 1, axis: 0) - |> Nx.squeeze() + |> Nx.squeeze(axes: [0]) |> Nx.take(indices, axis: 1) |> Nx.transpose() |> Nx.add(jll) diff --git a/lib/scholar/naive_bayes/complement.ex b/lib/scholar/naive_bayes/complement.ex index f2f32a40..31ef6c47 100644 --- a/lib/scholar/naive_bayes/complement.ex +++ b/lib/scholar/naive_bayes/complement.ex @@ -186,7 +186,7 @@ defmodule Scholar.NaiveBayes.Complement do sample_weights = Nx.tensor(sample_weights, type: x_type) {priors, opts} = Keyword.pop(opts, :priors, Nx.tensor(0.0, type: x_type)) - class_priors = Nx.tensor(priors) + class_priors = Nx.tensor(priors, type: x_type) {alpha, opts} = Keyword.pop!(opts, :alpha) alpha = Nx.tensor(alpha, type: x_type) @@ -385,7 +385,8 @@ defmodule Scholar.NaiveBayes.Complement do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: x_type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ diff --git a/lib/scholar/naive_bayes/multinomial.ex b/lib/scholar/naive_bayes/multinomial.ex index 82ec59dd..4ff0d95a 100644 --- a/lib/scholar/naive_bayes/multinomial.ex +++ b/lib/scholar/naive_bayes/multinomial.ex @@ -197,7 +197,7 @@ defmodule Scholar.NaiveBayes.Multinomial do priors_flag = opts[:class_priors] != nil {class_priors, opts} = Keyword.pop(opts, :class_priors, :nan) - class_priors = Nx.tensor(class_priors) + class_priors = Nx.tensor(class_priors, type: type) if priors_flag and Nx.size(class_priors) != num_classes do raise ArgumentError, @@ -271,7 +271,8 @@ defmodule Scholar.NaiveBayes.Multinomial do Nx.log(class_count) - Nx.log(Nx.sum(class_count)) true -> - Nx.broadcast(-Nx.log(num_classes), {num_classes}) + num_classes_t = Nx.tensor(1.0, type: type) * num_classes + Nx.broadcast(-Nx.log(num_classes_t), {num_classes}) end %__MODULE__{ diff --git a/test/scholar/naive_bayes/bernoulli_test.exs b/test/scholar/naive_bayes/bernoulli_test.exs index 77adaad5..aa2a624b 100644 --- a/test/scholar/naive_bayes/bernoulli_test.exs +++ b/test/scholar/naive_bayes/bernoulli_test.exs @@ -169,5 +169,90 @@ defmodule Scholar.NaiveBayes.BernoulliTest do expected_predictions = Nx.tensor([2, 1]) assert predictions == expected_predictions end + + test "applies binarize to the input at predict time, not only at fit time" do + x = Nx.iota({4, 3}) + y = Nx.tensor([1, 2, 0, 2]) + model = Bernoulli.fit(x, y, num_classes: 3) + + x_test = Nx.tensor([[6, 2, 4], [8, 5, 9]]) + + # Every entry of x_test is above the default threshold of 0.0, so both + # rows binarize to all-ones and must therefore score identically. + # Reference: sklearn.naive_bayes.BernoulliNB(binarize=0.0) on the same + # data returns these probabilities for both rows. + assert_all_close( + Bernoulli.predict_probability(model, x_test), + Nx.tensor([ + [0.23000899, 0.11500449, 0.65498656], + [0.23000899, 0.11500449, 0.65498656] + ]) + ) + + # Passing the already-binarized input must give the same answer. + x_test_binarized = Scholar.Preprocessing.Binarizer.fit_transform(x_test, threshold: 0.0) + model_no_binarize = Bernoulli.fit(x, y, num_classes: 3) + + assert_all_close( + Bernoulli.predict_probability(model_no_binarize, x_test_binarized), + Bernoulli.predict_probability(model, x_test) + ) + end + + test "respects a non-default binarize threshold at predict time" do + x = Nx.iota({4, 3}) + y = Nx.tensor([1, 2, 0, 2]) + model = Bernoulli.fit(x, y, num_classes: 3, binarize: 5.0) + + x_test = Nx.tensor([[6, 2, 4], [8, 5, 9]]) + + # With threshold 5.0 the two rows binarize differently ([1,0,0] and + # [1,0,1]), so unlike the default-threshold case they must not score + # identically. + probability = Bernoulli.predict_probability(model, x_test) + refute Nx.to_number(Nx.all_close(probability[0], probability[1])) == 1 + end + end + + describe "option validation" do + test "rejects a negative alpha instead of silently producing NaN" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]]) + y = Nx.tensor([0, 1, 1]) + + assert_raise NimbleOptions.ValidationError, fn -> + Bernoulli.fit(x, y, num_classes: 2, alpha: -1.0) + end + + assert_raise NimbleOptions.ValidationError, fn -> + Bernoulli.fit(x, y, num_classes: 2, alpha: [1.0, -2.0, 1.0]) + end + end + + test "keeps class_log_priors in the same type as the rest of the model" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]], type: :f64) + y = Nx.tensor([0, 1, 1]) + + for opts <- [ + [num_classes: 2], + [num_classes: 2, class_priors: [0.3, 0.7]], + [num_classes: 2, fit_priors: false] + ] do + model = Bernoulli.fit(x, y, opts) + + assert Nx.type(model.class_log_priors) == Nx.type(model.feature_log_probability), + "class_log_priors downcast for opts: #{inspect(opts)}" + end + end + + test "computes the uniform prior in the model type, not in f32" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]], type: :f64) + y = Nx.tensor([0, 1, 1]) + model = Bernoulli.fit(x, y, num_classes: 3, fit_priors: false) + + # -log(3) must be accurate to f64, not an f32 value widened to f64. + assert_all_close(model.class_log_priors[0], Nx.tensor(-:math.log(3), type: :f64), + atol: 1.0e-15 + ) + end end end diff --git a/test/scholar/naive_bayes/categorical_test.exs b/test/scholar/naive_bayes/categorical_test.exs index 179242aa..e46766e4 100644 --- a/test/scholar/naive_bayes/categorical_test.exs +++ b/test/scholar/naive_bayes/categorical_test.exs @@ -1027,4 +1027,90 @@ defmodule Scholar.NaiveBayes.CategoricalTest do ) end end + + describe "predict with num_features != num_classes" do + test "predicts when the feature count differs from the class count" do + # Every doctest and existing test happens to use num_features == + # num_classes, which is the only case where a jll accumulator sized + # from x instead of from the class axis works. + x = Nx.tensor([[1, 2, 2], [1, 2, 1], [2, 2, 0], [0, 1, 2], [2, 0, 1]]) + y = Nx.tensor([0, 1, 1, 0, 1]) + model = Categorical.fit(x, y, num_classes: 2) + + x_test = Nx.tensor([[1, 2, 2], [2, 0, 1], [0, 1, 0]]) + + assert Nx.shape(Categorical.predict_probability(model, x_test)) == {3, 2} + assert Categorical.predict(model, x_test, Nx.tensor([0, 1])) == Nx.tensor([0, 1, 0]) + + # Reference: sklearn.naive_bayes.CategoricalNB on the same data. + assert_all_close( + Categorical.predict_probability(model, x_test), + Nx.tensor([ + [0.6973365617433412, 0.3026634382566586], + [0.060150375939849614, 0.9398496240601504], + [0.6973365617433416, 0.3026634382566585] + ]), + atol: 1.0e-5 + ) + end + + test "check_dim validates against the feature count, not the class count" do + x = Nx.tensor([[1, 2, 2], [1, 2, 1], [2, 2, 0], [0, 1, 2], [2, 0, 1]]) + y = Nx.tensor([0, 1, 1, 0, 1]) + model = Categorical.fit(x, y, num_classes: 2) + + # 3 features, as used for fitting: must be accepted. + assert Nx.shape(Categorical.predict_probability(model, Nx.tensor([[1, 2, 2]]))) == {1, 2} + + # 2 features: must be rejected even though it matches num_classes. + assert_raise ArgumentError, fn -> + Categorical.predict_probability(model, Nx.tensor([[1, 2]])) + end + end + end + + describe "option validation" do + test "rejects a negative alpha instead of silently producing NaN" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]]) + y = Nx.tensor([0, 1, 1]) + + assert_raise NimbleOptions.ValidationError, fn -> + Categorical.fit(x, y, num_classes: 2, alpha: -1.0) + end + + assert_raise NimbleOptions.ValidationError, fn -> + Categorical.fit(x, y, num_classes: 2, alpha: [1.0, -2.0, 1.0]) + end + end + + test "min_categories actually sizes the category axis" do + x = Nx.tensor([[1, 2, 3], [1, 3, 4], [2, 2, 3], [1, 1, 3], [2, 1, 4]]) + y = Nx.tensor([0, 1, 2, 1, 0]) + + # Without the option the category count is inferred from the data (max 4 + # -> 5 categories). Asking for 7 must widen the axis to 7, matching + # sklearn's CategoricalNB(min_categories=...) semantics. + inferred = Categorical.fit(x, y, num_classes: 3) + assert Nx.axis_size(inferred.feature_count, 2) == 5 + + widened = Categorical.fit(x, y, num_classes: 3, min_categories: [7, 7, 7]) + assert Nx.axis_size(widened.feature_count, 2) == 7 + end + + test "keeps class_log_priors in the same type as the rest of the model" do + x = Nx.tensor([[1, 0, 1], [0, 1, 1], [1, 1, 0]]) + y = Nx.tensor([0, 1, 1]) + + for opts <- [ + [num_classes: 2], + [num_classes: 2, class_priors: [0.3, 0.7]], + [num_classes: 2, fit_priors: false] + ] do + model = Categorical.fit(x, y, opts) + + assert Nx.type(model.class_log_priors) == Nx.type(model.feature_log_probability), + "class_log_priors downcast for opts: #{inspect(opts)}" + end + end + end end