from typing import Tuple
import numpy as np
[docs]
def slope_intercept(peak_times: np.ndarray, longitudinals: np.ndarray, cleaning_mask) -> Tuple[np.ndarray, np.ndarray]:
"""Computes slope, intercept by fitting a line (LSE) on the `peak_times` of each pixels along the ellipsis major axis
This is essentially a less robust __batched__ numpy.linalg.lstsq
Parameters
----------
peak_times : np.ndarray
The time of maximum signal for each pixels as computed by the integration.
shape: ([N_batch,], N_pixels)
longitudinals : np.ndarray
The longitudinal coordinates (coordinate along the ellipsis major axis) of each pixels.
shape: ([N_batch,], N_pixels)
Returns
-------
Tuple[np.ndarray, np.ndarray]
Slope and intercept values for each waveform.
shape of each array: ([Batch,],)
"""
# Compute mean and sums only where the cleaning mask is True.
longitudinals_mean = longitudinals.mean(axis=-1, keepdims=True, where=cleaning_mask)
peak_times_mean = peak_times.mean(axis=-1, keepdims=True, where=cleaning_mask)
slope = (cleaning_mask * (longitudinals - longitudinals_mean) * (peak_times - peak_times_mean)).sum(axis=-1) / (
cleaning_mask * (longitudinals - longitudinals_mean) ** 2
).sum(axis=-1)
intercept = peak_times_mean.squeeze() - slope * longitudinals_mean.squeeze()
return slope, intercept