CS50AI - Lecture 2 - Bayesian Network updated source code

inference.py

import torch
from model import model

X = torch.tensor(
    [
        [
            -1,
            -1,
            1, # delayed
            -1,
        ]
    ]
)

X_masked = torch.masked.MaskedTensor(X, mask=(X != -1))

states = (
    ("rain", ["none", "light", "heavy"]),
    ("maintenance", ["yes", "no"]),
    ("train", ["on time", "delayed"]),
    ("appointment", ["attend", "miss"]),
)

# Calculate predictions
predictions = model.predict_proba(X_masked)

# Print predictions for each node
for (node_name, values), prediction in zip(states, predictions):
    if isinstance(prediction, str):
        print(f"{node_name}: {prediction}")
    else:
        print(f"{node_name}")
        for value, probability in zip(values, prediction[0]):
            print(f"    {value}: {probability:.4f}")

likelihood.py

import numpy
import torch
from model import model


rain_values = ["none", "light", "heavy"]
maintenance_values = ["yes", "no"]
train_values = ["on time", "delayed"]
appoinment_values = ["attend", "miss"]


probability = model.probability(
    torch.as_tensor(
        [
            [
                rain_values.index("none"),
                maintenance_values.index("no"),
                train_values.index("on time"),
                appoinment_values.index("attend"),
            ]
        ]
    )
)

print(probability)

model.py

from pomegranate import *

import numpy as np
from pomegranate.distributions import *
from pomegranate.bayesian_network import BayesianNetwork

rain = Categorical(
    [
        [0.7, 0.2, 0.1],
    ]
)

maintenance = ConditionalCategorical(
    [
        [
            [0.4, 0.6],
            [0.2, 0.8],
            [0.1, 0.9],
        ],
    ]
)

train = ConditionalCategorical(
    [
        [
            [
                [0.8, 0.2],
                [0.9, 0.1],
            ],
            [
                [0.6, 0.4],
                [0.7, 0.3],
            ],
            [
                [0.4, 0.6],
                [0.5, 0.5],
            ],
        ]
    ]
)


appointment = ConditionalCategorical(
    [
        [
            [0.9, 0.1],
            [0.6, 0.4],
        ],
    ]
)


# Create a Bayesian Network and add states
model = BayesianNetwork()
model.add_distributions([rain, maintenance, train, appointment])

# Add edges connecting nodes
model.add_edge(rain, maintenance)
model.add_edge(rain, train)
model.add_edge(maintenance, train)
model.add_edge(train, appointment)

sample.py

from pomegranate.distributions import ConditionalCategorical

from collections import Counter

from model import model

# Rejection sampling
# Compute distribution of Appointment given that train is delayed
N = 10000
data = []
for i in range(N):
    sample = model.sample(1)[0]
    # sample == "delayed"
    if sample[2] == 1.0:
        data.append("attend" if sample[3] == 0 else "miss")
print(Counter(data))