Refactor folder structure of the project

src/CMakeLists.txt

@@ -0,0 +1,2 @@
+add_library(BayesNet Network.cc Node.cc)
+target_link_libraries(BayesNet "${TORCH_LIBRARIES}")

src/Network.cc

@@ -0,0 +1,131 @@
+#include "Network.h"
+namespace bayesnet {
+    Network::Network() : laplaceSmoothing(1), root(nullptr), features(vector<string>()), className("") {}
+    Network::Network(int smoothing) : laplaceSmoothing(smoothing), root(nullptr), features(vector<string>()), className("") {}
+    Network::~Network()
+    {
+        for (auto& pair : nodes) {
+            delete pair.second;
+        }
+    }
+    void Network::addNode(string name, int numStates)
+    {
+        if (nodes.find(name) != nodes.end()) {
+            throw invalid_argument("Node " + name + " already exists");
+        }
+        nodes[name] = new Node(name, numStates);
+        if (root == nullptr) {
+            root = nodes[name];
+        }
+    }
+    void Network::setRoot(string name)
+    {
+        if (nodes.find(name) == nodes.end()) {
+            throw invalid_argument("Node " + name + " does not exist");
+        }
+        root = nodes[name];
+    }
+    Node* Network::getRoot()
+    {
+        return root;
+    }
+    bool Network::isCyclic(const string& nodeId, unordered_set<string>& visited, unordered_set<string>& recStack)
+    {
+        if (visited.find(nodeId) == visited.end()) // if node hasn't been visited yet
+        {
+            visited.insert(nodeId);
+            recStack.insert(nodeId);
+            for (Node* child : nodes[nodeId]->getChildren()) {
+                if (visited.find(child->getName()) == visited.end() && isCyclic(child->getName(), visited, recStack))
+                    return true;
+                else if (recStack.find(child->getName()) != recStack.end())
+                    return true;
+            }
+        }
+        recStack.erase(nodeId); // remove node from recursion stack before function ends
+        return false;
+    }
+    void Network::addEdge(const string parent, const string child)
+    {
+        if (nodes.find(parent) == nodes.end()) {
+            throw invalid_argument("Parent node " + parent + " does not exist");
+        }
+        if (nodes.find(child) == nodes.end()) {
+            throw invalid_argument("Child node " + child + " does not exist");
+        }
+        // Temporarily add edge to check for cycles
+        nodes[parent]->addChild(nodes[child]);
+        nodes[child]->addParent(nodes[parent]);
+        // temporarily add edge
+        unordered_set<string> visited;
+        unordered_set<string> recStack;
+        if (isCyclic(nodes[child]->getName(), visited, recStack)) // if adding this edge forms a cycle
+        {
+            // remove problematic edge
+            nodes[parent]->removeChild(nodes[child]);
+            nodes[child]->removeParent(nodes[parent]);
+            throw invalid_argument("Adding this edge forms a cycle in the graph.");
+        }
+
+    }
+    map<string, Node*>& Network::getNodes()
+    {
+        return nodes;
+    }
+    void Network::buildNetwork()
+    {
+        // Add features as nodes to the network
+        for (int i = 0; i < features.size(); ++i) {
+            addNode(features[i], *max_element(dataset[features[i]].begin(), dataset[features[i]].end()) + 1);
+        }
+        // Add class as node to the network
+        addNode(className, *max_element(dataset[className].begin(), dataset[className].end()) + 1);
+        // Add edges from class to features => naive Bayes
+        for (auto feature : features) {
+            addEdge(className, feature);
+        }
+        addEdge("petalwidth", "petallength");
+    }
+    void Network::fit(const vector<vector<int>>& dataset, const vector<int>& labels, const vector<string>& featureNames, const string& className)
+    {
+        features = featureNames;
+        this->className = className;
+        // Build dataset
+        for (int i = 0; i < featureNames.size(); ++i) {
+            this->dataset[featureNames[i]] = dataset[i];
+        }
+        this->dataset[className] = labels;
+        buildNetwork();
+        estimateParameters();
+    }
+
+    void Network::estimateParameters()
+    {
+        auto dimensions = vector<int64_t>();
+        for (auto [name, node] : nodes) {
+            // Get dimensions of the CPT
+            dimensions.clear();
+            dimensions.push_back(node->getNumStates());
+            for (auto father : node->getParents()) {
+                dimensions.push_back(father->getNumStates());
+            }
+            auto length = dimensions.size();
+            // Create a tensor of zeros with the dimensions of the CPT
+            torch::Tensor cpt = torch::zeros(dimensions, torch::kFloat) + laplaceSmoothing;
+            // Fill table with counts
+            for (int n_sample = 0; n_sample < dataset[name].size(); ++n_sample) {
+                torch::List<c10::optional<torch::Tensor>> coordinates;
+                coordinates.push_back(torch::tensor(dataset[name][n_sample]));
+                for (auto father : node->getParents()) {
+                    coordinates.push_back(torch::tensor(dataset[father->getName()][n_sample]));
+                }
+                // Increment the count of the corresponding coordinate
+                cpt.index_put_({ coordinates }, cpt.index({ coordinates }) + 1);
+            }
+            // Normalize the counts
+            cpt = cpt / cpt.sum(0);
+            // store thre resulting cpt in the node
+            node->setCPT(cpt);
+        }
+    }
+}

