appletree package

Subpackages

Submodules

appletree.component module

class appletree.component.Component(name: Optional[str] = None, llh_name: Optional[str] = None, **kwargs)[source]

Bases: object

Base class of component.

__init__(name: Optional[str] = None, llh_name: Optional[str] = None, **kwargs)[source]

Initialization.

Parameters
  • bins – bins to generate the histogram. * For irreg bins_type, bins must be bin edges of the two dimensions. * For meshgrid bins_type, bins are sent to jnp.histogramdd.

  • bins_type – binning scheme, can be either irreg or meshgrid.

_clip(result: list)[source]

Clip simulated result.

add_eps_to_hist: bool = True
compile()[source]

Hook for compiling simulation code.

deduce(*args, **kwargs)[source]

Hook for workflow deduction.

force_no_eff: bool = False
get_normalization(hist, parameters, batch_size=None)[source]

Return the normalization factor of the histogram.

implement_binning(mc, eff)[source]

Apply binning to MC data.

Parameters
  • mc – data from simulation.

  • eff – efficiency of each event, as the weight when making a histogram.

property lineage
property lineage_hash
multiple_simulations(key, batch_size, parameters, times, apply_eff=False)[source]

Simulate many times and move results to CPU because the memory limit of GPU.

multiple_simulations_compile(key, batch_size, parameters, times, apply_eff=False)[source]

Simulate many times after new compilation and move results to CPU because the memory limit of GPU.

needed_parameters: Set[str]
norm_type: str = ''
rate_name: str = ''
set_binning(**kwargs)[source]

Set binning of component.

simulate_hist(*args, **kwargs)[source]

Hook for simulation with histogram output.

class appletree.component.ComponentFixed(*args, **kwargs)[source]

Bases: Component

Component whose shape is fixed.

__init__(*args, **kwargs)[source]

Initialization.

deduce(data_names: Union[List[str], Tuple[str]] = ['cs1', 'cs2'])[source]

Deduce the needed parameters and make the fixed histogram.

property lineage
simulate()[source]

Fixed component does not need to simulate.

simulate_hist(parameters, *args, **kwargs)[source]

Return the fixed histogram.

simulate_weighted_data(parameters, *args, **kwargs)[source]

Simulate and return histogram.

class appletree.component.ComponentSim(*args, **kwargs)[source]

Bases: Component

Component that needs MC simulations.

__init__(*args, **kwargs)[source]

Initialization.

property code

Code of simulation function.

compile()[source]

Build simulation function and cache it to share._cached_functions.

deduce(data_names: Union[List[str], Tuple[str]] = ['cs1', 'cs2'], func_name: str = 'simulate', nodep_data_name: str = 'batch_size', force_no_eff: bool = False)[source]

Deduce workflow and code.

Parameters
  • data_names – data names that simulation will output.

  • func_name – name of the simulation function, used to cache it.

  • nodep_data_name – data_name without dependency will not be deduced.

  • force_no_eff – force to ignore the efficiency, used in yield prediction.

dependencies_deduce(data_names: Union[List[str], Tuple[str]] = ['cs1', 'cs2', 'eff'], dependencies: Optional[List[Dict]] = None, nodep_data_name: str = 'batch_size') list[source]

Deduce dependencies.

Parameters
  • data_names – data names that simulation will output.

  • dependencies – dependency tree.

  • nodep_data_name – data_name without dependency will not be deduced.

dependencies_simplify(dependencies)[source]

Simplify the dependencies.

flush_source_code(data_names: Union[List[str], Tuple[str]] = ['cs1', 'cs2', 'eff'], func_name: str = 'simulate', nodep_data_name: str = 'batch_size')[source]

Infer the simulation code from the dependency tree.

property lineage
new_component(llh_name: Optional[str] = None, pass_binning: bool = True)[source]

Generate new component with same binning, usually used on predicting yields.

register(plugin_class)[source]

Register a plugin to the component.

register_all(module)[source]

Register all plugins defined in module.

Can pass a list/tuple of modules to register all in each.

save_code(file_path)[source]

Save the code to file.

set_config(configs)[source]

Set new global configuration options.

Parameters

configs – dict, configuration file name or dictionary

show_config(data_names: Union[List[str], Tuple[str]] = ['cs1', 'cs2', 'eff'])[source]

Return configuration options that affect data_names.

Parameters

data_names – Data type name

simulate_hist(key, batch_size, parameters)[source]

Simulate and return histogram.

Parameters
  • key – key used for pseudorandom generator.

  • batch_size – number of events to be simulated.

  • parameters – a dictionary that contains all parameters needed in simulation.

simulate_weighted_data(key, batch_size, parameters)[source]

Simulate and return histogram.

