1"""Demo code to fit a 1d Gaussian model with soft histograms and jax.grad"""
2
3from collections import namedtuple
4
5from jax import jit as jjit
6from jax import numpy as jnp
7from jax import random as jran
8from jax import value_and_grad
9
10from diffsky.soft_histograms.signdhist_lomem import nnsig_ndhist
11
12GParams = namedtuple("GParams", ("mu", "sig"))
13DEFAULT_PARAMS = GParams(mu=-1.0, sig=1.0)
14
15NPTS = 20_000
16
17
18@jjit
19def mc_single_gaussian(params, ran_key):
20 """Draw a Monte Carlo realization of a Gaussian"""
21 xdata = jran.normal(ran_key, shape=(NPTS,)) * params.sig + params.mu
22 return xdata
23
24
25@jjit
26def mc_predict_hard_edged_xhist(params, xbins, ran_key):
27 """Predict histogram counts by applying jnp.histogram to
28 a Monte Carlo realization of a Gaussian"""
29 xdata = mc_single_gaussian(params, ran_key)
30 xhist, __ = jnp.histogram(xdata, bins=xbins)
31 return xhist
32
33
34@jjit
35def mc_predict_soft_xhist(params, xbins, ran_key):
36 """Predict histogram counts by applying a soft histogram to
37 a Monte Carlo realization of a Gaussian"""
38 xdata = mc_single_gaussian(params, ran_key)
39 n = xdata.shape[0]
40 xdata = xdata.reshape((n, 1))
41 xhist = soft_xhist(xdata, xbins)
42 return xhist
43
44
45@jjit
46def soft_xhist(xdata, xbins):
47 """Soft histogram function
48 This is a wrapper around diffsky.nnsig_ndhist for 1d data"""
49 nbins = xbins.shape[0]
50 xbins_lo = xbins[:-1].reshape((nbins - 1, 1))
51 xbins_hi = xbins[1:].reshape((nbins - 1, 1))
52 dx = jnp.diff(xbins).mean()
53 ndsig = jnp.zeros_like(xbins_lo) + dx / 2
54 xdata = xdata.reshape((-1, 1))
55 xhist = nnsig_ndhist(xdata, ndsig, xbins_lo, xbins_hi)
56 return xhist
57
58
59@jjit
60def _mae_kern(x, y):
61 """Mean absolute error"""
62 abs_diff = jnp.abs(y - x)
63 return jnp.mean(abs_diff)
64
65
66@jjit
67def hard_edged_xhist_loss(params, loss_data):
68 """Loss function based on a histogram with hard-edged bins"""
69 xhist_target, xbins, ran_key = loss_data
70 xhist_pred = mc_predict_hard_edged_xhist(params, xbins, ran_key)
71 loss = _mae_kern(xhist_pred, xhist_target)
72 return loss
73
74
75@jjit
76def soft_xhist_loss(params, loss_data):
77 """Loss function based on a soft histogram"""
78 xhist_target, xbins, ran_key = loss_data
79 xhist_pred = mc_predict_soft_xhist(params, xbins, ran_key)
80 loss = _mae_kern(xhist_pred, xhist_target)
81 return loss
82
83
84@jjit
85def param_update(params, grads, learning_rate):
86 """Update namedtuple params by taking a small step down the gradient"""
87 new_params = params._make(jnp.array(params) - jnp.array(grads) * learning_rate)
88 return new_params
89
90
91hard_edged_xhist_loss_and_grad = jjit(value_and_grad(hard_edged_xhist_loss, argnums=0))
92soft_xhist_loss_and_grad = jjit(value_and_grad(soft_xhist_loss, argnums=0))