|
import matplotlib.colors as colors |
|
import mplhep |
|
import numpy as np |
|
from matplotlib import pyplot as plt |
|
from matplotlib import rc_context |
|
from matplotlib.ticker import LogFormatterSciNotation |
|
from numpy import random |
|
from numpy.random import PCG64, SeedSequence |
|
from scipy.interpolate import LinearNDInterpolator |
|
|
|
# generate observed values (x, y) |
|
x_step = 250 |
|
x_range = np.arange(750, 3000 + x_step, step=x_step) |
|
y_step = 150 |
|
y_range = np.arange(700, 2000 + y_step, step=y_step) |
|
|
|
bit_generator = PCG64(SeedSequence(0)) |
|
rng = random.default_rng(bit_generator) |
|
|
|
x = [] |
|
y = [] |
|
for step in x_range: |
|
choose_n = rng.integers(low=0, high=y_range.size - 2) |
|
for value in y_range[: choose_n + 2]: |
|
x.append(step) |
|
y.append(value) |
|
|
|
x = np.asarray(x) |
|
y = np.asarray(y) |
|
|
|
# Generate uniform data on the interval [1e-3, 100.] |
|
# Uniform [a,b) = (b - a) * random_sample() + a |
|
uniform_range = [1e-3, 100.0] |
|
z = (uniform_range[1] - uniform_range[0]) * rng.random(x.size) + uniform_range[0] |
|
|
|
# Generate a 2D grid |
|
x_coords = np.linspace(min(x), max(x), 100) |
|
y_coords = np.linspace(min(y), max(y), 100) |
|
x_grid, y_grid = np.meshgrid(x_coords, y_coords) # 2D grid for interpolation |
|
|
|
# Interpolate with function of choice across the grid |
|
# between the known values |
|
interp_func = LinearNDInterpolator(list(zip(x, y)), z) |
|
interp_Z = interp_func(x_grid, y_grid) |
|
|
|
plt.style.use(mplhep.style.ATLAS) |
|
|
|
# c.f. https://github.com/scikit-hep/mplhep/issues/362 |
|
rc_options = {"ytick.minor.visible": False} # yticks |
|
with rc_context(rc_options): |
|
fig, ax = plt.subplots() |
|
|
|
# plot interpolated values |
|
real_valued_Z = interp_Z[~np.isnan(interp_Z)] |
|
step_exp = 0.25 |
|
levels_exp = np.arange( |
|
np.floor(np.log10(real_valued_Z.min()) - 1), |
|
np.ceil(np.log10(real_valued_Z.max()) + 1) + step_exp, |
|
step=step_exp, |
|
) |
|
levels = np.power(10, levels_exp) |
|
|
|
cs = ax.contourf( |
|
x_grid, |
|
y_grid, |
|
interp_Z, |
|
cmap="PuBu_r", |
|
levels=levels, |
|
norm=colors.LogNorm(vmin=levels.min(), vmax=levels.max()), |
|
) |
|
formatter = LogFormatterSciNotation(10, labelOnlyBase=False) |
|
cbar = fig.colorbar(cs, format=formatter) |
|
|
|
# plot observed values |
|
scatter_size = 15.0 |
|
ax.scatter(x, y, s=scatter_size, color="white", edgecolors="black") |
|
|
|
ax.set_xlabel(r"$x$") |
|
ax.set_ylabel(r"$y$") |
|
cbar.set_label(r"$z(x, y)$") |
|
|
|
file_types = ["png", "pdf"] |
|
for extension in file_types: |
|
fig.savefig(f"example.{extension}") |