Returns individual terms of inverse DFT of a function

get_fourier_components.py

def get_fourier_components(
    func, num_samples, domain_length, sort_by_descending_magnitude=False
):
    """
    Returns the N // 2 + 1 amplitude-phase terms of the inverse D series 
    of func(t), where funct t is evaluated on [0, T) in N evenly-spaced points.
    I.e each component
    
                     A[k] cos(2*pi * f[k] * t + phi[k])
    
    is returned in a 2D array for each frequency component and each time.
    
    Parameters
    ----------
    func : Callable[[numpy.ndarray], numpy.ndarray]
        A vectorized function whose fourier components are being evaluated
    
    num_samples : int
        The number of samples used, N
    
    domain_length : float
        The upper bound, T, on the sampling domain: [0, T)
    
    sort_by_descending_magnitude : bool, default=False
        If True, the components are returned in order of decreasing amplitude
    
    Returns
    -------
    numpy.ndarray, shape-(N // 2 + 1, N)
        Axis-0 carries the respective fourier components. I.e. element-k returns
        the shape-(N,) array vectorized over times:
                       A[k] cos(2*pi * f[k] * t + phi[k])   (where t is shape-(N,))
        Thus calling `out.sum(axis=0)` on the output returns `func(t)`
    """
    N = num_samples
    L = domain_length

    # evaluate on [0, N)
    t = np.arange(N) * (L / N)

    ck = np.fft.rfft(func(t))  # shape-(N,)

    amps = np.abs(ck)
    phases = np.arctan2(ck.imag, ck.real)
    times = t[:, None]
    freqs = np.arange(N // 2 + 1) / L

    # shape-(N, N) - time vs freq
    out = amps * np.cos((2 * np.pi * freqs) * times + phases) / N
    const = out[:, 0]
    out = out[:, 1:] * 2  # need to double count for n, N-n folding
    
    out = np.hstack([const[:, None], out]).T  # transpose: freq vs time

    if sort_by_descending_magnitude:
        out = out[np.argsort(amps), :][::-1]

    return out

plotting:

def plot_topk(out, topk,  function_label="true function", ax=None,):
    cumout = np.cumsum(out, axis=0)
    if ax is None:
        fig, ax = plt.subplots()
    else:
        fig = None

    for i in range(topk):
        ax.plot(
            t, out[i], c="blue", alpha=0.3, label="Fourier component" if i == 0 else ""
        )
    ax.plot(t, cumout[topk], lw=3, alpha=0.6, label=f"Sum of top-{topk} components")
    ax.plot(t, f(t), c="black", ls="--", lw=1, label=function_label)
    ax.set_xlabel(r"$t$")
    ax.legend()
    return fig, ax