def pdfflow_tester(pdf, members=None):
"""Test several pdfflow features:
- Check the single/many/all-pid signatures
- Checks the python and TF signatures
- Checks the output of the python and TF signatures are the same
- Checks the expected shape of the signatures is correct
"""
grid_size = 7
x = np.random.rand(grid_size)
q2 = 1000.0 * np.random.rand(grid_size)
xtf = float_me(x)
q2tf = float_me(q2)
# Check I can get just one pid
for i in range(-1, 2):
# as int
res_1 = pdf.py_xfxQ2(i, x, q2)
# as list
res_2 = pdf.py_xfxQ2([i], x, q2)
np.testing.assert_allclose(res_1, res_2)
# as tf objects
tfpid = int_me([i])
res_3 = pdf.xfxQ2(tfpid, xtf, q2tf)
np.testing.assert_allclose(res_2, res_3)
# Check shape
if members is None:
assert res_1.numpy().shape == (grid_size, )
else:
assert res_1.numpy().shape == (
members,
grid_size,
)
# Check I can get more than one pid
nfl_size = 6
fl_scheme = pdf.flavor_scheme.numpy()
nfl_total = fl_scheme.size
many_pid = np.random.choice(fl_scheme, nfl_size)
res_1 = pdf.py_xfxQ2(many_pid, x, q2)
res_2 = pdf.xfxQ2(int_me(many_pid), xtf, q2tf)
np.testing.assert_allclose(res_1, res_2)
# Check shape
if members is None:
assert res_1.numpy().shape == (grid_size, nfl_size)
else:
assert res_1.numpy().shape == (members, grid_size, nfl_size)
# Check I can actually get all PID
res_1 = pdf.py_xfxQ2_allpid(x, q2)
res_2 = pdf.xfxQ2_allpid(xtf, q2tf)
np.testing.assert_allclose(res_1, res_2)
# Check shape
if members is None:
assert res_1.numpy().shape == (grid_size, nfl_total)
else:
assert res_1.numpy().shape == (members, grid_size, nfl_total)
def alphas_first_subgrid(
shape,
a_q2,
log_q2min,
log_q2max,
padded_q2,
s_q2,
actual_padded,
):
"""
First subgrid interpolation
Calls
alphas_interpolate (basic interpolation) (0)
Parameters
----------
shape: tf.tensor of shape [None]
final output shape to scatter points into
For other parameters refer to subgrid.py:alphas_interpolate
Returns
----------
tf.tensor of shape `shape`
alphas interpolated values for each query point
"""
res = tf.zeros(shape, dtype=DTYPE)
# --------------------------------------------------------------------
# normal interpolation
stripe, f_idx = _condition_to_idx(a_q2 >= log_q2min, a_q2 < log_q2max)
if tf.math.equal(f_idx, 0) is not None:
in_q2 = tf.boolean_mask(a_q2, stripe)
ff_f = alphas_interpolate(in_q2, padded_q2, s_q2, actual_padded)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
# --------------------------------------------------------------------
# lowq2
stripe = a_q2 < log_q2min
f_idx = int_me(tf.where(stripe))
if tf.math.equal(f_idx, 0) is not None:
in_q2 = tf.boolean_mask(a_q2, stripe)
m = tf.math.log(actual_padded[2]/actual_padded[1])\
/(padded_q2[2] - padded_q2[1])
ff_f = actual_padded[1] * tf.math.pow(
tf.math.exp(in_q2) / tf.math.exp(padded_q2[1]), m)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
return res
def py_xfxQ2(self, pid, arr_x, arr_q2):
"""
Python interface for pdfflow
The input gets converted to the right tensorflow type
before calling the corresponding function: `xfxQ2`
returns PDF evaluated in each f(x,q2) for each pid
The output of the function is of shape
(members, number_of_points, flavours)
but note that dimensions of size 1 will be squeezed out.
Example
-------
>>> from pdfflow.pflow import mkPDFs
>>> from pdfflow.configflow import run_eager, float_me, int_me
>>> run_eager(True)
>>> pdf = mkPDFs("NNPDF31_nlo_as_0118", [0, 1, 4])
>>> pdf.py_xfxQ2(21, [0.4, 0.5], [15625.0, 15625.0])
<tf.Tensor: shape=(3, 2), dtype=float64, numpy=
array([[0.02977381, 0.00854525],
[0.03653673, 0.00929325],
[0.031387 , 0.00896622]])>
>>> pdf.py_xfxQ2([1,2], [0.4], [15625.0])
<tf.Tensor: shape=(3, 2), dtype=float64, numpy=
array([[0.05569674, 0.19323399],
[0.05352555, 0.18965438],
[0.04515956, 0.18704451]])>
Parameters
----------
pid: list(int)
list of PID to be computed
arr_x: np.array
grid on x where to compute the PDF
arr_q2: np.array
grid on q^2 where to compute the PDF
Returns
-------
pdf: tensor
grid of results ([members], [number of points], [flavour])
"""
# if pid is None, the mask is set to true everywhere
if pid is None:
return self.py_xfxQ2_allpid(arr_x, arr_q2)
tensor_pid = tf.reshape(int_me(pid), (-1, ))
a_x = float_me(arr_x)
a_q2 = float_me(arr_q2)
return self.xfxQ2(tensor_pid, a_x, a_q2)
def alphas_last_subgrid(
shape,
a_q2,
log_q2min,
log_q2max,
padded_q2,
s_q2,
actual_padded,
):
"""
Last subgrid interpolation.
Calls
alphas_interpolate: (0)
Parameters
----------
shape: tf.tensor of shape [None]
final output shape to scatter points into
For other parameters see :py:func:`pdfflow.alphas_region_interpolator.alphas_interpolate`
Returns
----------
tf.tensor, shape: `shape`
alphas interpolated values for each query point
"""
# Generate all conditions for all stripes
res = tf.zeros(shape, dtype=DTYPE)
# --------------------------------------------------------------------
# normal interpolation
stripe, f_idx = _condition_to_idx(a_q2 >= log_q2min, a_q2 <= log_q2max)
if tf.math.equal(f_idx, 0) is not None:
# Check whether there are any points in this region
# if there are, execute normal_interpolation
in_q2 = tf.boolean_mask(a_q2, stripe)
ff_f = alphas_interpolate(in_q2, padded_q2, s_q2, actual_padded)
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
# --------------------------------------------------------------------
# high q2
stripe = a_q2 > log_q2max
f_idx = int_me(tf.where(stripe))
if tf.math.equal(f_idx, 0) is not None:
ff_f = tf.ones_like(f_idx[:, 0], dtype=DTYPE) * actual_padded[-2]
res = tf.tensor_scatter_nd_update(res, f_idx, ff_f)
return res
def _condition_to_idx(cond1, cond2):
""" Take two boolean masks and returns the indexes in which both are true """
full_condition = tf.logical_and(cond1, cond2)
return full_condition, int_me(tf.where(full_condition))
def __init__(self,
grid,
i=0,
total=0,
compile_functions=True,
alpha_s=False):
name_sg = f"grid_{i}"
self.alpha_s = alpha_s
if alpha_s:
name_sg += "_alpha"
self.name_sg = name_sg
super().__init__(name=f"Parent_{name_sg}")
q2min = min(grid.q2)
q2max = max(grid.q2)
self.log_q2min = float_me(np.log(q2min))
self.log_q2max = float_me(np.log(q2max))
# Save grid shape information
self.s_q2 = int_me(grid.q2.size)
# Insert a padding at the beginning and the end
log_q2pad = np.pad(np.log(grid.q2), 1, mode="edge")
log_q2pad[0] *= 0.99
log_q2pad[-1] *= 1.01
self.padded_q2 = float_me(log_q2pad)
# Depending on whether it is an alphs_s grid or a pdf grid
# we might need to change some options
compilation_options = OPT.copy()
if alpha_s:
# the grid is sized (q.size), pad it with 0s
self.padded_grid = float_me(np.pad(grid.grid, (1, 1)))
if i == 0:
self.fn_interpolation = alphas_first_subgrid
elif i == (total - 1):
self.fn_interpolation = alphas_last_subgrid
else:
self.fn_interpolation = alphas_inner_subgrid
# Change the function signature to that of alpha_s
compilation_options[
"input_signature"] = ALPHAS_GRID_FUNCTION_SIGNATURE
else:
# If this is a pdf grid, save also the x information
xmin = min(grid.x)
xmax = max(grid.x)
self.log_xmin = float_me(np.log(xmin))
self.log_xmax = float_me(np.log(xmax))
self.s_x = int_me(grid.x.size)
log_xpad = np.pad(np.log(grid.x), 1, mode="edge")
log_xpad[0] *= 0.99
log_xpad[-1] *= 1.01
self.padded_x = float_me(log_xpad)
# Finally parse the grid
# the grid is sized (x.size * q.size, flavours)
reshaped_grid = grid.grid.reshape(grid.x.size, grid.q2.size, -1)
# and pad it with 0s in x and q
padded_grid = np.pad(reshaped_grid, ((1, 1), (1, 1), (0, 0)))
# flatten the x and q dimensions again and store it
self.padded_grid = float_me(padded_grid.reshape(
-1, grid.flav.size))
# Depending on the index of the grid, select with interpolation function should be run
if i == 0:
self.fn_interpolation = first_subgrid
elif i == (total - 1):
self.fn_interpolation = last_subgrid
else:
self.fn_interpolation = inner_subgrid
if compile_functions:
self.fn_interpolation = tf.function(self.fn_interpolation,
**compilation_options)
def __init__(self, dirname, fname, members, compilable=True):
if not compilable:
logger.warning("Running pdfflow in eager mode")
logger.warning("Setting eager mode will affect all of TF")
tf.config.experimental_run_functions_eagerly(True)
self.dirname = dirname
self.fname = fname
self.grids = []
info_file = os.path.join(self.dirname, self.fname, f"{fname}.info")
# Load the info file
with open(info_file, "r") as ifile:
self.info = yaml.load(ifile, Loader=yaml.FullLoader)
if members is None:
total_members = self.info.get("NumMembers", 1)
members = range(total_members)
if len(members) == 1:
logger.info("Loading member %d from %s", members[0], self.fname)
else:
logger.info("Loading %d members from %s", len(members), self.fname)
for member_int in members:
member = str(member_int).zfill(4)
filename = os.path.join(self.dirname, fname,
f"{fname}_{member}.dat")
logger.debug("Loading %s", filename)
grids = _load_data(filename)
subgrids = [
Subgrid(grid, i, len(grids)) for i, grid in enumerate(grids)
]
self.grids.append(subgrids)
self.members = members
# Get the flavour scheme from the info file
flavor_scheme = np.array(self.info.get("Flavors", None))
if flavor_scheme is None:
# fallback to getting the scheme from the first grid, as all grids should have the same number of flavours
# if there's a failure here it should be the grid fault so no need to check from our side?
flavor_scheme = grids[0].flav
# Look at the flavor_scheme and ensure that it is sorted
# save the whole thing in case it is not sorted
self.flavor_scheme = int_me(flavor_scheme)
flavor_scheme[flavor_scheme == PID_G.numpy()] = 0
if all(np.diff(flavor_scheme) == 1):
self.flavors_sorted = True
self.flavor_shift = -flavor_scheme[0]
else:
# TODO can't we rely on the PDF flavours to be sorted?
self.flavors_sorted = False
self.flavor_shift = 0
# Finalize by loading the alpha information form the .info file
alpha_grids = _load_alphas(self.info)
self.alphas_subgrids = [
Subgrid(grid, i, len(grids), alpha_s=True)
for i, grid in enumerate(alpha_grids)
]
import logging
import collections
import yaml
import subprocess as sp
import numpy as np
import os, sys
# import configflow before tf to set some tf options
from pdfflow.configflow import DTYPE, DTYPEINT, int_me, izero, float_me, find_pdf_path
import tensorflow as tf
from pdfflow.subgrid import Subgrid
# lhapdf gluon code
PID_G = int_me(21)
# expected input shapes to be found in this module
GRID_F = tf.TensorSpec(shape=[None], dtype=DTYPE)
GRID_I = tf.TensorSpec(shape=[None], dtype=DTYPEINT)
# instantiate logger
logger = logging.getLogger(__name__)
# create the Grid namedtuple
GridTuple = collections.namedtuple("Grid", ["x", "q2", "flav", "grid"])
AlphaTuple = collections.namedtuple("Alpha", ["q2", "grid"])
def _load_data(pdf_file):
"""
Reads pdf from file and retrieves a list of grids
Each grid is a tuple containing numpy arrays (x,Q2, flavours, pdf)
def test_int_me():
res = int_me(4)
assert res.dtype == DTYPEINT