appletree.component._add_component_extension(module, component, force=False)[source]

Add component to module.

appletree.component.add_component_extensions(module1, module2, force=False)[source]

Add components of module2 to module1.

appletree.config module

class appletree.config.Config(name: str, type: Union[type, tuple, list, str] = '<OMITTED>', default: Any = '<OMITTED>', help: str = '')[source]

Bases: object

Configuration option taken by a appletree plugin.

__init__(name: str, type: Union[type, tuple, list, str] = '<OMITTED>', default: Any = '<OMITTED>', help: str = '')[source]

Initialization.

Parameters
  • name – name of the map.

  • type – Excepted type of the option’s value.

  • default – Default value the option takes.

  • help – description of the map.

_resolve_cached_config(llh_name, transform=None)[source]

Look up config in cache, apply transform, resolve per-likelihood dict.

Parameters
  • llh_name – name of the likelihood component.

  • transform – optional callable applied to the default value before caching (e.g. get_file_path for Map configs).

Returns

The resolved scalar (or per-likelihood) config value.

build(llh_name: Optional[str] = None)[source]

Build configuration, set attributes to Config instance.

get_default()[source]

Get default value of configuration.

property lineage
property lineage_hash
llh_name: Optional[str] = None
required_parameter(llh_name=None)[source]
class appletree.config.Constant(name: str, type: Union[type, tuple, list, str] = '<OMITTED>', default: Any = '<OMITTED>', help: str = '')[source]

Bases: Config

Constant is a special config which takes only certain value.

build(llh_name: Optional[str] = None)[source]

Set value of Constant.

property lineage
value = None
class appletree.config.ConstantSet(name: str, type: Union[type, tuple, list, str] = '<OMITTED>', default: Any = '<OMITTED>', help: str = '')[source]

Bases: Config

ConstantSet is a special config which takes a set of values.

We will not specify any hard-coded distribution or function here. User should be careful with the actual function implemented. Fortunately, we only use these values as keyword arguments, so mismatch will be catched when running.

_sanity_check()[source]

Check if parameter set lengths are same.

build(llh_name: Optional[str] = None)[source]

Set value of ConstantSet.

property lineage
class appletree.config.Map(method='IDW', **kwargs)[source]

Bases: Config

Map is a special config that takes input files.

The method apply is dynamically assigned. When using points, the apply will be map_point, while using regular binning, the apply will be map_regbin. When using log-binning, we will first convert the positions to log space.

_get_regbin_interpolator(ndim)[source]

Return the regular-binning interpolator for the given dimensionality.

_validate_log_coords(*coord_arrays)[source]

Raise if any log-scaled coordinate is non-positive.

build(llh_name: Optional[str] = None)[source]

Cache the map to jnp.array.

build_point(data)[source]

Cache the map to jnp.array if bins_type is point.

build_regbin(data)[source]

Cache the map to jnp.array if bins_type is regbin.

coordinate_type may be a single string ("regbin" or "log_regbin") applied to every axis, or a list of such strings with one entry per axis for mixed linear/log grids (e.g. ["regbin", "log_regbin"]).

property lineage
map_point(pos)[source]
map_regbin(pos)[source]
pdf_to_cdf(x, pdf)[source]

Convert pdf map to cdf map.

preprocess(pos)[source]

Apply log10 to axes marked as log in _log_mask.

class appletree.config.SigmaMap(method='IDW', **kwargs)[source]

Bases: Config

Maps with uncertainty.

The value of a SigmaMap can be:
  • a list with four elements,

    which are the file names of median, lower, upper maps and the name of the scaler.

  • a list with three elements,

    which are the file names of median, lower and upper maps. The name of the scaler is the default one f”{self.name}_sigma”.

  • a string,

    which is the file name of the map for median, lower, upper.

In the first and second case, the name of the scaler will appear in Component.needed_parameters.

_resolve_sigma_value(configs, index, label)[source]

Extract the i-th sigma value from a list-or-string config.

Parameters
  • configs – list of file paths or a single file path string.

  • index – which sigma (0=median, 1=lower, 2=upper).

  • label – “configs” or “default configs” for error messages.

_update_sigma_cache(map_name, value)[source]

Update _cached_configs for a sigma sub-map, checking for conflicts.

apply(pos, parameters)[source]

Apply SigmaMap with sigma and position.

build(llh_name: Optional[str] = None)[source]

Read maps.

get_configs(llh_name=None)[source]

Get configs of SigmaMap.

property lineage
lower: Map
median: Map
required_parameter(llh_name=None)[source]

Get required parameter of SigmaMap.

upper: Map
appletree.config.takes_config(*configs)[source]

Decorator for plugin classes, to specify which configs it takes.

Parameters

configs – Config instances of configs this plugin takes.

