lib2828.py

from euxfel_bunch_pattern import indices_at_sase  # Installable from PyPI
import numpy as np
import string
from scipy.interpolate import interp1d
from pathlib import Path


# List of detector infos from external txt file:
# name, digitizer, channel, angle, sample_rate
detectors = np.genfromtxt(Path(__file__).parent / 'detectors.txt',
                          names=True, dtype=('|U5', '|U4', '|U3', '<f8', '<i4'))


def correct_adq_common_mode(trace, region, sym):
    """Baseline subtraction based on common mode.

    Since ADQ digitizers always interleave multiple ADCs per channel to sample
    a trace, regular baseline subtraction will cause an ADC-dependent common
    mode to appear. This correction directly uses a per-ADC baseline instead
    to perform include this in the substraction.

    Empirical testing has shown that a symmetry of 8 (should be 4 for
    non-interleaved) regularly shows better results in suppressing high
    frequency signals from the common mode. Going 16 or higher is not recommend
    in most cases, as additional artifacts appear.

    Parameters
    ----------
    trace : array_like
        Vector to correct, will be modified in-place.
    region : slice
        Region to use for computing baseline.
    sym : int
        Periodic symmetry of ADC common mode.

    Returns
    -------
    trace : array_like
        Corrected vector, same shape as trace.
    """
    trace = trace.astype(np.float32)

    for x in range(sym):
        trace[..., x::sym] -= trace[..., region][x::sym].mean()

    return trace


def separate_pulses(traces, ppt, adq_train_region, adq_pulse_region, adq_sample_rate=4, sase=3):
    """Separate train into seperate pulses using the pulse pattern table.

    Parameters
    ----------
    traces : array_like
        Array of traces to be split up
    ppt : array_like
        Pulse pattern table from time server device.
    adq_train_region : slice, optional
        Region containing actual signal for the entire train.
    adq_pulse_region : slice, optional
        Region after pulse separation containing actual signal for the pulse.
    adq_pulse_region : slice, optional
        Region after pulse separation containing actual signal for the pulse.
    adq_sample_rate : array_like or int, optional
        Sample rate for all digitizer channels in GS/s.
    sase : int, optional
        Which SASE the experiments run at, 3 by default.

    Returns
    -------
    traces : array_like
        Separated vector of trains.
    """
    traces = traces[..., adq_train_region]
    pulse_ids = indices_at_sase(ppt, sase=sase)
    num_pulses = len(pulse_ids)
    if num_pulses < 2:
        return traces[..., adq_pulse_region].reshape(traces.shape[0], 1, -1)
    else:
        # ADQ412s run at a clock factor of 440 relative to the PPT un-interleaved.
        pulse_spacing = (pulse_ids[1] - pulse_ids[0]) * 220 * adq_sample_rate
        traces_aligned = traces[..., :(num_pulses * pulse_spacing)]
        assert num_pulses * pulse_spacing == traces_aligned.shape[-1] # digitizer traces are too short for pulse pattern?
        return traces_aligned.reshape(*(traces_aligned.shape[:-1]), num_pulses, pulse_spacing)[..., adq_pulse_region]


def _f(t, a, b, c, d):
    return a / t**3 + b / t**2 + c / t + d

def _df(t, a, b, c, d):
    return -3 * a / t**4 - 2 * b / t**3 - c / t**2
    
def energy_calibration(traces, calib_params, sample_rate, valid_energies,
                       energy_nodes=None, model=_f, model_derivative=_df):
    """Calibrate time of flight traces to energy spectra.

    Parameters
    ----------
    traces : array_like
        Traces to be calibrated, shape (detectors, ..., samples).
    calib_params : array_like
        Calibration parameters passed to the model, first dimension needs
        to be number of function arguments.
    sample_rate : array_like or int, optional
        Sample rate for all digitizer channels in GS/s.
    valid_energies : slice
        Slice applied to input allowing to discard divergent values.
    energy_nodes : array_like, optional
        Energy values to be evaluated in interpolation, if you want.
    model : function, optional
        Mapping from time to energy, arguments (t, *args). Defaults to
        a / t**3 + b / t**2 + c / t + d.
    model_derivative : function, optional
        Derivative of `model` function.

    Returns
    -------
    energy : ndarray
        Energy values corresponding to time of flight.
    spectra : ndarray
        Reweighted spectra.
    resampled_spectra : ndarray
        Spectra resampled according to `energy_nodes`.
    time : ndarray
        Times of flight corresponding to samples.
    """
    # Convert digitizer channels into ToF [ns]
    time = np.arange(traces.shape[-1])[:, np.newaxis] / sample_rate * 1e-9
    # Calibrate ToF to energy
    energy = model(time, *calib_params)[valid_energies].T
    dE = model_derivative(time[:, None], *calib_params)[valid_energies].T
    spectra = traces[..., valid_energies] / dE

    if energy_nodes is not None:
        try:
            resampled_spectra = np.asarray([interp1d(e, t)(energy_nodes) for e, t in
                                            zip(energy, spectra)])
        except ValueError as e:
            print(e)
            resampled_spectra = np.full((*traces.shape[:-1], len(energy_nodes)), np.nan)
    else:
        resampled_spectra = None
    return energy, spectra, resampled_spectra, time[valid_energies]