def double_crystalball_integral(limits, params, model):
mu = params['mu']
sigma = params['sigma']
(lower,), (upper,) = limits.limits
lower = lower[0] # obs number 0
upper = upper[0]
limits_left = Space(limits.obs, (lower, mu))
limits_right = Space(limits.obs, (mu, upper))
params_left = dict(mu=mu, sigma=sigma, alpha=params["alphal"],
n=params["nl"])
params_right = dict(mu=mu, sigma=sigma, alpha=-params["alphar"],
n=params["nr"])
#
left = tf.cond(pred=tf.less(mu, lower), true_fn=lambda: z.constant(0.),
false_fn=lambda: crystalball_integral(limits_left, params_left, model))
right = tf.cond(pred=tf.greater(mu, upper), true_fn=lambda: z.constant(0.),
false_fn=lambda: crystalball_integral(limits_right, params_right, model))
integral = left + right
# integral = z.where(condition=tf.less(mu, lower),
# x=crystalball_integral(limits_left, params_left, model),
# y=crystalball_integral(limits_right, params_right, model))
return integral
def test_add():
param1 = zfit.Parameter("param1", 1.)
param2 = zfit.Parameter("param2", 2.)
pdfs = [0] * 4
pdfs[0] = Gauss(param1, 4, obs=obs1)
pdfs[1] = Gauss(param2, 5, obs=obs1)
pdfs[2] = Gauss(3, 6, obs=obs1)
pdfs[3] = Gauss(4, 7, obs=obs1)
datas = [0] * 4
datas[0] = z.constant(1.)
datas[1] = z.constant(2.)
datas[2] = z.constant(3.)
datas[3] = z.constant(4.)
ranges = [0] * 4
ranges[0] = (1, 4)
ranges[1] = Space(limits=(2, 5), obs=obs1)
ranges[2] = Space(limits=(3, 6), obs=obs1)
ranges[3] = Space(limits=(4, 7), obs=obs1)
constraint1 = zfit.constraint.nll_gaussian(params=param1,
observation=1.,
uncertainty=0.5)
constraint2 = zfit.constraint.nll_gaussian(params=param2,
observation=2.,
uncertainty=0.25)
merged_contraints = [constraint1, constraint2]
nll1 = UnbinnedNLL(model=pdfs[0],
data=datas[0],
fit_range=ranges[0],
constraints=constraint1)
nll2 = UnbinnedNLL(model=pdfs[1],
data=datas[1],
fit_range=ranges[1],
constraints=constraint2)
nll3 = UnbinnedNLL(model=[pdfs[2], pdfs[3]],
data=[datas[2], datas[3]],
fit_range=[ranges[2], ranges[3]])
simult_nll = nll1 + nll2 + nll3
assert simult_nll.model == pdfs
assert simult_nll.data == datas
ranges[0] = Space(
limits=ranges[0], obs='obs1',
axes=(0, )) # for comparison, Space can only compare with Space
ranges[1].coords._axes = (0, )
ranges[2].coords._axes = (0, )
ranges[3].coords._axes = (0, )
assert simult_nll.fit_range == ranges
def eval_constraint(constraints):
return z.reduce_sum([c.value() for c in constraints]).numpy()
assert eval_constraint(
simult_nll.constraints) == eval_constraint(merged_contraints)
def _exp_integral_func_shifting(lambd, lower, upper, model):
def raw_integral(x):
return z.exp(lambd * (model._shift_x(x))) / lambd # needed due to overflow in exp otherwise
lower_int = raw_integral(x=z.constant(lower))
upper_int = raw_integral(x=z.constant(upper))
integral = (upper_int - lower_int)
return integral
def double_crystalball_mu_integral_func(mu, sigma, alphal, nl, alphar, nr, lower, upper):
left = tf.cond(pred=tf.less(mu, lower), true_fn=lambda: z.constant(0.),
false_fn=lambda: crystalball_integral_func(mu=mu, sigma=sigma, alpha=alphal, n=nl,
lower=lower, upper=mu))
right = tf.cond(pred=tf.greater(mu, upper), true_fn=lambda: z.constant(0.),
false_fn=lambda: crystalball_integral_func(mu=mu, sigma=sigma, alpha=alphar, n=nr,
lower=mu, upper=upper))
integral = left + right
return integral
def _exp_integral_func_shifting(lambd, lower, upper, model):
def raw_integral(x):
return model._numerics_shifted_exp(
x=x,
lambda_=lambd) / lambd # needed due to overflow in exp otherwise
lower_int = raw_integral(x=z.constant(lower))
upper_int = raw_integral(x=z.constant(upper))
integral = (upper_int - lower_int)
return integral
def __call__(self, n_to_produce, limits, dtype):
importance_sampling_called[0] = True
import tensorflow_probability.python.distributions as tfd
n_to_produce = tf.cast(n_to_produce, dtype=tf.int32)
gaussian = tfd.TruncatedNormal(loc=z.constant(-1.), scale=z.constant(2.),
low=low, high=high)
sample = gaussian.sample(sample_shape=(n_to_produce, 1))
weights = gaussian.prob(sample)[:, 0]
thresholds = tf.random.uniform(shape=(n_to_produce,), dtype=dtype)
return sample, thresholds, weights, None, n_to_produce
def _exp_integral_from_any_to_any(limits, params, model):
lambda_ = params['lambda']
def raw_integral(x):
return model._numerics_shifted_exp(x=x, lambda_=lambda_) / lambda_ # needed due to overflow in exp otherwise
(lower,), (upper,) = limits.limits
if lower[0] == - upper[0] == np.inf:
raise NotImplementedError
lower_int = raw_integral(x=z.constant(lower))
upper_int = raw_integral(x=z.constant(upper))
return (upper_int - lower_int)[0]
def test_gradients(chunksize):
from numdifftools import Gradient
zfit.run.chunking.active = True
zfit.run.chunking.max_n_points = chunksize
initial1 = 1.0
initial2 = 2
param1 = zfit.Parameter("param1", initial1)
param2 = zfit.Parameter("param2", initial2)
gauss1 = Gauss(param1, 4, obs=obs1)
gauss1.set_norm_range((-5, 5))
gauss2 = Gauss(param2, 5, obs=obs1)
gauss2.set_norm_range((-5, 5))
data1 = zfit.Data.from_tensor(obs=obs1,
tensor=z.constant(1.0, shape=(100, )))
data1.set_data_range((-5, 5))
data2 = zfit.Data.from_tensor(obs=obs1,
tensor=z.constant(1.0, shape=(100, )))
data2.set_data_range((-5, 5))
nll = UnbinnedNLL(model=[gauss1, gauss2], data=[data1, data2])
def loss_func(values):
for val, param in zip(values, nll.get_cache_deps(only_floating=True)):
param.set_value(val)
return nll.value().numpy()
# theoretical, numerical = tf.test.compute_gradient(loss_func, list(params))
gradient1 = nll.gradient(params=param1)
gradient_func = Gradient(loss_func)
# gradient_func = lambda *args, **kwargs: list(gradient_func_numpy(*args, **kwargs))
assert gradient1[0].numpy() == pytest.approx(
gradient_func([param1.numpy()]))
param1.set_value(initial1)
param2.set_value(initial2)
params = [param2, param1]
gradient2 = nll.gradient(params=params)
both_gradients_true = list(
reversed(list(gradient_func([initial1, initial2
])))) # because param2, then param1
assert [g.numpy() for g in gradient2] == pytest.approx(both_gradients_true)
param1.set_value(initial1)
param2.set_value(initial2)
gradient3 = nll.gradient()
assert frozenset(g.numpy() for g in gradient3) == pytest.approx(
frozenset(both_gradients_true))
def gauss_4d(x, params):
norm = params[0]
mean = tf.stack(params[1:5])
sigma = tf.stack(params[5:9])
corr = tf.stack([[z.constant(1.), params[9], params[10], params[11]],
[params[9],
z.constant(1.), params[12], params[13]],
[params[10], params[12],
z.constant(1.), params[14]],
[params[11], params[13], params[14],
z.constant(1.)]])
cov = tf.einsum("i,ij,j->ij", sigma, corr, sigma)
invcov = tf.linalg.inv(cov)
return multivariate_gauss(x, norm, mean, invcov)
def create_loss():
with tf.compat.v1.variable_scope("func1"):
a_param = zfit.Parameter("variable_a15151",
1.5,
-1.,
20.,
step_size=z.constant(0.1))
b_param = zfit.Parameter("variable_b15151", 3.5)
c_param = zfit.Parameter("variable_c15151", -0.04)
obs1 = zfit.Space(obs='obs1', limits=(-2.4, 9.1))
# load params for sampling
a_param.load(true_a)
b_param.load(true_b)
c_param.load(true_c)
gauss1 = zfit.pdf.Gauss(mu=a_param, sigma=b_param, obs=obs1)
exp1 = zfit.pdf.Exponential(lambda_=c_param, obs=obs1)
sum_pdf1 = 0.9 * gauss1 + exp1
sampled_data = sum_pdf1.create_sampler(n=15000)
sampled_data.resample()
loss = zfit.loss.UnbinnedNLL(model=sum_pdf1,
data=sampled_data,
fit_range=obs1)
return loss, (a_param, b_param, c_param)
def func_integral_chebyshev2(limits, norm, params, model):
lower, upper = limits.limit1d
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)
lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)
# the integral of cheby2_ni is a cheby1_ni+1/(n+1). We add the (n+1) to the coeffs. The cheby1 shape makes
# the sum for us.
coeffs_cheby1 = {"c_0": z.constant(0.0, dtype=model.dtype)}
for name, coeff in params.items():
n_plus1 = int(name.split("_", 1)[-1]) + 1
coeffs_cheby1[f"c_{n_plus1}"] = coeff / z.convert_to_tensor(
n_plus1, dtype=model.dtype)
coeffs_cheby1 = convert_coeffs_dict_to_list(coeffs_cheby1)
def indefinite_integral(limits):
return chebyshev_shape(x=limits, coeffs=coeffs_cheby1)
integral = indefinite_integral(upper) - indefinite_integral(lower)
integral = znp.reshape(integral, newshape=())
integral *= 0.5 * model.space.area() # rescale back to whole width
return integral
def __init__(self, lam=None, obs: ztyping.ObsTypeInput = None, name: str = "Exponential", lambda_=None):
"""Exponential function exp(lambda * x).
The function is normalized over a finite range and therefore a pdf. So the PDF is precisely
defined as :math:`\\frac{ e^{\\lambda \\cdot x}}{ \\int_{lower}^{upper} e^{\\lambda \\cdot x} dx}`
Args:
lam: Accessed as parameter "lambda".
obs: The :py:class:`~zfit.Space` the pdf is defined in.
name: Name of the pdf.
dtype:
"""
if lambda_ is not None:
if lam is None:
lam = lambda_
else:
raise BreakingAPIChangeError("The 'lambda' parameter has been renamed from 'lambda_' to 'lam'.")
params = {'lambda': lam}
super().__init__(obs, name=name, params=params)
self._calc_numerics_data_shift = lambda: z.constant(0.)
if not self.space.has_limits:
warn_advanced_feature("Exponential pdf relies on a shift of the input towards 0 to keep the numerical "
f"stability high. The space {self.space} does not have limits set and no shift"
f" will occure. To set it manually, set _numerics_data_shift to the expected"
f" average values given to this function _in case you want things to be set_."
f"If this sounds unfamiliar, regard this as an error and use a normalization range.",
identifier='exp_shift')
self._set_numerics_data_shift(self.space)
def create_loss(n=15000):
a_param = zfit.Parameter("variable_a15151",
1.5,
-1.,
20.,
step_size=z.constant(0.1))
b_param = zfit.Parameter("variable_b15151", 3.5, 0, 20)
c_param = zfit.Parameter("variable_c15151", -0.04, -1, 0.)
obs1 = zfit.Space(obs='obs1', limits=(-2.4, 9.1))
# load params for sampling
a_param.set_value(true_a)
b_param.set_value(true_b)
c_param.set_value(true_c)
gauss1 = zfit.pdf.Gauss(mu=a_param, sigma=b_param, obs=obs1)
exp1 = zfit.pdf.Exponential(lambda_=c_param, obs=obs1)
sum_pdf1 = zfit.pdf.SumPDF((gauss1, exp1), 0.7)
sampled_data = sum_pdf1.create_sampler(n=n)
sampled_data.resample()
loss = zfit.loss.UnbinnedNLL(model=sum_pdf1, data=sampled_data)
return loss, (a_param, b_param, c_param)
def func_integral_hermite(limits, norm, params, model):
lower, upper = limits.limit1d
lower_rescaled = model._polynomials_rescale(lower)
upper_rescaled = model._polynomials_rescale(upper)
lower = z.convert_to_tensor(lower_rescaled)
upper = z.convert_to_tensor(upper_rescaled)
# the integral of hermite is a hermite_ni. We add the ni to the coeffs.
coeffs = {"c_0": z.constant(0.0, dtype=model.dtype)}
for name, coeff in params.items():
ip1_coeff = int(name.split("_", 1)[-1]) + 1
coeffs[f"c_{ip1_coeff}"] = coeff / z.convert_to_tensor(
ip1_coeff * 2.0, dtype=model.dtype)
coeffs = convert_coeffs_dict_to_list(coeffs)
def indefinite_integral(limits):
return hermite_shape(x=limits, coeffs=coeffs)
integral = indefinite_integral(upper) - indefinite_integral(lower)
integral = znp.reshape(integral, newshape=())
integral *= 0.5 * model.space.area() # rescale back to whole width
return integral
def test_normalization(obs1, pdf_factory):
import numpy as np
import tensorflow as tf
import zfit
from zfit import z
test_yield = 1524.3
dist = pdf_factory()
samples = tf.cast(np.random.uniform(low=low, high=high, size=100000),
dtype=tf.float64)
small_samples = tf.cast(np.random.uniform(low=low, high=high, size=10),
dtype=tf.float64)
with dist.set_norm_range(zfit.Space(obs1, limits=(low, high))):
probs = dist.pdf(samples)
probs_small = dist.pdf(small_samples)
log_probs = dist.log_pdf(small_samples)
probs, log_probs = probs.numpy(), log_probs.numpy()
probs = np.average(probs) * (high - low)
assert probs == pytest.approx(1., rel=0.05)
assert log_probs == pytest.approx(tf.math.log(probs_small).numpy(),
rel=0.05)
dist = dist.create_extended(z.constant(test_yield))
probs = dist.pdf(samples)
probs_extended = dist.ext_pdf(samples)
result = probs.numpy()
result = np.average(result) * (high - low)
result_ext = np.average(probs_extended) * (high - low)
assert result == pytest.approx(1, rel=0.05)
assert result_ext == pytest.approx(test_yield, rel=0.05)
def create_loss(n=15000, weights=None):
avalue = 1.5
a_param = zfit.Parameter("variable_a15151", avalue, -1., 20.,
step_size=z.constant(0.1))
a_param.init_val = avalue
bvalue = 3.5
b_param = zfit.Parameter("variable_b15151", bvalue, 0, 20)
b_param.init_val = bvalue
cvalue = -0.04
c_param = zfit.Parameter("variable_c15151", cvalue, -1, 0.)
c_param.init_val = cvalue
obs1 = zfit.Space(obs='obs1', limits=(-2.4, 9.1))
# load params for sampling
a_param.set_value(true_a)
b_param.set_value(true_b)
c_param.set_value(true_c)
gauss1 = zfit.pdf.Gauss(mu=a_param, sigma=b_param, obs=obs1)
exp1 = zfit.pdf.Exponential(lam=c_param, obs=obs1)
sum_pdf1 = zfit.pdf.SumPDF((gauss1, exp1), 0.7)
sampled_data = sum_pdf1.create_sampler(n=n)
sampled_data.resample()
if weights is not None:
sampled_data.set_weights(weights)
loss = zfit.loss.UnbinnedNLL(model=sum_pdf1, data=sampled_data)
return loss, (a_param, b_param, c_param)
def __init__(
self,
obs,
coeffs: list,
apply_scaling: bool = True,
coeff0: tf.Tensor | None = None,
name: str = "Polynomial",
): # noqa
"""Base class to create 1 dimensional recursive polynomials that can be rescaled. Overwrite _poly_func.
Args:
coeffs: Coefficients for each polynomial. Used to calculate the degree.
apply_scaling: Rescale the data so that the actual limits represent (-1, 1).
.. math::
x_{n+1} = recurrence(x_{n}, x_{n-1}, n)
"""
# 0th coefficient set to 1 by default
coeff0 = (z.constant(1.0)
if coeff0 is None else tf.cast(coeff0, dtype=ztypes.float))
coeffs = convert_to_container(coeffs).copy()
coeffs.insert(0, coeff0)
params = {f"c_{i}": coeff for i, coeff in enumerate(coeffs)}
self._degree = len(coeffs) - 1 # 1 coeff -> 0th degree
self._do_scale = apply_scaling
if apply_scaling and not (isinstance(obs, Space)
and obs.n_limits == 1):
raise ValueError(
"obs need to be a Space with exactly one limit if rescaling is requested."
)
super().__init__(obs=obs, name=name, params=params)
def test_simple_loss():
true_a = 1.0
true_b = 4.0
true_c = -0.3
truevals = true_a, true_b, true_c
a_param = zfit.Parameter("variable_a15151loss",
1.5,
-1.0,
20.0,
step_size=z.constant(0.1))
b_param = zfit.Parameter("variable_b15151loss", 3.5)
c_param = zfit.Parameter("variable_c15151loss", -0.23)
param_list = [a_param, b_param, c_param]
def loss_func(params):
a_param, b_param, c_param = params
probs = (z.convert_to_tensor((a_param - true_a)**2 +
(b_param - true_b)**2 +
(c_param - true_c)**4) + 0.42)
return tf.reduce_sum(input_tensor=tf.math.log(probs))
with pytest.raises(ValueError):
_ = zfit.loss.SimpleLoss(func=loss_func, params=param_list)
loss_func.errordef = 1
loss_deps = zfit.loss.SimpleLoss(func=loss_func, params=param_list)
# loss = zfit.loss.SimpleLoss(func=loss_func)
loss = zfit.loss.SimpleLoss(func=loss_func, params=param_list)
loss2 = zfit.loss.SimpleLoss(func=loss_func, params=truevals)
assert loss_deps.get_cache_deps() == set(param_list)
assert set(loss_deps.get_params()) == set(param_list)
loss_tensor = loss_func(param_list)
loss_value_np = loss_tensor.numpy()
assert loss.value().numpy() == pytest.approx(loss_value_np)
assert loss_deps.value().numpy() == pytest.approx(loss_value_np)
with pytest.raises(IntentionAmbiguousError):
_ = loss + loss_deps
minimizer = zfit.minimize.Minuit()
result = minimizer.minimize(loss=loss)
assert result.valid
assert true_a == pytest.approx(result.params[a_param]["value"], rel=0.03)
assert true_b == pytest.approx(result.params[b_param]["value"], rel=0.06)
assert true_c == pytest.approx(result.params[c_param]["value"], rel=0.5)
zfit.param.set_values(param_list, np.array(zfit.run(param_list)) + 0.6)
result2 = minimizer.minimize(loss=loss2)
assert result2.valid
params = list(result2.params)
assert true_a == pytest.approx(result2.params[params[0]]["value"],
rel=0.03)
assert true_b == pytest.approx(result2.params[params[1]]["value"],
rel=0.06)
assert true_c == pytest.approx(result2.params[params[2]]["value"], rel=0.5)
def true_nll_gaussian(x, mu, sigma):
x = convert_to_container(x, container=tuple)
mu = convert_to_container(mu, container=tuple)
sigma = convert_to_container(sigma, container=tuple)
constraint = z.constant(0.0)
if not len(x) == len(mu) == len(sigma):
raise ValueError("params, mu and sigma have to have the same length.")
for x_, mean, sig in zip(x, mu, sigma):
constraint += z.reduce_sum(z.square(x_ - mean) / (2.0 * z.square(sig)))
return constraint
def test_extended_unbinned_nll():
test_values = z.constant(test_values_np)
test_values = zfit.Data.from_tensor(obs=obs1, tensor=test_values)
gaussian3, mu3, sigma3, yield3 = create_gauss3ext()
nll = zfit.loss.ExtendedUnbinnedNLL(model=gaussian3,
data=test_values,
fit_range=(-20, 20))
assert {mu3, sigma3, yield3} == nll.get_params()
minimizer = Minuit()
status = minimizer.minimize(loss=nll)
params = status.params
assert params[mu3]['value'] == pytest.approx(np.mean(test_values_np), rel=0.007)
assert params[sigma3]['value'] == pytest.approx(np.std(test_values_np), rel=0.007)
assert params[yield3]['value'] == pytest.approx(yield_true, rel=0.007)
def dependents_func():
a = zfit.pdf.Gauss(var1, var2, obs='obs1').sample(n=500,
limits=(-5, 5)) * 5.
b = z.constant(3.) + 4 * var1
c = 5. * b * var3
d = b * var2 + a
e = d * 3.
return a, b, c, d, e
assert get_dependents_auto(
e, [a, b, c, d, var1, var2, var3]) == [a, b, d, var1, var2]
assert get_dependents_auto(e, [var1, var2, var3]) == [var1, var2]
assert get_dependents_auto(
c, [a, b, d, var1, var2, var3]) == [b, var1, var3]
return True
def test_extended_unbinned_nll(size):
if size is None:
test_values = z.constant(test_values_np)
size = test_values.shape[0]
else:
test_values = create_test_values(size)
test_values = zfit.Data.from_tensor(obs=obs1, tensor=test_values)
gaussian3, mu3, sigma3, yield3 = create_gauss3ext()
nll = zfit.loss.ExtendedUnbinnedNLL(model=gaussian3,
data=test_values,
fit_range=(-20, 20))
assert {mu3, sigma3, yield3} == nll.get_params()
minimizer = Minuit(tol=1e-4)
status = minimizer.minimize(loss=nll)
params = status.params
assert params[mu3]['value'] == pytest.approx(zfit.run(tf.math.reduce_mean(test_values)), rel=0.05)
assert params[sigma3]['value'] == pytest.approx(zfit.run(tf.math.reduce_std(test_values)), rel=0.05)
assert params[yield3]['value'] == pytest.approx(size, rel=0.005)
def test_simple_loss():
true_a = 1.
true_b = 4.
true_c = -0.3
a_param = zfit.Parameter("variable_a15151loss",
1.5,
-1.,
20.,
step_size=z.constant(0.1))
b_param = zfit.Parameter("variable_b15151loss", 3.5)
c_param = zfit.Parameter("variable_c15151loss", -0.23)
param_list = [a_param, b_param, c_param]
def loss_func():
probs = z.convert_to_tensor((a_param - true_a)**2 +
(b_param - true_b)**2 +
(c_param - true_c)**4) + 0.42
return tf.reduce_sum(input_tensor=tf.math.log(probs))
loss_deps = zfit.loss.SimpleLoss(func=loss_func, deps=param_list)
# loss = zfit.loss.SimpleLoss(func=loss_func)
loss = zfit.loss.SimpleLoss(func=loss_func, deps=param_list)
assert loss_deps.get_cache_deps() == set(param_list)
assert set(loss_deps.get_params()) == set(param_list)
loss_tensor = loss_func()
loss_value_np = loss_tensor.numpy()
assert loss.value().numpy() == loss_value_np
assert loss_deps.value().numpy() == loss_value_np
with pytest.raises(IntentionAmbiguousError):
_ = loss + loss_deps
minimizer = zfit.minimize.Minuit()
result = minimizer.minimize(loss=loss)
assert true_a == pytest.approx(result.params[a_param]['value'], rel=0.03)
assert true_b == pytest.approx(result.params[b_param]['value'], rel=0.06)
assert true_c == pytest.approx(result.params[c_param]['value'], rel=0.5)
def do_integrate():
return tf.while_loop(
cond=all_calculated,
body=body,
loop_vars=[
z.constant(0.), # integral
initial_points[:-1], # lower
initial_points[1:], # upper
0 # n_iter
],
# here we specify the shape of the loop_vars: since they change (of the second and third),
# we need to specify them, with None as "shape is not fixed". For the integral as well as for
# the number of iterations, this is a scalar with shape ()
shape_invariants=[
tf.TensorShape(()),
tf.TensorShape((None, )),
tf.TensorShape((None, )),
tf.TensorShape(()),
],
maximum_iterations=n_iter_max,
)
def test_wrapped_func():
rnd = z.random.uniform(shape=(10, ))
result = wrapped_numpy_func(rnd, z.constant(3.0))
np.testing.assert_allclose(rnd * np.sqrt(3), result)
def model(self, x):
d = z.constant(0.)
for i in self.params:
d += gauss_2d(x, i[0], i[1], i[2], i[3], i[4], i[5])
return d
from zfit import z
def integrate(func, lower, upper):
func_upper = func(upper)
func_lower = func(lower)
uncertainty = tf.abs(func_upper - func_lower) # can be improved of course
integrals = (func_lower + func_upper) / 2 * (upper - lower)
return integrals, uncertainty
# func = lambda x: tf.where(tf.less(x, 0.1),
# tf.sin(x * 100),
# tf.sin(x))
func = lambda x: tf.sin(x) + tf.cos(x * 100) # example func to integrate
lower, upper = z.constant(0.), z.constant(math.pi)
n_iter_max = 32 # maximum iteration: if we have a discontinuous function, we won't reacht the precision requested
# so we should break
def body(integral, lower, upper, n_iter):
integrals, uncertainties = integrate(func, lower, upper)
uncertainties_too_large = tf.greater(uncertainties, 1e-5)
# if we reached the max number of iterations, we take the values anyway, so the uncertainties are just "too large",
# or need to be redone if we did not yet reach the max iterations
uncertainties_too_large = tf.logical_and(uncertainties_too_large,
n_iter < n_iter_max)
too_large_indices = tf.where(uncertainties_too_large)[:, 0]
def chunked_average(func, x, num_batches, batch_size, space, mc_sampler):
lower, upper = space.limits
fake_resource_var = tf.Variable("fake_hack_ResVar_for_custom_gradient",
initializer=z.constant(4242.))
fake_x = z.constant(42.) * fake_resource_var
@tf.custom_gradient
def dummy_func(fake_x): # to make working with custom_gradient
if x is not None:
raise WorkInProgressError("partial not yet implemented")
def body(batch_num, mean):
if mc_sampler == tfp.mcmc.sample_halton_sequence:
start_idx = batch_num * batch_size
end_idx = start_idx + batch_size
indices = tf.range(start_idx, end_idx, dtype=tf.int32)
sample = mc_sampler(space.n_obs,
sequence_indices=indices,
dtype=ztypes.float,
randomized=False)
else:
sample = mc_sampler(shape=(batch_size, space.n_obs),
dtype=ztypes.float)
sample = tf.guarantee_const(sample)
sample = (np.array(upper[0]) -
np.array(lower[0])) * sample + lower[0]
sample = tf.transpose(a=sample)
sample = func(sample)
sample = tf.guarantee_const(sample)
batch_mean = tf.reduce_mean(input_tensor=sample)
batch_mean = tf.guarantee_const(batch_mean)
err_weight = 1 / tf.cast(batch_num + 1, dtype=tf.float64)
# err_weight /= err_weight + 1
# print_op = tf.print(batch_mean)
do_print = False
if do_print:
tf.print(batch_num + 1, mean, err_weight * (batch_mean - mean))
return batch_num + 1, mean + err_weight * (batch_mean - mean)
cond = lambda batch_num, _: batch_num < num_batches
initial_mean = tf.convert_to_tensor(value=0, dtype=ztypes.float)
_, final_mean = tf.while_loop(cond=cond,
body=body,
loop_vars=(0, initial_mean),
parallel_iterations=1,
swap_memory=False,
back_prop=False,
maximum_iterations=num_batches)
def dummy_grad_with_var(dy, variables=None):
raise WorkInProgressError("Who called me? Mayou36")
if variables is None:
raise WorkInProgressError(
"Is this needed? Why? It's not a NN. Please make an issue."
)
def dummy_grad_func(x):
values = func(x)
if variables:
gradients = tf.gradients(ys=values,
xs=variables,
grad_ys=dy)
else:
gradients = None
return gradients
return chunked_average(func=dummy_grad_func,
x=x,
num_batches=num_batches,
batch_size=batch_size,
space=space,
mc_sampler=mc_sampler)
def dummy_grad_without_var(dy):
return dummy_grad_with_var(dy=dy, variables=None)
do_print = False
if do_print:
tf.print("Total mean calculated = ", final_mean)
return final_mean, dummy_grad_with_var
try:
return dummy_func(fake_x)
except TypeError:
return dummy_func(fake_x)