appletree.context module

class appletree.context.Context(instruct, par_config=None)[source]

Bases: object

Combine all likelihood (e.g. Rn220, Ar37), handle MCMC and post-fitting analysis.

__getitem__(keys)[source]

Get likelihood in context.

__init__(instruct, par_config=None)[source]

Create an appletree context.

Parameters

instruct – dict or str, instruct file name or dictionary.

_dump_meta(batch_size, metadata=None)[source]

Save parameters name as attributes.

_sanity_check()[source]

Check if needed parameters are provided.

continue_fitting(context=None, iteration=500, batch_size=1000000, moves=None)[source]

Continue a fitting of another context.

Parameters
  • context – appletree context.

  • iteration – int, number of steps to generate.

dump_post_parameters(file_name)[source]

Dump max posterior parameter in .json file.

fitting(nwalkers=200, iteration=500, batch_size=1000000, moves=None)[source]

Fitting posterior distribution of needed parameters.

Parameters
  • nwalkers – int, number of walkers in the ensemble.

  • iteration – int, number of steps to generate.

classmethod from_backend(backend_h5_file_name, moves=None)[source]

Initialize context from a backend_h5 file.

get_all_post_parameters(**kwargs)[source]

Return all posterior parameters.

get_num_events_accepted(parameters, batch_size=1000000)[source]

Get number of events in the histogram under given parameters.

Parameters
  • batch_size – int of number of simulated events.

  • parameters – dict of parameters used in simulation.

get_parameter_config(par_config)[source]

Get configuration for parameter manager.

Parameters

par_config – str, parameters configuration file.

get_post_parameters(which='mpe')[source]

Get parameters from the backend.

Parameters
  • which – str, ‘mpe’, ‘random’ or ‘median’. ‘mpe’ is the maximum posterior estimate,

  • a (i.e. the parameter set with the highest posterior value. 'random' returns) –

  • medians. (random parameter set from the posterior distribution. 'median' is the marginal) –

get_template(likelihood_name: str, component_name: str, batch_size: int = 1000000, seed: Optional[int] = None)[source]

Get parameters correspondes to max posterior.

Parameters
  • likelihood_name – name of Likelihood.

  • component_name – name of Component.

  • batch_size – int of number of simulated events.

  • seed – random seed.

property lineage
property lineage_hash
log_posterior(parameters, batch_size=1000000)[source]

Get log likelihood of given parameters.

Parameters
  • batch_size – int of number of simulated events.

  • parameters – dict of parameters used in simulation.

pre_fitting(nwalkers=100, read_only=True, reset=False, batch_size=1000000, moves=None)[source]

Prepare for fitting, initialize backend and sampler.

print_context_summary(short=True)[source]

Print summary of the context.

register_all_likelihood(config)[source]

Create all appletree likelihoods.

Parameters

config – dict, configuration file name or dictionary.

register_component(likelihood_name, component_cls, component_name, file_name=None)[source]

Register component to likelihood.

Parameters
  • likelihood_name – name of Likelihood.

  • component_cls – class of Component.

  • component_name – name of Component.

register_likelihood(likelihood_name, likelihood_config)[source]

Create an appletree likelihood.

Parameters
  • likelihood_name – name of Likelihood.

  • likelihood_config – dict of likelihood configuration.

update_parameter_config(likelihoods)[source]
update_url_base(url_base)[source]

Update url_base in appletree.share.

appletree.hist module

appletree.hist.make_hist_irreg_bin_1d(sample, bins, weights)[source]

Make a histogram with irregular binning.

Parameters
  • sample – array with shape N.

  • bins – array with shape M.

  • weights – array with shape (N,).

appletree.hist.make_hist_irreg_bin_2d(sample, bins_x, bins_y, weights)[source]

Make a histogram with irregular binning.

Parameters
  • sample – array with shape (N, 2).

  • bins_x – array with shape (M1, ).

  • bins_y – array with shape (M1-1, M2).

  • weights – array with shape (N,).

appletree.hist.make_hist_mesh_grid(sample, bins=10, weights=None)[source]

Same as jnp.histogramdd.

appletree.interpolation module

appletree.interpolation._L2_dist2(pos1, pos2)[source]

Calculate L2 distance between pos1 and pos2.

Parameters
  • pos1 – array with shape (N, D).

  • pos2 – array with shape (M, D).

Returns

L2 distance squared with shape (N, M).

appletree.interpolation.curve_interpolator(pos, ref_pos, ref_val)[source]

Inverse distance weighting average as interpolation using KNN (K=2) for 1D map.

Parameters
  • pos – array with shape (N,), as the points to be interpolated.

  • ref_pos – array with shape (M,), as the reference points.

  • ref_val – array with shape (M,), as the reference values.

