12922237403

Committed 23 Jan 2025 04:24AM UTC coverage: 73.077% (-6.2%) from 79.264%

Build # 12922237403

Build Type

push

github

Committed by

avivajpeyi

Commit Message

feat: add jax as optional backend

Run Details

140 of 231 branches covered (60.61%)

Branch coverage included in aggregate %.

49 of 140 new or added lines in 15 files covered. (35.0%)

21 existing lines in 3 files now uncovered.

753 of 991 relevant lines covered (75.98%)

0.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

95.45

/src/pywavelet/transforms/phi_computer.py

from ..backend import xp, PI, betainc, ifft



def phitilde_vec(
    omega: xp.ndarray, Nf: int, d: float = 4.0
) -> xp.ndarray:
    """Compute phi_tilde(omega_i) array, nx is filter steepness, defaults to 4.

    Eq 11 of https://arxiv.org/pdf/2009.00043.pdf (Cornish et al. 2020)

    phi(omega_i) =
        1/sqrt(2π∆F) if |omega_i| < A
        1/sqrt(2π∆F) cos(nu_d π/2 * |omega|-A / B) if A < |omega_i| < A + B

    Where nu_d = normalized incomplete beta function



    Parameters
    ----------
    ω : xp.ndarray
        Array of angular frequencies
    Nf : int
        Number of frequency bins
    d : float, optional
        Number of standard deviations for the gaussian wavelet, by default 4.

    Returns
    -------
    xp.ndarray
        Array of phi_tilde(omega_i) values

    """
    dF = 1.0 / (2 * Nf)  # NOTE: missing 1/dt?
    dOmega = 2 * PI * dF  # Near Eq 10 # 2 pi times DF
    inverse_sqrt_dOmega = 1.0 / xp.sqrt(dOmega)

    A = dOmega / 4
    B = dOmega - 2 * A  # Cannot have B \leq 0.
    if B <= 0:
        raise ValueError("B must be greater than 0")

    phi = xp.zeros(omega.size)
    mask = (A <= xp.abs(omega)) & (xp.abs(omega) < A + B)  # Minor changes
    vd = (PI / 2.0) * __nu_d(omega[mask], A, B, d=d)  # different from paper
    phi[mask] = inverse_sqrt_dOmega * xp.cos(vd)
    phi[xp.abs(omega) < A] = inverse_sqrt_dOmega
    return phi


def __nu_d(
    omega: xp.ndarray, A: float, B: float, d: float = 4.0
) -> xp.ndarray:
    """Compute the normalized incomplete beta function.

    Parameters
    ----------
    ω : xp.ndarray
        Array of angular frequencies
    A : float
        Lower bound for the beta function
    B : float
        Upper bound for the beta function
    d : float, optional
        Number of standard deviations for the gaussian wavelet, by default 4.

    Returns
    -------
    xp.ndarray
        Array of ν_d values

    scipy.special.betainc
    https://docs.scipy.org/doc/scipy-1.7.1/reference/reference/generated/scipy.special.betainc.html

    """
    x = (xp.abs(omega) - A) / B
    return betainc(d, d, x) / betainc(d, d, 1)


