power_level.py

import numpy as np
import scipy as sp

def tcp_siggen(samp_rate, freq, amp, noise_var, N):
    t = np.arange(start=0, stop=N / samp_rate, step=1 / samp_rate, dtype=np.complex64)
    # Create the phase of the sin wave at each samples
    phase = 2 * np.pi * freq * t
    # create a complex sin wave
    carrier = amp * np.exp(phase * 1j)
    # generate complex white(ish) gaussian noise with the right variance
    noise_dev = np.sqrt(noise_var)
    noise = (noise_dev / np.sqrt(2)) * (np.random.normal(0, 1, N) + 1j * np.random.normal(0, 1, N))
    # add the sin wave and the noise together
    return np.add(carrier, noise, dtype=np.complex64)


Fs = 100000.0
N = int(Fs)
noise_variance = 100.0

# Make choices at random
amp = np.random.choice([200, 400, 800])
freq = Fs / 32
gen = tcp_siggen(samp_rate=Fs, freq=freq, amp=amp, noise_var=noise_variance, N=N)


#  ------- SOLUTION --------

import scipy.signal as signal
f, Pxx_den = signal.periodogram(gen, Fs)
peak_freq = f[np.argmax(Pxx_den)]
print(f"peak freq: {peak_freq * 2 * np.pi} rad/s")


def unilateral_ztransform(gen, z):
    return np.sum(gen * (z.reshape((-1, 1)) ** -np.arange(len(gen))), axis=-1)


def unilateral_discrete_stransform(gen, s):
    z = np.exp(s / Fs)
    return unilateral_ztransform(gen, z) / Fs


def finite_difference_coefficients(k, xbar, x):
    """
    based on the maths in Finite Difference Methods for Ordinary and Partial Differential Equations
    :param k: the order of the derivative
    :param xbar: the point at which the derivative is to be evaluated
    :param x: the set of points at which the function is known
    :return: the coefficients of the finite difference approximation
    """
    x = np.array(x)
    n = len(x)
    factorial_values = 1 / sp.special.factorial(np.arange(n))
    matrix = np.einsum('i,ji->ij', factorial_values, np.vander(x - xbar, n, increasing=True), dtype=np.complex128)

    b = np.zeros(len(x))
    b[k] = 1
    return sp.linalg.solve(matrix, b)

def nth_derivative(f, n: int, step_size=1e-7):
    """
    the nth derivative of f numerically around x using finite differences
    does not evaluate f at x
    """

    def diff(x):
        if n == 0:
            return f(x + step_size) / step_size
        xbar = x
        # sample a circle of complex oints around x but not including x
        number_of_samples = 2 * n + 1
        x = np.array(
            [xbar + step_size * np.exp(2j * np.pi * i / number_of_samples) for i in range(0, number_of_samples)])
        return np.dot(finite_difference_coefficients(n, xbar, x), f(x)) / (step_size ** n)

    return diff


def find_residue(f, pole, order):
    return nth_derivative(lambda z: ((z - pole) ** order * f(z)), order - 1)(pole) / np.math.factorial(order - 1)


residue = np.real(find_residue(lambda s: unilateral_discrete_stransform(gen, s), peak_freq * 1j * 2 * np.pi, 1))
print(f"residue of s-transform at peak frequency (amplitude): {residue}")

t = np.arange(start=0, stop=N / Fs, step=1 / Fs, dtype=np.complex64)
predicted_noise_variance = np.var(gen) - np.var(residue * np.exp(2 * np.pi * peak_freq * t * 1j))
print(f"predicted noise variance: {predicted_noise_variance}")

predicted_snr = 10 * np.log10(residue * residue / predicted_noise_variance)
snr = 10 * np.log10(amp * amp / noise_variance)
print(f"predicted snr: {predicted_snr}")