Returns

interpolated values with shape (N,),

weighted by the inverse of the distance to k nearest neighbors.

appletree.interpolation.map_interpolator_knn(pos, ref_pos, ref_val, k=3)[source]

Inverse distance weighting average as interpolation using KNN.

Parameters
  • pos – array with shape (N, D), as the points to be interpolated.

  • ref_pos – array with shape (M, D), as the reference points.

  • ref_val – array with shape (M,), as the reference values.

Returns

interpolated values with shape (N,),

weighted by the inverse of the distance to k nearest neighbors.

appletree.interpolation.map_interpolator_linear_1d(pos, ref_pos, ref_val)[source]

Linear 1D interpolation.

Copied to prevent misuse of other arguments of jnp.interp.
Args:

pos: array with shape (N,), as the points to be interpolated. ref_pos: array with shape (M,), as the reference points. ref_val: array with shape (M,), as the reference values.

appletree.interpolation.map_interpolator_nearest_neighbor_1d(pos, ref_pos, ref_val)[source]

Nearest neighbor 1D interpolation.

Parameters
  • pos – array with shape (N,), as the points to be interpolated.

  • ref_pos – array with shape (M,), as the reference points.

  • ref_val – array with shape (M,), as the reference values.

appletree.interpolation.map_interpolator_regular_binning_1d(pos, ref_pos_lowers, ref_pos_uppers, ref_val)[source]

Inverse distance weighting average as 1D interpolation using KNN(K=2). A uniform mesh grid binning is assumed.

Parameters
  • pos – array with shape (N,), positions at which the interp is calculated.

  • ref_pos_lowers – array with shape (1,), the lower edges of the binning on each dimension.

  • ref_pos_uppers – array with shape (1,), the upper edges of the binning on each dimension.

  • ref_val – array with shape (M1,), map values.

appletree.interpolation.map_interpolator_regular_binning_2d(pos, ref_pos_lowers, ref_pos_uppers, ref_val)[source]

Inverse distance weighting average as 2D interpolation using KNN(K=4). A uniform mesh grid binning is assumed.

Parameters
  • pos – array with shape (N, 2), positions at which the interp is calculated.

  • ref_pos_lowers – array with shape (2,), the lower edges of the binning on each dimension.

  • ref_pos_uppers – array with shape (2,), the upper edges of the binning on each dimension.

  • ref_val – array with shape (M1, M2), map values.

appletree.interpolation.map_interpolator_regular_binning_3d(pos, ref_pos_lowers, ref_pos_uppers, ref_val)[source]

Inverse distance weighting average as 3D interpolation using KNN(K=8). A uniform mesh grid binning is assumed.

Parameters
  • pos – array with shape (N, 3), positions at which the interp is calculated.

  • ref_pos_lowers – array with shape (3,), the lower edges of the binning on each dimension.

  • ref_pos_uppers – array with shape (3,), the upper edges of the binning on each dimension.

  • ref_val – array with shape (M1, M2, M3), map values.

appletree.interpolation.map_interpolator_regular_binning_nearest_neighbor_2d(pos, ref_pos_lowers, ref_pos_uppers, ref_val)[source]

Nearest neighbor 2D interpolation.

A uniform mesh grid binning is assumed.
Args:

pos: array with shape (N, 2), positions at which the interp is calculated. ref_pos_lowers: array with shape (2,), the lower edges of the binning on each dimension. ref_pos_uppers: array with shape (2,), the upper edges of the binning on each dimension. ref_val: array with shape (M1, M2), map values.

appletree.interpolation.map_interpolator_regular_binning_nearest_neighbor_3d(pos, ref_pos_lowers, ref_pos_uppers, ref_val)[source]

Nearest neighbor 3D interpolation.

A uniform mesh grid binning is assumed.
Args:

pos: array with shape (N, 3), positions at which the interp is calculated. ref_pos_lowers: array with shape (3,), the lower edges of the binning on each dimension. ref_pos_uppers: array with shape (3,), the upper edges of the binning on each dimension. ref_val: array with shape (M1, M2, M3), map values.

appletree.likelihood module

class appletree.likelihood.Likelihood(name: Optional[str] = None, **config)[source]

Bases: object

Combine all components (e.g. ER, AC, Wall), and calculate log posterior likelihood.

__getitem__(keys)[source]

Get component in likelihood.

__init__(name: Optional[str] = None, **config)[source]

Create an appletree likelihood.

Parameters

config – Dictionary with configuration options that will be applied, should include: * data_file_name: the data used in fitting, usually calibration data * bins_type: either meshgrid or equiprob * bins_on: observables where we will perform inference on, usually [cs1, cs2] * x_clip, y_clip: ROI of the fitting, should be list of upper and lower boundary * parameter_alias: alias of parameters that will be renamed from key to value

