import numpy as np

def _adddims(a, b):
    n = max(a.ndim, b.ndim)
    a = np.expand_dims(a, tuple(range(n - a.ndim)))
    b = np.expand_dims(b, tuple(range(n - b.ndim)))
    return a, b

def _yz(y, z, t, yout):
    """
    Shared implementation of peaksampl and sumpeaks.
    """
    y = np.asarray(y)
    z = np.asarray(z)
    t = np.asarray(t)
    
    y = np.pad(y, [(0, 0)] * (y.ndim - 1) + [(1, 1)], constant_values=yout)
    offset = np.argmax(np.abs(y), axis=-1)
    ampl = np.take_along_axis(y, np.expand_dims(offset, -1), -1)
    ampl = np.squeeze(ampl, -1)

    indices = t[..., :, None] - t[..., None, :] + offset[..., None, None]
    indices = np.minimum(indices, y.shape[-1] - 1)
    indices = np.maximum(indices, 0)
    
    N = t.shape[-1]
    indices = indices.reshape(indices.shape[:-2] + (N * N,))
    n = max(y.ndim, indices.ndim)
    y, indices = _adddims(y, indices)
    y = np.take_along_axis(y, indices, -1)
    eps = np.finfo(float).eps * N * N * ampl
    y[..., ::N + 1] += np.expand_dims(eps, -1)
    y = y.reshape(y.shape[:-1] + (N, N))
    
    z = z[..., None]
    y, z = _adddims(y, z)
    
    return y, z

def peaksampl(y, z, t, yout=0):
    """
    Get peak amplitudes given their sum.
    
    This assumes that the position of the signals is given by peaks positions
    even when they are summed.
    
    Parameters
    ----------
    y : array (..., M,)
        The single signal shape.
    z : array (..., N,)
        The peak height in the sum of the signals for each peak.
    t : int array (..., N,)
        The indices of the peaks in the sum.
    yout : scalar
        The value of the signal outside the provided values, default 0.
    
    Return
    ------
    a : array (..., N),
        The amplitudes such that z_i = sum_j a_j * y[t_i - t_j].
        Broadcasted along non-last axis.
    """
    y, z = _yz(y, z, t, yout)
    a = np.linalg.solve(y, z)
    return np.squeeze(a, -1)

def sumpeaks(y, a, t, yout=0):
    """
    Compute the peak heights of a sum of signals.
    
    This assumes that the position of the peaks is given by the signal
    positions even when they are summed.
    
    Parameters
    ----------
    y : array (..., M,)
        The single signal shape.
    a : array (..., N,)
        The amplitudes of the signals (`y` is multiplied by `a`).
    t : int array (..., N,)
        The indices of the position of the signals.
    yout : scalar
        The value of the signal outside the provided values, default 0.
    
    Return
    ------
    z : array (..., N,)
        The peak height in the sum of the signals for each signal. Broadcasted
        along non-last axis.
    """
    y, a = _yz(y, a, t, yout)
    z = np.matmul(y, a)
    return np.squeeze(z, axis=-1)

if __name__ == '__main__':
    from matplotlib import pyplot as plt
    from scipy import signal
    
    y = np.exp(-np.linspace(0, 10, 1000) / 10)
    i = np.arange(1, 1000)
    t0 = np.array([10, 340, 523])
    a0 = np.array([3, 2, 1])
    indices = i - t0[:, None]
    z = np.take(y, indices, mode='clip') * a0[:, None]
    z = np.where((indices < 0) | (indices >= len(y)), 0, z)
    z = np.sum(z, axis=0)
    t, = signal.argrelmax(z)
    assert len(t) == len(t0)
    a = peaksampl(y, z[t], t)
    h = sumpeaks(y, a, t)
    
    fig, ax = plt.subplots(num='peaksampl', clear=True)
    
    ax.plot(z, color='#f55')
    ax.vlines(t0, 0, a0, color='gray', zorder=3)
    ax.vlines(t, 0, a, linestyle='--', zorder=3)
    ax.plot(t, h, 'ok')
    
    ax.grid('major', linestyle='--')
    
    fig.tight_layout()
    fig.show()

[(13, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (14, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (15, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (17, 'numpy.pad', 'np.pad', True, 'import numpy as np\n'), (20, 'numpy.squeeze', 'np.squeeze', True, 'import numpy as np\n'), (23, 'numpy.minimum', 'np.minimum', True, 'import numpy as np\n'), (24, 'numpy.maximum', 'np.maximum', True, 'import numpy as np\n'), (30, 'numpy.take_along_axis', 'np.take_along_axis', True, 'import numpy as np\n'), (32, 'numpy.expand_dims', 'np.expand_dims', True, 'import numpy as np\n'), (65, 'numpy.linalg.solve', 'np.linalg.solve', True, 'import numpy as np\n'), (66, 'numpy.squeeze', 'np.squeeze', True, 'import numpy as np\n'), (93, 'numpy.matmul', 'np.matmul', True, 'import numpy as np\n'), (94, 'numpy.squeeze', 'np.squeeze', True, 'import numpy as np\n'), (101, 'numpy.arange', 'np.arange', True, 'import numpy as np\n'), (102, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (103, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (107, 'numpy.sum', 'np.sum', True, 'import numpy as np\n'), (108, 'scipy.signal.argrelmax', 'signal.argrelmax', False, 'from scipy import signal\n'), (113, 'matplotlib.pyplot.subplots', 'plt.subplots', True, 'from matplotlib import pyplot as plt\n'), (18, 'numpy.abs', 'np.abs', True, 'import numpy as np\n'), (19, 'numpy.expand_dims', 'np.expand_dims', True, 'import numpy as np\n'), (105, 'numpy.take', 'np.take', True, 'import numpy as np\n'), (100, 'numpy.linspace', 'np.linspace', True, 'import numpy as np\n'), (31, 'numpy.finfo', 'np.finfo', True, 'import numpy as np\n')]