cs50ai_bayesian_network_pgmpy.py

# based on edx harvard cs50ai uncertainty lesson source code
from pgmpy.models import BayesianNetwork
from pgmpy.factors.discrete import (
    TabularCPD,
)
from pgmpy.inference import VariableElimination


model = BayesianNetwork(
    [
        ("rain", "maintenance"),
        ("maintenance", "train"),
        ("rain", "train"),
        ("train", "appointment"),
    ]
)

rain_cpd = TabularCPD(
    variable="rain",
    variable_card=3,
    values=[
        [0.7],
        [0.2],
        [0.1],
    ],
)

maintenance_cpd = TabularCPD(
    variable="maintenance",
    variable_card=2,
    values=[
        [0.4, 0.2, 0.1],
        [0.6, 0.8, 0.9],
    ],
    evidence=["rain"],
    evidence_card=[3],
)

train_cpd = TabularCPD(
    variable="train",
    variable_card=2,
    values=[
        [0.8, 0.9, 0.6, 0.7, 0.4, 0.5],
        [0.2, 0.1, 0.4, 0.3, 0.6, 0.5],
    ],
    evidence=["rain", "maintenance"],
    evidence_card=[3, 2],
)


appointment_cpd = TabularCPD(
    variable="appointment",
    variable_card=2,
    values=[
        [0.9, 0.6],
        [0.1, 0.4],
    ],
    evidence=["train"],
    evidence_card=[2],
)

model.add_cpds(
    rain_cpd,
    maintenance_cpd,
    train_cpd,
    appointment_cpd,
)

model_infer = VariableElimination(model)

prob = model_infer.query(variables=["rain"], evidence={"train": 1})
print(prob)

# +---------+-------------+
# | rain    |   phi(rain) |
# +=========+=============+
# | rain(0) |      0.4601 |
# +---------+-------------+
# | rain(1) |      0.3005 |
# +---------+-------------+
# | rain(2) |      0.2394 |
# +---------+-------------+