fourier.py

#! /usr/bin/env python3
###############################################################################
#
# Copyright 2024
# ASTRON (Netherlands Institute for Radio Astronomy) <http://www.astron.nl/>
# P.O.Box 2, 7990 AA Dwingeloo, The Netherlands
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

###############################################################################

# Author: Eric Kooistra
# Purpose: Fourier transform functions for DSP
# Description:
#
# References:
# [1] dsp_study_erko.txt

import numpy as np
from .utilities import c_rtol, c_atol, ceil_pow2, is_even


###############################################################################
# Time domain interpolation using DFT
###############################################################################

def fourier_interpolate(HFfilter, Ncoefs):
    """Use Fourier interpolation to create final FIR filter coefs.

    HF contains filter transfer function for N points, in order 0 to fs. The
    interpolation inserts Ncoefs - N zeros and then performs IFFT to get the
    interpolated impulse response.

    Use phase correction dependent on interpolation factor Q to have fractional
    time shift of hInterpolated, to make it symmetrical. Similar as done in
    pfs_coeff_final.m and pfir_coeff.m. Use upper = conj(lower), because that
    is easier than using upper from HFfilter.

    Reference: LYONS 13.27, 13.28
    """
    N = len(HFfilter)
    K = N // 2

    # Interpolation factor Q can be fractional, for example:
    # - Ncoefs = 1024 * 16
    #   . firls: N = 1024 + 1  --> Q = 15.98439
    #   . remez: N = 1024 --> Q = 16
    Q = Ncoefs / N

    # Apply phase correction (= time shift) to make the impulse response
    # exactly symmetric
    f = np.arange(N) / Ncoefs
    p = np.exp(-1j * np.pi * (Q - 1) * f)
    HF = HFfilter * p

    HFextended = np.zeros(Ncoefs, dtype=np.complex_)
    if is_even(N):
        # Copy DC and K - 1 positive frequencies in lower part (K values)
        HFextended[0:K] = HF[0:K]  # starts with DC at [0]
        # Create the upper part of the spectrum from the K - 1 positive
        # frequencies and fs/2 in lower part (K values)
        HFflip = np.conjugate(np.flip(HF[0:K + 1]))  # contains DC at [0]
        HFextended[-K:] = HFflip[0:K]  # omit DC at [K]
        # K + K = N values, because N is even and K = N // 2
    else:
        # Copy DC and K positive frequencies in lower part (K + 1 values)
        HFextended[0:K + 1] = HF[0:K + 1]  # starts with DC at [0]
        # Create the upper part of the spectrum from the K positive
        # frequencies in the lower part (K values)
        HFflip = np.conjugate(np.flip(HF[0:K + 1]))  # contains DC at [0]
        HFextended[-K:] = HFflip[0:K]  # omit DC at [K]
        # K + 1 + K = N values, because N is odd and K = N // 2
    hInterpolated = np.fft.ifft(HFextended)
    if np.allclose(hInterpolated.imag, np.zeros(Ncoefs), rtol=c_rtol, atol=c_atol):
        print('hInterpolated.imag ~= 0')
    else:
        print('WARNING: hInterpolated.imag != 0 (max(abs) = %e)' % np.max(np.abs(hInterpolated.imag)))
    return hInterpolated.real


###############################################################################
# Discrete Time Fourier Transform (DTFT)
###############################################################################

def dtft(coefs, Ndtft=None, zeroCenter=True, fftShift=True):
    """Calculate DTFT of filter impulse response or window.

    Use DFT with Ndtft points, to have frequency resolution of 2 pi / Ndtft.
    Define h by extending coefs with Ndtft - Ncoefs zeros, where Ncoefs =
    len(coefs). This DFT approaches the DTFT using bandlimited interpolation,
    similar to INTERP() which interpolates by factor L = Ndtft / Ncoefs [JOS1].

    Input:
    . coefs: filter impulse response or window coefficients
    . Ndtft: number of points in DFT to calculate DTFT
    . zeroCenter: when True zero center h to have even function that aligns
      with cos() in DTFT, for zero-phase argument (+-1 --> 0, +-pi). Else
      apply h as causal function.
    . fftShift: when True fft shift to have -0.5 to +0.5 frequency axis,
      else use 0 to 1.0.
    Return:
    . h: zero padded coefs
    . fn: normalized frequency axis for HF, fs = 1
    . HF: dtft(h), the frequency transfer function of h
    """
    c_interpolate = 10
    Ncoefs = len(coefs)
    if Ndtft is None:
        Ndtft = ceil_pow2(Ncoefs * c_interpolate)
    # Time series, causal with coefs at left in h
    h = np.concatenate((coefs, np.zeros([Ndtft - Ncoefs])))
    if zeroCenter:
        # Zero center h to try to make it even
        h = np.roll(h, -(Ncoefs // 2))
    # DFT
    HF = np.fft.fft(h)
    # Normalized frequency axis, fs = 1, ws = 2 pi
    fn = np.arange(0, 1, 1 / Ndtft)  # fn = 0,pos, neg
    if fftShift:
        # FFT shift to center HF, fn = neg, 0,pos
        fn = fn - 0.5
        HF = np.roll(HF, Ndtft // 2)
    return h, fn, HF


###############################################################################
# Estimate f for certain gain or estimate gain for certain frequency
###############################################################################

def estimate_gain_at_frequency(f, HF, freq):
    """Find gain = abs(HF) at frequency in f that is closest to freq.

    Input:
    . f: frequency axis as f = fn * fs for fn from dtft() with fftShift=True
    . HF: spectrum as from dtft() with fftShift=True
    . freq : frequency to look for in f
    Return:
    . fIndex: index of frequency in f that is closest to freq
    . fValue: frequency at fIndex
    . fGain : gain = abs(HF) at fIndex
    """
    fDiff = np.abs(f - freq)
    fIndex = fDiff.argmin()
    fValue = f[fIndex]
    fGain = np.abs(HF[fIndex])
    return fIndex, fValue, fGain


def estimate_frequency_at_gain(f, HF, gain):
    """Find positive frequency in f that has gain = abs(HF) closest to gain.

    Look only in positive frequencies and from highest frequency to lowest
    frequency to avoid band pass ripple gain variation in LPF.

    Input:
    . f: frequency axis as f = fn * fs for fn from dtft() with fftShift=True
    . HF: spectrum as from dtft() with fftShift=True
    . gain : gain to look for in HF
    Return:
    . fIndex: index of frequency in f that has abs(HF) closest to gain
    . fValue: frequency at fIndex
    . fGain : gain = abs(HF) at fIndex
    """
    # Look only in positive frequencies
    pos = len(f[f <= 0])
    HFpos = HF[pos:]
    HFgain = np.abs(HFpos)
    HFdiff = np.abs(HFgain - gain)
    # Flip to avoid detections in LPF band pass
    HFdiff = np.flipud(HFdiff)
    fIndex = pos + len(HFpos) - 1 - HFdiff.argmin()
    fValue = f[fIndex]
    fGain = np.abs(HF[fIndex])
    return fIndex, fValue, fGain


###############################################################################
# Radix FFT
###############################################################################

# def radix_fft():
#     """Calculate DFT using smaller DFTs
#
#     Use DFTs with radices equal to the greatest common dividers (gcd) of the
#     input DFT size.
#
#     https://stackoverflow.com/questions/9940192/how-to-calculate-a-large-size-fft-using-smaller-sized-ffts
#     """