实现了感知机模型，K最近邻模型，朴素贝叶斯模型

.gitignore

@@ -0,0 +1,2 @@
+/target
+.idea/

Cargo.lock

@@ -0,0 +1,89 @@
+# This file is automatically @generated by Cargo.
+# It is not intended for manual editing.
+version = 4
+
+[[package]]
+name = "autocfg"
+version = "1.5.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53"
+
+[[package]]
+name = "matrixmultiply"
+version = "0.3.11"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3f607c237553f086e7043417a51df26b2eb899d3caff94e6a67592ff992fedc7"
+dependencies = [
+ "autocfg",
+ "rawpointer",
+]
+
+[[package]]
+name = "ml_learn"
+version = "0.1.0"
+dependencies = [
+ "ndarray",
+]
+
+[[package]]
+name = "ndarray"
+version = "0.17.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "portable-atomic",
+ "portable-atomic-util",
+ "rawpointer",
+]
+
+[[package]]
+name = "num-complex"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.47"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7ce2d95d4b3734dc35aa2f45e1aa22cd416814592a4f9d9205e11affd5b8e10b"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
+[[package]]
+name = "portable-atomic"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "05c8b63e8d9609db387f0324918f81d68fe27748f084ef092fb35954d0539a85"
+
+[[package]]
+name = "portable-atomic-util"
+version = "0.2.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "10ab3eb7f3becc3a1cbc4f2c6f20267996cfc1a6467a873763411b136a122715"
+dependencies = [
+ "portable-atomic",
+]
+
+[[package]]
+name = "rawpointer"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3"

Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "ml_learn"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+ndarray = "0.17.2"

README.md

@@ -0,0 +1,8 @@
+一个练手的机器学习项目，目前包含以下几种模型：
+
+- perceptron(基础感知机模型)
+- KNN(K最近邻模型)
+- naive_bayes(朴素贝叶斯模型)包含以下几种:
+    - gaussian_naive_bayes(高斯朴素贝叶斯)
+    - poisson_naive_bayes(泊松朴素贝叶斯)
+    - berboulli_naive_bayes(伯努利朴素贝叶斯)

src/lib.rs

@@ -0,0 +1,6 @@
+pub mod model;
+
+pub use model::model::{Model, TestDataSet, TestOutput, TrainDataset};
+pub use model::knn::knn::KNNModel;
+pub use model::naive_bayes::naive_bayes::NaiveBayesClassifier;
+pub use model::perceptron::perceptron::Perceptron;

src/main.rs

@@ -0,0 +1,5 @@
+
+
+fn main() {
+    println!("Hello, world!");
+}

src/model.rs

@@ -0,0 +1,5 @@
+pub mod perceptron;
+mod utils;
+pub mod knn;
+pub mod naive_bayes;
+pub mod model;

src/model/knn.rs

@@ -0,0 +1 @@
+pub mod knn;

src/model/knn/knn.rs

