run_model.py

import gefry3 # My model code
import pymc

import numpy as np
import seaborn as sb

from scipy.stats import multivariate_normal

# P is an object that encompasses the geometry and problem, can be
# called to compute detector response
P = gefry3.read_input('g3_deck_varxs.yml')


# The actual location and intensity (used for starting the chains)
S0 = np.array([158., 98.])
I0 = 3.214e9
BG = 300

NS = int(1e4) # Number of samples

XMIN, YMIN, XMAX, YMAX = P.domain.all.bounds
IMIN, IMAX = 1e9, 5e9

# Relative perturbation used for all cross sections
XS_DELTA = 0.25

DWELL = np.array([i.dwell for i in P.detectors])

# Call P at the nominal values to get the real response
nominal = P(
    S0,
    I0,
    P.interstitial_material,
    P.materials,
)

nominal += BG * DWELL

# Generate the data
data = np.random.poisson(nominal)
C = np.diag(data)

# Model factory builds the dict of variables to pass to PyMC
def model_factory():
    x = pymc.Uniform("x", value=S0[0], lower=XMIN, upper=XMAX)
    y = pymc.Uniform("y", value=S0[1], lower=YMIN, upper=YMAX)
    I = pymc.Uniform("I", value=I0, lower=IMIN, upper=IMAX)
 
    # interstitial_xs = pymc.Normal(
        # "Sigma_inter",
        # P.interstitial_material.Sigma_T,
        # (P.interstitial_material.Sigma_T * XS_DELTA) ** 2,
        # value=P.interstitial_material.Sigma_T,
        # observed=True,
    # )

    s_i_xs = P.interstitial_material.Sigma_T
    interstitial_xs = pymc.Uniform(
        "Sigma_inter",
        s_i_xs * (1 - XS_DELTA),
        s_i_xs * (1 + XS_DELTA),
        value=s_i_xs,
        observed=True,
    )

    mu_xs = np.array([M.Sigma_T for M in P.materials])
    # var_xs = np.diag([(M * XS_DELTA) ** 2. for M in mu_xs])

    # building_xs = pymc.MvNormalCov(
        # "Sigma",
        # mu_xs,
        # var_xs,
        # value=mu_xs,
        # observed=True,
    # )

    building_xs = pymc.Uniform(
        "Sigma",
        mu_xs * (1 - XS_DELTA),
        mu_xs * (1 + XS_DELTA),
        value=mu_xs,
        observed=True,
    )

    @pymc.deterministic(plot=False)
    def model_pred(x=x, y=y, I=I, interstitial_xs=interstitial_xs, building_xs=building_xs):
        inter_mat = gefry3.Material(1.0, interstitial_xs)
        building_mats = [gefry3.Material(1.0, s) for s in building_xs]

        return P(
            [x, y],
            I,
            inter_mat,
            building_mats,
        )

    background = pymc.Poisson(
        "b",
        DWELL * BG,
        value=DWELL * BG,
        observed=True,
        plot=False,
    ) 
    
    @pymc.stochastic(plot=False, observed=True)
    def observed_response(value=nominal, model_pred=model_pred, background=background):
        resp = model_pred + background
        return multivariate_normal.logpdf(resp, mean=data, cov=C)

    return {
        "x": x,
        "y": y,
        "I": I,
        "interstitial_xs": interstitial_xs,
        "building_xs": building_xs,
        "model_pred": model_pred,
        "background": background,
        "observed_response": observed_response,
    } 

mvars = model_factory()
M = pymc.MCMC(mvars)

# This sets AdaptiveMetropolis for all variables
# Note that observed variables won't actually use the sampler, but
# it's fine to set it on everything. It just ignores ones that aren't
# being sampled.
M.use_step_method(
    pymc.AdaptiveMetropolis,
    [mvars[i] for i in mvars]
)

M.sample(NS)

res = np.vstack([M.trace(z)[:] for z in ["x", "y", "I"]])
np.savetxt("out_25pct.dat", res.T)

print("\n\n==== Results ====\n")
print("x: {}".format(np.mean(M.trace("x")[:])))
print("y: {}".format(np.mean(M.trace("y")[:])))
print("I: {}".format(np.mean(M.trace("I")[:])))