_make_data_hist(use_meshgrid=False)[source]

Create the data histogram from self.data and self._bins.

_print_components(indent: str = '    ', short: bool = True)[source]

Print component details shared by Likelihood and LikelihoodLit.

_resolve_bin_edges(bin_spec, clip_range, config, warn_key)[source]

Resolve a single axis bin specification into bin edges.

Parameters
  • bin_spec – int (number of bins) or array-like (explicit edges).

  • clip_range – tuple of (lo, hi) for linspace when bin_spec is int.

  • config – the full config dict.

  • warn_key – config key (e.g. “x_clip”) to check for spurious warning.

_resolve_explicit_bins(config)[source]

Resolve explicit bin specs (meshgrid or irreg) into a tuple of arrays.

_sanity_check()[source]

Check equality between number of bins group and observables.

_validate_irreg_bins_2d()[source]

Validate irregular 2D bins: x-bins length and y-bins uniformity.

get_log_likelihood(*args, **kwargs)[source]
get_num_events_accepted(*args, **kwargs)[source]
property lineage
property lineage_hash
print_likelihood_summary(indent: str = '    ', short: bool = True)[source]

Print likelihood summary: components, bins, file names.

Parameters
  • indent – str of indent.

  • short – bool, whether only print short summary.

register_component(component_cls: Type[Component], component_name: str, file_name: Optional[str] = None)[source]

Create an appletree likelihood.

Parameters
  • component_cls – class of Component.

  • component_name – name of Component.

  • file_name – file used in ComponentFixed.

replace_alias(parameters)[source]

Replace alias in parameters from key to value.

Note that the value will be popped out.

set_binning(config)[source]

Set binning of likelihood.

simulate_weighted_data(*args, **kwargs)[source]
class appletree.likelihood.LikelihoodLit(name: Optional[str] = None, **config)[source]

Bases: Likelihood

Using literature constraint to build LLH.

The idea is to simulate light and charge yields directly with given energy distribution. And then fit the result with provided literature measurement points. The energy distribution will always be twohalfnorm(TwoHalfNorm), norm or band, which is specified by

__init__(name: Optional[str] = None, **config)[source]

Create an appletree likelihood.

Parameters

config – Dictionary with configuration options that will be applied.

_sanity_check()[source]

Check sanities of supported distribution and dimension.

get_log_likelihood(*args, **kwargs)[source]
property lineage
print_likelihood_summary(indent: str = '    ', short: bool = True)[source]

Print likelihood summary: components, bins, file names.

Parameters
  • indent – str of indent.

  • short – bool, whether only print short summary.

register_component(*args, **kwargs)[source]

Create an appletree likelihood.

Parameters
  • component_cls – class of Component.

  • component_name – name of Component.

  • file_name – file used in ComponentFixed.

appletree.likelihood.need_replacing_alias(func)[source]

Decorator to replace alias in parameters from key to value.

appletree.parameter module

class appletree.parameter.Parameter(parameter_config)[source]

Bases: object

Parameter handler to update parameters and calculate prior.

__getitem__(keys)[source]

__getitem__, keys can be str/list/set.

__init__(parameter_config)[source]

Initialization.

Parameters

parameter_config – can be either * str: the json file name where the config is stored. * dict: config dictionary.

check_parameter_exist(keys, return_not_exist=False)[source]

Check whether the keys exist in parameters.

Parameters
  • keys – Parameter names. Can be a single str, or a list of str.

  • return_not_exist – If False, function will return a bool if all keys exist. If True, function will additionally return the list of the not existing keys.

get_all_parameter()[source]

Return all parameters as a dict.

get_parameter(keys)[source]

Return parameter values.

Parameters

keys – Parameter names. Can be a single str, or a list of str.

init_parameter(seed=None)[source]

Initializing parameters by sampling prior. sampling from the. initial guess. If the prior is free, then.

Parameters

seed – integer, sent to np.random.seed(seed)

property log_prior

Return log prior.

If any parameter is out of allowed_range return -np.inf.

property parameter_fit

Return sorted list of parameters name waiting for fitting.

property parameter_fit_array

Return non-fixed parameters, ordered by self._parameter_fit.

sample_init()[source]

Samping parameters from initial guess clipped by the allowed_range and set self._parameter_dict.

sample_prior()[source]

Sampling parameters from prior and set self._parameter_dict.

If the prior is free, then sampling from the initial guess.

set_parameter(keys, vals=None)[source]

Set parameter values.

Parameters
  • keys – Parameter names. Can be either * str: vals must be int or float. * list: vals must have the same length. * dict: vals will be overwritten as keys.values().

  • vals – Values to be set.