@@ -0,0 +1,81 @@
+use std::collections::HashMap;
+use crate::model::model::{Model, TestDataSet, TestOutput, TrainDataset};
+use crate::model::utils::l2_distance;
+
+pub struct KdTree;
+
+pub struct BaseModel {
+    k: usize,
+    input: Vec<(Vec<f64>, String)>,
+}
+
+pub enum KNNModel {
+    KdTree(KdTree),
+    BaseKNN(BaseModel),
+}
+
+impl KNNModel {
+    pub fn new(model_class: &str, k: usize) -> Self {
+        match model_class {
+            "base" => Self::BaseKNN(BaseModel { k, input: Vec::new() }),
+            _ => todo!()
+        }
+    }
+}
+
+impl Model for KNNModel {
+    fn fit(&mut self, input: &TrainDataset) {
+        match self {
+            Self::BaseKNN(model) => {
+                // KNN 的"训练"只是把数据存下来
+                model.input = input.input_data.rows()
+                    .into_iter()
+                    .map(|row| row.to_vec())
+                    .zip(input.output_data.iter().cloned())
+                    .collect();
+            }
+            _ => todo!()
+        }
+    }
+
+    fn predict(&self, input: &TestDataSet) -> TestOutput {
+        match self {
+            Self::BaseKNN(model) => {
+                let output_data = input.input_data.rows()
+                    .into_iter()
+                    .map(|row| {
+                        let row = row.to_vec();
+                        // 1. 算所有距离
+                        let mut neighbors: Vec<(f64, &String)> = model.input.iter()
+                            .map(|(x, label)| (l2_distance(x, &row), label))
+                            .collect();
+                        // 2. 按距离升序排（f64 用 partial_cmp 处理 NaN）
+                        neighbors.sort_by(|a, b| {
+                            a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal)
+                        });
+                        // 3. 取前 k 个投票，找票数最多的
+                        let mut votes: HashMap<&String, usize> = HashMap::new();
+                        for (_, label) in neighbors.iter().take(model.k) {
+                            *votes.entry(label).or_insert(0) += 1;
+                        }
+                        votes.into_iter()
+                            .max_by_key(|&(_, count)| count)
+                            .map(|(label, _)| label.clone())
+                            .unwrap_or_default()
+                    })
+                    .collect();
+                TestOutput { output_data }
+            }
+            _ => todo!()
+        }
+    }
+
+    fn score(&self, input: &TrainDataset) -> f64 {
+        let prediction = self.predict(&TestDataSet { input_data: input.input_data.clone() });
+        let correct = prediction.output_data.iter()
+            .zip(input.output_data.iter())
+            .filter(|(p, t)| p == t)
+            .count();
+        correct as f64 / input.output_data.len() as f64
+    }
+}

src/model/model.rs

@@ -0,0 +1,19 @@
+use ndarray::Array2;
+
+pub struct TrainDataset {
+    pub input_data: Array2<f64>,
+    pub output_data: Vec<String>,
+}
+
+pub struct TestDataSet {
+    pub input_data: Array2<f64>,
+}
+
+pub struct TestOutput {
+    pub output_data: Vec<String>,
+}
+pub trait Model {
+    fn fit(&mut self, input: &TrainDataset);
+    fn predict(&self, input: &TestDataSet) -> TestOutput;
+    fn score(&self, input: &TrainDataset) -> f64;
+}

src/model/naive_bayes.rs

@@ -0,0 +1,4 @@
+pub mod gaussian_naive_bayes;
+pub mod naive_bayes;
+pub mod poisson_naive_bayes;
+pub mod bernoulli_naive_bayes;

src/model/naive_bayes/bernoulli_naive_bayes.rs

@@ -0,0 +1,29 @@
+use super::naive_bayes::NaiveBayes;
+use ndarray::{ArrayView1, Array1};
+
+pub struct BernoulliNaiveBayesLayer {
+    p: Array1<f64>,
+}
+
+impl NaiveBayes for BernoulliNaiveBayesLayer {
+    fn fit(input: ArrayView1<f64>, output: &Vec<usize>, output_space_size: usize) -> Self {
+        let mut freq: Vec<(f64, usize)> = vec![(0.0, 0); output_space_size];
+        for (i, o) in input.iter().zip(output.iter()) {
+            freq[*o].0 += i;
+            freq[*o].1 += 1;
+        }
+        let p: Array1<f64> = freq.iter().map(|(sum, freq)| sum / *freq as f64).collect();
+        Self {
+            p
+        }
+    }
+    fn predict(&self, input: f64) -> Array1<f64> {
+        self.p.iter().map(|p| bernoulli_pmf(input, p)).collect()
+    }
+}
+
+// 伯努利分布：input 只能取 0 或 1
+fn bernoulli_pmf(input: f64, p: &f64) -> f64 {
+    let k = input as i32;
+    p.powi(k) * (1.0 - p).powi(1 - k)
+}

