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