appletree.plugin module

class appletree.plugin.Plugin(llh_name: Optional[str] = None)[source]

Bases: object

The smallest simulation unit.

__call__(*args, **kwargs)[source]

Calls self.simulate.

__init__(llh_name: Optional[str] = None)[source]

Initialization.

depends_on: List[str] = []
property lineage
property lineage_hash
parameters: Tuple = ()
provides: List[str] = []
sanity_check()[source]

Check the consistency between depends_on, provides and in(out)put of self.simulate

simulate(*args, **kwargs)[source]

The main simulation function.

Parameters
  • key – a jnp.array with length 2, used to generate random variables. See https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html

  • parameters – a dictionary with key being parameters’ names. Plugin will get values of self.parameters from this dictionary.

  • args – other args following key and parameters must be in the order of self.depends_on.

Returns

key and output simulated variables, ordered by self.provides. key will be updated if it’s used inside self.simulate to generate random variables.

takes_config = immutabledict({})
appletree.plugin._add_plugin_extension(module, plugin, force=False)[source]

Add plugin to module.

appletree.plugin.add_plugin_extensions(module1, module2, force=False)[source]

Add plugins of module2 to module1.

appletree.randgen module

appletree.randgen.bernoulli(key, p, shape=())[source]

Bernoulli random sampler.

Parameters
  • key – seed for random generator.

  • p – <jnp.array>-like probability in bernoulli distribution.

  • shape – output shape. If not given, output has shape jnp.shape(lam).

Returns

an updated seed, random variables.

appletree.randgen.binomial(key, p, n, shape=())[source]

Binomial random sampler.

Parameters
  • key – seed for random generator.

  • p – <jnp.array>-like probability in binomial distribution.

  • n – <jnp.array>-like count in binomial distribution.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(p), jnp.shape(n)).

  • always_use_normal – If true, then Norm(np, sqrt(npq)) is always used. Otherwise if n * p < 5, use the inversion method instead.

Returns

an updated seed, random variables.

appletree.randgen.binomial_hybrid(key, p, n, shape=(), threshold=128)[source]

Hybrid binomial random sampler: Bernoulli sum + Cornish-Fisher expansion.

For small n (n <= threshold): uses exact Bernoulli sum (no approximation error). For large n (n > threshold): uses Cornish-Fisher corrected normal approximation.

This hybrid approach is faster than exact binomial while maintaining exact accuracy for small n where the normal approximation is least accurate.

Parameters
  • key – seed for random generator.

  • p – <jnp.array>-like probability in binomial distribution.

  • n – <jnp.array>-like count in binomial distribution.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(p), jnp.shape(n)).

  • threshold – use Bernoulli sum for n <= threshold, CF for n > threshold.

Returns

an updated seed, random variables.

appletree.randgen.gamma(key, alpha, beta, shape=())[source]

Gamma distribution random sampler.

Parameters
  • key – seed for random generator.

  • alpha – <jnp.array>-like shape in gamma distribution.

  • beta – <jnp.array>-like rate in normal distribution.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(alpha), jnp.shape(beta)).

Returns

an updated seed, random variables.

appletree.randgen.generalized_poisson(key, lam, eta, shape=())[source]

Generalized Poisson Distribution(GPD) random sampler.

Parameters
  • key – seed for random generator.

  • lam – <jnp.array>-like expectation(location parameter) in GPD.

  • eta – <jnp.array>-like scale parameter in GPD, within [0, 1).

  • shape – output shape. If not given, output has shape jnp.shape(lam).

Returns

an updated seed, random variables.

References

  1. https://gist.github.com/danmackinlay/00e957b11c488539bd3e2a3804922b9d

  2. https://search.r-project.org/CRAN/refmans/LaplacesDemon/html/dist.Generalized.Poisson.html # noqa

appletree.randgen.get_key(seed=None)[source]

Generate a key for jax.random.

appletree.randgen.negative_binomial(key, p, n, shape=())[source]

Negative binomial distribution random sampler.

Using Gamma–Poisson mixture.
Args:

key: seed for random generator. p: <jnp.array>-like probability of a single success in negative binomial distribution. n: <jnp.array>-like number of successes in negative binomial distribution. shape: output shape.

If not given, output has shape jnp.broadcast_shapes(jnp.shape(p), jnp.shape(n)).

Returns:

an updated seed, random variables.

References:
  1. https://en.wikipedia.org/wiki/Negative_binomial_distribution#Gamma%E2%80%93Poisson_mixture # noqa

  2. https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.nbinom.html

appletree.randgen.normal(key, mean, std, shape=())[source]

Normal distribution random sampler.