src/model/naive_bayes/gaussian_naive_bayes.rs

@@ -0,0 +1,34 @@
+use super::naive_bayes::NaiveBayes;
+use ndarray::{ArrayView1, Array1};
+pub struct GaussianNaiveBayesLayer {
+    mean_value: Array1<f64>,
+    variance: Array1<f64>,
+}
+
+impl NaiveBayes for GaussianNaiveBayesLayer {
+    fn fit(input: ArrayView1<f64>, output: &Vec<usize>, output_space_size: usize) -> GaussianNaiveBayesLayer {
+        let mut mean_value: Vec<(f64, usize)> = vec![(0.0,0); output_space_size];
+        for (i,o) in input.iter().zip(output.iter()){
+            mean_value[*o].0 += i;
+            mean_value[*o].1 += 1;
+        }
+        let mean_value: Array1<f64> = mean_value.iter().map(|(sum, freq)| sum / *freq as f64).collect();
+        let mut variance: Vec<(f64, usize)> = vec![(0.0,0); output_space_size];
+        for (i, o) in input.iter().zip(output.iter()) {
+            variance[*o].0 += (i - mean_value[*o]) * (i - mean_value[*o]);
+            variance[*o].1 += 1;
+        }
+        let variance: Array1<f64> = variance.iter().map(|(sum, freq)| sum / *freq as f64).collect();
+        Self {
+            mean_value,
+            variance
+        }
+    }
+    fn predict(&self, input: f64) -> Array1<f64> {
+        self.mean_value.iter().zip(self.variance.iter()).map(|(m, v)| gaussian_pdf(input,m,v)).collect()
+    }
+}
+fn gaussian_pdf(input: f64, mean_value: &f64, variance: &f64) -> f64 {
+    let exp_part = - (input - mean_value).powi(2) / (2.0 * variance);
+    exp_part.exp() / (2.0 * std::f64::consts::PI * variance).sqrt()
+}

src/model/naive_bayes/naive_bayes.rs

