Update CHANGELOG

Reviewed-on: #32
Added

- Add a new hyperparameter to the BoostAODE class, alphablock, to control the way α is computed, with the last model or with the ensmble built so far. Default value is false.
- Add a new hyperparameter to the SPODE class, parent, to set the root node of the model. If no value is set the root parameter of the constructor is used.
- Add a new hyperparameter to the TAN class, parent, to set the root node of the model. If not set the first feature is used as root.

CHANGELOG.md

@@ -13,6 +13,14 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add a new hyperparameter to the SPODE class, *parent*, to set the root node of the model. If no value is set the root parameter of the constructor is used.
 - Add a new hyperparameter to the TAN class, *parent*, to set the root node of the model. If not set the first feature is used as root.
 
+### Internal
+
+- Optimize ComputeCPT method in the Node class.
+
+### Changed
+
+- Hyperparameter *maxTolerance* in the BoostAODE class is now in [1, 6] range (it was in [1, 4] range before).
+
 ## [1.0.6] 2024-11-23
 
 ### Fixed

bayesnet/BaseClassifier.h

@@ -28,8 +28,8 @@ namespace bayesnet {
         status_t virtual getStatus() const = 0;
         float virtual score(std::vector<std::vector<int>>& X, std::vector<int>& y) = 0;
         float virtual score(torch::Tensor& X, torch::Tensor& y) = 0;
-        int virtual getNumberOfNodes()const = 0;
+        int virtual getNumberOfNodes() const = 0;
-        int virtual getNumberOfEdges()const = 0;
+        int virtual getNumberOfEdges() const = 0;
         int virtual getNumberOfStates() const = 0;
         int virtual getClassNumStates() const = 0;
         std::vector<std::string> virtual show() const = 0;
@@ -37,7 +37,7 @@ namespace bayesnet {
         virtual std::string getVersion() = 0;
         std::vector<std::string> virtual topological_order() = 0;
         std::vector<std::string> virtual getNotes() const = 0;
-        std::string virtual dump_cpt()const = 0;
+        std::string virtual dump_cpt() const = 0;
         virtual void setHyperparameters(const nlohmann::json& hyperparameters) = 0;
         std::vector<std::string>& getValidHyperparameters() { return validHyperparameters; }
     protected:

bayesnet/ensembles/Boost.cc

@@ -48,8 +48,8 @@ namespace bayesnet {
         }
         if (hyperparameters.contains("maxTolerance")) {
             maxTolerance = hyperparameters["maxTolerance"];
-            if (maxTolerance < 1 || maxTolerance > 4)
+            if (maxTolerance < 1 || maxTolerance > 6)
-                throw std::invalid_argument("Invalid maxTolerance value, must be greater in [1, 4]");
+                throw std::invalid_argument("Invalid maxTolerance value, must be greater in [1, 6]");
             hyperparameters.erase("maxTolerance");
         }
         if (hyperparameters.contains("predict_voting")) {

bayesnet/ensembles/Ensemble.h

@@ -33,7 +33,12 @@ namespace bayesnet {
         }
         std::string dump_cpt() const override
         {
-            return "";
+            std::string output;
+            for (auto& model : models) {
+                output += model->dump_cpt();
+                output += std::string(80, '-') + "\n";
+            }
+            return output;
         }
     protected:
         torch::Tensor predict_average_voting(torch::Tensor& X);

bayesnet/ensembles/WA2DE.cc

@@ -0,0 +1,267 @@
+// ***************************************************************
+// SPDX-FileCopyrightText: Copyright 2024 Ricardo Montañana Gómez
+// SPDX-FileType: SOURCE
+// SPDX-License-Identifier: MIT
+// ***************************************************************
+#include "WA2DE.h"
+namespace bayesnet {
+    WA2DE::WA2DE(bool predict_voting)
+        : num_classes_(0), num_attributes_(0), total_count_(0.0), weighted_a2de_(false), smoothing_factor_(1.0)
+    {
+        validHyperparameters = { "predict_voting" };
+        std::cout << "WA2DE classifier created.\n";
+    }
+
+    void bayesnet::WA2DE::setHyperparameters(const nlohmann::json& hyperparameters_)
+    {
+        auto hyperparameters = hyperparameters_;
+        if (hyperparameters.contains("predict_voting")) {
+            predict_voting = hyperparameters["predict_voting"];
+            hyperparameters.erase("predict_voting");
+        }
+        Classifier::setHyperparameters(hyperparameters);
+    }
+
+
+    void WA2DE::buildModel(const torch::Tensor& weights)
+    {
+        for (int c = 0; c < num_classes_; ++c) {
+            class_counts_[c] += 1e-4; // Laplace smoothing
+        }
+        for (int a = 0; a < num_attributes_; ++a) {
+            for (int v = 0; v < attribute_cardinalities_[a]; ++v) {
+                for (int c = 0; c < num_classes_; ++c) {
+                    freq_attr_class_[a][v][c] =
+                        (freq_attr_class_[a][v][c] + 1.0) / (class_counts_[c] + attribute_cardinalities_[a]);
+                }
+            }
+        }
+
+        for (int sp = 0; sp < num_attributes_; ++sp) {
+            for (int spv = 0; spv < attribute_cardinalities_[sp]; ++spv) {
+                for (int ch = 0; ch < num_attributes_; ++ch) {
+                    if (sp != ch) {
+                        for (int chv = 0; chv < attribute_cardinalities_[ch]; ++chv) {
+                            for (int c = 0; c < num_classes_; ++c) {
+                                freq_pair_class_[sp][spv][ch][chv][c] =
+                                    (freq_pair_class_[sp][spv][ch][chv][c] + 1.0) /
+                                    (class_counts_[c] + attribute_cardinalities_[sp] * attribute_cardinalities_[ch]);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+        std::cout << "Model probabilities computed.\n";
+    }
+    void WA2DE::trainModel(const torch::Tensor& weights, const Smoothing_t smoothing)
+    {
+        auto data = dataset.clone();
+        auto labels = data[-1];
+        // Remove class row from data
+        data = data.index({ at::indexing::Slice(0, -1) });
+        std::cout << "Training A2DE model...\n";
+        std::cout << "Data: " << data.sizes() << std::endl;
+        std::cout << "Labels: " << labels.sizes() << std::endl;
+        std::cout << std::string(80, '-') << std::endl;
+        if (data.dim() != 2 || labels.dim() != 1) {
+            throw std::invalid_argument("Invalid input dimensions.");
+        }
+        num_attributes_ = data.size(0);
+        num_classes_ = labels.max().item<int>() + 1;
+        total_count_ = data.size(1);
+        std::cout << "Number of attributes: " << num_attributes_ << std::endl;
+        std::cout << "Number of classes: " << num_classes_ << std::endl;
+        std::cout << "Total count: " << total_count_ << std::endl;
+
+        // Compute cardinalities
+        attribute_cardinalities_.clear();
+        for (int i = 0; i < num_attributes_; ++i) {
+            attribute_cardinalities_.push_back(data[i].max().item<int>() + 1);
+        }
+        std::cout << "Attribute cardinalities: ";
+        for (int i = 0; i < num_attributes_; ++i) {
+            std::cout << attribute_cardinalities_[i] << " ";
+        }
+        std::cout << std::endl;
+        // output the map of states
+        std::cout << "States: ";
+        for (int i = 0; i < states.size() - 1; i++) {
+            std::cout << features[i] << " " << states[features[i]].size() << std::endl;
+        }
+
+        // Resize storage
+        class_counts_.resize(num_classes_, 0.0);
+        freq_attr_class_.resize(num_attributes_);
+        freq_pair_class_.resize(num_attributes_);
+
+        for (int i = 0; i < num_attributes_; ++i) {
+            freq_attr_class_[i].resize(attribute_cardinalities_[i], std::vector<double>(num_classes_, 0.0));
+            freq_pair_class_[i].resize(attribute_cardinalities_[i]); // Ensure first level exists
+            for (int j = 0; j < attribute_cardinalities_[i]; ++j) {
+                freq_pair_class_[i][j].resize(num_attributes_); // Ensure second level exists
+                for (int k = 0; k < num_attributes_; ++k) {
+                    if (i != k) {
+                        freq_pair_class_[i][j][k].resize(attribute_cardinalities_[k]); // Ensure third level exists
+                        for (int l = 0; l < attribute_cardinalities_[k]; ++l) {
+                            freq_pair_class_[i][j][k][l].resize(num_classes_, 0.0); // Finally, initialize with 0.0
+                        }
+                    }
+                }
+            }
+        }
+        // Count frequencies
+        auto data_cpu = data.to(torch::kCPU);
+        auto labels_cpu = labels.to(torch::kCPU);
+        int32_t* data_ptr = data_cpu.data_ptr<int32_t>();
+        int32_t* labels_ptr = labels_cpu.data_ptr<int32_t>();
+
+        for (int i = 0; i < total_count_; ++i) {
+            int class_label = labels_ptr[i];
+            class_counts_[class_label] += 1.0;
+
+            std::vector<int> attr_values(num_attributes_);
+            for (int a = 0; a < num_attributes_; ++a) {
+                attr_values[a] = toIntValue(a, data_ptr[i * num_attributes_ + a]);
+                freq_attr_class_[a][attr_values[a]][class_label] += 1.0;
+            }
+
+            // Pairwise counts
+            for (int sp = 0; sp < num_attributes_; ++sp) {
+                for (int ch = 0; ch < num_attributes_; ++ch) {
+                    if (sp != ch) {
+                        freq_pair_class_[sp][attr_values[sp]][ch][attr_values[ch]][class_label] += 1.0;
+                    }
+                }
+            }
+        }
+        std::cout << "Verifying Frequency Counts:\n";
+        for (int c = 0; c < num_classes_; ++c) {
+            std::cout << "Class " << c << " Count: " << class_counts_[c] << std::endl;
+        }
+
+        for (int a = 0; a < num_attributes_; ++a) {
+            for (int v = 0; v < attribute_cardinalities_[a]; ++v) {
+                std::cout << "P(A[" << a << "]=" << v << "|C): ";
+                for (int c = 0; c < num_classes_; ++c) {
+                    std::cout << freq_attr_class_[a][v][c] << " ";
+                }
+                std::cout << std::endl;
+            }
+        }
+
+    }
+
+    torch::Tensor WA2DE::computeProbabilities(const torch::Tensor& data) const
+    {
+        int M = data.size(1);
+        auto output = torch::zeros({ M, num_classes_ }, torch::kF64);
+
+        auto data_cpu = data.to(torch::kCPU);
+        int32_t* data_ptr = data_cpu.data_ptr<int32_t>();
+
+        for (int i = 0; i < M; ++i) {
+            std::vector<int> attr_values(num_attributes_);
+            for (int a = 0; a < num_attributes_; ++a) {
+                attr_values[a] = toIntValue(a, data_ptr[i * num_attributes_ + a]);
+            }
+
+            std::vector<double> log_prob(num_classes_, 0.0);
+            for (int c = 0; c < num_classes_; ++c) {
+                log_prob[c] = std::log((class_counts_[c] + smoothing_factor_) / (total_count_ + num_classes_ * smoothing_factor_));
+
+                double sum_log = 0.0;
+                for (int sp = 0; sp < num_attributes_; ++sp) {
+                    double sp_log = log_prob[c];
+                    for (int ch = 0; ch < num_attributes_; ++ch) {
+                        if (sp == ch) continue;
+                        double num = freq_pair_class_[sp][attr_values[sp]][ch][attr_values[ch]][c] + smoothing_factor_;
+                        double denom = class_counts_[c] + attribute_cardinalities_[sp] * attribute_cardinalities_[ch] * smoothing_factor_;
+                        sp_log += std::log(num / denom);
+                    }
+                    sum_log += std::exp(sp_log);
+                }
+                log_prob[c] = std::log(sum_log / num_attributes_);
+            }
+
+            double max_log = *std::max_element(log_prob.begin(), log_prob.end());
+            double sum_exp = 0.0;
+            for (int c = 0; c < num_classes_; ++c) {
+                sum_exp += std::exp(log_prob[c] - max_log);
+            }
+            double log_sum_exp = max_log + std::log(sum_exp);
+
+            for (int c = 0; c < num_classes_; ++c) {
+                output[i][c] = std::exp(log_prob[c] - log_sum_exp);
+            }
+        }
+
+        return output.to(torch::kF32);
+    }
+    int WA2DE::toIntValue(int attributeIndex, float value) const
+    {
+        int v = static_cast<int>(value);
+        return std::max(0, std::min(v, attribute_cardinalities_[attributeIndex] - 1));
+    }
+    torch::Tensor WA2DE::AODEConditionalProb(const torch::Tensor& data)
+    {
+        int M = data.size(1);  // Number of test samples
+        torch::Tensor output = torch::zeros({ M, num_classes_ }, torch::kF32);
+
+        auto data_cpu = data.to(torch::kCPU);
+        int32_t* data_ptr = data_cpu.data_ptr<int32_t>();
+
+        for (int i = 0; i < M; ++i) {
+            std::vector<int> attr_values(num_attributes_);
+            for (int a = 0; a < num_attributes_; ++a) {
+                attr_values[a] = toIntValue(a, data_ptr[i * num_attributes_ + a]);
+            }
+
+            std::vector<double> log_prob(num_classes_, 0.0);
+            for (int c = 0; c < num_classes_; ++c) {
+                log_prob[c] = std::log(class_counts_[c] / total_count_);
+
+                double sum_log = 0.0;
+                for (int sp = 0; sp < num_attributes_; ++sp) {
+                    double sp_log = log_prob[c];
+                    for (int ch = 0; ch < num_attributes_; ++ch) {
+                        if (sp == ch) continue;
+                        double prob = freq_pair_class_[sp][attr_values[sp]][ch][attr_values[ch]][c];
+                        sp_log += std::log(prob);
+                    }
+                    sum_log += std::exp(sp_log);
+                }
+                log_prob[c] = std::log(sum_log / num_attributes_);
+            }
+
+            double max_log = *std::max_element(log_prob.begin(), log_prob.end());
+            double sum_exp = 0.0;
+            for (int c = 0; c < num_classes_; ++c) {
+                sum_exp += std::exp(log_prob[c] - max_log);
+            }
+            double log_sum_exp = max_log + std::log(sum_exp);
+
+            for (int c = 0; c < num_classes_; ++c) {
+                output[i][c] = std::exp(log_prob[c] - log_sum_exp);
+            }
+        }
+
+        return output;
+    }
+
+    double WA2DE::score(const torch::Tensor& X, const torch::Tensor& y)
+    {
+        torch::Tensor preds = AODEConditionalProb(X);
+        torch::Tensor pred_labels = preds.argmax(1);
+
+        auto correct = pred_labels.eq(y).sum().item<int>();
+        auto total = y.size(0);
+
+        return static_cast<double>(correct) / total;
+    }
+
+    std::vector<std::string> WA2DE::graph(const std::string& title) const
+    {
+        return { title, "Graph visualization not implemented." };
+    }
+}

bayesnet/ensembles/WA2DE.h

@@ -0,0 +1,53 @@
+// ***************************************************************
+// SPDX-FileCopyrightText: Copyright 2024 Ricardo Montañana Gómez
+// SPDX-FileType: SOURCE
+// SPDX-License-Identifier: MIT
+// ***************************************************************
+#ifndef WA2DE_H
+#define WA2DE_H
+#include "Ensemble.h"
+#include <torch/torch.h>
+#include <vector>
+#include <map>
+#include <nlohmann/json.hpp>
+namespace bayesnet {
+    /**
+     * Geoffrey I. Webb's A2DE (Averaged 2-Dependence Estimators) classifier
+     * Implements the A2DE algorithm as an ensemble of SPODE models.
+     */
+    class WA2DE : public Ensemble {
+    public:
+        explicit WA2DE(bool predict_voting = false);
+        virtual ~WA2DE() {};
+
+        // Override method to set hyperparameters
+        void setHyperparameters(const nlohmann::json& hyperparameters) override;
+
+        // Graph visualization function
+        std::vector<std::string> graph(const std::string& title = "A2DE") const override;
+        torch::Tensor computeProbabilities(const torch::Tensor& data) const;
+        double score(const torch::Tensor& X, const torch::Tensor& y);
+    protected:
+        // Model-building function
+        void buildModel(const torch::Tensor& weights) override;
+
+    private:
+        int num_classes_;                // Number of classes
+        int num_attributes_;             // Number of attributes
+        std::vector<int> attribute_cardinalities_; // Cardinalities of attributes
+
+        // Frequency counts (similar to Java implementation)
+        std::vector<double> class_counts_;  // Class frequency
+        std::vector<std::vector<std::vector<double>>> freq_attr_class_; // P(A | C)
+        std::vector<std::vector<std::vector<std::vector<std::vector<double>>>>> freq_pair_class_; // P(A_i, A_j | C)
+
+        double total_count_; // Total instance count
+
+        bool weighted_a2de_; // Whether to use weighted A2DE
+        double smoothing_factor_; // Smoothing parameter (default: Laplace)
+        torch::Tensor AODEConditionalProb(const torch::Tensor& data);
+        void trainModel(const torch::Tensor& data, const Smoothing_t smoothing);
+        int toIntValue(int attributeIndex, float value) const;
+    };
+}
+#endif

bayesnet/network/Node.cc

@@ -93,36 +93,42 @@ namespace bayesnet {
     void Node::computeCPT(const torch::Tensor& dataset, const std::vector<std::string>& features, const double smoothing, const torch::Tensor& weights)
     {
         dimensions.clear();
+        dimensions.reserve(parents.size() + 1);
         // Get dimensions of the CPT
         dimensions.push_back(numStates);
-        transform(parents.begin(), parents.end(), back_inserter(dimensions), [](const auto& parent) { return parent->getNumStates(); });
+        for (const auto& parent : parents) {
-        // Create a tensor of zeros with the dimensions of the CPT
+            dimensions.push_back(parent->getNumStates());
-        cpTable = torch::zeros(dimensions, torch::kDouble) + smoothing;
+        }
-        // Fill table with counts
+        //transform(parents.begin(), parents.end(), back_inserter(dimensions), [](const auto& parent) { return parent->getNumStates(); });
-        auto pos = find(features.begin(), features.end(), name);
+        // Create a tensor initialized with smoothing
-        if (pos == features.end()) {
+        cpTable = torch::full(dimensions, smoothing, torch::kDouble);
-            throw std::logic_error("Feature " + name + " not found in dataset");
+        // Create a map for quick feature index lookup
+        std::unordered_map<std::string, int> featureIndexMap;
+        for (size_t i = 0; i < features.size(); ++i) {
+            featureIndexMap[features[i]] = i;
+        }
+        // Fill table with counts
+        // Get the index of this node's feature
+        int name_index = featureIndexMap[name];
+        // Get parent indices in dataset
+        std::vector<int> parent_indices;
+        parent_indices.reserve(parents.size());
+        for (const auto& parent : parents) {
+            parent_indices.push_back(featureIndexMap[parent->getName()]);
         }
-        int name_index = pos - features.begin();
         c10::List<c10::optional<at::Tensor>> coordinates;
         for (int n_sample = 0; n_sample < dataset.size(1); ++n_sample) {
             coordinates.clear();
             auto sample = dataset.index({ "...", n_sample });
             coordinates.push_back(sample[name_index]);
-            for (auto parent : parents) {
+            for (size_t i = 0; i < parent_indices.size(); ++i) {
-                pos = find(features.begin(), features.end(), parent->getName());
+                coordinates.push_back(sample[parent_indices[i]]);
-                if (pos == features.end()) {
-                    throw std::logic_error("Feature parent " + parent->getName() + " not found in dataset");
-                }
-                int parent_index = pos - features.begin();
-                coordinates.push_back(sample[parent_index]);
             }
             // Increment the count of the corresponding coordinate
             cpTable.index_put_({ coordinates }, weights.index({ n_sample }), true);
         }
-        // Normalize the counts
+        // Normalize the counts (dividing each row by the sum of the row)
-        // Divide each row by the sum of the row
+        cpTable /= cpTable.sum(0, true);
-        cpTable = cpTable / cpTable.sum(0);
     }
     double Node::getFactorValue(std::map<std::string, int>& evidence)
     {

sample/CMakeLists.txt

@@ -18,7 +18,7 @@ include_directories(
     ../tests/lib/Files
     lib/json/include
     /usr/local/include
-    ${FImdlp_INCLUDE_DIRS}
+    /usr/local/include/fimdlp/
 )
 
 add_executable(bayesnet_sample sample.cc) 

sample/sample.cc

@@ -60,7 +60,21 @@ int main(int argc, char* argv[])
     auto clf = bayesnet::BoostAODE(false); // false for not using voting in predict
     std::cout << "Library version: " << clf.getVersion() << std::endl;
     tie(X, y, features, className, states) = loadDataset(file_name, true);
-    clf.fit(X, y, features, className, states, bayesnet::Smoothing_t::LAPLACE);
+    torch::Tensor weights = torch::full({ X.size(1) }, 15, torch::kDouble);
+    torch::Tensor dataset;
+    try {
+        auto yresized = torch::transpose(y.view({ y.size(0), 1 }), 0, 1);
+        dataset = torch::cat({ X, yresized }, 0);
+    }
+    catch (const std::exception& e) {
+        std::stringstream oss;
+        oss << "* Error in X and y dimensions *\n";
+        oss << "X dimensions: " << dataset.sizes() << "\n";
+        oss << "y dimensions: " << y.sizes();
+        throw std::runtime_error(oss.str());
+    }
+    //Classifier& fit(torch::Tensor& dataset, const std::vector<std::string>& features, const std::string& className, std::map<std::string, std::vector<int>>& states, const torch::Tensor& weights, const Smoothing_t smoothing) override;
+    clf.fit(dataset, features, className, states, weights, bayesnet::Smoothing_t::LAPLACE);
     auto score = clf.score(X, y);
     std::cout << "File: " << file_name << " Model: BoostAODE score: " << score << std::endl;
     return 0;

tests/CMakeLists.txt

@@ -9,11 +9,12 @@ if(ENABLE_TESTING)
     )
     file(GLOB_RECURSE BayesNet_SOURCES "${BayesNet_SOURCE_DIR}/bayesnet/*.cc")
     add_executable(TestBayesNet TestBayesNetwork.cc TestBayesNode.cc TestBayesClassifier.cc 
-        TestBayesModels.cc TestBayesMetrics.cc TestFeatureSelection.cc TestBoostAODE.cc TestA2DE.cc 
+        TestBayesModels.cc TestBayesMetrics.cc TestFeatureSelection.cc TestBoostAODE.cc TestA2DE.cc TestWA2DE.cc
         TestUtils.cc TestBayesEnsemble.cc TestModulesVersions.cc TestBoostA2DE.cc TestMST.cc ${BayesNet_SOURCES})
     target_link_libraries(TestBayesNet PUBLIC "${TORCH_LIBRARIES}" fimdlp PRIVATE Catch2::Catch2WithMain)
     add_test(NAME BayesNetworkTest COMMAND TestBayesNet)
     add_test(NAME A2DE COMMAND TestBayesNet "[A2DE]")
+    add_test(NAME WA2DE COMMAND TestBayesNet "[WA2DE]")
     add_test(NAME BoostA2DE COMMAND TestBayesNet "[BoostA2DE]")
     add_test(NAME BoostAODE COMMAND TestBayesNet "[BoostAODE]")
     add_test(NAME Classifier COMMAND TestBayesNet "[Classifier]")

tests/TestBoostA2DE.cc

@@ -123,7 +123,7 @@ TEST_CASE("Oddities2", "[BoostA2DE]")
         { { "order", "duck" } },
         { { "select_features", "duck" } },
         { { "maxTolerance", 0 } },
-        { { "maxTolerance", 5 } },
+        { { "maxTolerance", 7 } },
     };
     for (const auto& hyper : bad_hyper.items()) {
         INFO("BoostA2DE hyper: " + hyper.value().dump());

tests/TestBoostAODE.cc

@@ -118,7 +118,7 @@ TEST_CASE("Oddities", "[BoostAODE]")
         { { "order", "duck" } },
         { { "select_features", "duck" } },
         { { "maxTolerance", 0 } },
-        { { "maxTolerance", 5 } },
+        { { "maxTolerance", 7 } },
     };
     for (const auto& hyper : bad_hyper.items()) {
         INFO("BoostAODE hyper: " << hyper.value().dump());

tests/TestWA2DE.cc

@@ -0,0 +1,31 @@
+// ***************************************************************
+// SPDX-FileCopyrightText: Copyright 2024 Ricardo Montañana Gómez
+// SPDX-FileType: SOURCE
+// SPDX-License-Identifier: MIT
+// ***************************************************************
+
+#include <type_traits>
+#include <catch2/catch_test_macros.hpp>
+#include <catch2/catch_approx.hpp>
+#include <catch2/generators/catch_generators.hpp>
+#include "bayesnet/ensembles/WA2DE.h"
+#include "TestUtils.h"
+
+
+TEST_CASE("Fit and Score", "[WA2DE]")
+{
+    auto raw = RawDatasets("iris", true);
+    auto clf = bayesnet::WA2DE();
+    clf.fit(raw.Xv, raw.yv, raw.features, raw.className, raw.states, raw.smoothing);
+    REQUIRE(clf.score(raw.Xt, raw.yt) == Catch::Approx(0.831776).epsilon(raw.epsilon));
+}
+TEST_CASE("Test graph", "[WA2DE]")
+{
+    auto raw = RawDatasets("iris", true);
+    auto clf = bayesnet::WA2DE();
+    clf.fit(raw.Xv, raw.yv, raw.features, raw.className, raw.states, raw.smoothing);
+    auto graph = clf.graph("BayesNet WA2DE");
+    REQUIRE(graph.size() == 2);
+    REQUIRE(graph[0] == "BayesNet WA2DE");
+    REQUIRE(graph[1] == "Graph visualization not implemented.");
+}