src/Network.h

@@ -0,0 +1,32 @@
+#ifndef NETWORK_H
+#define NETWORK_H
+#include "Node.h"
+#include <map>
+#include <vector>
+
+
+namespace bayesnet {
+    class Network {
+    private:
+        map<string, Node*> nodes;
+        map<string, vector<int>> dataset;
+        Node* root;
+        vector<string> features;
+        string className;
+        int laplaceSmoothing;
+        bool isCyclic(const std::string&, std::unordered_set<std::string>&, std::unordered_set<std::string>&);
+    public:
+        Network();
+        Network(int);
+        ~Network();
+        void addNode(string, int);
+        void addEdge(const string, const string);
+        map<string, Node*>& getNodes();
+        void fit(const vector<vector<int>>&, const vector<int>&, const vector<string>&, const string&);
+        void estimateParameters();
+        void buildNetwork();
+        void setRoot(string);
+        Node* getRoot();
+    };
+}
+#endif

src/Node.cc

@@ -0,0 +1,52 @@
+#include "Node.h"
+
+namespace bayesnet {
+    int Node::next_id = 0;
+
+    Node::Node(const std::string& name, int numStates)
+        : id(next_id++), name(name), numStates(numStates), cpt(torch::Tensor()), parents(vector<Node*>()), children(vector<Node*>())
+    {
+    }
+
+    string Node::getName() const
+    {
+        return name;
+    }
+
+    void Node::addParent(Node* parent)
+    {
+        parents.push_back(parent);
+    }
+    void Node::removeParent(Node* parent)
+    {
+        parents.erase(std::remove(parents.begin(), parents.end(), parent), parents.end());
+    }
+    void Node::removeChild(Node* child)
+    {
+        children.erase(std::remove(children.begin(), children.end(), child), children.end());
+    }
+    void Node::addChild(Node* child)
+    {
+        children.push_back(child);
+    }
+    vector<Node*>& Node::getParents()
+    {
+        return parents;
+    }
+    vector<Node*>& Node::getChildren()
+    {
+        return children;
+    }
+    int Node::getNumStates() const
+    {
+        return numStates;
+    }
+    torch::Tensor& Node::getCPT()
+    {
+        return cpt;
+    }
+    void Node::setCPT(const torch::Tensor& cpt)
+    {
+        this->cpt = cpt;
+    }
+}

src/Node.h

@@ -0,0 +1,33 @@
+#ifndef NODE_H
+#define NODE_H
+#include <torch/torch.h>
+#include <vector>
+#include <string>
+namespace bayesnet {
+    using namespace std;
+    class Node {
+    private:
+        static int next_id;
+        const int id;
+        string name;
+        vector<Node*> parents;
+        vector<Node*> children;
+        torch::Tensor cpTable;
+        int numStates;
+        torch::Tensor cpt;
+    public:
+        Node(const std::string&, int);
+        void addParent(Node*);
+        void addChild(Node*);
+        void removeParent(Node*);
+        void removeChild(Node*);
+        string getName() const;
+        vector<Node*>& getParents();
+        vector<Node*>& getChildren();
+        torch::Tensor& getCPT();
+        void setCPT(const torch::Tensor&);
+        int getNumStates() const;
+        int getId() const { return id; }
+    };
+}
+#endif