@@ -0,0 +1,136 @@
+use crate::model::model::{Model, TestOutput, TrainDataset, TestDataSet};
+use ndarray::{ArrayView1,Array1};
+use std::collections::HashMap;
+use super::gaussian_naive_bayes::GaussianNaiveBayesLayer;
+use super::poisson_naive_bayes::PoissonNaiveBayesLayer;
+use super::bernoulli_naive_bayes::BernoulliNaiveBayesLayer;
+
+pub trait NaiveBayes {
+    fn fit(input: ArrayView1<f64>, output: &Vec<usize>, output_space_size: usize) -> Self;
+    fn predict(&self, input: f64) -> Array1<f64>;
+}
+pub enum NaiveBayesLayer {
+    GaussianNaiveBayes(Option<GaussianNaiveBayesLayer>),
+    PoissonNaiveBayes(Option<PoissonNaiveBayesLayer>),
+    BernoulliNaiveBayes(Option<BernoulliNaiveBayesLayer>),
+}
+
+impl NaiveBayesLayer {
+    fn predict(&self, input: f64) -> Array1<f64> {
+        match self {
+            NaiveBayesLayer::GaussianNaiveBayes(Some(layer)) => layer.predict(input),
+            NaiveBayesLayer::GaussianNaiveBayes(None) => panic!("layer not fitted"),
+            NaiveBayesLayer::PoissonNaiveBayes(Some(layer)) => layer.predict(input),
+            NaiveBayesLayer::PoissonNaiveBayes(None) => panic!("layer not fitted"),
+            NaiveBayesLayer::BernoulliNaiveBayes(Some(layer)) => layer.predict(input),
+            NaiveBayesLayer::BernoulliNaiveBayes(None) => panic!("layer not fitted"),
+        }
+    }
+}
+
+pub struct NaiveBayesClassifier{
+    pub feature_size: usize,
+    pub feature_layer: Vec<NaiveBayesLayer>,
+    pub output_map: HashMap<String, usize>,
+    pub index_to_label: Vec<String>,
+    pub prior_probability: Array1<f64>,
+}
+
+impl NaiveBayesClassifier {
+    pub fn new() -> Self {
+        Self{
+            feature_size: 0,
+            feature_layer: Vec::new(),
+            output_map: HashMap::new(),
+            index_to_label: Vec::new(),
+            prior_probability: Array1::<f64>::zeros(1)
+        }
+    }
+    pub fn create_classifier(&mut self, layer_names: Vec<&str>) {
+        self.feature_size = layer_names.len();
+        for layer_name in layer_names {
+            match layer_name {
+                "GNB" => self.feature_layer.push(NaiveBayesLayer::GaussianNaiveBayes(None)),
+                "PoissonNB" => self.feature_layer.push(NaiveBayesLayer::PoissonNaiveBayes(None)),
+                "BernoulliNB" => self.feature_layer.push(NaiveBayesLayer::BernoulliNaiveBayes(None)),
+                _ => panic!("UNKNOWN LAYER NAME")
+            }
+        }
+    }
+}
+
+impl Model for NaiveBayesClassifier {
+    fn fit(&mut self, input: &TrainDataset) {
+        let output = index_output(self, input);
+        let output_space_size = self.output_map.len();
+        for i in 0..self.feature_size{
+            let layer: NaiveBayesLayer = match self.feature_layer[i] {
+                NaiveBayesLayer::GaussianNaiveBayes(_) => {
+                    NaiveBayesLayer::GaussianNaiveBayes(Some(GaussianNaiveBayesLayer::fit(input.input_data.column(i), &output, output_space_size)))
+                }
+                NaiveBayesLayer::PoissonNaiveBayes(_) => {
+                    NaiveBayesLayer::PoissonNaiveBayes(Some(PoissonNaiveBayesLayer::fit(input.input_data.column(i), &output, output_space_size)))
+                }
+                NaiveBayesLayer::BernoulliNaiveBayes(_) => {
+                    NaiveBayesLayer::BernoulliNaiveBayes(Some(BernoulliNaiveBayesLayer::fit(input.input_data.column(i), &output, output_space_size)))
+                }
+            };
+            self.feature_layer[i] = layer;
+        }
+    }
+    fn predict(&self, input: &TestDataSet) -> TestOutput {
+        let mut output: Vec<String> = Vec::with_capacity(input.input_data.nrows());
+        for item in input.input_data.rows() {
+            let mut result: Array1<f64> = self.prior_probability.clone();
+            for i in 0..item.len(){
+                result = result * self.feature_layer[i].predict(item[i]);
+            }
+            output.push(self.index_to_label[get_max(result)].clone());
+        }
+        TestOutput { output_data: output }
+    }
+    fn score(&self, input: &TrainDataset) -> f64 {
+        let prediction = self.predict(&TestDataSet { input_data: input.input_data.clone() });
+        let correct = prediction.output_data.iter()
+            .zip(input.output_data.iter())
+            .filter(|(p, t)| p == t)
+            .count();
+        correct as f64 / input.output_data.len() as f64
+    }
+}
+
+fn index_output(model: &mut NaiveBayesClassifier,dataset: &TrainDataset) -> Vec<usize> {
+    let mut statistics: HashMap<String, usize> = HashMap::new();
+    let mut i: usize = 0;
+    let mut output: Vec<usize> = Vec::new();
+    for sample in dataset.output_data.iter() {
+        match model.output_map.get(sample) {
+            Some(index) => output.push(*index),
+            None => {
+                model.output_map.insert(sample.to_string(),i);
+                model.index_to_label.push(sample.to_string());
+                output.push(i);
+                i += 1;
+            }
+        }
+        *statistics.entry(sample.to_string()).or_insert(0) += 1;
+    }
+    let size = dataset.output_data.len();
+    let mut prior_probability = vec![0.0; model.output_map.len()];
+    for (key, freq) in &statistics {
+        if let Some(index) = model.output_map.get(key) {
+            prior_probability[*index] = *freq as f64 / size as f64;
+        }
+    }
+    model.prior_probability = Array1::<f64>::from(prior_probability);
+    output
+}
+
+fn get_max(arr: Array1<f64>) -> usize {
+    let mut max: usize = 0;
+    for i in 1..arr.len( ){
+        if arr[i] > arr[max] {max = i;}
+    }
+    max
+}
+

