Example of context

[1]:

import numpy as np
import jax.numpy as jnp
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
import multihist as mh

import GOFevaluation

import appletree as apt
from appletree.utils import get_file_path

Using aptext package from https://github.com/XENONnT/applefiles

[2]:

apt.set_gpu_memory_usage(0.2)

Define context

[3]:

# Load configuration file
config = get_file_path("rn220.json")

# Initialize context
tree = apt.Context(config)

[4]:

# To see all the likelihoods

tree.print_context_summary(short=True)


========================================
LIKELIHOOD rn220_llh

----------------------------------------
BINNING

    bins_type: equiprob
    bins_on: ['cs1', 'cs2']

----------------------------------------
DATA

    file_name: data_Rn220.csv
    data_rate: 2000.0

----------------------------------------
MODEL

    COMPONENT 0: rn220_er
        type: simulation
        rate_par: rn220_er_rate
        pars: {'nex_ni_ratio', 'py4', 's2_cut_acc_sigma', 'fano', 'drift_velocity', 'p_dpe', 's1_cut_acc_sigma', 's1_eff_3f_sigma', 'gas_gain', 'py2', 'elife_sigma', 'rn220_er_rate', 'w', 'g1', 'g2', 'rf0', 's2_threshold', 'field', 'py1', 'rf1', 'py3', 'py0'}

    COMPONENT 1: rn220_ac
        type: fixed
        file_name: AC_Rn220.pkl
        rate_par: rn220_ac_rate
        pars: {'rn220_ac_rate'}

----------------------------------------

========================================

Fit, run the MCMC

[5]:

result = tree.fitting(nwalkers=100, iteration=10)

With no backend

100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [02:54<00:00, 17.42s/it]

[6]:

logp = tree.sampler.get_log_prob()

for _logp in logp.T:
    plt.plot(_logp, lw=0.1)

plt.xlabel("iteration")
plt.ylabel("log posterior")
plt.show()

Generate templates

[7]:

cs1, cs2, eff = tree.get_template("rn220_llh", "rn220_er")

[8]:

h, be = jnp.histogramdd(
    jnp.asarray([cs1, cs2]).T,
    bins=(jnp.linspace(0, 100, 101), jnp.logspace(2.5, 4.1, 81)),
    weights=eff,
)

h = mh.Histdd.from_histogram(np.array(h), be, axis_names=["cs1", "cs2"])
h.plot(norm=LogNorm())
plt.scatter(*tree["rn220_llh"].data.T, color="r", s=2.0)
plt.yscale("log")
plt.show()

[9]:

parameters = tree.get_post_parameters()
key = apt.randgen.get_key()
batch_size = int(1e6)

[10]:

key, result = tree["rn220_llh"]["rn220_er"].simulate(key, batch_size, parameters)

[11]:

data_sample = np.array(result[:-1]).T
x_clip = tree["rn220_llh"]._config["x_clip"]
y_clip = tree["rn220_llh"]._config["y_clip"]

n, bin_edges = GOFevaluation.utils.equiprobable_histogram(
    tree["rn220_llh"].data,
    data_sample,
    tree["rn220_llh"]._config["bins"],
    order=[0, 1],
    plot=True,
    nevents_expected=len(tree["rn220_llh"].data),
    plot_xlim=x_clip,
    plot_ylim=y_clip,
    plot_mode="sigma_deviation",
)
plt.yscale("log")

/home/xudc/GOFevaluation/GOFevaluation/utils.py:214: UserWarning: reference_sample contains ties, this might cause problems!
  check_for_ties(reference_sample)

Number of events in equal-probability binning

[12]:

key, h = tree["rn220_llh"]["rn220_er"].simulate_hist(key, batch_size, parameters)

[13]:

apt.utils.plot_irreg_histogram_2d(*tree["rn220_llh"]["rn220_er"].bins, h, density=False)
plt.yscale("log")

plt.show()