# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import numpy as np
from becquerel import rebin, Spectrum
import matplotlib.pyplot as plt
from becquerel import SpectrumPlotter as sp
# %matplotlib inline
np.random.seed(0)

def plot_rebin(bins, edges, lim=None):
    for b, e in zip(bins, edges):
        plt.plot(*sp.bin_edges_and_heights_to_steps(e, b))
        plt.fill_between(*sp.bin_edges_and_heights_to_steps(e, b), alpha=0.5)
    if lim:
        plt.gca().set_ylim(lim)

def plot_rebin_spec(specs, lim=None):
    fig, ax = plt.subplots()
    for s in specs:
        s.plot(ax=ax)
        s.fill_between(ax=ax, alpha=0.5)
    if lim:
        ax.set_ylim(lim)


# -

# ## Rebinning module
#
# Rebin can take in numpy arrays (both 1D and 2D, with the shape (number of spectra, number of bins) directly:

in_bins_1d = np.array([1,2,2,3,1,0])
in_edges = np.array([0,2,4,6,8,10,12])
out_edges = np.array([0,1,3,6,7,9])  # The edges does not have to be equally spaced

# ### Rebinning methods
#
# there are two rebinning methods, `interpolation` and `listmode`  
# - `interpolation`: assuming that the distribution within each bin is flat (or as specified in the `slopes` argument), put the proportion of counts into each of the corresponding new bins. Counts may become fractional after the rebinning, and this does not preserve Poisson statistics. (TODO: add how the `slopes` argument works)
# - `listmode`: carries out a stochastic rebinning, by randomly sampling energies within each bin and binning all these energies into the new bin structure

out_bins_1d = rebin(in_bins_1d, in_edges, out_edges, method='interpolation', slopes=None)
plot_rebin([in_bins_1d, out_bins_1d], [in_edges, out_edges], lim=(0,4.5))

out_bins_1d = rebin(in_bins_1d, in_edges, out_edges, method='listmode', slopes=None)
plot_rebin([in_bins_1d, out_bins_1d], [in_edges, out_edges], lim=(0,4.5))

# ### 2D Rebinning
# You can rebin multiple spectra at the same time, too. The example uses the same input bin edges as the 1D example.

in_bins_2d = np.array([[1,2,2,3,1,0], [1,1,5,2,6,3]])
plot_rebin(in_bins_2d, [in_edges]*2, lim=(0,10.5))

out_bins_2d = rebin(in_bins_2d, in_edges, out_edges, method='interpolation', slopes=None)
plot_rebin(out_bins_2d, [out_edges]*2, lim=(0,10.5))

out_bins_2d = rebin(in_bins_2d, in_edges, out_edges, method='listmode', slopes=None)
plot_rebin(out_bins_2d, [out_edges]*2, lim=(0,10.5))

# ## Spectrum rebinning
#
# It is possible to directly rebin a spectrum object

in_spec = Spectrum(in_bins_1d, bin_edges_kev=in_edges)
out_spec = in_spec.rebin(out_edges)
plot_rebin_spec([in_spec, out_spec])

# There is also a `rebin_like` method, that takes another spectrum as argument and rebins according to that spectrum.

out_spec_like = in_spec.rebin_like(out_spec)
plot_rebin_spec([in_spec, out_spec])

# ## Usage with other packages than numpy
#
# The rebin module can also be used with other packages such as `xarray`. But there is a bit more code involved to apply the rebinning alogrithm in parallel as a `ufunc`. Also this doesn't perform any tests on the inputs. Lets first create a `xarray` data set. For this the `xarray` package needs to be installed.

# +
import xarray as xr
import becquerel as bq

in_bins_2d = xr.DataArray([[1,2,2,3,1,0], [1,1,5,2,6,3]], dims=['spec', 'bin'])
# -

# Lets redo the 'interpolation' example in 2d with the newely created `DataArray`.

# +
out_bins_2d = xr.apply_ufunc(
    bq.core.rebin._rebin_interpolation,
    in_bins_2d, in_edges, out_edges[:-1], np.zeros(6),
    input_core_dims=[["bin"], [], [], []],
    output_core_dims=[["out_bin"]],
    dask="parallelized",
    output_dtypes=[float],
)

plot_rebin(out_bins_2d, [out_edges]*2, lim=(0,10.5))
# -

# It is important to note that the last element in the `out_edges` array was omitted in the input. This is necessary to used the `_rebin_interpolate` and `_rebin_listmode` methods directly.
#
# Simlarly, also the 'listmode' binning can be applied. The differences to the above sample originates from the randomness of the 'listmode' binning method.

# +
out_bins_2d = xr.apply_ufunc(
    bq.core.rebin._rebin_listmode,
    in_bins_2d, in_edges, out_edges[:-1],
    input_core_dims=[["bin"], [], []],
    output_core_dims=[["out_bin"]],
    dask="parallelized",
    output_dtypes=[int],
)

plot_rebin(out_bins_2d, [out_edges]*2, lim=(0,10.5))
# -



StarcoderdataJupyterScriptsDedupFiltered