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()

[((290, 303), 'numpy.asarray', 'np.asarray', (['y'], {}), '(y)\n', (300, 303), True, 'import numpy as np\n'), ((312, 325), 'numpy.asarray', 'np.asarray', (['z'], {}), '(z)\n', (322, 325), True, 'import numpy as np\n'), ((334, 347), 'numpy.asarray', 'np.asarray', (['t'], {}), '(t)\n', (344, 347), True, 'import numpy as np\n'), ((361, 428), 'numpy.pad', 'np.pad', (['y', '([(0, 0)] * (y.ndim - 1) + [(1, 1)])'], {'constant_values': 'yout'}), '(y, [(0, 0)] * (y.ndim - 1) + [(1, 1)], constant_values=yout)\n', (367, 428), True, 'import numpy as np\n'), ((548, 568), 'numpy.squeeze', 'np.squeeze', (['ampl', '(-1)'], {}), '(ampl, -1)\n', (558, 568), True, 'import numpy as np\n'), ((658, 694), 'numpy.minimum', 'np.minimum', (['indices', '(y.shape[-1] - 1)'], {}), '(indices, y.shape[-1] - 1)\n', (668, 694), True, 'import numpy as np\n'), ((709, 731), 'numpy.maximum', 'np.maximum', (['indices', '(0)'], {}), '(indices, 0)\n', (719, 731), True, 'import numpy as np\n'), ((898, 932), 'numpy.take_along_axis', 'np.take_along_axis', (['y', 'indices', '(-1)'], {}), '(y, indices, -1)\n', (916, 932), True, 'import numpy as np\n'), ((1001, 1024), 'numpy.expand_dims', 'np.expand_dims', (['eps', '(-1)'], {}), '(eps, -1)\n', (1015, 1024), True, 'import numpy as np\n'), ((1882, 1903), 'numpy.linalg.solve', 'np.linalg.solve', (['y', 'z'], {}), '(y, z)\n', (1897, 1903), True, 'import numpy as np\n'), ((1915, 1932), 'numpy.squeeze', 'np.squeeze', (['a', '(-1)'], {}), '(a, -1)\n', (1925, 1932), True, 'import numpy as np\n'), ((2697, 2712), 'numpy.matmul', 'np.matmul', (['y', 'a'], {}), '(y, a)\n', (2706, 2712), True, 'import numpy as np\n'), ((2724, 2746), 'numpy.squeeze', 'np.squeeze', (['z'], {'axis': '(-1)'}), '(z, axis=-1)\n', (2734, 2746), True, 'import numpy as np\n'), ((2905, 2923), 'numpy.arange', 'np.arange', (['(1)', '(1000)'], {}), '(1, 1000)\n', (2914, 2923), True, 'import numpy as np\n'), ((2933, 2957), 'numpy.array', 'np.array', (['[10, 340, 523]'], {}), '([10, 340, 523])\n', (2941, 2957), True, 'import numpy as np\n'), ((2967, 2986), 'numpy.array', 'np.array', (['[3, 2, 1]'], {}), '([3, 2, 1])\n', (2975, 2986), True, 'import numpy as np\n'), ((3140, 3157), 'numpy.sum', 'np.sum', (['z'], {'axis': '(0)'}), '(z, axis=0)\n', (3146, 3157), True, 'import numpy as np\n'), ((3167, 3186), 'scipy.signal.argrelmax', 'signal.argrelmax', (['z'], {}), '(z)\n', (3183, 3186), False, 'from scipy import signal\n'), ((3291, 3332), 'matplotlib.pyplot.subplots', 'plt.subplots', ([], {'num': '"""peaksampl"""', 'clear': '(True)'}), "(num='peaksampl', clear=True)\n", (3303, 3332), True, 'from matplotlib import pyplot as plt\n'), ((452, 461), 'numpy.abs', 'np.abs', (['y'], {}), '(y)\n', (458, 461), True, 'import numpy as np\n'), ((505, 531), 'numpy.expand_dims', 'np.expand_dims', (['offset', '(-1)'], {}), '(offset, -1)\n', (519, 531), True, 'import numpy as np\n'), ((3025, 3057), 'numpy.take', 'np.take', (['y', 'indices'], {'mode': '"""clip"""'}), "(y, indices, mode='clip')\n", (3032, 3057), True, 'import numpy as np\n'), ((2866, 2890), 'numpy.linspace', 'np.linspace', (['(0)', '(10)', '(1000)'], {}), '(0, 10, 1000)\n', (2877, 2890), True, 'import numpy as np\n'), ((943, 958), 'numpy.finfo', 'np.finfo', (['float'], {}), '(float)\n', (951, 958), True, 'import numpy as np\n')]