137 lines
5.0 KiB
Python
137 lines
5.0 KiB
Python
import numpy as np
|
|
'''
|
|
Functions for eta interpolation
|
|
1. build_xy_lut_from_hist(H)
|
|
2. bilinear_xy_lookup(etaX, etaY, X_tab, Y_tab, x_centers, y_centers)
|
|
3. map_xy_no_interp(etaX, etaY, X_tab, Y_tab, x_edges, y_edges)
|
|
'''
|
|
def build_xy_lut_Rosenblatt(H):
|
|
"""
|
|
Generate two 2D tables X_tab, Y_tab from a 2D histogram H
|
|
such that (X_tab[i,j], Y_tab[i,j]) is the (x,y) value at the center of bin (i,j).
|
|
Input:
|
|
H: 2D histogram of charge centroid (eta_x, eta_y) distribution, shape (nBinX, nBinY), ranging from 0 to 1.
|
|
Output:
|
|
X_tab, Y_tab: two 2D tables of interpolated (x,y) values in position space, shape (nBinX, nBinY), ranging from 0 to 1.
|
|
"""
|
|
H = np.asarray(H, dtype=float)
|
|
nBinX, nBinY = H.shape
|
|
etaX_edges = np.linspace(0, 1, nBinX+1)
|
|
etaY_edges = np.linspace(0, 1, nBinY+1)
|
|
|
|
# slight smoothing to avoid zero-count rows/columns causing flat CDF
|
|
eps=1e-9
|
|
alpha = max(eps, 1e-12 * max(1.0, H.sum()))
|
|
Hs = H + alpha
|
|
|
|
# marginal 1D CDF F_EtaX
|
|
pEtaX = Hs.sum(axis=1) # Marginal PDF (nBinX,)
|
|
F_EtaX_edges = np.concatenate([[0.0], np.cumsum(pEtaX)])# CDF at bin edges (nBinX+1,)
|
|
F_EtaX_edges = F_EtaX_edges / F_EtaX_edges[-1]
|
|
|
|
# conditional CDF F_{Y|X} (a CDF in y direction for each etaX bin)
|
|
F_EtaY_cond_edges = np.zeros((nBinX, nBinY + 1), float)
|
|
for i in range(nBinX):
|
|
cy = np.concatenate([[0.0], np.cumsum(Hs[i, :])])
|
|
F_EtaY_cond_edges[i, :] = cy / cy[-1]
|
|
|
|
# calculate (u,v) at the center of each original bin
|
|
Xc_1d = 0.5 * (F_EtaX_edges[:-1] + F_EtaX_edges[1:])
|
|
Yc_2d = 0.5 * (F_EtaY_cond_edges[:, :-1] + F_EtaY_cond_edges[:, 1:])
|
|
|
|
X_tab = np.repeat(Xc_1d[:, None], nBinY, axis=1) # (nBinX, nBinY)
|
|
Y_tab = Yc_2d # (nBinX, nBinY)
|
|
|
|
return (X_tab, Y_tab)
|
|
|
|
def build_xy_lut_DoubleCDF(H):
|
|
"""
|
|
Generate two 2D tables X_tab, Y_tab from a 2D histogram H
|
|
such that (X_tab[i,j], Y_tab[i,j]) is the (x,y) value at the center of bin (i,j).
|
|
Input:
|
|
H: 2D histogram of charge centroid (eta_x, eta_y) distribution, shape (nBinX, nBinY), ranging from 0 to 1.
|
|
Output:
|
|
X_tab, Y_tab: two 2D tables of interpolated (x,y) values in position space, shape (nBinX, nBinY), ranging from 0 to 1.
|
|
"""
|
|
H = np.asarray(H, dtype=float)
|
|
nBinX, nBinY = H.shape
|
|
etaX_edges = np.linspace(0, 1, nBinX+1)
|
|
etaY_edges = np.linspace(0, 1, nBinY+1)
|
|
|
|
# slight smoothing to avoid zero-count rows/columns causing flat CDF
|
|
eps=1e-9
|
|
alpha = max(eps, 1e-12 * max(1.0, H.sum()))
|
|
Hs = H + alpha
|
|
|
|
# marginal 1D CDF F_EtaX
|
|
pEtaX = Hs.sum(axis=1) # Marginal PDF (nBinX,)
|
|
F_EtaX_edges = np.concatenate([[0.0], np.cumsum(pEtaX)])# CDF at bin edges (nBinX+1,)
|
|
F_EtaX_edges = F_EtaX_edges / F_EtaX_edges[-1]
|
|
|
|
# marginal 1D CDF F_EtaY
|
|
pEtaY = Hs.sum(axis=0) # Marginal PDF (nBinY,)
|
|
F_EtaY_edges = np.concatenate([[0.0], np.cumsum(pEtaY)])# CDF at bin edges (nBinY+1,)
|
|
F_EtaY_edges = F_EtaY_edges / F_EtaY_edges[-1]
|
|
|
|
# calculate (u,v) at the center of each original bin
|
|
Xc_1d = 0.5 * (F_EtaX_edges[:-1] + F_EtaX_edges[1:])
|
|
Yc_1d = 0.5 * (F_EtaY_edges[:-1] + F_EtaY_edges[1:])
|
|
|
|
X_tab = np.repeat(Xc_1d[:, None], nBinY, axis=1) # (nBinX, nBinY)
|
|
Y_tab = np.repeat(Yc_1d[None, :], nBinX, axis=0) # (nBinX, nBinY)
|
|
|
|
return (X_tab, Y_tab)
|
|
|
|
def bilinear_xy_lookup(etaX, etaY, X_tab, Y_tab):
|
|
"""
|
|
bilinear interpolation, from (eta_x, eta_y) to (u,v):
|
|
"""
|
|
x = np.asarray(etaX); y = np.asarray(etaY)
|
|
nBinX, nBinY = X_tab.shape
|
|
|
|
### bin centers
|
|
x_centers = (np.arange(nBinX) + 0.5) / nBinX
|
|
y_centers = (np.arange(nBinY) + 0.5) / nBinY
|
|
|
|
# search bin
|
|
ix = np.clip(np.searchsorted(x_centers, x) - 1, 0, nBinX-2)
|
|
iy = np.clip(np.searchsorted(y_centers, y) - 1, 0, nBinY-2)
|
|
|
|
# four corners
|
|
x0 = x_centers[ix]; x1 = x_centers[ix+1]
|
|
y0 = y_centers[iy]; y1 = y_centers[iy+1]
|
|
tx = np.where(x1>x0, (x - x0)/(x1 - x0), 0.0)
|
|
ty = np.where(y1>y0, (y - y0)/(y1 - y0), 0.0)
|
|
|
|
# values at corners
|
|
X00 = X_tab[ix, iy ]
|
|
X10 = X_tab[ix+1, iy ]
|
|
X01 = X_tab[ix, iy+1]
|
|
X11 = X_tab[ix+1, iy+1]
|
|
Y00 = Y_tab[ix, iy ]
|
|
Y10 = Y_tab[ix+1, iy ]
|
|
Y01 = Y_tab[ix, iy+1]
|
|
Y11 = Y_tab[ix+1, iy+1]
|
|
|
|
# bilinear interpolation
|
|
X0 = (1-tx)*X00 + tx*X10
|
|
X1 = (1-tx)*X01 + tx*X11
|
|
Y0 = (1-tx)*Y00 + tx*Y10
|
|
Y1 = (1-tx)*Y01 + tx*Y11
|
|
X = (1-ty)*X0 + ty*X1
|
|
Y = (1-ty)*Y0 + ty*Y1
|
|
return X, Y
|
|
|
|
def xy_lookup(etaX, etaY, X_tab, Y_tab):
|
|
"""
|
|
mapping without interpolation.
|
|
"""
|
|
x = np.asarray(etaX)
|
|
y = np.asarray(etaY)
|
|
nBinX, nBinY = X_tab.shape
|
|
|
|
# Compute nearest bin index directly
|
|
ix = np.clip(np.rint(x * nBinX - 0.5).astype(int), 0, nBinX - 1)
|
|
iy = np.clip(np.rint(y * nBinY - 0.5).astype(int), 0, nBinY - 1)
|
|
|
|
return X_tab[ix, iy], Y_tab[ix, iy] |