80 lines
2.6 KiB
C++
80 lines
2.6 KiB
C++
#include "Network.h"
|
|
namespace bayesnet {
|
|
Network::~Network()
|
|
{
|
|
for (auto& pair : nodes) {
|
|
delete pair.second;
|
|
}
|
|
}
|
|
void Network::addNode(std::string name, int numStates)
|
|
{
|
|
nodes[name] = new Node(name, numStates);
|
|
}
|
|
void Network::addEdge(const std::string parent, const std::string child)
|
|
{
|
|
if (nodes.find(parent) == nodes.end()) {
|
|
throw std::invalid_argument("Parent node " + parent + " does not exist");
|
|
}
|
|
if (nodes.find(child) == nodes.end()) {
|
|
throw std::invalid_argument("Child node " + child + " does not exist");
|
|
}
|
|
nodes[parent]->addChild(nodes[child]);
|
|
nodes[child]->addParent(nodes[parent]);
|
|
}
|
|
std::map<std::string, Node*>& Network::getNodes()
|
|
{
|
|
return nodes;
|
|
}
|
|
void Network::fit(const std::vector<std::vector<int>>& dataset, const int smoothing)
|
|
{
|
|
auto jointCounts = [](const std::vector<std::vector<int>>& data, const std::vector<int>& indices, int numStates) {
|
|
int size = indices.size();
|
|
std::vector<int64_t> sizes(size, numStates);
|
|
torch::Tensor counts = torch::zeros(sizes, torch::kLong);
|
|
|
|
for (const auto& row : data) {
|
|
int idx = 0;
|
|
for (int i = 0; i < size; ++i) {
|
|
idx = idx * numStates + row[indices[i]];
|
|
}
|
|
counts.view({ -1 }).add_(idx, 1);
|
|
}
|
|
|
|
return counts;
|
|
};
|
|
|
|
auto marginalCounts = [](const torch::Tensor& jointCounts) {
|
|
return jointCounts.sum(-1);
|
|
};
|
|
|
|
for (auto& pair : nodes) {
|
|
Node* node = pair.second;
|
|
|
|
std::vector<int> indices;
|
|
for (const auto& parent : node->getParents()) {
|
|
indices.push_back(nodes[parent->getName()]->getId());
|
|
}
|
|
indices.push_back(node->getId());
|
|
|
|
for (auto& child : node->getChildren()) {
|
|
torch::Tensor counts = jointCounts(dataset, indices, node->getNumStates()) + smoothing;
|
|
torch::Tensor parentCounts = marginalCounts(counts);
|
|
parentCounts = parentCounts.unsqueeze(-1);
|
|
|
|
torch::Tensor cpt = counts.to(torch::kDouble) / parentCounts.to(torch::kDouble);
|
|
setCPD(node->getCPDKey(child), cpt);
|
|
}
|
|
}
|
|
}
|
|
|
|
torch::Tensor& Network::getCPD(const std::string& key)
|
|
{
|
|
return cpds[key];
|
|
}
|
|
|
|
void Network::setCPD(const std::string& key, const torch::Tensor& cpt)
|
|
{
|
|
cpds[key] = cpt;
|
|
}
|
|
}
|