Binning demonstration on locally generated fake data

In this example, we generate a table with random data simulating a single event dataset. We showcase the binning method, first on a simple single table using the bin_partition method and then in the distributed method bin_dataframe, using daks dataframes. The first method is never really called directly, as it is simply the function called by the bin_dataframe on each partition of the dask dataframe.

import dask
import numpy as np
import pandas as pd
import dask.dataframe

import matplotlib.pyplot as plt

from sed.binning import bin_partition, bin_dataframe

%matplotlib widget

/opt/hostedtoolcache/Python/3.12.12/x64/lib/python3.12/site-packages/dask/dataframe/__init__.py:42: FutureWarning:
Dask dataframe query planning is disabled because dask-expr is not installed.

You can install it with `pip install dask[dataframe]` or `conda install dask`.
This will raise in a future version.

  warnings.warn(msg, FutureWarning)

Generate Fake Data

n_pts = 100000
cols = ["posx", "posy", "energy"]
df = pd.DataFrame(np.random.randn(n_pts, len(cols)), columns=cols)
df

	posx	posy	energy
0	-1.048649	-0.627300	0.144712
1	0.213310	-0.731747	0.036102
2	0.276974	-1.231643	0.179607
3	0.360123	1.132908	-1.192721
4	-0.509570	-0.278248	-0.027518
...	...	...	...
99995	0.005450	1.435372	0.414865
99996	-0.389768	-0.042776	0.740735
99997	-0.207318	-0.734152	0.893433
99998	-0.711687	0.001476	-1.303670
99999	-0.066542	2.021989	-1.145597

100000 rows × 3 columns

Define the binning range

binAxes = ["posx", "posy", "energy"]
nBins = [120, 120, 120]
binRanges = [(-2, 2), (-2, 2), (-2, 2)]
coords = {ax: np.linspace(r[0], r[1], n) for ax, r, n in zip(binAxes, binRanges, nBins)}

Compute the binning along the pandas dataframe

%%time
res = bin_partition(
    part=df,
    bins=nBins,
    axes=binAxes,
    ranges=binRanges,
    hist_mode="numba",
)

CPU times: user 1.18 s, sys: 70.9 ms, total: 1.25 s
Wall time: 1.25 s

fig, axs = plt.subplots(1, 3, figsize=(6, 1.875), constrained_layout=True)
for i in range(3):
    axs[i].imshow(res.sum(i))

Transform to dask dataframe

ddf = dask.dataframe.from_pandas(df, npartitions=50)
ddf

Dask DataFrame Structure:

	posx	posy	energy
npartitions=50
0	float64	float64	float64
2000	...	...	...
...	...	...	...
98000	...	...	...
99999	...	...	...

Dask Name: from_pandas, 1 graph layer

Compute distributed binning on the partitioned dask dataframe

In this example, the small dataset does not give significant improvement over the pandas implementation, at least using this number of partitions. A single partition would be faster (you can try…) but we use multiple for demonstration purposes.

%%time
res = bin_dataframe(
    df=ddf,
    bins=nBins,
    axes=binAxes,
    ranges=binRanges,
    hist_mode="numba",
)

CPU times: user 659 ms, sys: 148 ms, total: 807 ms
Wall time: 711 ms

fig, axs = plt.subplots(1, 3, figsize=(6, 1.875), constrained_layout=True)
for dim, ax in zip(binAxes, axs):
    res.sum(dim).plot(ax=ax)