Example of component

import os
from time import time

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

import appletree as apt
from appletree.utils import get_file_path

Using Normal as an approximation of Binomial
Using aptext package from https://github.com/XENONnT/applefiles

# constrain the GPU memory usage

apt.set_gpu_memory_usage(0.2)

Define component

ComponentSim

# The components is associated with bins, so first we load bins
data = pd.read_csv(get_file_path("data_Rn220.csv"))
bins_cs1, bins_cs2 = apt.utils.get_equiprob_bins_2d(
    data[["cs1", "cs2"]].to_numpy(),
    [15, 15],
    order=[0, 1],
    x_clip=[0, 100],
    y_clip=[1e2, 1e4],
    which_np=jnp,
)

# Initialize component
er = apt.ERBand(bins=[bins_cs1, bins_cs2], bins_type="irreg")

# Deduce the workflow(datastructure)
er.deduce(data_names=["cs1", "cs2"], func_name="simulate")  # 'eff'(efficiency) is always simulated
er.rate_name = "er_rate"  # also we have to specify a normalization factor of the component

# Compile ER script
# This is meta-programing because  appletree can generate codes dynamically
er.compile()

# For reference, this is the compiled code, the function is stored in appletree.share._cached_functions
print(er.code)

from functools import partial
from jax import jit
from appletree.plugins import PositionSpectra
from appletree.plugins import UniformEnergySpectra
from appletree.plugins import RecombFluct
from appletree.plugins import mTI
from appletree.plugins import Quanta
from appletree.plugins import TrueRecombER
from appletree.plugins import IonizationER
from appletree.plugins import DriftLoss
from appletree.plugins import RecombinationER
from appletree.plugins import ElectronDrifted
from appletree.plugins import PositionRecon
from appletree.plugins import S2Correction
from appletree.plugins import S2PE
from appletree.plugins import S2
from appletree.plugins import S1Correction
from appletree.plugins import PhotonDetection
from appletree.plugins import S1PE
from appletree.plugins import S1
from appletree.plugins import S2CutAccept
from appletree.plugins import S1CutAccept
from appletree.plugins import S1ReconEff
from appletree.plugins import S2Threshold
from appletree.plugins import Eff
from appletree.plugins import cS2
from appletree.plugins import cS1
PositionSpectra_ERBand = PositionSpectra('ERBand_llh')
UniformEnergySpectra_ERBand = UniformEnergySpectra('ERBand_llh')
RecombFluct_ERBand = RecombFluct('ERBand_llh')
mTI_ERBand = mTI('ERBand_llh')
Quanta_ERBand = Quanta('ERBand_llh')
TrueRecombER_ERBand = TrueRecombER('ERBand_llh')
IonizationER_ERBand = IonizationER('ERBand_llh')
DriftLoss_ERBand = DriftLoss('ERBand_llh')
RecombinationER_ERBand = RecombinationER('ERBand_llh')
ElectronDrifted_ERBand = ElectronDrifted('ERBand_llh')
PositionRecon_ERBand = PositionRecon('ERBand_llh')
S2Correction_ERBand = S2Correction('ERBand_llh')
S2PE_ERBand = S2PE('ERBand_llh')
S2_ERBand = S2('ERBand_llh')
S1Correction_ERBand = S1Correction('ERBand_llh')
PhotonDetection_ERBand = PhotonDetection('ERBand_llh')
S1PE_ERBand = S1PE('ERBand_llh')
S1_ERBand = S1('ERBand_llh')
S2CutAccept_ERBand = S2CutAccept('ERBand_llh')
S1CutAccept_ERBand = S1CutAccept('ERBand_llh')
S1ReconEff_ERBand = S1ReconEff('ERBand_llh')
S2Threshold_ERBand = S2Threshold('ERBand_llh')
Eff_ERBand = Eff('ERBand_llh')
cS2_ERBand = cS2('ERBand_llh')
cS1_ERBand = cS1('ERBand_llh')

