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, 8, 13, 21), 'numpy.asarray', 'np.asarray', ({(13, 19, 13, 20): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((14, 8, 14, 21), 'numpy.asarray', 'np.asarray', ({(14, 19, 14, 20): 'z'}, {}), '(z)', True, 'import numpy as np\n'), ((15, 8, 15, 21), 'numpy.asarray', 'np.asarray', ({(15, 19, 15, 20): 't'}, {}), '(t)', True, 'import numpy as np\n'), ((17, 8, 17, 75), 'numpy.pad', 'np.pad', (), '', True, 'import numpy as np\n'), ((20, 11, 20, 31), 'numpy.squeeze', 'np.squeeze', ({(20, 22, 20, 26): 'ampl', (20, 28, 20, 30): '-1'}, {}), '(ampl, -1)', True, 'import numpy as np\n'), ((23, 14, 23, 50), 'numpy.minimum', 'np.minimum', ({(23, 25, 23, 32): 'indices', (23, 34, 23, 49): 'y.shape[-1] - 1'}, {}), '(indices, y.shape[-1] - 1)', True, 'import numpy as np\n'), ((24, 14, 24, 36), 'numpy.maximum', 'np.maximum', ({(24, 25, 24, 32): 'indices', (24, 34, 24, 35): '0'}, {}), '(indices, 0)', True, 'import numpy as np\n'), ((30, 8, 30, 42), 'numpy.take_along_axis', 'np.take_along_axis', ({(30, 27, 30, 28): 'y', (30, 30, 30, 37): 'indices', (30, 39, 30, 41): '-1'}, {}), '(y, indices, -1)', True, 'import numpy as np\n'), ((32, 23, 32, 46), 'numpy.expand_dims', 'np.expand_dims', ({(32, 38, 32, 41): 'eps', (32, 43, 32, 45): '(-1)'}, {}), '(eps, -1)', True, 'import numpy as np\n'), ((65, 8, 65, 29), 'numpy.linalg.solve', 'np.linalg.solve', ({(65, 24, 65, 25): 'y', (65, 27, 65, 28): 'z'}, {}), '(y, z)', True, 'import numpy as np\n'), ((66, 11, 66, 28), 'numpy.squeeze', 'np.squeeze', ({(66, 22, 66, 23): 'a', (66, 25, 66, 27): '(-1)'}, {}), '(a, -1)', True, 'import numpy as np\n'), ((93, 8, 93, 23), 'numpy.matmul', 'np.matmul', ({(93, 18, 93, 19): 'y', (93, 21, 93, 22): 'a'}, {}), '(y, a)', True, 'import numpy as np\n'), ((94, 11, 94, 33), 'numpy.squeeze', 'np.squeeze', (), '', True, 'import numpy as np\n'), ((101, 8, 101, 26), 'numpy.arange', 'np.arange', ({(101, 18, 101, 19): '1', (101, 21, 101, 25): '1000'}, {}), '(1, 1000)', True, 'import numpy as np\n'), ((102, 9, 102, 33), 'numpy.array', 'np.array', ({(102, 18, 102, 32): '[10, 340, 523]'}, {}), '([10, 340, 523])', True, 'import numpy as np\n'), ((103, 9, 103, 28), 'numpy.array', 'np.array', ({(103, 18, 103, 27): '[3, 2, 1]'}, {}), '([3, 2, 1])', True, 'import numpy as np\n'), ((107, 8, 107, 25), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((108, 9, 108, 28), 'scipy.signal.argrelmax', 'signal.argrelmax', ({(108, 26, 108, 27): 'z'}, {}), '(z)', False, 'from scipy import signal\n'), ((113, 14, 113, 55), 'matplotlib.pyplot.subplots', 'plt.subplots', (), '', True, 'from matplotlib import pyplot as plt\n'), ((18, 23, 18, 32), 'numpy.abs', 'np.abs', ({(18, 30, 18, 31): 'y'}, {}), '(y)', True, 'import numpy as np\n'), ((19, 33, 19, 59), 'numpy.expand_dims', 'np.expand_dims', ({(19, 48, 19, 54): 'offset', (19, 56, 19, 58): '-1'}, {}), '(offset, -1)', True, 'import numpy as np\n'), ((105, 8, 105, 40), 'numpy.take', 'np.take', (), '', True, 'import numpy as np\n'), ((100, 16, 100, 40), 'numpy.linspace', 'np.linspace', ({(100, 28, 100, 29): '0', (100, 31, 100, 33): '10', (100, 35, 100, 39): '1000'}, {}), '(0, 10, 1000)', True, 'import numpy as np\n'), ((31, 10, 31, 25), 'numpy.finfo', 'np.finfo', ({(31, 19, 31, 24): 'float'}, {}), '(float)', True, 'import numpy as np\n')]