def phitilde_vec_norm(Nf: int, Nt: int,  d: float) -> xp.ndarray:
    """Normalize phitilde for inverse frequency domain transform."""

    # Calculate the frequency values
    ND = Nf * Nt
    omegas = 2 * xp.pi / ND * xp.arange(0, Nt // 2 + 1)

    # Calculate the unnormalized phitilde (u_phit)
    u_phit = phitilde_vec(omegas, Nf,  d)

    # Normalize the phitilde
    normalising_factor = PI ** (-1 / 2)  # Ollie's normalising factor

    # Notes: this is the overall normalising factor that is different from Cornish's paper
    # It is the only way I can force this code to be consistent with our work in the
    # frequency domain. First note that

    # old normalising factor -- This factor is absolutely ridiculous. Why!?
    # Matt_normalising_factor = np.sqrt(
    #     (2 * np.sum(u_phit[1:] ** 2) + u_phit[0] ** 2) * 2 * PI / ND
    # )
    # Matt_normalising_factor /= PI**(3/2)/PI

    # The expression above is equal to np.pi**(-1/2) after working through the maths.
    # I have pulled (2/Nf) from __init__.py (from freq to wavelet) into the normalsiing
    # factor here. I thnk it's cleaner to have ONE normalising constant. Avoids confusion
    # and it is much easier to track.

    # TODO: understand the following:
    # (2 * np.sum(u_phit[1:] ** 2) + u_phit[0] ** 2) = 0.5 * Nt / dOmega
    # Matt_normalising_factor is equal to 1/sqrt(pi)... why is this computed?
    # in such a stupid way?

    return u_phit / (normalising_factor)


def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> xp.ndarray:
    """get time domain phi as fourier transform of phitilde_vec"""
    insDOM = 1.0 / xp.sqrt(PI / Nf)
    K = q * 2 * Nf
    half_K = q * Nf  # xp.int64(K/2)

    dom = 2 * PI / K  # max frequency is K/2*dom = pi/dt = OM

    DX = xp.zeros(K, dtype=xp.complex128)

    # zero frequency
    DX[0] = insDOM

    DX = DX.copy()
    # postive frequencies
    DX[1 : half_K + 1] = phitilde_vec(
        dom * xp.arange(1, half_K + 1), Nf, d
    )
    # negative frequencies
    DX[half_K + 1 :] = phitilde_vec(
        -dom * xp.arange(half_K - 1, 0, -1), Nf,  d
    )
    DX = K * ifft(DX, K)

    phi = xp.zeros(K)
    phi[0:half_K] = xp.real(DX[half_K:K])
    phi[half_K:] = xp.real(DX[0:half_K])

    nrm = xp.sqrt(K / dom)  # *xp.linalg.norm(phi)

    fac = xp.sqrt(2.0) / nrm
    phi *= fac
    return phi

1	from ..backend import xp, PI, betainc, ifft	1✔
2
3
4		1✔
5	def phitilde_vec(
6	omega: xp.ndarray, Nf: int, d: float = 4.0
7	) -> xp.ndarray:
8	"""Compute phi_tilde(omega_i) array, nx is filter steepness, defaults to 4.
9
10	Eq 11 of https://arxiv.org/pdf/2009.00043.pdf (Cornish et al. 2020)
11
12	phi(omega_i) =
13	1/sqrt(2π∆F) if \|omega_i\| < A
14	1/sqrt(2π∆F) cos(nu_d π/2 * \|omega\|-A / B) if A < \|omega_i\| < A + B
15
16	Where nu_d = normalized incomplete beta function
17
18
19
20	Parameters
21	----------
22	ω : xp.ndarray
23	Array of angular frequencies
24	Nf : int
25	Number of frequency bins
26	d : float, optional
27	Number of standard deviations for the gaussian wavelet, by default 4.
28
29	Returns
30	-------
31	xp.ndarray
32	Array of phi_tilde(omega_i) values	1✔
33		1✔
34	"""	1✔
35	dF = 1.0 / (2 * Nf) # NOTE: missing 1/dt?
36	dOmega = 2 * PI * dF # Near Eq 10 # 2 pi times DF	1✔
37	inverse_sqrt_dOmega = 1.0 / xp.sqrt(dOmega)	1✔
38		1!
UNCOV 39	A = dOmega / 4	×
40	B = dOmega - 2 * A # Cannot have B \leq 0.
41	if B <= 0:	1✔
42	raise ValueError("B must be greater than 0")	1✔
43		1✔
44	phi = xp.zeros(omega.size)	1✔
45	mask = (A <= xp.abs(omega)) & (xp.abs(omega) < A + B) # Minor changes	1✔
46	vd = (PI / 2.0) * __nu_d(omega[mask], A, B, d=d) # different from paper	1✔
47	phi[mask] = inverse_sqrt_dOmega * xp.cos(vd)
48	phi[xp.abs(omega) < A] = inverse_sqrt_dOmega
49	return phi	1✔
50
51
52	def __nu_d(
53	omega: xp.ndarray, A: float, B: float, d: float = 4.0
54	) -> xp.ndarray:
55	"""Compute the normalized incomplete beta function.
56
57	Parameters
58	----------
59	ω : xp.ndarray
60	Array of angular frequencies
61	A : float
62	Lower bound for the beta function
63	B : float
64	Upper bound for the beta function
65	d : float, optional
66	Number of standard deviations for the gaussian wavelet, by default 4.
67
68	Returns
69	-------
70	xp.ndarray
71	Array of ν_d values
72
73	scipy.special.betainc
74	https://docs.scipy.org/doc/scipy-1.7.1/reference/reference/generated/scipy.special.betainc.html	1✔
75		1✔
76	"""
77	x = (xp.abs(omega) - A) / B
78	return betainc(d, d, x) / betainc(d, d, 1)	1✔
79
80
81	def phitilde_vec_norm(Nf: int, Nt: int, d: float) -> xp.ndarray:
82	"""Normalize phitilde for inverse frequency domain transform."""	1✔
83		1✔
84	# Calculate the frequency values
85	ND = Nf * Nt
86	omegas = 2 * xp.pi / ND * xp.arange(0, Nt // 2 + 1)	1✔
87
88	# Calculate the unnormalized phitilde (u_phit)
89	u_phit = phitilde_vec(omegas, Nf, d)	1✔
90
91	# Normalize the phitilde
92	normalising_factor = PI ** (-1 / 2) # Ollie's normalising factor
93
94	# Notes: this is the overall normalising factor that is different from Cornish's paper
95	# It is the only way I can force this code to be consistent with our work in the
96	# frequency domain. First note that
97
98	# old normalising factor -- This factor is absolutely ridiculous. Why!?
99	# Matt_normalising_factor = np.sqrt(
100	# (2 * np.sum(u_phit[1:] 2) + u_phit[0] 2) * 2 * PI / ND
101	# )
102	# Matt_normalising_factor /= PI**(3/2)/PI
103
104	# The expression above is equal to np.pi**(-1/2) after working through the maths.
105	# I have pulled (2/Nf) from __init__.py (from freq to wavelet) into the normalsiing
106	# factor here. I thnk it's cleaner to have ONE normalising constant. Avoids confusion
107	# and it is much easier to track.
108
109	# TODO: understand the following:
110	# (2 * np.sum(u_phit[1:] 2) + u_phit[0] 2) = 0.5 * Nt / dOmega
111	# Matt_normalising_factor is equal to 1/sqrt(pi)... why is this computed?	1✔
112	# in such a stupid way?
113
114	return u_phit / (normalising_factor)	1✔
115
116		1✔
117	def phi_vec(Nf: int, d: float = 4.0, q: int = 16) -> xp.ndarray:	1✔
118	"""get time domain phi as fourier transform of phitilde_vec"""	1✔
119	insDOM = 1.0 / xp.sqrt(PI / Nf)
120	K = q * 2 * Nf	1✔
121	half_K = q * Nf # xp.int64(K/2)
122		1✔
123	dom = 2 * PI / K # max frequency is K/2*dom = pi/dt = OM
124
125	DX = xp.zeros(K, dtype=xp.complex128)	1✔
126
127	# zero frequency	1✔
128	DX[0] = insDOM
129		1✔
130	DX = DX.copy()
131	# postive frequencies	1✔
132	DX[1 : half_K + 1] = phitilde_vec(	1✔
133	dom * xp.arange(1, half_K + 1), Nf, d
134	)	1✔
135	# negative frequencies	1✔
136	DX[half_K + 1 :] = phitilde_vec(	1✔
137	-dom * xp.arange(half_K - 1, 0, -1), Nf, d
138	)	1✔
139	DX = K * ifft(DX, K)
140		1✔
141	phi = xp.zeros(K)	1✔
142	phi[0:half_K] = xp.real(DX[half_K:K])	1✔
143	phi[half_K:] = xp.real(DX[0:half_K])
144
145	nrm = xp.sqrt(K / dom) # *xp.linalg.norm(phi)
146
147	fac = xp.sqrt(2.0) / nrm
148	phi *= fac
149	return phi

avivajpeyi / pywavelet / 12922237403

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous