Add parent hyperparameter to TAN & SPODE

bayesnet/classifiers/SPODE.cc

@@ -8,14 +8,29 @@
 
 namespace bayesnet {
 
-    SPODE::SPODE(int root) : Classifier(Network()), root(root) {}
+    SPODE::SPODE(int root) : Classifier(Network()), root(root)
+    {
+        validHyperparameters = { "parent" };
+    }
 
+    void SPODE::setHyperparameters(const nlohmann::json& hyperparameters_)
+    {
+        auto hyperparameters = hyperparameters_;
+        if (hyperparameters.contains("parent")) {
+            root = hyperparameters["parent"];
+            hyperparameters.erase("parent");
+        }
+        Classifier::setHyperparameters(hyperparameters);
+    }
     void SPODE::buildModel(const torch::Tensor& weights)
     {
         // 0. Add all nodes to the model
         addNodes();
         // 1. Add edges from the class node to all other nodes
         // 2. Add edges from the root node to all other nodes
+        if (root >= static_cast<int>(features.size())) {
+            throw std::invalid_argument("The parent node is not in the dataset");
+        }
         for (int i = 0; i < static_cast<int>(features.size()); ++i) {
             model.addEdge(className, features[i]);
             if (i != root) {

bayesnet/classifiers/SPODE.h

@@ -10,14 +10,15 @@
 
 namespace bayesnet {
     class SPODE : public Classifier {
-    private:
-        int root;
-    protected:
-        void buildModel(const torch::Tensor& weights) override;
     public:
         explicit SPODE(int root);
         virtual ~SPODE() = default;
+        void setHyperparameters(const nlohmann::json& hyperparameters_) override;
         std::vector<std::string> graph(const std::string& name = "SPODE") const override;
+    protected:
+        void buildModel(const torch::Tensor& weights) override;
+    private:
+        int root;
     };
 }
 #endif

bayesnet/classifiers/TAN.cc

@@ -7,8 +7,20 @@
 #include "TAN.h"
 
 namespace bayesnet {
-    TAN::TAN() : Classifier(Network()) {}
+    TAN::TAN() : Classifier(Network())
+    {
+        validHyperparameters = { "parent" };
+    }
 
+    void TAN::setHyperparameters(const nlohmann::json& hyperparameters_)
+    {
+        auto hyperparameters = hyperparameters_;
+        if (hyperparameters.contains("parent")) {
+            parent = hyperparameters["parent"];
+            hyperparameters.erase("parent");
+        }
+        Classifier::setHyperparameters(hyperparameters);
+    }
     void TAN::buildModel(const torch::Tensor& weights)
     {
         // 0. Add all nodes to the model
@@ -23,7 +35,10 @@ namespace bayesnet {
             mi.push_back({ i, mi_value });
         }
         sort(mi.begin(), mi.end(), [](const auto& left, const auto& right) {return left.second < right.second;});
-        auto root = mi[mi.size() - 1].first;
+        auto root = parent == -1 ? mi[mi.size() - 1].first : parent;
+        if (root >= static_cast<int>(features.size())) {
+            throw std::invalid_argument("The parent node is not in the dataset");
+        }
         // 2. Compute mutual information between each feature and the class
         auto weights_matrix = metrics.conditionalEdge(weights);
         // 3. Compute the maximum spanning tree

bayesnet/classifiers/TAN.h

@@ -9,13 +9,15 @@
 #include "Classifier.h"
 namespace bayesnet {
     class TAN : public Classifier {
-    private:
-    protected:
-        void buildModel(const torch::Tensor& weights) override;
     public:
         TAN();
         virtual ~TAN() = default;
+        void setHyperparameters(const nlohmann::json& hyperparameters_) override;
         std::vector<std::string> graph(const std::string& name = "TAN") const override;
+    protected:
+        void buildModel(const torch::Tensor& weights) override;
+    private:
+        int parent = -1;
     };
 }
 #endif