1"""Demo code to fit a 2d 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
12DGParams = namedtuple("DGParams", ("mu0", "sig0", "mu1", "sig1", "frac0"))
13DEFAULT_PARAMS = DGParams(mu0=-1.0, sig0=0.5, mu1=1.0, sig1=1.0, frac0=0.75)
14
15NPTS = 20_000
16
17
18@jjit
19def mc_double_gaussian(params, ran_key):
20 """Draw a stochastic Monte Carlo realization of a double Gaussian"""
21 u_key, n0_key, n1_key = jran.split(ran_key, 3)
22 uran = jran.uniform(u_key, minval=0, maxval=1, shape=(NPTS,))
23 n0 = jran.normal(n0_key, shape=(NPTS,)) * params.sig0 + params.mu0
24 n1 = jran.normal(n1_key, shape=(NPTS,)) * params.sig1 + params.mu1
25 mc_0 = uran < params.frac0
26 xdata = jnp.where(mc_0, n0, n1)
27 return xdata
28
29
30@jjit
31def predict_soft_xhist_mc(params, xbins, ran_key):
32 """Predict histogram counts by applying soft histogram to
33 a stochastic Monte Carlo realization of a double Gaussian"""
34 xdata = mc_double_gaussian(params, ran_key)
35 xhist = soft_xhist(xdata, xbins)
36 return xhist
37
38
39@jjit
40def predict_soft_xhist_weighted(params, xbins, ran_key):
41 """Predict histogram counts by applying soft histogram to
42 a PDF-weighted Monte Carlo realization of a double Gaussian"""
43 n0_key, n1_key = jran.split(ran_key, 2)
44 n0 = jran.normal(n0_key, shape=(NPTS,)) * params.sig0 + params.mu0
45 n1 = jran.normal(n1_key, shape=(NPTS,)) * params.sig1 + params.mu1
46 xhist0 = soft_xhist(n0, xbins)
47 xhist1 = soft_xhist(n1, xbins)
48 xhist = params.frac0 * xhist0 + (1.0 - params.frac0) * xhist1
49 return xhist
50
51
52@jjit
53def soft_xhist(xdata, xbins):
54 """Soft histogram function
55 This is a wrapper around diffsky.nnsig_ndhist for 1d data"""
56 nbins = xbins.shape[0]
57 xbins_lo = xbins[:-1].reshape((nbins - 1, 1))
58 xbins_hi = xbins[1:].reshape((nbins - 1, 1))
59 dx = jnp.diff(xbins).mean()
60 ndsig = jnp.zeros_like(xbins_lo) + dx / 2
61 xdata = xdata.reshape((-1, 1))
62 xhist = nnsig_ndhist(xdata, ndsig, xbins_lo, xbins_hi)
63 return xhist
64
65
66@jjit
67def _mae_kern(x, y):
68 """Mean absolute error"""
69 abs_diff = jnp.abs(y - x)
70 return jnp.mean(abs_diff)
71
72
73@jjit
74def weighted_mae_loss(params, loss_data):
75 """Loss function based on a PDF-weighted soft histogram"""
76 xhist_target, xbins, ran_key = loss_data
77 xhist_pred = predict_soft_xhist_weighted(params, xbins, ran_key)
78 loss = _mae_kern(xhist_pred, xhist_target)
79 return loss
80
81
82@jjit
83def mc_mae_loss(params, loss_data):
84 """Loss function based on a stochastic Monte Carlo with a soft histogram"""
85 xhist_target, xbins, ran_key = loss_data
86 xhist_pred = predict_soft_xhist_mc(params, xbins, ran_key)
87 loss = _mae_kern(xhist_pred, xhist_target)
88 return loss
89
90
91@jjit
92def param_update(params, grads, learning_rate):
93 """Update namedtuple params by taking a small step down the gradient"""
94 new_params = params._make(jnp.array(params) - jnp.array(grads) * learning_rate)
95 return new_params
96
97
98weighted_mae_loss_and_grad = jjit(value_and_grad(weighted_mae_loss, argnums=0))
99mc_mae_loss_and_grad = jjit(value_and_grad(mc_mae_loss, argnums=0))