Parameters
  • key – seed for random generator.

  • mean – <jnp.array>-like mean in normal distribution.

  • std – <jnp.array>-like std in normal distribution.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(mean), jnp.shape(std)).

Returns

an updated seed, random variables.

appletree.randgen.poisson(key, lam, shape=())[source]

Poisson random sampler.

Parameters
  • key – seed for random generator.

  • lam – <jnp.array>-like expectation in poisson distribution.

  • shape – output shape. If not given, output has shape jnp.shape(lam).

Returns

an updated seed, random variables.

appletree.randgen.skewnormal(key, a, loc, scale, shape=())[source]

Skew-normal distribution random sampler.

Parameters
  • key – seed for random generator.

  • a – <jnp.array>-like skewness in skewnormal distribution.

  • loc – <jnp.array>-like loc in skewnormal distribution.

  • scale – <jnp.array>-like scale in skewnormal distribution.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(loc), jnp.shape(scale)).

Returns

an updated seed, random variables.

References

1. Dariush Ghorbanzadeh, Luan Jaupi, Philippe Durand. A Method to Simulate the Skew Normal Distribution. Applied Mathematics, 2014, 5 (13), pp.2073-2076. ff10.4236/am.2014.513201ff. ffhal02467997

appletree.randgen.truncate_normal(key, mean, std, vmin=None, vmax=None, shape=())[source]

Truncated normal distribution random sampler.

Parameters
  • key – seed for random generator.

  • mean – <jnp.array>-like mean in normal distribution.

  • std – <jnp.array>-like std in normal distribution.

  • vmin – <jnp.array>-like min value to clip. By default it’s None. vmin and vmax cannot be both None.

  • vmax – <jnp.array>-like max value to clip. By default it’s None. vmin and vmax cannot be both None.

  • shape – parameter passed to normal(…, shape=shape)

Returns

an updated seed, random variables.

appletree.randgen.twohalfnorm(key, mu, sigma_pos, sigma_neg, shape=())[source]

JAX version of TwoHalfNorm.rvs.

Parameters
  • key – seed for random generator.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(mean), jnp.shape(std)).

Returns

an updated seed, random variables.

appletree.randgen.uniform(key, vmin, vmax, shape=())[source]

Uniform random sampler.

Parameters
  • key – seed for random generator.

  • vmin – <jnp.array>-like min in uniform distribution.

  • vmax – <jnp.array>-like max in uniform distribution.

  • shape – output shape. If not given, output has shape jnp.broadcast_shapes(jnp.shape(vmin), jnp.shape(vmax)).

Returns

an updated seed, random variables.

appletree.randgen.uniform_key_vectorized(key)[source]

Uniform(0,1) distribution sampler, vectorized by key. Note: key won’t be updated!

Parameters

key – seed for random generator, with shape (N, 2)

Returns

random varibles with shape (N,)

appletree.share module

class appletree.share.RecordingDict(*args, **kwargs)[source]

Bases: dict

clear() None.  Remove all items from D.[source]
class appletree.share.StaticValueDict[source]

Bases: dict

appletree.share.clear_cache()[source]
appletree.share.set_global_config(configs)[source]

Set new global configuration options.

Parameters

configs – dict, configuration file name or dictionary

appletree.utils module

appletree.utils._add_extension(module, subclass, base, force=False)[source]

Add subclass to module Skip the class when it is base class.

It is no allowed to assign a class which has same name to an already assigned class. We do not allowed class name covering! Please change the name of your class when Error shows itself.

appletree.utils._files(package: Union[module, str]) Traversable

Get a Traversable resource from a package

appletree.utils._get_abspath(file_name)[source]

Get the abspath of the file.

Raise FileNotFoundError when not found in any subfolder

appletree.utils._package_path(sub_directory)[source]

Get the abs path of the requested sub folder.

appletree.utils.add_deps_to_graph_tree(component, graph_tree, data_names: list = ['cs1', 'cs2', 'eff'], _seen=None)[source]

Recursively add nodes to graph base on plugin.deps.

Parameters
  • context – Context instance.

  • graph_tree – Digraph instance.

  • data_names – Data type name.

  • _seen – list or None, the seen data_name should not be plot.

appletree.utils.add_extensions(module1, module2, base, force=False)[source]

Add subclasses of module2 to module1.

When ComponentSim compiles the dependency tree, it will search in the appletree.plugins module for Plugin(as attributes). When building Likelihood, it will also search for corresponding Component(s) specified in the instructions(e.g. NRBand).

So we need to assign the attributes before compilation. These plugins are mostly user defined.

appletree.utils.add_plugins_to_graph_tree(component, graph_tree, data_names: list = ['cs1', 'cs2', 'eff'], _seen=None, with_data_names=False)[source]

Recursively add nodes to graph base on plugin.deps.

