Skip to content

Instantly share code, notes, and snippets.

@avivajpeyi

avivajpeyi/wdm_transform.ipynb

Last active April 28, 2025 19:51
Show Gist options
  • Save avivajpeyi/b0a99b3f54b6841fb37e60af312b58a4 to your computer and use it in GitHub Desktop.
Save avivajpeyi/b0a99b3f54b6841fb37e60af312b58a4 to your computer and use it in GitHub Desktop.
Download ZIP
wdm_transform.ipynb
Display the source blob
Display the rendered blob
Raw
wdm_transform.ipynb
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"provenance": [],
"authorship_tag": "ABX9TyOk7xD8r533f+eImNBxoHDY",
"include_colab_link": true
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"language_info": {
"name": "python"
}
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/avivajpeyi/b0a99b3f54b6841fb37e60af312b58a4/wdm_transform.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"source": [
"# Comparing JAX and Numba WDM Transforms\n",
"\n",
"\n",
"## Common functions"
],
"metadata": {
"id": "H6YcKEr5mHIV"
}
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {
"id": "0gJu4ONiQ6sB",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "4ad76cba-8b79-4422-ddc6-deb4939d91e8"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Running JAX on cpu\n"
]
}
],
"source": [
"\"\"\" COMMON FUNCS \"\"\"\n",
"import numpy as np\n",
"from numpy import fft\n",
"from scipy.special import betainc\n",
"from scipy.signal import chirp\n",
"from typing import List, Tuple\n",
"import jax\n",
"jax.config.update('jax_enable_x64', False)\n",
"\n",
"DEVICE = jax.devices()[0].device_kind\n",
"print(f\"Running JAX on {DEVICE}\")\n",
"\n",
"PI = np.pi\n",
"FREQ_RANGE = [20, 100]\n",
"\n",
"def phitilde_vec_norm(Nf: int, Nt: int, dt: float, d: float) -> np.ndarray:\n",
" ND = Nf * Nt\n",
" omegas = 2 * np.pi / ND * np.arange(0, Nt // 2 + 1)\n",
" u_phit = phitilde_vec(omegas, Nf, dt, d)\n",
" normalising_factor = np.pi ** (-1 / 2) # Ollie's normalising factor\n",
" return u_phit / (normalising_factor)\n",
"\n",
"def phitilde_vec(omega: np.ndarray, Nf: int, dt: float, d: float = 4.0) -> np.ndarray:\n",
" \"\"\"Compute phi_tilde(omega_i) array.\"\"\"\n",
" dF = 1.0 / (2 * Nf * dt)\n",
" dOmega = 2 * PI * dF\n",
" inverse_sqrt_dOmega = 1.0 / np.sqrt(dOmega)\n",
" A, B = dOmega / 4, 3 * dOmega / 4\n",
" if B <= 0:\n",
" raise ValueError(\"B must be greater than 0\")\n",
" phi = np.full_like(omega, inverse_sqrt_dOmega)\n",
" mask = (A <= np.abs(omega)) & (np.abs(omega) < A + B)\n",
" phi[mask] *= np.cos((PI / 2.0) * _nu_d(omega[mask], A, B, d))\n",
" return phi\n",
"\n",
"def _nu_d(omega: np.ndarray, A: float, B: float, d: float = 4.0) -> np.ndarray:\n",
" \"\"\"Compute the normalized incomplete beta function.\"\"\"\n",
" x = (np.abs(omega) - A) / B\n",
" return betainc(d, d, x)\n",
"\n",
"\n",
"def simulate_data(Nf):\n",
" # assert Nf is power of 2\n",
" assert Nf & (Nf - 1) == 0, \"Nf must be a power of 2\"\n",
" fs = 512\n",
" dt = 1 / fs\n",
" Nt = Nf\n",
" mult = 16\n",
" nx = 4.0\n",
" ND = Nt * Nf\n",
" t = np.arange(0, ND) * dt\n",
" y = chirp(t, f0=FREQ_RANGE[0], f1=FREQ_RANGE[1], t1=t[-1], method=\"quadratic\")\n",
" phif = phitilde_vec_norm(Nf, Nt, dt=dt, d=nx)\n",
" yf = fft.fft(y)[:ND//2+1]\n",
"\n",
"\n",
"\n",
" plt.plot(phif)\n",
"\n",
"\n",
" return yf, phif\n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"source": [
"## NUMBA Transform"
],
"metadata": {
"id": "4JhI5f_A3p2V"
}
},
{
"cell_type": "code",
"source": [
"\"\"\"NUMBA VERSION\"\"\"\n",
"import numpy as np\n",
"from numba import njit\n",
"from numpy import fft\n",
"\n",
"\n",
"def transform_wavelet_freq_helper_numba(\n",
" data: np.ndarray, Nf: int, Nt: int, phif: np.ndarray\n",
") -> np.ndarray:\n",
" \"\"\"helper to do the wavelet transform using the fast wavelet domain transform\"\"\"\n",
" wave = np.zeros((Nt, Nf)) # wavelet wavepacket transform of the signal\n",
"\n",
" DX = np.zeros(Nt, dtype=np.complex128)\n",
" freq_strain = data.copy() # Convert\n",
" for f_bin in range(0, Nf + 1):\n",
" __fill_wave_1_numba(f_bin, Nt, Nf, DX, freq_strain, phif)\n",
" DX_trans = fft.ifft(\n",
" DX, Nt\n",
" ) # A fix because numba doesn't support np.fft\n",
" __fill_wave_2_numba(f_bin, DX_trans, wave, Nt, Nf)\n",
"\n",
" return wave\n",
"\n",
"\n",
"@njit()\n",
"def __fill_wave_1_numba(\n",
" f_bin: int,\n",
" Nt: int,\n",
" Nf: int,\n",
" DX: np.ndarray,\n",
" data: np.ndarray,\n",
" phif: np.ndarray,\n",
") -> None:\n",
" \"\"\"helper for assigning DX in the main loop\"\"\"\n",
" i_base = Nt // 2\n",
" jj_base = f_bin * Nt // 2\n",
"\n",
" if f_bin == 0 or f_bin == Nf:\n",
" # NOTE this term appears to be needed to recover correct constant (at least for m=0), but was previously missing\n",
" DX[Nt // 2] = phif[0] * data[f_bin * Nt // 2] / 2.0\n",
" else:\n",
" DX[Nt // 2] = phif[0] * data[f_bin * Nt // 2]\n",
"\n",
" for jj in range(jj_base + 1 - Nt // 2, jj_base + Nt // 2):\n",
" j = np.abs(jj - jj_base)\n",
" i = i_base - jj_base + jj\n",
" if f_bin == Nf and jj > jj_base:\n",
" DX[i] = 0.0\n",
" elif f_bin == 0 and jj < jj_base:\n",
" DX[i] = 0.0\n",
" elif j == 0:\n",
" continue\n",
" else:\n",
" DX[i] = phif[j] * data[jj]\n",
"\n",
"\n",
"@njit()\n",
"def __fill_wave_2_numba(\n",
" f_bin: int, DX_trans: np.ndarray, wave: np.ndarray, Nt: int, Nf: int\n",
") -> None:\n",
" if f_bin == 0:\n",
" # half of lowest and highest frequency bin pixels are redundant, so store them in even and odd components of f_bin=0 respectively\n",
" for n in range(0, Nt, 2):\n",
" wave[n, 0] = np.real(DX_trans[n] * np.sqrt(2))\n",
" elif f_bin == Nf:\n",
" for n in range(0, Nt, 2):\n",
" wave[n + 1, 0] = np.real(DX_trans[n] * np.sqrt(2))\n",
" else:\n",
" for n in range(0, Nt):\n",
" if f_bin % 2:\n",
" if (n + f_bin) % 2:\n",
" wave[n, f_bin] = -np.imag(DX_trans[n])\n",
" else:\n",
" wave[n, f_bin] = np.real(DX_trans[n])\n",
" else:\n",
" if (n + f_bin) % 2:\n",
" wave[n, f_bin] = np.imag(DX_trans[n])\n",
" else:\n",
" wave[n, f_bin] = np.real(DX_trans[n])\n",
"\n",
"\n",
"\n"
],
"metadata": {
"id": "RkJzTw0t3pPw"
},
"execution_count": 2,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"## JAX WDM transform"
],
"metadata": {
"id": "u9mBzg9Vmhgm"
}
},
{
"cell_type": "code",
"source": [
"\"\"\"JAX VERSION\"\"\"\n",
"import jax\n",
"import jax.numpy as jnp\n",
"from functools import partial\n",
"from jax import jit\n",
"from jax.numpy.fft import ifft\n",
"\n",
"@partial(jit, static_argnames=('Nf', 'Nt'))\n",
"def transform_wavelet_freq_helper_JAX(\n",
" data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray\n",
") -> jnp.ndarray:\n",
" wave = jnp.zeros((Nt, Nf))\n",
" f_bins = jnp.arange(Nf)\n",
"\n",
" i_base = Nt // 2\n",
" jj_base = f_bins * Nt // 2\n",
"\n",
" initial_values = jnp.where(\n",
" (f_bins == 0) | (f_bins == Nf),\n",
" phif[0] * data[f_bins * Nt // 2] / 2.0,\n",
" phif[0] * data[f_bins * Nt // 2]\n",
" )\n",
"\n",
" DX = jnp.zeros((Nf, Nt), dtype=jnp.complex64)\n",
" DX = DX.at[:, Nt // 2].set(initial_values)\n",
"\n",
" j_range = jnp.arange(1 - Nt // 2, Nt // 2)\n",
" j = jnp.abs(j_range)\n",
" i = i_base + j_range\n",
"\n",
" cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)\n",
" cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)\n",
" cond3 = j[None, :] == 0\n",
"\n",
" jj = jj_base[:, None] + j_range[None, :]\n",
" val = jnp.where(cond1 | cond2, 0.0, phif[j] * data[jj])\n",
" DX = DX.at[:, i].set(jnp.where(cond3, DX[:, i], val))\n",
"\n",
" # Vectorized ifft\n",
" DX_trans = ifft(DX, axis=1)\n",
"\n",
" # Vectorized __fill_wave_2_jax\n",
" n_range = jnp.arange(Nt)\n",
" cond1 = (n_range[:, None] + f_bins[None, :]) % 2 == 1\n",
" cond2 = f_bins % 2 == 1\n",
"\n",
" real_part = jnp.where(cond2, -jnp.imag(DX_trans), jnp.real(DX_trans))\n",
" imag_part = jnp.where(cond2, jnp.real(DX_trans), jnp.imag(DX_trans))\n",
"\n",
" wave = jnp.where(cond1, imag_part, real_part)\n",
"\n",
" # Special cases for f_bin 0 and Nf\n",
" wave = wave.at[::2, 0].set(jnp.real(DX_trans[0, ::2] * jnp.sqrt(2)))\n",
" wave = wave.at[1::2, -1].set(jnp.real(DX_trans[-1, ::2] * jnp.sqrt(2)))\n",
"\n",
" return wave.T\n"
],
"metadata": {
"id": "Cf4wt-8dX9Xb"
},
"execution_count": 6,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"## Cupy version"
],
"metadata": {
"id": "b_5FzcvgLHr2"
}
},
{
"cell_type": "code",
"source": [
"try:\n",
" import cupy as cp\n",
" cupy_available = True\n",
"except:\n",
" cupy_available = False\n",
" cp = np\n",
"\n",
"\n",
"\n",
"def transform_wavelet_freq_helper_CuPy(data: cp.ndarray, Nf: int, Nt: int, phif: cp.ndarray) -> cp.ndarray:\n",
" if not cupy_available:\n",
" return None\n",
"\n",
" wave = cp.zeros((Nt, Nf), dtype=cp.complex64)\n",
" f_bins = cp.arange(Nf)\n",
"\n",
" # Base indices\n",
" i_base = Nt // 2\n",
" jj_base = f_bins * Nt // 2\n",
"\n",
" # Initial values\n",
" initial_values = cp.where(\n",
" (f_bins == 0) | (f_bins == Nf),\n",
" phif[0] * data[f_bins * Nt // 2] / 2.0,\n",
" phif[0] * data[f_bins * Nt // 2]\n",
" )\n",
"\n",
" DX = cp.zeros((Nf, Nt), dtype=cp.complex64)\n",
" DX[:, Nt // 2] = initial_values\n",
"\n",
" j_range = cp.arange(1 - Nt // 2, Nt // 2)\n",
" j = cp.abs(j_range)\n",
" i = i_base + j_range\n",
"\n",
" cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)\n",
" cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)\n",
" cond3 = j[None, :] == 0\n",
"\n",
" jj = jj_base[:, None] + j_range[None, :]\n",
" val = cp.where(cond1 | cond2, 0.0, phif[j] * data[jj])\n",
" DX[:, i] = cp.where(cond3, DX[:, i], val)\n",
"\n",
" # Vectorized ifft using CuPy\n",
" DX_trans = cp.fft.ifft(DX, axis=1)\n",
"\n",
" # Vectorized __fill_wave_2\n",
" n_range = cp.arange(Nt)\n",
" cond1 = (n_range[:, None] + f_bins[None, :]) % 2 == 1\n",
" cond2 = f_bins % 2 == 1\n",
"\n",
" real_part = cp.where(cond2, -cp.imag(DX_trans), cp.real(DX_trans))\n",
" imag_part = cp.where(cond2, cp.real(DX_trans), cp.imag(DX_trans))\n",
"\n",
" wave = cp.where(cond1, imag_part, real_part)\n",
"\n",
" # Special cases for f_bin 0 and Nf\n",
" wave[::2, 0] = cp.real(DX_trans[0, ::2] * cp.sqrt(2))\n",
" wave[1::2, -1] = cp.real(DX_trans[-1, ::2] * cp.sqrt(2))\n",
"\n",
" return wave.T"
],
"metadata": {
"id": "EWQ8ugU6LGwe"
},
"execution_count": 7,
"outputs": []
},
{
"cell_type": "code",
"source": [
"\n",
"KERNEL_CODE = r'''\n",
"extern \"C\" __global__\n",
"void wavelet_kernel(\n",
" const float* data, const float* phif, int Nt, int Nf, float* DX_real, float* DX_imag) {\n",
"\n",
" int f = blockIdx.x * blockDim.x + threadIdx.x; // Frequency index\n",
" int t = blockIdx.y * blockDim.y + threadIdx.y; // Time index\n",
"\n",
" if (f < Nf && t < Nt) {\n",
" int i_base = Nt / 2;\n",
" int jj_base = f * Nt / 2;\n",
"\n",
" // Initial values logic\n",
" if (t == Nt / 2) {\n",
" float init_val;\n",
" if (f == 0 || f == Nf - 1) {\n",
" init_val = phif[0] * data[jj_base] / 2.0;\n",
" } else {\n",
" init_val = phif[0] * data[jj_base];\n",
" }\n",
" DX_real[f * Nt + t] = init_val;\n",
" DX_imag[f * Nt + t] = 0.0f; // Initialize imaginary part to zero\n",
" }\n",
" }\n",
"}\n",
"'''\n",
"\n",
"\n",
"if cupy_available:\n",
" wavelet_kernel = cp.RawKernel(KERNEL_CODE, 'wavelet_kernel')\n",
"\n",
"import cupy as cp\n",
"\n",
"# Define the custom CUDA kernel\n",
"\n",
"\n",
"def transform_wavelet_freq_helper_CUDA(data: cp.ndarray, Nf: int, Nt: int, phif: cp.ndarray) -> cp.ndarray:\n",
" if not cupy_available:\n",
" return None\n",
" # Prepare output arrays\n",
" wave = cp.zeros((Nt, Nf), dtype=cp.complex64)\n",
"\n",
" # Prepare DX as separate real and imaginary parts\n",
" DX_real = cp.zeros((Nf, Nt), dtype=cp.float32)\n",
" DX_imag = cp.zeros((Nf, Nt), dtype=cp.float32)\n",
"\n",
" # Launch the CUDA kernel to populate DX\n",
" threads_per_block = (16, 16)\n",
" blocks_per_grid = ((Nf + threads_per_block[0] - 1) // threads_per_block[0],\n",
" (Nt + threads_per_block[1] - 1) // threads_per_block[1])\n",
"\n",
" # Unpack the blocks_per_grid tuple\n",
" wavelet_kernel(\n",
" blocks_per_grid, threads_per_block,\n",
" (data.data.ptr, phif.data.ptr, Nt, Nf, DX_real.data.ptr, DX_imag.data.ptr)\n",
" )\n",
"\n",
" # Perform inverse FFT on DX using CuPy's optimized FFT\n",
" DX = DX_real + 1j * DX_imag\n",
" DX_trans = cp.fft.ifft(DX, axis=1)\n",
"\n",
" # Fill the wave array using CuPy\n",
" f_bins = cp.arange(Nf)\n",
" n_range = cp.arange(Nt)\n",
"\n",
" # Condition logic (this part may not require CUDA kernel)\n",
" cond1 = (n_range[:, None] + f_bins[None, :]) % 2 == 1\n",
" cond2 = f_bins % 2 == 1\n",
"\n",
" real_part = cp.where(cond2, -cp.imag(DX_trans), cp.real(DX_trans))\n",
" imag_part = cp.where(cond2, cp.real(DX_trans), cp.imag(DX_trans))\n",
"\n",
" wave = cp.where(cond1, imag_part, real_part)\n",
"\n",
" # Special cases for f_bin 0 and Nf\n",
" wave[::2, 0] = cp.real(DX_trans[0, ::2] * cp.sqrt(2))\n",
" wave[1::2, -1] = cp.real(DX_trans[-1, ::2] * cp.sqrt(2))\n",
"\n",
" return wave.T\n",
"\n"
],
"metadata": {
"id": "TNZzHj2zMa2f"
},
"execution_count": 8,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import matplotlib.pyplot as plt\n",
"import jax\n",
"import jax.numpy as jnp\n",
"\n",
"# collect plotting data\n",
"Nf = Nt = 2 ** 6\n",
"yf, phif = simulate_data(Nf)\n",
"wave = transform_wavelet_freq_helper_numba(yf, Nf, Nt, phif)\n",
"jax_yf = jnp.array(yf)\n",
"jax_phif = jnp.array(phif)\n",
"wave_jax = transform_wavelet_freq_helper_JAX(jax_yf, Nf, Nt, jax_phif)\n",
"nplots = 3\n",
"\n",
"print(f\"Phif shape {phif.shape}\")\n",
"print(f\"data shape {jax_yf.shape}\")\n",
"\n",
"\n",
"\n",
"def xn_i(m:int, data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray) -> jnp.ndarray:\n",
"\n",
"\n",
" return jnp.fft.ifft(jnp.roll(data, - (m* Nt//2))[:Nt//2] * phif)\n",
"\n",
"\n",
"def transform_conv(data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray) -> jnp.ndarray:\n",
"\n",
" # fftshifft data to match up with Neil Eq 17\n",
" data_shifted = jnp.fft.fftshift(data)\n",
"\n",
" wave= jax.lax.fori_loop(\n",
" 0, Nt,\n",
" lambda i, accum: jnp.stack((accum, xn_i(i, data, Nf, Nt, phif)), 1),\n",
" jnp.zeros((Nt,1), dtype=jnp.complex64)\n",
" )\n",
" return wave[:,1:]\n",
"\n",
"\n",
"# @partial(jit, static_argnames=('Nf', 'Nt'))\n",
"# def transform_conv(\n",
"# data: jnp.ndarray, Nf: int, Nt: int, phif: jnp.ndarray\n",
"# ) -> jnp.ndarray:\n",
"\n",
"\n",
"\n",
"\n",
"# wave = jnp.zeros((Nt, Nf))\n",
"# return\n",
"\n",
"\n",
"\n",
"\n",
" # i_base = Nt // 2\n",
" # jj_base = f_bins * Nt // 2\n",
"\n",
" # initial_values = jnp.where(\n",
" # (f_bins == 0) | (f_bins == Nf),\n",
" # phif[0] * data[f_bins * Nt // 2] / 2.0,\n",
" # phif[0] * data[f_bins * Nt // 2]\n",
" # )\n",
"\n",
" # DX = jnp.zeros((Nf, Nt), dtype=jnp.complex64)\n",
" # DX = DX.at[:, Nt // 2].set(initial_values)\n",
"\n",
" # j_range = jnp.arange(1 - Nt // 2, Nt // 2)\n",
" # j = jnp.abs(j_range)\n",
" # i = i_base + j_range\n",
"\n",
" # cond1 = (f_bins[:, None] == Nf) & (j_range[None, :] > 0)\n",
" # cond2 = (f_bins[:, None] == 0) & (j_range[None, :] < 0)\n",
" # cond3 = j[None, :] == 0\n",
"\n",
" # jj = jj_base[:, None] + j_range[None, :]\n",
" # val = jnp.where(cond1 | cond2, 0.0, phif[j] * data[jj])\n",
" # DX = DX.at[:, i].set(jnp.where(cond3, DX[:, i], val))\n",
"\n",
" # # Vectorized ifft\n",
" # DX_trans = ifft(DX, axis=1)\n",
"\n",
" # # Vectorized __fill_wave_2_jax\n",
" # n_range = jnp.arange(Nt)\n",
" # cond1 = (n_range[:, None] + f_bins[None, :]) % 2 == 1\n",
" # cond2 = f_bins % 2 == 1\n",
"\n",
" # real_part = jnp.where(cond2, -jnp.imag(DX_trans), jnp.real(DX_trans))\n",
" # imag_part = jnp.where(cond2, jnp.real(DX_trans), jnp.imag(DX_trans))\n",
"\n",
" # wave = jnp.where(cond1, imag_part, real_part)\n",
"\n",
" # # Special cases for f_bin 0 and Nf\n",
" # wave = wave.at[::2, 0].set(jnp.real(DX_trans[0, ::2] * jnp.sqrt(2)))\n",
" # wave = wave.at[1::2, -1].set(jnp.real(DX_trans[-1, ::2] * jnp.sqrt(2)))\n",
"\n",
" # return wave.T\n",
"\n",
"\n",
"\n",
"wave_conv = transform_conv(jax_yf, Nf, Nt, jax_phif)\n",
"\n",
"# if cupy_available:\n",
"# cupy_yf = cp.array(yf)\n",
"# cupy_phif = cp.array(phif)\n",
"# wave_cupy = transform_wavelet_freq_helper_CuPy(cupy_yf, Nf, Nt, cupy_phif).get()\n",
"# wave_CUDA = transform_wavelet_freq_helper_CUDA(cupy_yf, Nf, Nt, cupy_phif).get()\n",
"# nplots += 2\n",
"\n",
"# render plot\n",
"fig, ax = plt.subplots(1, nplots, figsize=(5, 5), sharex=True, sharey=True)\n",
"ax[0].imshow(np.abs(np.rot90(wave)))\n",
"ax[0].set_title(\"Numba\")\n",
"ax[1].imshow(np.abs(np.rot90(wave_jax)))\n",
"ax[1].set_title(\"Jax\")\n",
"ax[2].imshow(np.abs(np.rot90(wave_conv)))\n",
"ax[2].set_title(\"JaxCONV\")\n",
"# if cupy_available:\n",
"# ax[2].imshow(np.abs(np.rot90(wave_cupy)))\n",
"# ax[2].set_title(\"Cupy\")\n",
"# ax[3].imshow(np.abs(np.rot90(wave_CUDA)))\n",
"# ax[3].set_title(\"CUDA\")\n",
"\n",
"for a in ax:\n",
" a.set_xticks([])\n",
" a.set_yticks([])\n",
" # a.set_ylim(120, 70)\n",
"plt.tight_layout()\n",
"plt.show()"
],
"metadata": {
"id": "Tx1yiDeFh7Yl",
"colab": {
"base_uri": "https://localhost:8080/",
"height": 919
},
"outputId": "3e427137-57b7-4bfb-88e1-ef2de5522bc2"
},
"execution_count": 28,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Phif shape (33,)\n",
"data shape (2049,)\n"
]
},
{
"output_type": "error",
"ename": "TypeError",
"evalue": "mul got incompatible shapes for broadcasting: (32,), (33,).",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mTypeError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m<ipython-input-28-9015245b2ced>\u001b[0m in \u001b[0;36m<cell line: 0>\u001b[0;34m()\u001b[0m\n\u001b[1;32m 91\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 92\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 93\u001b[0;31m \u001b[0mwave_conv\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtransform_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjax_yf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mjax_phif\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 94\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 95\u001b[0m \u001b[0;31m# if cupy_available:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m<ipython-input-28-9015245b2ced>\u001b[0m in \u001b[0;36mtransform_conv\u001b[0;34m(data, Nf, Nt, phif)\u001b[0m\n\u001b[1;32m 24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 25\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mtransform_conv\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNf\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mint\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mphif\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m->\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mndarray\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 26\u001b[0;31m wave= jax.lax.fori_loop(\n\u001b[0m\u001b[1;32m 27\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 28\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maccum\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maccum\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mxn_i\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mphif\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
" \u001b[0;31m[... skipping hidden 13 frame]\u001b[0m\n",
"\u001b[0;32m<ipython-input-28-9015245b2ced>\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(i, accum)\u001b[0m\n\u001b[1;32m 26\u001b[0m wave= jax.lax.fori_loop(\n\u001b[1;32m 27\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 28\u001b[0;31m \u001b[0;32mlambda\u001b[0m \u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0maccum\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0maccum\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mxn_i\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mphif\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 29\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mzeros\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mNt\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcomplex64\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 30\u001b[0m )\n",
"\u001b[0;32m<ipython-input-28-9015245b2ced>\u001b[0m in \u001b[0;36mxn_i\u001b[0;34m(m, data, Nf, Nt, phif)\u001b[0m\n\u001b[1;32m 20\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 21\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 22\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfft\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mifft\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mroll\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m-\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mm\u001b[0m\u001b[0;34m*\u001b[0m \u001b[0mNt\u001b[0m\u001b[0;34m//\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0mNt\u001b[0m\u001b[0;34m//\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m*\u001b[0m \u001b[0mphif\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 23\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 24\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/jax/_src/numpy/array_methods.py\u001b[0m in \u001b[0;36mop\u001b[0;34m(self, *args)\u001b[0m\n\u001b[1;32m 1058\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_forward_operator_to_aval\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1059\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1060\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mgetattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34mf\"_{name}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1061\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mop\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1062\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/jax/_src/numpy/array_methods.py\u001b[0m in \u001b[0;36mdeferring_binary_op\u001b[0;34m(self, other)\u001b[0m\n\u001b[1;32m 577\u001b[0m \u001b[0margs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mswap\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mother\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 578\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mother\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_accepted_binop_types\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 579\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mbinary_op\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 580\u001b[0m \u001b[0;31m# Note: don't use isinstance here, because we don't want to raise for\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 581\u001b[0m \u001b[0;31m# subclasses, e.g. NamedTuple objects that may override operators.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/jax/_src/numpy/ufunc_api.py\u001b[0m in \u001b[0;36m__call__\u001b[0;34m(self, out, where, *args)\u001b[0m\n\u001b[1;32m 178\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mNotImplementedError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34mf\"where argument of {self}\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 179\u001b[0m \u001b[0mcall\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m__static_props\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'call'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_vectorized\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 180\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 181\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 182\u001b[0m \u001b[0;34m@\u001b[0m\u001b[0mpartial\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mstatic_argnames\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'self'\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
" \u001b[0;31m[... skipping hidden 13 frame]\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/jax/_src/numpy/ufuncs.py\u001b[0m in \u001b[0;36mmultiply\u001b[0;34m(x, y)\u001b[0m\n\u001b[1;32m 1254\u001b[0m \"\"\"\n\u001b[1;32m 1255\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpromote_args\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"multiply\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1256\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmul\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mbool\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0mlax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbitwise_and\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0my\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1257\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1258\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
" \u001b[0;31m[... skipping hidden 9 frame]\u001b[0m\n",
"\u001b[0;32m/usr/local/lib/python3.11/dist-packages/jax/_src/lax/lax.py\u001b[0m in \u001b[0;36m_try_broadcast_shapes\u001b[0;34m(name, *shapes)\u001b[0m\n\u001b[1;32m 130\u001b[0m \u001b[0mresult_shape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mappend\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnon_1s\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 131\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 132\u001b[0;31m raise TypeError(f'{name} got incompatible shapes for broadcasting: '\n\u001b[0m\u001b[1;32m 133\u001b[0m f'{\", \".join(map(str, map(tuple, shapes)))}.')\n\u001b[1;32m 134\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mtuple\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mresult_shape\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mTypeError\u001b[0m: mul got incompatible shapes for broadcasting: (32,), (33,)."
]
},
{
"output_type": "display_data",
"data": {
"text/plain": [
"<Figure size 640x480 with 1 Axes>"
],
"image/png": "iVBORw0KGgoAAAANSUhEUgAAAjUAAAGdCAYAAADqsoKGAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjAsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvlHJYcgAAAAlwSFlzAAAPYQAAD2EBqD+naQAAL21JREFUeJzt3X9wFHWC//9XJjAJhvwgxCQkBgLJSkCFkQzMps7LhkswUep0gbsLtyiYo0BFQBkPj8gtiNTWsGsdFxUkdXfrrhVY4bxFt0QuqKOgHCNoMJVdL2QXVi9GSCBQZEhcB8j05w++jt9ZEsmgEvLm+ajqKqb7/avf1VXzovs9nSjLsiwBAAAMcLb+HgAAAMC3gVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADDCoP4ewJUSDAZ19OhRxcfHKyoqqr+HAwAA+sCyLJ05c0YZGRmy2b7+Xsw1E2qOHj2qrKys/h4GAAC4DJ9++qluuOGGry1zzYSa+Ph4SRcmJSEhoZ9HAwAA+sLv9ysrKyv0Pf51rplQ8+Ujp4SEBEINAAADTF+WjrBQGAAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEuK9Rs3LhR2dnZio2Nlcvl0oEDB3otu337djmdTiUlJSkuLk4Oh0M1NTVhZaKionrcnnrqqVCZU6dOac6cOUpISFBSUpLmz5+vzs7Oyxk+AAAwUMShZtu2bXK73Vq9erUOHjyoiRMnqrS0VMePH++xfHJyslauXCmfz6eGhgZVVFSooqJCu3btCpU5duxY2Pb8888rKipKs2bNCpWZM2eOPvroI73xxhvasWOH3nnnHS1cuPAyThkAAJgoyrIsK5IKLpdLkydP1oYNGyRJwWBQWVlZWrJkiVasWNGnNiZNmqTp06dr7dq1PR7/4Q9/qDNnzsjr9UqSGhsbNX78eL3//vtyOp2SpNraWt15551qaWlRRkbGJfv0+/1KTExUR0eHEhIS+jROAADQvyL5/o7oTs3Zs2dVV1enkpKSrxqw2VRSUiKfz3fJ+pZlyev1qqmpSYWFhT2WaWtr02uvvab58+eH9vl8PiUlJYUCjSSVlJTIZrNp//79PbYTCATk9/vDNgAAYK6IQk17e7u6u7uVlpYWtj8tLU2tra291uvo6NDQoUNlt9s1ffp0Pfvss5o2bVqPZV944QXFx8dr5syZoX2tra1KTU0NKzdo0CAlJyf32q/H41FiYmJoy8rK6utpAgCAAeiK/PopPj5e9fX1ev/99/WTn/xEbrdbu3fv7rHs888/rzlz5ig2NvYb9VlZWamOjo7Q9umnn36j9gAAwNVtUCSFU1JSFB0drba2trD9bW1tSk9P77WezWZTbm6uJMnhcKixsVEej0dFRUVh5d599101NTVp27ZtYfvT09MvWoh8/vx5nTp1qtd+Y2JiFBMT09dTAwAAA1xEd2rsdrvy8/NDC3ilCwuFvV6vCgoK+txOMBhUIBC4aP/Pf/5z5efna+LEiWH7CwoKdPr0adXV1YX2vfXWWwoGg3K5XJGcAgAAMFREd2okye12a968eXI6nZoyZYqqqqrU1dWliooKSdLcuXOVmZkpj8cj6cLaFqfTqZycHAUCAe3cuVM1NTXatGlTWLt+v18vvfSS/uVf/uWiPseNG6eysjItWLBA1dXVOnfunBYvXqzZs2f36ZdPAADAfBGHmvLycp04cUKrVq1Sa2urHA6HamtrQ4uHm5ubZbN9dQOoq6tLixYtUktLi4YMGaK8vDxt3rxZ5eXlYe1u3bpVlmXp7//+73vsd8uWLVq8eLGKi4tls9k0a9YsPfPMM5EOHwAAGCri99QMVLynBgCAgec7e08NAADA1YpQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARLivUbNy4UdnZ2YqNjZXL5dKBAwd6Lbt9+3Y5nU4lJSUpLi5ODodDNTU1F5VrbGzUXXfdpcTERMXFxWny5Mlqbm4OHS8qKlJUVFTY9sADD1zO8AEAgIEGRVph27Ztcrvdqq6ulsvlUlVVlUpLS9XU1KTU1NSLyicnJ2vlypXKy8uT3W7Xjh07VFFRodTUVJWWlkqSjhw5ottuu03z58/XmjVrlJCQoI8++kixsbFhbS1YsEBPPvlk6PN1110X6fABAIChoizLsiKp4HK5NHnyZG3YsEGSFAwGlZWVpSVLlmjFihV9amPSpEmaPn261q5dK0maPXu2Bg8e3OMdnC8VFRXJ4XCoqqoqkuGG+P1+JSYmqqOjQwkJCZfVBgAAuLIi+f6O6PHT2bNnVVdXp5KSkq8asNlUUlIin893yfqWZcnr9aqpqUmFhYWSLoSi1157TTfeeKNKS0uVmpoql8ulV1555aL6W7ZsUUpKim6++WZVVlbq888/77WvQCAgv98ftgEAAHNFFGra29vV3d2ttLS0sP1paWlqbW3ttV5HR4eGDh0qu92u6dOn69lnn9W0adMkScePH1dnZ6fWrVunsrIyvf7665oxY4ZmzpypPXv2hNr40Y9+pM2bN+vtt99WZWWlampqdM899/Tap8fjUWJiYmjLysqK5FQBAMAAE/GamssRHx+v+vp6dXZ2yuv1yu12a8yYMSoqKlIwGJQk3X333Vq2bJkkyeFwaN++faqurtYPfvADSdLChQtD7d1yyy0aMWKEiouLdeTIEeXk5FzUZ2Vlpdxud+iz3+8n2AAAYLCIQk1KSoqio6PV1tYWtr+trU3p6em91rPZbMrNzZV0IbA0NjbK4/GoqKhIKSkpGjRokMaPHx9WZ9y4cdq7d2+vbbpcLknS4cOHeww1MTExiomJ6fO5AQCAgS2ix092u135+fnyer2hfcFgUF6vVwUFBX1uJxgMKhAIhNqcPHmympqawsr8/ve/16hRo3pto76+XpI0YsSICM4AAACYKuLHT263W/PmzZPT6dSUKVNUVVWlrq4uVVRUSJLmzp2rzMxMeTweSRfWtjidTuXk5CgQCGjnzp2qqanRpk2bQm0uX75c5eXlKiws1NSpU1VbW6tXX31Vu3fvlnThJ9+/+tWvdOedd2r48OFqaGjQsmXLVFhYqAkTJnwL0wAAAAa6iENNeXm5Tpw4oVWrVqm1tVUOh0O1tbWhxcPNzc2y2b66AdTV1aVFixappaVFQ4YMUV5enjZv3qzy8vJQmRkzZqi6uloej0dLly7V2LFj9etf/1q33XabpAt3c958881QgMrKytKsWbP0z//8z9/0/AEAgCEifk/NQMV7agAAGHi+s/fUAAAAXK0INQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwwmWFmo0bNyo7O1uxsbFyuVw6cOBAr2W3b98up9OppKQkxcXFyeFwqKam5qJyjY2Nuuuuu5SYmKi4uDhNnjxZzc3NoeNffPGFHnroIQ0fPlxDhw7VrFmz1NbWdjnDBwAABoo41Gzbtk1ut1urV6/WwYMHNXHiRJWWlur48eM9lk9OTtbKlSvl8/nU0NCgiooKVVRUaNeuXaEyR44c0W233aa8vDzt3r1bDQ0N+vGPf6zY2NhQmWXLlunVV1/VSy+9pD179ujo0aOaOXPmZZwyAAAwUZRlWVYkFVwulyZPnqwNGzZIkoLBoLKysrRkyRKtWLGiT21MmjRJ06dP19q1ayVJs2fP1uDBg3u8gyNJHR0duv766/WrX/1Kf/M3fyNJOnTokMaNGyefz6fvf//7l+zT7/crMTFRHR0dSkhI6NM4AQBA/4rk+zuiOzVnz55VXV2dSkpKvmrAZlNJSYl8Pt8l61uWJa/Xq6amJhUWFkq6EIpee+013XjjjSotLVVqaqpcLpdeeeWVUL26ujqdO3curN+8vDyNHDmy134DgYD8fn/YBgAAzBVRqGlvb1d3d7fS0tLC9qelpam1tbXXeh0dHRo6dKjsdrumT5+uZ599VtOmTZMkHT9+XJ2dnVq3bp3Kysr0+uuva8aMGZo5c6b27NkjSWptbZXdbldSUlKf+/V4PEpMTAxtWVlZkZwqAAAYYAZdiU7i4+NVX1+vzs5Oeb1eud1ujRkzRkVFRQoGg5Kku+++W8uWLZMkORwO7du3T9XV1frBD35wWX1WVlbK7XaHPvv9foINAAAGiyjUpKSkKDo6+qJfHbW1tSk9Pb3XejabTbm5uZIuBJbGxkZ5PB4VFRUpJSVFgwYN0vjx48PqjBs3Tnv37pUkpaen6+zZszp9+nTY3Zqv6zcmJkYxMTGRnB4AABjAInr8ZLfblZ+fL6/XG9oXDAbl9XpVUFDQ53aCwaACgUCozcmTJ6upqSmszO9//3uNGjVKkpSfn6/BgweH9dvU1KTm5uaI+gUAAOaK+PGT2+3WvHnz5HQ6NWXKFFVVVamrq0sVFRWSpLlz5yozM1Mej0fShbUtTqdTOTk5CgQC2rlzp2pqarRp06ZQm8uXL1d5ebkKCws1depU1dbW6tVXX9Xu3bslSYmJiZo/f77cbreSk5OVkJCgJUuWqKCgoE+/fAIAAOaLONSUl5frxIkTWrVqlVpbW+VwOFRbWxtaPNzc3Cyb7asbQF1dXVq0aJFaWlo0ZMgQ5eXlafPmzSovLw+VmTFjhqqrq+XxeLR06VKNHTtWv/71r3XbbbeFyvzrv/6rbDabZs2apUAgoNLSUj333HPf5NwBAIBBIn5PzUDFe2oAABh4vrP31AAAAFytCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGOGyQs3GjRuVnZ2t2NhYuVwuHThwoNey27dvl9PpVFJSkuLi4uRwOFRTUxNW5r777lNUVFTYVlZWFlYmOzv7ojLr1q27nOEDAAADDYq0wrZt2+R2u1VdXS2Xy6WqqiqVlpaqqalJqampF5VPTk7WypUrlZeXJ7vdrh07dqiiokKpqakqLS0NlSsrK9MvfvGL0OeYmJiL2nryySe1YMGC0Of4+PhIhw8AAAwVcahZv369FixYoIqKCklSdXW1XnvtNT3//PNasWLFReWLiorCPj/88MN64YUXtHfv3rBQExMTo/T09K/tOz4+/pJlAADAtSmix09nz55VXV2dSkpKvmrAZlNJSYl8Pt8l61uWJa/Xq6amJhUWFoYd2717t1JTUzV27Fg9+OCDOnny5EX1161bp+HDh+vWW2/VU089pfPnz/faVyAQkN/vD9sAAIC5IrpT097eru7ubqWlpYXtT0tL06FDh3qt19HRoczMTAUCAUVHR+u5557TtGnTQsfLyso0c+ZMjR49WkeOHNHjjz+uO+64Qz6fT9HR0ZKkpUuXatKkSUpOTta+fftUWVmpY8eOaf369T326fF4tGbNmkhODwAADGBRlmVZfS189OhRZWZmat++fSooKAjtf+yxx7Rnzx7t37+/x3rBYFB//OMf1dnZKa/Xq7Vr1+qVV1656NHUl/74xz8qJydHb775poqLi3ss8/zzz+v+++9XZ2dnj+tvAoGAAoFA6LPf71dWVpY6OjqUkJDQ11MGAAD9yO/3KzExsU/f3xHdqUlJSVF0dLTa2trC9re1tX3tWhebzabc3FxJksPhUGNjozweT6+hZsyYMUpJSdHhw4d7DTUul0vnz5/XJ598orFjx150PCYmpsewAwAAzBTRmhq73a78/Hx5vd7QvmAwKK/XG3bn5lKCwWDYXZQ/19LSopMnT2rEiBG9lqmvr5fNZuvxF1cAAODaE/Gvn9xut+bNmyen06kpU6aoqqpKXV1doV9DzZ07V5mZmfJ4PJIurG1xOp3KyclRIBDQzp07VVNTo02bNkmSOjs7tWbNGs2aNUvp6ek6cuSIHnvsMeXm5oZ+HeXz+bR//35NnTpV8fHx8vl8WrZsme655x4NGzbs25oLAAAwgEUcasrLy3XixAmtWrVKra2tcjgcqq2tDS0ebm5uls321Q2grq4uLVq0SC0tLRoyZIjy8vK0efNmlZeXS5Kio6PV0NCgF154QadPn1ZGRoZuv/12rV27NvT4KCYmRlu3btUTTzyhQCCg0aNHa9myZXK73d/GHAAAAANEtFB4IItkoREAALg6RPL9zd9+AgAARiDUAAAAIxBqAACAEQg1AADACBH/+gnhLMvSn8519/cwAAC4KgwZHK2oqKh+6ZtQ8w396Vy3xq/a1d/DAADgqvC/T5bqOnv/xAsePwEAACNwp+YbGjI4Wv/7ZGl/DwMAgKvCkMHR/dY3oeYbioqK6rfbbAAA4Cs8fgIAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEuK9Rs3LhR2dnZio2Nlcvl0oEDB3otu337djmdTiUlJSkuLk4Oh0M1NTVhZe677z5FRUWFbWVlZWFlTp06pTlz5ighIUFJSUmaP3++Ojs7L2f4AADAQBGHmm3btsntdmv16tU6ePCgJk6cqNLSUh0/frzH8snJyVq5cqV8Pp8aGhpUUVGhiooK7dq1K6xcWVmZjh07FtpefPHFsONz5szRRx99pDfeeEM7duzQO++8o4ULF0Y6fAAAYKgoy7KsSCq4XC5NnjxZGzZskCQFg0FlZWVpyZIlWrFiRZ/amDRpkqZPn661a9dKunCn5vTp03rllVd6LN/Y2Kjx48fr/fffl9PplCTV1tbqzjvvVEtLizIyMi7Zp9/vV2Jiojo6OpSQkNCncQIAgP4Vyfd3RHdqzp49q7q6OpWUlHzVgM2mkpIS+Xy+S9a3LEter1dNTU0qLCwMO7Z7926lpqZq7NixevDBB3Xy5MnQMZ/Pp6SkpFCgkaSSkhLZbDbt37+/x74CgYD8fn/YBgAAzDUoksLt7e3q7u5WWlpa2P60tDQdOnSo13odHR3KzMxUIBBQdHS0nnvuOU2bNi10vKysTDNnztTo0aN15MgRPf7447rjjjvk8/kUHR2t1tZWpaamhg980CAlJyertbW1xz49Ho/WrFkTyekBAIABLKJQc7ni4+NVX1+vzs5Oeb1eud1ujRkzRkVFRZKk2bNnh8recsstmjBhgnJycrR7924VFxdfVp+VlZVyu92hz36/X1lZWd/oPAAAwNUrolCTkpKi6OhotbW1he1va2tTenp6r/VsNptyc3MlSQ6HQ42NjfJ4PKFQ8+fGjBmjlJQUHT58WMXFxUpPT79oIfL58+d16tSpXvuNiYlRTExMBGcHAAAGsojW1NjtduXn58vr9Yb2BYNBeb1eFRQU9LmdYDCoQCDQ6/GWlhadPHlSI0aMkCQVFBTo9OnTqqurC5V56623FAwG5XK5IjkFAABgqIgfP7ndbs2bN09Op1NTpkxRVVWVurq6VFFRIUmaO3euMjMz5fF4JF1Y2+J0OpWTk6NAIKCdO3eqpqZGmzZtkiR1dnZqzZo1mjVrltLT03XkyBE99thjys3NVWlpqSRp3LhxKisr04IFC1RdXa1z585p8eLFmj17dp9++QQAAMwXcagpLy/XiRMntGrVKrW2tsrhcKi2tja0eLi5uVk221c3gLq6urRo0SK1tLRoyJAhysvL0+bNm1VeXi5Jio6OVkNDg1544QWdPn1aGRkZuv3227V27dqwx0dbtmzR4sWLVVxcLJvNplmzZumZZ575pucPAAAMEfF7agYq3lMDAMDA8529pwYAAOBqRagBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiXFWo2btyo7OxsxcbGyuVy6cCBA72W3b59u5xOp5KSkhQXFyeHw6Gamppeyz/wwAOKiopSVVVV2P7s7GxFRUWFbevWrbuc4QMAAAMNirTCtm3b5Ha7VV1dLZfLpaqqKpWWlqqpqUmpqakXlU9OTtbKlSuVl5cnu92uHTt2qKKiQqmpqSotLQ0r+/LLL+u9995TRkZGj30/+eSTWrBgQehzfHx8pMMHAACGivhOzfr167VgwQJVVFRo/Pjxqq6u1nXXXafnn3++x/JFRUWaMWOGxo0bp5ycHD388MOaMGGC9u7dG1bus88+05IlS7RlyxYNHjy4x7bi4+OVnp4e2uLi4iIdPgAAMFREoebs2bOqq6tTSUnJVw3YbCopKZHP57tkfcuy5PV61dTUpMLCwtD+YDCoe++9V8uXL9dNN93Ua/1169Zp+PDhuvXWW/XUU0/p/PnzvZYNBALy+/1hGwAAMFdEj5/a29vV3d2ttLS0sP1paWk6dOhQr/U6OjqUmZmpQCCg6OhoPffcc5o2bVro+E9/+lMNGjRIS5cu7bWNpUuXatKkSUpOTta+fftUWVmpY8eOaf369T2W93g8WrNmTSSnBwAABrCI19Rcjvj4eNXX16uzs1Ner1dut1tjxoxRUVGR6urq9PTTT+vgwYOKiorqtQ232x3694QJE2S323X//ffL4/EoJibmovKVlZVhdfx+v7Kysr7dEwMAAFeNiEJNSkqKoqOj1dbWFra/ra1N6enpvdaz2WzKzc2VJDkcDjU2Nsrj8aioqEjvvvuujh8/rpEjR4bKd3d369FHH1VVVZU++eSTHtt0uVw6f/68PvnkE40dO/ai4zExMT2GHQAAYKaI1tTY7Xbl5+fL6/WG9gWDQXm9XhUUFPS5nWAwqEAgIEm699571dDQoPr6+tCWkZGh5cuXa9euXb22UV9fL5vN1uMvrgAAwLUn4sdPbrdb8+bNk9Pp1JQpU1RVVaWuri5VVFRIkubOnavMzEx5PB5JF9a2OJ1O5eTkKBAIaOfOnaqpqdGmTZskScOHD9fw4cPD+hg8eLDS09NDd2B8Pp/279+vqVOnKj4+Xj6fT8uWLdM999yjYcOGfaMJAAAAZog41JSXl+vEiRNatWqVWltb5XA4VFtbG1o83NzcLJvtqxtAXV1dWrRokVpaWjRkyBDl5eVp8+bNKi8v73OfMTEx2rp1q5544gkFAgGNHj1ay5YtC1szAwAArm1RlmVZ/T2IK8Hv9ysxMVEdHR1KSEjo7+EAAIA+iOT7m7/9BAAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARLivUbNy4UdnZ2YqNjZXL5dKBAwd6Lbt9+3Y5nU4lJSUpLi5ODodDNTU1vZZ/4IEHFBUVpaqqqrD9p06d0pw5c5SQkKCkpCTNnz9fnZ2dlzN8AABgoIhDzbZt2+R2u7V69WodPHhQEydOVGlpqY4fP95j+eTkZK1cuVI+n08NDQ2qqKhQRUWFdu3adVHZl19+We+9954yMjIuOjZnzhx99NFHeuONN7Rjxw698847WrhwYaTDBwAAhoqyLMuKpILL5dLkyZO1YcMGSVIwGFRWVpaWLFmiFStW9KmNSZMmafr06Vq7dm1o32effSaXy6Vdu3Zp+vTpeuSRR/TII49IkhobGzV+/Hi9//77cjqdkqTa2lrdeeedamlp6TEE/Tm/36/ExER1dHQoISEhklMGAAD9JJLv74ju1Jw9e1Z1dXUqKSn5qgGbTSUlJfL5fJesb1mWvF6vmpqaVFhYGNofDAZ17733avny5brpppsuqufz+ZSUlBQKNJJUUlIim82m/fv399hXIBCQ3+8P2wAAgLkiCjXt7e3q7u5WWlpa2P60tDS1trb2Wq+jo0NDhw6V3W7X9OnT9eyzz2ratGmh4z/96U81aNAgLV26tMf6ra2tSk1NDds3aNAgJScn99qvx+NRYmJiaMvKyurraQIAgAFo0JXoJD4+XvX19ers7JTX65Xb7daYMWNUVFSkuro6Pf300zp48KCioqK+tT4rKyvldrtDn/1+P8EGAACDRRRqUlJSFB0drba2trD9bW1tSk9P77WezWZTbm6uJMnhcKixsVEej0dFRUV69913dfz4cY0cOTJUvru7W48++qiqqqr0ySefKD09/aKFyOfPn9epU6d67TcmJkYxMTGRnB4AABjAInr8ZLfblZ+fL6/XG9oXDAbl9XpVUFDQ53aCwaACgYAk6d5771VDQ4Pq6+tDW0ZGhpYvXx76hVRBQYFOnz6turq6UBtvvfWWgsGgXC5XJKcAAAAMFfHjJ7fbrXnz5snpdGrKlCmqqqpSV1eXKioqJElz585VZmamPB6PpAtrW5xOp3JychQIBLRz507V1NRo06ZNkqThw4dr+PDhYX0MHjxY6enpGjt2rCRp3LhxKisr04IFC1RdXa1z585p8eLFmj17dp9++QQAAMwXcagpLy/XiRMntGrVKrW2tsrhcKi2tja0eLi5uVk221c3gLq6urRo0SK1tLRoyJAhysvL0+bNm1VeXh5Rv1u2bNHixYtVXFwsm82mWbNm6Zlnnol0+AAAwFARv6dmoOI9NQAADDzf2XtqAAAArlaEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABiBUAMAAIxAqAEAAEYg1AAAACMQagAAgBEINQAAwAiEGgAAYARCDQAAMAKhBgAAGIFQAwAAjECoAQAARiDUAAAAIxBqAACAEQg1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMcFmhZuPGjcrOzlZsbKxcLpcOHDjQa9nt27fL6XQqKSlJcXFxcjgcqqmpCSvzxBNPKC8vT3FxcRo2bJhKSkq0f//+sDLZ2dmKiooK29atW3c5wwcAAAaKONRs27ZNbrdbq1ev1sGDBzVx4kSVlpbq+PHjPZZPTk7WypUr5fP51NDQoIqKClVUVGjXrl2hMjfeeKM2bNig3/72t9q7d6+ys7N1++2368SJE2FtPfnkkzp27FhoW7JkSaTDBwAAhoqyLMuKpILL5dLkyZO1YcMGSVIwGFRWVpaWLFmiFStW9KmNSZMmafr06Vq7dm2Px/1+vxITE/Xmm2+quLhY0oU7NY888ogeeeSRSIZ7UZsdHR1KSEi4rDYAAMCVFcn3d0R3as6ePau6ujqVlJR81YDNppKSEvl8vkvWtyxLXq9XTU1NKiws7LWPf/u3f1NiYqImTpwYdmzdunUaPny4br31Vj311FM6f/58r30FAgH5/f6wDQAAmGtQJIXb29vV3d2ttLS0sP1paWk6dOhQr/U6OjqUmZmpQCCg6OhoPffcc5o2bVpYmR07dmj27Nn6/PPPNWLECL3xxhtKSUkJHV+6dKkmTZqk5ORk7du3T5WVlTp27JjWr1/fY58ej0dr1qyJ5PQAAMAAFlGouVzx8fGqr69XZ2envF6v3G63xowZo6KiolCZqVOnqr6+Xu3t7fr3f/93/d3f/Z3279+v1NRUSZLb7Q6VnTBhgux2u+6//355PB7FxMRc1GdlZWVYHb/fr6ysrO/uJAEAQL+KKNSkpKQoOjpabW1tYfvb2tqUnp7eaz2bzabc3FxJksPhUGNjozweT1ioiYuLU25urnJzc/X9739f3/ve9/Tzn/9clZWVPbbpcrl0/vx5ffLJJxo7duxFx2NiYnoMOwAAwEwRramx2+3Kz8+X1+sN7QsGg/J6vSooKOhzO8FgUIFA4BuVqa+vl81mC93JAQAA17aIHz+53W7NmzdPTqdTU6ZMUVVVlbq6ulRRUSFJmjt3rjIzM+XxeCRdWNvidDqVk5OjQCCgnTt3qqamRps2bZIkdXV16Sc/+YnuuusujRgxQu3t7dq4caM+++wz/e3f/q0kyefzaf/+/Zo6dari4+Pl8/m0bNky3XPPPRo2bNi3NRcAAGAAizjUlJeX68SJE1q1apVaW1vlcDhUW1sbWjzc3Nwsm+2rG0BdXV1atGiRWlpaNGTIEOXl5Wnz5s0qLy+XJEVHR+vQoUN64YUX1N7eruHDh2vy5Ml69913ddNNN0m68Chp69ateuKJJxQIBDR69GgtW7YsbM0MAAC4tkX8npqBivfUAAAw8ETy/X1Ffv10Nfgyu/G+GgAABo4vv7f7cg/mmgk1Z86ckSR+1g0AwAB05swZJSYmfm2Za+bxUzAY1NGjRxUfH6+oqKhvte0v34Hz6aef8mjr/4d56R1z0zPmpXfMTc+Yl56ZNC+WZenMmTPKyMgIW7Pbk2vmTo3NZtMNN9zwnfaRkJAw4C+e7wLz0jvmpmfMS++Ym54xLz0zZV4udYfmSxH/lW4AAICrEaEGAAAYgVDzLYiJidHq1av5swx/hnnpHXPTM+ald8xNz5iXnl2r83LNLBQGAABm404NAAAwAqEGAAAYgVADAACMQKgBAABGINR8Qxs3blR2drZiY2Plcrl04MCB/h5Sv3viiScUFRUVtuXl5fX3sK64d955R3/913+tjIwMRUVF6ZVXXgk7blmWVq1apREjRmjIkCEqKSnRH/7wh/4Z7BV2qbm57777LrqGysrK+mewV5DH49HkyZMVHx+v1NRU/fCHP1RTU1NYmS+++EIPPfSQhg8frqFDh2rWrFlqa2vrpxFfGX2Zl6KioouumQceeKCfRnzlbNq0SRMmTAi9ZK+goED//d//HTp+rV0vhJpvYNu2bXK73Vq9erUOHjyoiRMnqrS0VMePH+/vofW7m266SceOHQtte/fu7e8hXXFdXV2aOHGiNm7c2OPxn/3sZ3rmmWdUXV2t/fv3Ky4uTqWlpfriiy+u8EivvEvNjSSVlZWFXUMvvvjiFRxh/9izZ48eeughvffee3rjjTd07tw53X777erq6gqVWbZsmV599VW99NJL2rNnj44ePaqZM2f246i/e32ZF0lasGBB2DXzs5/9rJ9GfOXccMMNWrdunerq6vTBBx/or/7qr3T33Xfro48+knQNXi8WLtuUKVOshx56KPS5u7vbysjIsDweTz+Oqv+tXr3amjhxYn8P46oiyXr55ZdDn4PBoJWenm499dRToX2nT5+2YmJirBdffLEfRth//nxuLMuy5s2bZ9199939Mp6ryfHjxy1J1p49eyzLunCNDB482HrppZdCZRobGy1Jls/n669hXnF/Pi+WZVk/+MEPrIcffrj/BnUVGTZsmPUf//Ef1+T1wp2ay3T27FnV1dWppKQktM9ms6mkpEQ+n68fR3Z1+MMf/qCMjAyNGTNGc+bMUXNzc38P6ary8ccfq7W1Nez6SUxMlMvl4vr5/+zevVupqakaO3asHnzwQZ08ebK/h3TFdXR0SJKSk5MlSXV1dTp37lzYdZOXl6eRI0deU9fNn8/Ll7Zs2aKUlBTdfPPNqqys1Oeff94fw+s33d3d2rp1q7q6ulRQUHBNXi/XzB+0/La1t7eru7tbaWlpYfvT0tJ06NChfhrV1cHlcumXv/ylxo4dq2PHjmnNmjX6y7/8S/3ud79TfHx8fw/vqtDa2ipJPV4/Xx67lpWVlWnmzJkaPXq0jhw5oscff1x33HGHfD6foqOj+3t4V0QwGNQjjzyiv/iLv9DNN98s6cJ1Y7fblZSUFFb2WrpuepoXSfrRj36kUaNGKSMjQw0NDfqnf/onNTU1afv27f042ivjt7/9rQoKCvTFF19o6NChevnllzV+/HjV19dfc9cLoQbfujvuuCP07wkTJsjlcmnUqFH6z//8T82fP78fR4aBYvbs2aF/33LLLZowYYJycnK0e/duFRcX9+PIrpyHHnpIv/vd767J9Whfp7d5WbhwYejft9xyi0aMGKHi4mIdOXJEOTk5V3qYV9TYsWNVX1+vjo4O/dd//ZfmzZunPXv29Pew+gWPny5TSkqKoqOjL1pF3tbWpvT09H4a1dUpKSlJN954ow4fPtzfQ7lqfHmNcP30zZgxY5SSknLNXEOLFy/Wjh079Pbbb+uGG24I7U9PT9fZs2d1+vTpsPLXynXT27z0xOVySdI1cc3Y7Xbl5uYqPz9fHo9HEydO1NNPP31NXi+Emstkt9uVn58vr9cb2hcMBuX1elVQUNCPI7v6dHZ26siRIxoxYkR/D+WqMXr0aKWnp4ddP36/X/v37+f66UFLS4tOnjxp/DVkWZYWL16sl19+WW+99ZZGjx4ddjw/P1+DBw8Ou26amprU3Nxs9HVzqXnpSX19vSQZf830JBgMKhAIXJvXS3+vVB7Itm7dasXExFi//OUvrf/93/+1Fi5caCUlJVmtra39PbR+9eijj1q7d++2Pv74Y+t//ud/rJKSEislJcU6fvx4fw/tijpz5oz14YcfWh9++KElyVq/fr314YcfWv/3f/9nWZZlrVu3zkpKSrJ+85vfWA0NDdbdd99tjR492vrTn/7UzyP/7n3d3Jw5c8b6x3/8R8vn81kff/yx9eabb1qTJk2yvve971lffPFFfw/9O/Xggw9aiYmJ1u7du61jx46Fts8//zxU5oEHHrBGjhxpvfXWW9YHH3xgFRQUWAUFBf046u/epebl8OHD1pNPPml98MEH1scff2z95je/scaMGWMVFhb288i/eytWrLD27Nljffzxx1ZDQ4O1YsUKKyoqynr99dcty7r2rhdCzTf07LPPWiNHjrTsdrs1ZcoU67333uvvIfW78vJya8SIEZbdbrcyMzOt8vJy6/Dhw/09rCvu7bfftiRdtM2bN8+yrAs/6/7xj39spaWlWTExMVZxcbHV1NTUv4O+Qr5ubj7//HPr9ttvt66//npr8ODB1qhRo6wFCxZcE/9Z6GlOJFm/+MUvQmX+9Kc/WYsWLbKGDRtmXXfdddaMGTOsY8eO9d+gr4BLzUtzc7NVWFhoJScnWzExMVZubq61fPlyq6Ojo38HfgX8wz/8gzVq1CjLbrdb119/vVVcXBwKNJZ17V0vUZZlWVfuvhAAAMB3gzU1AADACIQaAABgBEINAAAwAqEGAAAYgVADAACMQKgBAABGINQAAAAjEGoAAIARCDUAAMAIhBoAAGAEQg0AADACoQYAABjh/wEjh930hRAXSgAAAABJRU5ErkJggg==\n"
},
"metadata": {}
}
]
},
{
"cell_type": "code",
"source": [
"phif"
],
"metadata": {
"id": "sYZcibSW7mbF",
"outputId": "e33bacc9-e926-40b2-f338-841e2496e38f",
"colab": {
"base_uri": "https://localhost:8080/"
}
},
"execution_count": 26,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"array([0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339])"
]
},
"metadata": {},
"execution_count": 26
}
]
},
{
"cell_type": "code",
"source": [
"1 - Nt // 2, Nt // 2"
],
"metadata": {
"id": "7BP4UO8O5qh-",
"outputId": "b3dba92b-5693-449e-bf2f-58bad8779e6c",
"colab": {
"base_uri": "https://localhost:8080/"
}
},
"execution_count": 22,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"(-31, 32)"
]
},
"metadata": {},
"execution_count": 22
}
]
},
{
"cell_type": "code",
"source": [],
"metadata": {
"id": "yoeuSRb46J4_"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"phif"
],
"metadata": {
"id": "5wv3FDxY7d8n",
"outputId": "2e66e3b1-7f09-438c-a1f6-50c25d72a289",
"colab": {
"base_uri": "https://localhost:8080/"
}
},
"execution_count": 24,
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"array([0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339, 0.35355339, 0.35355339,\n",
" 0.35355339, 0.35355339, 0.35355339])"
]
},
"metadata": {},
"execution_count": 24
}
]
},
{
"cell_type": "markdown",
"source": [
"## Timing comparisons"
],
"metadata": {
"id": "v8MIBbNHmndU"
}
},
{
"cell_type": "code",
"source": [
"%timeit transform_wavelet_freq_helper_numba(yf, Nf, Nt, phif)"
],
"metadata": {
"id": "yT8KZlsZjFZq",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "4d79e0d4-9276-44d6-aa1f-6bffddc0c690"
},
"execution_count": 13,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"872 µs ± 14.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"%timeit transform_wavelet_freq_helper_JAX(jax_yf, Nf, Nt, jax_phif)"
],
"metadata": {
"id": "X3ObJ_BrjIHC"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"%timeit transform_wavelet_freq_helper_CuPy(cupy_yf, Nf, Nt, cupy_phif)"
],
"metadata": {
"id": "UdHonCg3OaT9"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import time\n",
"from tqdm.auto import tqdm\n",
"import pandas as pd\n",
"import numpy as np\n",
"import jax.numpy as jnp\n",
"\n",
"\n",
"def time_function(func, *args):\n",
" try:\n",
" t0 = time.time()\n",
" func(*args)\n",
" return time.time() - t0\n",
" except Exception:\n",
" return np.nan\n",
"\n",
"\n",
"def run_numba_timings(Nf=2**6, Nt=None, n_rep=10):\n",
" \"\"\"Run timing for Numba implementation.\"\"\"\n",
" Nt = Nt or Nf\n",
" yf, phif = simulate_data(Nf)\n",
" # Warmup\n",
" transform_wavelet_freq_helper_numba(yf, Nf, Nt, phif)\n",
"\n",
" times = np.zeros(n_rep)\n",
" for i in range(n_rep):\n",
" times[i] = time_function(transform_wavelet_freq_helper_numba, yf, Nf, Nt, phif)\n",
" return times\n",
"\n",
"def run_jax_timings(Nf=2**6, Nt=None, n_rep=10):\n",
" \"\"\"Run timing for JAX implementation.\"\"\"\n",
" Nt = Nt or Nf\n",
" yf, phif = simulate_data(Nf)\n",
" jax_yf, jax_phif = jnp.array(yf), jnp.array(phif)\n",
" # Warmup\n",
" transform_wavelet_freq_helper_JAX(jax_yf, Nf, Nt, jax_phif)\n",
"\n",
" times = np.zeros(n_rep)\n",
" for i in range(n_rep):\n",
" times[i] = time_function(transform_wavelet_freq_helper_JAX, jax_yf, Nf, Nt, jax_phif)\n",
" return times\n",
"\n",
"def run_cupy_timings(Nf=2**6, Nt=None, n_rep=10):\n",
" \"\"\"Run timing for CuPy implementation.\"\"\"\n",
" Nt = Nt or Nf\n",
"\n",
" if (Nf >= 8192 # Cupy seems to run out of mem here.\n",
" or not cupy_available):\n",
" return np.array([np.nan] * n_rep)\n",
"\n",
" yf, phif = simulate_data(Nf)\n",
" cupy_yf, cupy_phif = cp.array(yf), cp.array(phif)\n",
" # Warmup\n",
" transform_wavelet_freq_helper_CuPy(cupy_yf, Nf, Nt, cupy_phif)\n",
"\n",
" times = np.zeros(n_rep)\n",
" for i in range(n_rep):\n",
" times[i] = time_function(transform_wavelet_freq_helper_CuPy, cupy_yf, Nf, Nt, cupy_phif)\n",
" return times\n",
"\n",
"def run_experiments(min_n=5, max_n=10, n_rep=10):\n",
" \"\"\"\n",
" Run experiments for different sizes and collect runtime data.\n",
"\n",
" Parameters:\n",
" - min_n: Minimum power of 2 for the size of Nf.\n",
" - max_n: Maximum power of 2 for the size of Nf.\n",
" - n_rep: Number of repetitions for each size.\n",
"\n",
" Returns:\n",
" - pd.DataFrame: DataFrame containing runtimes for different methods and sizes.\n",
" \"\"\"\n",
" jax_label = f\"jax[{DEVICE}]_times\"\n",
" data = {\n",
" \"Ns\": [],\n",
" \"numba_times\": [],\n",
" jax_label: [],\n",
" \"cupy_times\": []\n",
" }\n",
"\n",
" Nfs = [2**i for i in range(min_n, max_n)]\n",
"\n",
" pbar = tqdm(Nfs, desc=\"Running experiments\")\n",
" for Nf in pbar:\n",
" pbar.set_postfix({\"Nf\": Nf})\n",
" numba_times = run_numba_timings(Nf=Nf, n_rep=n_rep)\n",
" jax_times = run_jax_timings(Nf=Nf, n_rep=n_rep)\n",
" cupy_times = run_cupy_timings(Nf=Nf, n_rep=n_rep)\n",
"\n",
" data[\"Ns\"].extend([Nf] * n_rep)\n",
" data[\"numba_times\"].extend(numba_times)\n",
" data[jax_label].extend(jax_times)\n",
" data[\"cupy_times\"].extend(cupy_times)\n",
"\n",
" return pd.DataFrame(data)\n",
"\n",
"\n",
"\n",
"# Running the experiments and collecting data\n",
"data = run_experiments(max_n=14, n_rep=5)\n",
"data.to_csv(\"runtimes.csv\", index=False)"
],
"metadata": {
"id": "rXA0cHe5gahJ"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import numpy as np\n",
"import os\n",
"import matplotlib.pyplot as plt\n",
"import pandas as pd\n",
"\n",
"def plot_runtime_vs_size(data, column_label='numba', color='blue', ax=None):\n",
" \"\"\"\n",
" Plots the runtime vs size graph with error bands for a given column in the dataset.\n",
"\n",
" Parameters:\n",
" - data: DataFrame containing the data to be plotted.\n",
" - column_label: Label for the column prefix (e.g., 'numba', 'jax', 'cupy') to plot.\n",
" - color: Color for the plot lines and fill area (default: 'blue').\n",
" - ax: Optional matplotlib axis to plot on. If None, a new figure and axis are created.\n",
"\n",
" Returns:\n",
" - ax: The axis with the plot.\n",
" \"\"\"\n",
"\n",
" df = data.copy()\n",
" grouped = df.groupby('Ns').agg(\n",
" mean=(f'{column_label}_times', 'mean'),\n",
" sem=(f'{column_label}_times', 'sem')\n",
" ).reset_index()\n",
" grouped['exponent'] = np.log2(grouped['Ns']).astype(int)\n",
"\n",
" # Create an axis if not provided\n",
" if ax is None:\n",
" fig, ax = plt.subplots(figsize=(4, 3))\n",
"\n",
"\n",
" # Plotting\n",
" ax.fill_between(grouped['exponent'],\n",
" grouped['mean'] - grouped['sem'],\n",
" grouped['mean'] + grouped['sem'],\n",
" color=color, alpha=0.3, label=f'{column_label.capitalize()}', lw=0)\n",
" ax.plot(grouped['exponent'], grouped['mean'], color=color, marker='')\n",
"\n",
" # Set labels, legend\n",
" ax.set_xlabel('$N_fxN_t$', fontsize=12)\n",
" ax.set_ylabel('Runtime [s]', fontsize=12)\n",
" ax.legend(fontsize=10, frameon=False)\n",
"\n",
" ax.grid(True, which=\"both\", linestyle=\"-\", alpha=0.2)\n",
" fmt = '$2^{e}x2^{e}$'\n",
" ax.set_xticks(grouped['exponent'])\n",
" ax.set_xticklabels([fmt.replace(\"e\", str(exp)) for exp in grouped['exponent']])\n",
" ax.tick_params(axis='x', rotation=45)\n",
" ax.set_yscale('log')\n",
" plt.tight_layout()\n",
" return ax\n",
"\n",
"if 'runtimes_gpu.csv' in os.listdir():\n",
" df_gpu = pd.read_csv(\"runtimes_gpu.csv\")\n",
" df_tpu = pd.read_csv(\"runtimes_tpu.csv\")\n",
" df = df_gpu.copy()\n",
" df['jax[TPU v2]_times'] = df_tpu['jax[TPU v2]_times']\n",
"else:\n",
" df = pd.read_csv(\"runtimes.csv\")\n",
"fig, ax = plt.subplots(figsize=(4, 3))\n",
"labels = [i.split(\"_times\")[0] for i in df.columns.values if 'times' in i]\n",
"for i, l in enumerate(labels):\n",
" plot_runtime_vs_size(df, column_label=l, color=f\"C{i}\", ax=ax)"
],
"metadata": {
"id": "bNoJIMuHFftq"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"Note: we've used lower precision data for JAX + Cupy (this can be changed)."
],
"metadata": {
"id": "m5UPtRxf32cz"
}
}
]
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment