sprinkler.py

"""
This is a version of the Sprinkler Bayesian Network example from
https://healthyalgorithms.com/2011/11/23/causal-modeling-in-python-bayesian-networks-in-pymc/
translated from pymc to pymc3.
"""

from pymc3 import Model, Bernoulli, sample, traceplot
from pymc3.math import switch
from matplotlib.pyplot import savefig

Nsamples = 10000

with Model() as sprinkler_model:
    Rain = Bernoulli('Rain', .2)
    pSprinkler = switch( Rain, 0.01, 0.4)
    Sprinkler = Bernoulli('Sprinkler', pSprinkler)
    pWetGrass = switch( Rain,
                        switch( Sprinkler, 0.99, 0.8 ),
                        # need to apply some smoothing:
                        switch( Sprinkler, .9, 1/Nsamples ) )
    WetGrass = Bernoulli('WetGrass', pWetGrass, observed=1 )
    trace = sample(Nsamples)
    traceplot(trace)
    savefig('trace_{}.png'.format(Nsamples))