Parameters
  • context – Context instance.

  • graph_tree – Digraph instance.

  • data_names – Data type name.

  • _seen – list or None, the seen data_name should not be plot.

  • with_data_names – bool, whether plot even with messy data_names

appletree.utils.add_spaces(x)[source]

Add four spaces to every line in x This is needed to make html raw blocks in rst format correctly.

appletree.utils.calculate_sha256(file_path)[source]

Get sha256 hash of the file.

appletree.utils.check_unused_configs()[source]

Check if there are unused configs.

appletree.utils.dump_lineage(file_path, entity)[source]

Dump lineage of whatever level into .json file.

appletree.utils.errors_to_two_half_norm_sigmas(errors)[source]

This function solves the sigmas for a two-half-norm distribution, such that the 16 and 84 percentile corresponds to the given errors.

In the two-half-norm distribution, the positive and negative errors are assumed to be the std of the glued normal distributions. While we interpret the 16 and 84 percentile as the input errors, thus we need to solve the sigmas for the two-half-norm distribution. The solution is determined by the following conditions: - The 16 percentile of the two-half-norm distribution should be the negative error. - The 84 percentile of the two-half-norm distribution should be the positive error. - The mode of the two-half-norm distribution should be 0.

appletree.utils.exporter(export_self=False)[source]

Export utility modified from https://stackoverflow.com/a/41895194 Returns export decorator, __all__ list

appletree.utils.get_equiprob_bins_1d(data, n_partitions, clip=(-inf, inf), integer=False, left=True, which_np=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/appletree/envs/latest/lib/python3.10/site-packages/numpy/__init__.py'>)[source]

Get 2D equiprobable binning edges.

Parameters
  • data – array with shape N.

  • n_partitions – M1 which is the number of bins.

  • clip – lower and upper binning edges on the 0th dimension. Data outside the clip will be dropped.

  • dropped. (Data outside the y_clip will be) –

  • integer – bool, whether the corresponding dimension is integer.

  • left – bool, whether start searching for bin edges from the left side.

  • which_np – can be numpy or jax.numpy, determining the returned array type.

appletree.utils.get_equiprob_bins_2d(data, n_partitions, order=(0, 1), x_clip=(-inf, inf), y_clip=(-inf, inf), integer=[False, False], left=True, which_np=<module 'numpy' from '/home/docs/checkouts/readthedocs.org/user_builds/appletree/envs/latest/lib/python3.10/site-packages/numpy/__init__.py'>)[source]

Get 2D equiprobable binning edges.

Parameters
  • data – array with shape (N, 2).

  • n_partitions – [M1, M2] where M1 M2 are the number of bins on each dimension.

  • x_clip – lower and upper binning edges on the 0th dimension. Data outside the x_clip will be dropped.

  • y_clip – lower and upper binning edges on the 1st dimension.

  • integer – list of bool with length 2, whether the corresponding dimension is integer.

  • left – bool, whether start searching for bin edges from the left side.

  • dropped. (Data outside the y_clip will be) –

  • which_np – can be numpy or jax.numpy, determining the returned array type.

appletree.utils.get_file_path(fname)[source]

Find the full path to the resource file.

Try 5 methods in the following order.
  • fname begin with ‘/’, return absolute path

  • url_base begin with ‘/’, return url_base + name

  • can get file from _get_abspath, return appletree internal file path

  • can be found in local installed ntauxfiles, return ntauxfiles absolute path

  • can be downloaded from MongoDB, download and return cached path

appletree.utils.get_platform() str[source]

Show the platform we are using, e.g. ‘cpu’, ‘gpu’, or ‘tpu’.

appletree.utils.load_data(file_name: str)[source]

Load data from file.

The suffix can be “.csv”, “.pkl”.

appletree.utils.load_json(file_name: str)[source]

Load data from json file.

appletree.utils.plot_irreg_histogram_2d(bins_x, bins_y, hist, **kwargs)[source]

Plot histogram defined by irregular binning.

Parameters
  • bins_x – array with shape (M1,).

  • bins_y – array with shape (M1-1, M2).

  • hist – array with shape (M1-1, M2-1).

  • density – boolean.

appletree.utils.set_gpu_memory_usage(fraction=0.3)[source]

Set GPU memory usage.

See more on https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html

appletree.utils.timeit(indent='')[source]

Use timeit as a decorator.

It will print out the running time of the decorated function.

appletree.utils.tree_to_svg(graph_tree, save_as='data_types', view=True)[source]

Where to save this node.

Parameters
  • graph_tree – Digraph instance.

  • save_as – str, file name.

  • view – bool, Open the rendered result with the default application.

appletree.utils.use_xenon_plot_style()[source]

Set matplotlib plot style.

Module contents