def main() :
nev = int(sys.argv[1])
outfile = sys.argv[2]
chunk_size = 1000000 # Events will be generated in chunks of this size
atfi.set_seed(nev)
n = 0 # Current tuple size
arrays = []
while(True) :
# Create Dalitz plot sample
unfiltered_sample = observables_phase_space.unfiltered_sample(chunk_size) # Unfiltered array
sample = observables_phase_space.filter(unfiltered_sample)
size = sample.shape[0]
print(f"Filtered chunk size = {size}")
# Generate final state momenta from Dalitz plot and run through selection
rnd = tf.random.uniform([size, random_array_size], dtype = atfi.fptype() ) # Auxiliary random array
array = atfi.stack(generate_selection( true_cuts, rnd, constant_cuts = True ), axis = 1)
arrays += [ array ]
# Increment counters and check if we are done
size = array.shape[0]
n += size
if n > nev : break
print(f"Selected size = {n}, last = {size}")
tfr.write_tuple(outfile, atfi.concat(arrays, axis = 0)[:nev,:], observables_toys)
def main():
nev = 1000000
outfile = "toy_tuple.root"
atfi.set_seed(nev + 1)
chunk_size = 1000000 # Events will be generated in chunks of this size
bounds = {i[0]: (i[2], i[3])
for i in parameters_list
} # Bounds and exponential factor for generation of cuts
branches = generated_variables + [i[0] for i in parameters_list]
n = 0 # Current tuple size
arrays = []
while (True):
# Generate final state momenta from Dalitz plot and run through selection
rnd = tf.random.uniform([chunk_size, random_array_size + len(bounds)],
dtype=atfi.fptype()) # Auxiliary random array
array = atfi.stack(generate_candidates_and_cuts(rnd), axis=1)
arrays += [array]
# Increment counters and check if we are done
size = array.shape[0]
n += size
if n > nev: break
print(f"Selected size = {n}, last = {size}")
tfr.write_tuple(outfile, atfi.concat(arrays, axis=0)[:nev, :], branches)
def run_toymc(pdf, phsp, size, maximum, chunk=200000, seed=None):
"""
Create toy MC sample. To save memory, the sample is generated in "chunks" of a fixed size
pdf : Function returning PDF graph for a given sample as an agrument
phsp : phase space
size : size of the target data sample (if >0) or number of chunks (if <0)
maximum : maximum PDF value for accept-reject method
chunk : chunk size
seed : initial random seed. Not initalised if None
"""
length, nchunk, curr_maximum = 0, 0, maximum
dim = phsp.dimensionality()
data = tf.Variable(np.empty((0, dim)),
shape=(None, dim),
dtype=atfi.fptype())
if seed is not None:
atfi.set_seed(seed)
def condition(length, size, nchunk):
return length < size or nchunk < -size
@atfi.function
def pdf_vals(chunk, curr_maximum):
d = accept_reject_sample(
pdf, phsp.filter(phsp.unfiltered_sample(chunk, curr_maximum)))
return d, pdf(d)
print(type(length), type(size), type(nchunk))
while condition(length, size, nchunk):
d, v = pdf_vals(chunk, curr_maximum)
over_maximum = v[v > curr_maximum]
if len(over_maximum) > 0:
new_maximum = tf.reduce_max(over_maximum) * 1.5
print(
f' Updating maximum: {curr_maximum} -> {new_maximum}. Starting over.'
)
length, nchunk, curr_maximum = 0, 0, new_maximum
data = tf.Variable(np.empty((0, dim)),
shape=(None, dim),
dtype=atfi.fptype())
continue
data = tf.concat([data, d], axis=0)
length += len(d)
nchunk += 1
print(f' Chunk {nchunk}, size={len(d)}, total length={length}')
return data[:size] if size > 0 else data
def main():
nev = 2000000
outfile = "toy_tuple.root"
atfi.set_seed(nev + 1)
chunk_size = 1000000 # Events will be generated in chunks of this size
bounds = {i[0]: (i[2], i[3])
for i in parameters_list
} # Bounds and exponential factor for generation of cuts
branches = observables_toys + [i[0] for i in parameters_list]
n = 0 # Current tuple size
arrays = []
while (True):
# Create Dalitz plot sample
unfiltered_sample = observables_phase_space.unfiltered_sample(
chunk_size) # Unfiltered array
sample = observables_phase_space.filter(unfiltered_sample)
size = sample.shape[0]
print(f"Filtered chunk size = {size}")
# Generate final state momenta from Dalitz plot and run through selection
rnd = tf.random.uniform([size, random_array_size],
dtype=atfi.fptype()) # Auxiliary random array
array = atfi.stack(selection_with_random_cuts(sample, rnd), axis=1)
arrays += [array]
# Increment counters and check if we are done
size = array.shape[0]
n += size
if n > nev: break
print(f"Selected size = {n}, last = {size}")
tfr.write_tuple(outfile, atfi.concat(arrays, axis=0)[:nev, :], branches)
def estimate_density(
phsp,
data,
ranges,
labels,
weight=None,
transform=None,
transform_ranges=None,
learning_rate=0.001,
training_epochs=100000,
norm_size=1000000,
print_step=50,
display_step=500,
weight_penalty=1.,
n_hidden=[32, 8],
initfile="init.npy",
outfile="train",
seed=1,
fig=None,
axes=None,
units=None,
model=None,
regularisation=None,
):
n_input = len(ranges)
bins = n_input * [50]
atfi.set_seed(seed)
if model is None:
try:
init_w = np.load(initfile, allow_pickle=True)
except:
init_w = None
if isinstance(init_w, np.ndarray):
print("Loading saved weights")
(weights, biases) = init_weights_biases(init_w[2:])
else:
print("Creating random weights")
(weights, biases) = create_weights_biases(n_input, n_hidden)
if not transform_ranges: transform_ranges = ranges
def fitmodel(x):
if transform: x2 = transform(x)
else: x2 = x
# to make sure PDF is always strictly positive
return multilayer_perceptron(x2, transform_ranges, weights,
biases) + 1e-20
parameters = [weights, biases]
else:
fitmodel, parameters = model
@tf.function
def unbinned_nll(pdf, integral):
return -tf.reduce_sum(atfi.log(pdf / integral))
@tf.function
def integral(pdf):
return tf.reduce_mean(pdf)
opt = tf.keras.optimizers.Adam(learning_rate=learning_rate)
data_sample = phsp.filter(data)
if isinstance(norm_size, int):
norm_sample = phsp.uniform_sample(norm_size)
else:
norm_sample = norm_size
print("Normalisation sample size = ", len(norm_sample))
print(norm_sample)
print("Data sample size = ", len(data_sample))
print(data_sample)
if model is None:
# Define loss and optimizer
if regularisation is None:
@tf.function
def nll():
return unbinned_nll(
fitmodel(data_sample), integral(fitmodel(norm_sample))
) + l2_regularisation(weights) * weight_penalty
else:
@tf.function
def nll():
return unbinned_nll(
fitmodel(data_sample), integral(
fitmodel(norm_sample))) + regularisation(weights)
else:
@tf.function
def nll():
return unbinned_nll(fitmodel(data_sample),
integral(fitmodel(norm_sample)))
# Training cycle
best_cost = 1e10
display = atfp.MultidimDisplay(data_sample,
norm_sample,
bins,
ranges,
labels,
fig,
axes,
units=units)
plt.ion()
plt.show()
plt.pause(0.1)
for epoch in range(training_epochs):
if epoch % display_step == 0 and fig:
w = fitmodel(norm_sample)
display.draw(w)
plt.tight_layout(pad=0., w_pad=0., h_pad=0.)
plt.draw()
plt.pause(0.1)
plt.savefig(outfile + ".pdf")
if epoch % print_step == 0:
c = nll()
s = "Epoch %d, cost %.9f" % (epoch + 1, c)
print(s)
if c < best_cost:
best_cost = c
w = fitmodel(norm_sample)
scale = 1. / np.mean(w)
if model is None:
weights = [w.numpy() for w in parameters[0]]
biases = [b.numpy() for b in parameters[1]]
np.save(outfile,
[scale, transform_ranges, weights, biases])
else:
np.save(outfile, [scale, transform_ranges] +
[p.numpy() for p in parameters()])
f = open(outfile + ".txt", "w")
f.write(s + "\n")
f.close()
opt.minimize(nll, parameters)
print("Optimization Finished!")
dp_phi_r,
dp_theta_d,
dp_phi_d,
1,
spin,
pol_lb,
-1,
0,
pol_p,
0,
))
density += atfi.density(ampl)
return density
atfi.set_seed(2)
# Produce normalisation sample (rectangular 2D grid of points)
norm_sample = phsp.rectangular_grid_sample(norm_grid, norm_grid)
print("Normalisation sample size = ", norm_sample.shape)
print(norm_sample)
# Calculate maximum of the PDF for accept-reject toy MC generation
maximum = tft.maximum_estimator(model, phsp, 100000) * 1.5
print("Maximum = ", maximum)
# Create toy MC data sample
data_sample = tft.run_toymc(model,
phsp,
toy_sample,
maximum,
atfi.backend_tf()
npoints = 100000000
@atfi.function
def func(x):
return 1.0 / math.sqrt(math.pi) * atfi.exp(-(x**2))
@atfi.function
def integral(x):
return atfi.reduce_sum(func(x)) / npoints * 10.0
atfi.set_seed(1)
x = atfi.random_uniform((npoints, 1), -5.0, 5.0)
strategy = tf.distribute.MirroredStrategy()
global_batch_size = x.shape[0]
tf_dataset = tf.data.Dataset.from_tensor_slices(x).batch(global_batch_size)
print(tf_dataset)
dist_dataset = strategy.experimental_distribute_dataset(tf_dataset)
print(dist_dataset)
dist_values = iter(dist_dataset).get_next()
print(dist_values)
@tf.function
def distributed_integral(x):
observables_bounds = exp_phase_space.bounds()
# Density model as a multilayer perceptron
def model(x, pars):
# Constant vectors of fit parameters (the same for every data point)
vec = tf.reshape(
tf.concat([tf.constant(ndim * [0.], dtype=atfi.fptype()), pars],
axis=0), [1, ndim + len(pars)])
# Input tensor for MLP, 5 first variables are data, the rest are constant optimisable parameters
x2 = tf.pad(x, [[0, 0], [0, len(pars)]], 'CONSTANT') + vec
return scale * tfn.multilayer_perceptron(x2, ranges, weights, biases)
# Initialise random seeds
atfi.set_seed(seed)
# Declare fit parameters
pars = [
tfo.FitParameter(par[0], (par[2][0] + par[2][1]) / 2., par[2][0],
par[2][1]) for par in parameters_list
]
# Unbinned negative log likelihood
@atfi.function
def nll(pars):
parslist = [pars[i[0]] for i in parameters_list]
return atfl.unbinned_nll(model(data_sample, parslist),
atfl.integral(model(norm_sample, parslist)))
def run_toymc(pdf, phsp, size, maximum, chunk=200000, seed=None, components=True):
"""
Create toy MC sample. To save memory, the sample is generated in "chunks" of a fixed size
pdf : Function returning PDF graph for a given sample as an agrument
phsp : phase space
size : size of the target data sample (if >0) or number of chunks (if <0)
maximum : maximum PDF value for accept-reject method
chunk : size of chunk (initial number of events before rejection, generated in parallel)
seed : initial random seed. Not initalised if None
components : if True, generate weights for components
(if the argument "switches" is available in pdf)
"""
import inspect
length, nchunk, curr_maximum = 0, 0, maximum
if seed is not None:
atfi.set_seed(seed)
def condition(length, size, nchunk):
return length < size or nchunk < -size
@tf.function
def pdf_vals(chunk, curr_maximum):
d = accept_reject_sample(
pdf, phsp.filter(phsp.unfiltered_sample(chunk, curr_maximum))
)
return d, pdf(d)
args, varargs, keywords, defaults = inspect.getargspec(pdf)
num_switches = 0
if defaults:
default_dict = dict(zip(args[-len(defaults) :], defaults))
if "switches" in default_dict:
num_switches = len(default_dict["switches"])
# dim = phsp.dimensionality()
# data = tf.Variable(np.empty((0, dim)), shape=(None, dim), dtype = atfi.fptype())
data = None
# @tf.function
def pdf_components(d):
result = []
for i in range(num_switches):
switches = num_switches * [0]
switches[i] = 1
result += [pdf(d, switches=tuple(switches))]
return result
while condition(length, size, nchunk):
d, v = pdf_vals(chunk, curr_maximum)
over_maximum = v[v > curr_maximum]
if len(over_maximum) > 0:
new_maximum = atfi.reduce_max(over_maximum) * 1.5
print(
f" Updating maximum: {curr_maximum} -> {new_maximum}. Starting over."
)
length, nchunk, curr_maximum = 0, 0, new_maximum
# data = tf.Variable(np.empty((0, dim)), shape=(None, dim), dtype = atfi.fptype())
data = None
continue
if components and num_switches > 0:
vs = pdf_components(d)
wd = tf.stack([i / v for i in vs], axis=1)
d = tf.concat([d, wd], axis=1)
if data is not None:
data = tf.concat([data, d], axis=0)
else:
data = d
length += len(d)
nchunk += 1
print(f" Chunk {nchunk}, size={len(d)}, total length={length}")
return data[:size] if size > 0 else data