src/model/naive_bayes/poisson_naive_bayes.rs

@@ -0,0 +1,27 @@
+use super::naive_bayes::NaiveBayes;
+use ndarray::{ArrayView1, Array1};
+use crate::model::utils::product;
+
+pub struct PoissonNaiveBayesLayer {
+    lambda: Array1<f64>,
+}
+
+impl NaiveBayes for PoissonNaiveBayesLayer {
+    fn fit(input: ArrayView1<f64>, output: &Vec<usize>, output_space_size: usize) -> Self {
+        let mut lambda: Vec<(f64, usize)> = vec![(0.0,0); output_space_size];
+        for (i,o) in input.iter().zip(output.iter()){
+            lambda[*o].0 += i;
+            lambda[*o].1 += 1;
+        }
+        let lambda: Array1<f64> = lambda.iter().map(|(sum, freq)| sum / *freq as f64).collect();
+        Self {
+            lambda
+        }
+    }
+    fn predict(&self, input: f64) -> Array1<f64> {
+        self.lambda.iter().map(|l| poisson_pmf(input, l)).collect()
+    }
+}
+fn poisson_pmf(input: f64, lambda: &f64) -> f64 {
+    lambda.powf(input) * (-lambda).exp() / product(input)
+}

src/model/perceptron.rs

@@ -0,0 +1 @@
+pub mod perceptron;

src/model/perceptron/perceptron.rs

@@ -0,0 +1,76 @@
+use crate::model::model::{Model, TestDataSet, TestOutput, TrainDataset};
+use crate::model::utils::dot;
+
+pub struct Perceptron {
+    weight: Vec<f64>,
+    b: f64,
+    learning_rate: f64,
+    max_iterations: usize,
+    // 标签映射：labels[0] 视为正类（+1），labels[1] 视为负类（-1）
+    labels: Vec<String>,
+}
+
+impl Perceptron {
+    pub fn new(learning_rate: f64, max_iterations: usize) -> Self {
+        Self {
+            weight: Vec::new(),
+            b: 0.0,
+            learning_rate,
+            max_iterations,
+            labels: Vec::new(),
+        }
+    }
+}
+
+impl Model for Perceptron {
+    fn fit(&mut self, input: &TrainDataset) {
+        self.weight = vec![0.0; input.input_data.ncols()];
+        self.b = 0.0;
+        for label in &input.output_data {
+            if !self.labels.contains(label) {
+                self.labels.push(label.clone());
+            }
+        }
+        assert_eq!(self.labels.len(), 2, "感知机只支持二分类");
+
+        let xs: Vec<Vec<f64>> = input.input_data.rows()
+            .into_iter()
+            .map(|row| row.to_vec())
+            .collect();
+
+        for _ in 0..self.max_iterations {
+            let mut mistake = false;
+            for (x, label) in xs.iter().zip(input.output_data.iter()) {
+                let y = if *label == self.labels[0] { 1.0 } else { -1.0 };
+                if (dot(&self.weight, x) + self.b) * y <= 0.0 {
+                    for (w, xi) in self.weight.iter_mut().zip(x.iter()) {
+                        *w += self.learning_rate * y * xi;
+                    }
+                    self.b += self.learning_rate * y;
+                    mistake = true;
+                }
+            }
+            if !mistake { break; }
+        }
+    }
+
+    fn predict(&self, input: &TestDataSet) -> TestOutput {
+        let output_data = input.input_data.rows()
+            .into_iter()
+            .map(|row| {
+                let result = dot(&self.weight, &row.to_vec()) + self.b;
+                if result >= 0.0 { self.labels[0].clone() } else { self.labels[1].clone() }
+            })
+            .collect();
+        TestOutput { output_data }
+    }
+
+    fn score(&self, input: &TrainDataset) -> f64 {
+        let prediction = self.predict(&TestDataSet { input_data: input.input_data.clone() });
+        let correct = prediction.output_data.iter()
+            .zip(input.output_data.iter())
+            .filter(|(p, t)| p == t)
+            .count();
+        correct as f64 / input.output_data.len() as f64
+    }
+}