@partial(jit, static_argnums=(1, ))
def simulate(key, batch_size, parameters):
    key, x, y, z = PositionSpectra_ERBand(key, parameters, batch_size)
    key, energy = UniformEnergySpectra_ERBand(key, parameters, batch_size)
    key, recomb_std = RecombFluct_ERBand(key, parameters, energy)
    key, recomb_mean = mTI_ERBand(key, parameters, energy)
    key, num_quanta = Quanta_ERBand(key, parameters, energy)
    key, recomb = TrueRecombER_ERBand(key, parameters, recomb_mean, recomb_std)
    key, num_ion = IonizationER_ERBand(key, parameters, num_quanta)
    key, drift_survive_prob = DriftLoss_ERBand(key, parameters, z)
    key, num_photon, num_electron = RecombinationER_ERBand(key, parameters, num_quanta, num_ion, recomb)
    key, num_electron_drifted = ElectronDrifted_ERBand(key, parameters, num_electron, drift_survive_prob)
    key, rec_x, rec_y, rec_z, rec_r = PositionRecon_ERBand(key, parameters, x, y, z, num_electron_drifted)
    key, s2_correction = S2Correction_ERBand(key, parameters, rec_x, rec_y)
    key, num_s2_pe = S2PE_ERBand(key, parameters, num_electron_drifted, s2_correction)
    key, s2_area = S2_ERBand(key, parameters, num_s2_pe)
    key, s1_correction = S1Correction_ERBand(key, parameters, rec_x, rec_y, rec_z)
    key, num_s1_phd = PhotonDetection_ERBand(key, parameters, num_photon, s1_correction)
    key, num_s1_pe = S1PE_ERBand(key, parameters, num_s1_phd)
    key, s1_area = S1_ERBand(key, parameters, num_s1_phd, num_s1_pe)
    key, cut_acc_s2 = S2CutAccept_ERBand(key, parameters, s2_area)
    key, cut_acc_s1 = S1CutAccept_ERBand(key, parameters, s1_area)
    key, acc_s1_recon_eff = S1ReconEff_ERBand(key, parameters, num_s1_phd)
    key, acc_s2_threshold = S2Threshold_ERBand(key, parameters, s2_area)
    key, eff = Eff_ERBand(key, parameters, acc_s2_threshold, acc_s1_recon_eff, cut_acc_s1, cut_acc_s2)
    key, cs2 = cS2_ERBand(key, parameters, s2_area, s2_correction, drift_survive_prob)
    key, cs1 = cS1_ERBand(key, parameters, s1_area, s1_correction)
    return key, [cs1, cs2, eff]

ComponentFixed

# Initialize component, not based on simulation, but the input template
ac = apt.AC(bins=[bins_cs1, bins_cs2], bins_type="irreg", file_name="AC_Rn220.pkl")
ac.rate_name = "ac_rate"

# Do not forget to deduce
ac.deduce()

# Of course these initialization process looks messy, but we will do the initialization automatically in context

Load parameters

# Of course we have to load parameters(and their priors) in simulation

par_manager = apt.Parameter(get_file_path("er.json"))
par_manager.sample_init()
parameters = par_manager.get_all_parameter()

Simulation

# Really do the simulation

batch_size = int(1e6)
key = apt.randgen.get_key(seed=137)

key, (cs1, cs2, eff) = er.simulate(key, batch_size, parameters)

# Just to show the histogram

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.yscale("log")
plt.show()

Simulation and make equiprob hist

# Actually `simulate_hist` is just a wrapper of `simulate`

batch_size = int(1e6)
key = apt.randgen.get_key(seed=137)
key, h = er.simulate_hist(key, batch_size, parameters)

apt.utils.plot_irreg_histogram_2d(bins_cs1, bins_cs2, h, density=False)
plt.yscale("log")
plt.ylim(5e2, 1e4)
plt.show()

h = ac.simulate_hist(parameters)

apt.utils.plot_irreg_histogram_2d(bins_cs1, bins_cs2, h, density=False)
plt.yscale("log")
plt.ylim(5e2, 1e4)
plt.show()

Speed test

@apt.utils.timeit
def test(key, batch_size, parameters):
    return er.simulate_hist(key, batch_size, parameters)

@apt.utils.timeit
def benchmark():
    key = apt.randgen.get_key()
    for _ in range(100):
        key, _ = test(key, int(1e6), parameters)

benchmark()

 Function <benchmark> starts.
 Function <test> starts.
 Function <test> ends! Time cost = 8.48 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.53 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.02 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.70 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.23 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.32 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.84 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.98 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.15 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.27 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.68 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.86 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.78 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.88 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.76 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.63 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.32 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.68 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.80 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.42 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.34 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.96 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.27 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.65 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.82 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.28 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.91 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.11 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.19 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.61 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.04 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.81 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.02 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.60 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.13 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.87 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.91 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.58 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.89 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.76 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.29 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.63 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.72 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.47 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.44 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.00 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.77 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.67 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.90 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.52 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.80 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.01 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.44 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.75 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.30 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.66 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.84 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.30 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.73 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.04 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.72 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.51 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.50 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.19 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.97 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.05 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.44 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.41 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.43 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.71 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.42 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.94 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.88 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 7.84 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.60 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.21 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.22 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.73 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.78 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.65 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 3.48 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.77 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.53 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.89 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.75 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.65 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.67 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 6.39 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 8.08 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.54 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 9.05 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 5.62 msec.
 Function <test> starts.
 Function <test> ends! Time cost = 4.42 msec.
 Function <benchmark> ends! Time cost = 582.11 msec.