def ubar0(p, i):
"""Compute the dirac spinor ubar0"""
zeros = tf.zeros_like(p[0], dtype=DTYPE)
czeros = tf.complex(zeros, zeros)
ones = tf.ones_like(p[0], dtype=DTYPE)
# case 1) py == 0
rz = p[3] / (p[0] + epsilon)
theta1 = tf.where(rz > 0, zeros, rz)
theta1 = tf.where(rz < 0, np.pi * ones, theta1)
phi1 = zeros
# case 2) py != 0
rrr = rz
rrr = tf.where(rz < -1, -ones, rrr)
rrr = tf.where(rz > 1, ones, rrr)
theta2 = tf.acos(rrr)
rrr = p[1] / p[0] / tf.sin(theta2)
rrr = tf.where(rrr < -1, -ones, rrr)
rrr = tf.where(rrr > 1, ones, rrr)
phi2 = tf.acos(rrr)
ry = p[2] / p[0]
phi2 = tf.where(ry < 0, -phi2, phi2)
# combine
theta = tf.where(p[1] == 0, theta1, theta2)
phi = tf.where(p[1] == 0, phi1, phi2)
prefact = tf.complex(float_me(np.sqrt(2)), zeros) * tf.sqrt(
tf.complex(p[0], zeros))
if i == -1:
a = tf.complex(tf.sin(theta / 2), zeros)
b = tf.complex(tf.abs(tf.cos(theta / 2)), zeros)
ubar0_0 = prefact * a * tf.complex(tf.cos(phi), tf.sin(phi))
ubar0_1 = -prefact * b
ubar0_2 = czeros
ubar0_3 = czeros
else:
a = tf.complex(tf.cos(theta / 2), zeros)
b = tf.complex(tf.sin(theta / 2), zeros)
ubar0_0 = czeros
ubar0_1 = czeros
ubar0_2 = prefact * a
ubar0_3 = prefact * b * tf.complex(tf.cos(phi), -tf.sin(phi))
return tf.stack([ubar0_0, ubar0_1, ubar0_2, ubar0_3])
def py_consume_array_into_indices(input_arr, indices, result_size):
"""
Python interface wrapper for ``consume_array_into_indices``.
It casts the possible python-object input into the correct tensorflow types.
"""
return consume_array_into_indices(float_me(input_arr), int_me(indices), int_me(result_size))
import time import numpy as np import subprocess as sp from pdfflow.configflow import DTYPEINT as pint import tensorflow as tf from vegasflow.vflow import vegas_wrapper from pdfflow.pflow import mkPDF from vegasflow.configflow import DTYPE, DTYPEINT, float_me from pdfflow.configflow import DTYPE as pfloat # MC integration setup dim = 3 ncalls = np.int32(5e6) n_iter = 5 pdfset = "NNPDF31_nlo_as_0118/0" epsilon = float_me(1e-7) # Physics setup # top mass mt = tf.constant(173.2, dtype=DTYPE) # center of mass energy sqrts = tf.constant(8000, dtype=DTYPE) # minimum allowed center of mass energy sqrtsmin = tf.constant(173.2, dtype=DTYPE) # W-boson mass mw = tf.constant(80.419, dtype=DTYPE) # gaw gaw = tf.constant(2.1054, dtype=DTYPE) # GF gf = tf.constant(1.16639e-5, dtype=DTYPE) # load pdf
def refine_grid_per_dimension(t_res_sq, subdivisions):
"""
Modifies the boundaries for the vegas grid for a given dimension
Parameters
----------
`t_res_sq`: tensor
array of results squared per bin
`subdivision`: tensor
current boundaries for the grid
Returns
-------
`new_divisions`: tensor
array with the new boundaries of the grid
"""
# Define some constants
paddings = int_me([[1, 1]])
tmp_meaner = tf.fill([BINS_MAX - 2], float_me(3.0))
meaner = tf.pad(tmp_meaner, paddings, constant_values=2.0)
# Pad the vector of results
res_padded = tf.pad(t_res_sq, paddings)
# First we need to smear out the array of results squared
smeared_tensor_tmp = res_padded[1:-1] + res_padded[2:] + res_padded[:-2]
smeared_tensor = tf.maximum(smeared_tensor_tmp / meaner, float_me(1e-30))
# Now we refine the grid according to
# journal of comp phys, 27, 192-203 (1978) G.P. Lepage
sum_t = tf.reduce_sum(smeared_tensor)
log_t = tf.math.log(smeared_tensor)
aux_t = (1.0 - smeared_tensor / sum_t) / (tf.math.log(sum_t) - log_t)
wei_t = tf.pow(aux_t, ALPHA)
ave_t = tf.reduce_sum(wei_t) / BINS_MAX
###### Auxiliary functions for the while loop
@tf.function
def while_check(bin_weight, *args):
"""Checks whether the bin has enough weight
to beat the average"""
return bin_weight < ave_t
@tf.function(input_signature=[
tf.TensorSpec(shape=[], dtype=DTYPE),
tf.TensorSpec(shape=[], dtype=DTYPEINT),
tf.TensorSpec(shape=[], dtype=DTYPE),
tf.TensorSpec(shape=[], dtype=DTYPE),
])
def while_body(bin_weight, n_bin, cur, prev):
"""Fills the bin weight until it surpassed the avg
once it's done, returns the limits of the last bin"""
n_bin += 1
bin_weight += wei_t[n_bin]
prev = cur
cur = subdivisions[n_bin + 1]
return bin_weight, n_bin, cur, prev
###########################
# And now resize all bins
new_bins = [fzero]
# Auxiliary variables
bin_weight = fzero
n_bin = -1
cur = fzero
prev = fzero
for _ in range(BINS_MAX - 1):
bin_weight, n_bin, cur, prev = tf.while_loop(
while_check,
while_body,
(bin_weight, n_bin, cur, prev),
parallel_iterations=1,
)
bin_weight -= ave_t
delta = (cur - prev) * bin_weight / wei_t[n_bin]
new_bins.append(cur - delta)
new_bins.append(fone)
new_divisions = tf.stack(new_bins)
return new_divisions
`vegas_wrapper`
"""
from vegasflow.configflow import DTYPE, DTYPEINT, fone, fzero, float_me, ione, int_me
import json
import numpy as np
import tensorflow as tf
from vegasflow.configflow import BINS_MAX, ALPHA
from vegasflow.monte_carlo import MonteCarloFlow, wrapper, sampler
from vegasflow.utils import consume_array_into_indices
import logging
logger = logging.getLogger(__name__)
FBINS = float_me(BINS_MAX)
# Auxiliary functions for Vegas
@tf.function(input_signature=[
tf.TensorSpec(shape=[None, None], dtype=DTYPE),
tf.TensorSpec(shape=[None, BINS_MAX + 1], dtype=DTYPE),
])
def _generate_random_array(rnds, divisions):
"""
Generates the Vegas random array for any number of events
Parameters
----------
rnds: array shaped (None, n_dim)
Random numbers used as an input for Vegas
def xjac(self):
"""The default jacobian is 1 / total number of events"""
return float_me([1.0 / self.n_events])
def test_float_me():
res = float_me(4.0)
assert res.dtype == DTYPE