src/model/utils.rs

@@ -0,0 +1,24 @@
+pub fn dot(a: &[f64], b: &[f64]) -> f64 {
+    assert_eq!(a.len(), b.len(), "长度不等");
+    a.iter()
+        .zip(b.iter())
+        .map(|(x, y)| x * y)
+        .sum()
+}
+
+pub fn l2_distance(a: &[f64], b: &[f64]) -> f64 {
+    assert_eq!(a.len(), b.len(), "长度不等");
+    a.iter()
+        .zip(b.iter())
+        .map(|(x, y)| (x - y).powi(2))
+        .sum::<f64>()
+        .sqrt()
+}
+
+pub fn product(input: f64) -> f64 {
+    let mut prod: u64 = 1;
+    for i in 2..=input as u64{
+        prod = prod * i;
+    }
+    prod as f64
+}

tests/model_test.rs

@@ -0,0 +1,146 @@
+use ml_learn::{KNNModel, Model, NaiveBayesClassifier, Perceptron, TestDataSet, TrainDataset};
+use ndarray::Array2;
+
+#[test]
+fn perceptron_fit_and_predict() {
+    // 线性可分数据：x1 + x2 > 2 为正类
+    let input_data = Array2::from_shape_vec((6, 2), vec![
+        2.0, 2.0, 3.0, 1.0, 2.0, 3.0,   // "pos"
+        0.0, 0.0, 1.0, 0.0, 0.0, 1.0,   // "neg"
+    ]).unwrap();
+    let output_data = vec![
+        "pos".to_string(), "pos".to_string(), "pos".to_string(),
+        "neg".to_string(), "neg".to_string(), "neg".to_string(),
+    ];
+    let train = TrainDataset { input_data, output_data };
+
+    let mut model = Perceptron::new(1.0, 1000);
+    model.fit(&train);
+    assert_eq!(model.score(&train), 1.0);
+
+    let test = TestDataSet {
+        input_data: Array2::from_shape_vec((2, 2), vec![4.0, 4.0, -1.0, 0.0]).unwrap(),
+    };
+    let result = model.predict(&test);
+    assert_eq!(result.output_data, vec!["pos".to_string(), "neg".to_string()]);
+}
+
+#[test]
+fn knn_fit_and_predict() {
+    let input_data = Array2::from_shape_vec((6, 2), vec![
+        1.0, 1.0, 1.2, 0.8, 0.8, 1.2,   // "A"
+        5.0, 5.0, 5.2, 4.8, 4.8, 5.2,   // "B"
+    ]).unwrap();
+    let output_data = vec![
+        "A".to_string(), "A".to_string(), "A".to_string(),
+        "B".to_string(), "B".to_string(), "B".to_string(),
+    ];
+    let train = TrainDataset { input_data, output_data };
+
+    let mut model = KNNModel::new("base", 3);
+    model.fit(&train);
+    assert_eq!(model.score(&train), 1.0);
+
+    let test = TestDataSet {
+        input_data: Array2::from_shape_vec((2, 2), vec![1.1, 0.9, 5.1, 5.1]).unwrap(),
+    };
+    let result = model.predict(&test);
+    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
+}
+
+#[test]
+fn gaussian_naive_bayes_fit_and_predict() {
+    // 两个类别，单一连续特征，分布明显分开
+    let input_data = Array2::from_shape_vec((6, 1), vec![
+        1.0, 1.2, 0.8,   // "A"
+        5.0, 5.2, 4.8,   // "B"
+    ]).unwrap();
+    let output_data = vec![
+        "A".to_string(), "A".to_string(), "A".to_string(),
+        "B".to_string(), "B".to_string(), "B".to_string(),
+    ];
+    let train = TrainDataset { input_data, output_data };
+
+    let mut clf = NaiveBayesClassifier::new();
+    clf.create_classifier(vec!["GNB"]);
+    clf.fit(&train);
+    assert_eq!(clf.score(&train), 1.0);
+
+    let test = TestDataSet {
+        input_data: Array2::from_shape_vec((2, 1), vec![1.1, 5.1]).unwrap(),
+    };
+    let result = clf.predict(&test);
+    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
+}
+
+#[test]
+fn poisson_naive_bayes_fit_and_predict() {
+    // 计数特征：A 类 λ≈1，B 类 λ≈5
+    let input_data = Array2::from_shape_vec((6, 1), vec![
+        0.0, 1.0, 2.0,   // "A"
+        4.0, 5.0, 6.0,   // "B"
+    ]).unwrap();
+    let output_data = vec![
+        "A".to_string(), "A".to_string(), "A".to_string(),
+        "B".to_string(), "B".to_string(), "B".to_string(),
+    ];
+    let train = TrainDataset { input_data, output_data };
+
+    let mut clf = NaiveBayesClassifier::new();
+    clf.create_classifier(vec!["PoissonNB"]);
+    clf.fit(&train);
+    assert_eq!(clf.score(&train), 1.0);
+
+    let test = TestDataSet {
+        input_data: Array2::from_shape_vec((2, 1), vec![1.0, 5.0]).unwrap(),
+    };
+    let result = clf.predict(&test);
+    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
+}
+
+#[test]
+fn bernoulli_naive_bayes_fit_and_predict() {
+    // 0/1 特征：A 类大概率为 1，B 类大概率为 0
+    let input_data = Array2::from_shape_vec((6, 1), vec![
+        1.0, 1.0, 1.0,   // "A"
+        0.0, 0.0, 0.0,   // "B"
+    ]).unwrap();
+    let output_data = vec![
+        "A".to_string(), "A".to_string(), "A".to_string(),
+        "B".to_string(), "B".to_string(), "B".to_string(),
+    ];
+    let train = TrainDataset { input_data, output_data };
+
+    let mut clf = NaiveBayesClassifier::new();
+    clf.create_classifier(vec!["BernoulliNB"]);
+    clf.fit(&train);
+
+    let test = TestDataSet {
+        input_data: Array2::from_shape_vec((2, 1), vec![1.0, 0.0]).unwrap(),
+    };
+    let result = clf.predict(&test);
+    assert_eq!(result.output_data, vec!["A".to_string(), "B".to_string()]);
+}
+
+#[test]
+fn naive_bayes_mixed_layers() {
+    // 两个特征列分别用不同层：连续特征 GNB + 计数特征 PoissonNB
+    let input_data = Array2::from_shape_vec((6, 2), vec![
+        1.0, 0.0,
+        1.2, 1.0,
+        0.8, 2.0,   // "A"
+        5.0, 4.0,
+        5.2, 5.0,
+        4.8, 6.0,   // "B"
+    ]).unwrap();
+    let output_data = vec![
+        "A".to_string(), "A".to_string(), "A".to_string(),
+        "B".to_string(), "B".to_string(), "B".to_string(),
+    ];
+    let train = TrainDataset { input_data, output_data };
+
+    let mut clf = NaiveBayesClassifier::new();
+    clf.create_classifier(vec!["GNB", "PoissonNB"]);
+    clf.fit(&train);
+    assert_eq!(clf.score(&train), 1.0);
+}