From 3cfe70c6557d547afc234c00956755e27bf2c07f Mon Sep 17 00:00:00 2001 From: thomaspmurphy Date: Mon, 15 Jun 2026 14:41:41 +0100 Subject: [PATCH 1/6] feat: implement Chirp Z-Transform (CZT) and zoom_fft Adds `NxSignal.czt/2` (Bluestein's algorithm, O((N+M)log(N+M))) and `NxSignal.zoom_fft/4` for evaluating the z-transform on an arbitrary spiral contour or zoomed frequency band. Includes tests and a Livebook guide. --- guides/chirp_z.livemd | 307 +++++++++++++++++++++++++++++++++ lib/nx_signal.ex | 151 ++++++++++++++++ test/nx_signal_test.exs | 81 +++++++++ test/support/nx_signal_case.ex | 3 +- 4 files changed, 541 insertions(+), 1 deletion(-) create mode 100644 guides/chirp_z.livemd diff --git a/guides/chirp_z.livemd b/guides/chirp_z.livemd new file mode 100644 index 0000000..6be74c3 --- /dev/null +++ b/guides/chirp_z.livemd @@ -0,0 +1,307 @@ +# Chirp Z-Transform and Zoom FFT + +```elixir +Mix.install([ + {:nx_signal, path: __DIR__ |> Path.join("..") |> Path.expand()}, + {:kino_vega_lite, "~> 0.1"}, + {:tucan, "~> 0.5"} +]) +``` + +## What is the Chirp Z-Transform? + +The standard Discrete Fourier Transform (DFT) evaluates the z-transform of a +signal at $N$ equally-spaced points on the **unit circle** in the complex plane: + +$$ +z_k = e^{j2\pi k/N}, \quad k = 0, \ldots, N-1 +$$ + +The **Chirp Z-Transform (CZT)** generalises this. Instead of the unit circle, +it evaluates the z-transform on an arbitrary spiral contour defined by a +starting point $A$ and a step ratio $W$: + +$$ +z_k = A \cdot W^{-k}, \quad k = 0, \ldots, M-1 +$$ + +$$ +X[k] = \sum_{n=0}^{N-1} x[n] \cdot A^{-n} \cdot W^{nk} +$$ + +This unlocks two powerful capabilities: + +* **$M \neq N$**: compute any number of output points from any length input. +* **Zoom FFT**: concentrate all $M$ output bins over a narrow frequency band, + giving arbitrarily fine frequency resolution without increasing the signal + length. + +The CZT is computed efficiently using Bluestein's algorithm, which rewrites +the sum as a convolution and evaluates it via FFT in +$O((N + M) \log(N + M))$ time — the same asymptotic cost as the FFT. + +## CZT as a generalised DFT + +With default parameters ($A = 1$, $W = e^{-j2\pi/N}$, $M = N$), +`NxSignal.czt/2` is identical to `Nx.fft/1`. + +```elixir +fs = 200 +duration = 1.0 +n = trunc(fs * duration) + +t = Nx.linspace(0, duration, n: n, endpoint: false, type: {:f, 32}) + +# Signal: two tones at 10 Hz and 35 Hz +x = + Nx.add( + Nx.sin(Nx.multiply(t, 2 * :math.pi() * 10)), + Nx.multiply(0.5, Nx.sin(Nx.multiply(t, 2 * :math.pi() * 35))) + ) + +fft_result = Nx.fft(Nx.as_type(x, {:c, 64})) +czt_result = NxSignal.czt(x) + +max_diff = + Nx.subtract(Nx.abs(czt_result), Nx.abs(fft_result)) + |> Nx.abs() + |> Nx.reduce_max() + +IO.inspect(Nx.to_number(max_diff), label: "max |CZT| - |FFT| difference") +``` + +The difference is at the level of floating-point rounding error, confirming +the two are mathematically equivalent in the default case. + +The signal itself, 50 ms of the two tones, shows their superposition clearly: + +```elixir +n_show = trunc(fs * 0.05) + +time_data = + Enum.zip(Nx.to_flat_list(t[0..(n_show - 1)]), Nx.to_flat_list(x[0..(n_show - 1)])) + |> Enum.map(fn {ti, xi} -> %{time: ti, amplitude: xi} end) + +Tucan.lineplot(time_data, "time", "amplitude") +|> Tucan.Axes.set_x_title("Time (s)") +|> Tucan.Axes.set_y_title("Amplitude") +|> Tucan.set_title("Signal 10 Hz + 0.5 × 35 Hz (first 50 ms)") +|> Tucan.set_width(680) +|> Tucan.set_height(180) +``` + +Now let's look at the full spectrum: + +```elixir +half = div(n, 2) +freqs = NxSignal.fft_frequencies(fs, fft_length: n)[0..half] |> Nx.to_flat_list() +amps = Nx.abs(fft_result)[0..half] |> Nx.to_flat_list() + +spectrum_data = + Enum.zip(freqs, amps) + |> Enum.map(fn {f, a} -> %{frequency: f, amplitude: a} end) + +Tucan.bar(spectrum_data, "frequency", "amplitude", tooltip: true) +|> VegaLite.encode_field(:x, "frequency", type: :quantitative, title: "Frequency (Hz)") +|> Tucan.Axes.set_y_title("Amplitude") +|> Tucan.set_title("DFT spectrum: 10 Hz and 35 Hz tones") +|> Tucan.set_width(680) +|> Tucan.set_height(260) +``` + +## Zoom FFT: resolving closely-spaced tones + +This is where the CZT really shines. + +Suppose we have two tones at 49 Hz and 51 Hz, sampled at 200 Hz for 1 second. +With a 200-point DFT the bin spacing is $200/200 = 1\,\text{Hz}$, so in +principle we have just enough resolution to separate them. In practice, spectral +leakage smears the peaks and makes the separation unclear. + +With `zoom_fft`, we can point all $M$ output bins at the narrow band +$[45, 55]\,\text{Hz}$, giving a bin spacing of just $10/M\,\text{Hz}$ regardless of the signal length. + +```elixir +fs = 200 +duration = 1.0 +n = trunc(fs * duration) + +t = Nx.linspace(0, duration, n: n, endpoint: false, type: {:f, 64}) + +f_a = 49.0 +f_b = 51.0 + +x2 = + Nx.add( + Nx.cos(Nx.multiply(t, 2 * :math.pi() * f_a)), + Nx.cos(Nx.multiply(t, 2 * :math.pi() * f_b)) + ) +``` + +The standard DFT and the Zoom FFT side by side (the resolution gain is +immediately visible): + +```elixir +# Standard DFT — 40 to 60 Hz region +fft_result2 = Nx.fft(x2) +half2 = div(n, 2) +freqs_full = NxSignal.fft_frequencies(fs, fft_length: n, type: {:f, 64})[0..half2] +amps_full = Nx.abs(fft_result2)[0..half2] + +dft_data = + Enum.zip(Nx.to_flat_list(freqs_full), Nx.to_flat_list(amps_full)) + |> Enum.filter(fn {f, _} -> f >= 40 and f <= 60 end) + |> Enum.map(fn {f, a} -> %{frequency: f, amplitude: a, method: "Standard DFT (1 Hz bins)"} end) + +# Zoom FFT — 512 bins over 45–55 Hz +m = 512 +f1_norm = 45.0 / fs +f2_norm = 55.0 / fs + +zoom_result = NxSignal.zoom_fft(x2, f1_norm, f2_norm, m: m) +zoom_amps = Nx.abs(zoom_result) + +zoom_data = + Nx.linspace(45.0, 55.0, n: m, endpoint: false, type: {:f, 64}) + |> Nx.to_flat_list() + |> Enum.zip(Nx.to_flat_list(zoom_amps)) + |> Enum.map(fn {f, a} -> %{frequency: f, amplitude: a, method: "Zoom FFT (#{10.0 / m |> Float.round(3)} Hz bins)"} end) + +comparison_data = dft_data ++ zoom_data + +Tucan.lineplot(comparison_data, "frequency", "amplitude") +|> Tucan.color_by("method") +|> Tucan.Axes.set_x_title("Frequency (Hz)") +|> Tucan.Axes.set_y_title("Amplitude") +|> Tucan.set_title("DFT vs Zoom FFT — resolving 49 Hz and 51 Hz tones") +|> Tucan.set_width(680) +|> Tucan.set_height(260) +``` + +Let's confirm by finding the top-2 frequency bins in the zoom output: + +```elixir +{_vals, indices} = Nx.top_k(zoom_amps, k: 2) + +peak_freqs = + indices + |> Nx.to_flat_list() + |> Enum.map(fn i -> 45.0 + i * (10.0 / m) end) + |> Enum.sort() + +IO.inspect(peak_freqs, label: "Peak frequencies (Hz)") +``` + +## The contour in the z-plane + +To understand *why* the Zoom FFT has higher resolution, it helps to see what +the two transforms are actually doing in the complex plane. + +The DFT evaluates the z-transform at $N$ equally-spaced points around the +**full** unit circle. The Zoom FFT evaluates at $M$ points along a **short arc** +of the unit circle — the arc corresponding to the frequency band $[f_1, f_2]$. + +With the same $M$ and $N$, the arc points are spaced $\frac{f_2 - f_1}{M}$ +apart instead of $\frac{f_s}{N}$ apart, giving $\frac{(f_2-f_1)/M}{f_s/N}$ +times finer resolution. + +```elixir +n_dft = 64 # number of DFT evaluation points to display +m_zoom = 64 # same number of zoom points, to make spacing difference visible + +# DFT: N equally-spaced points on the full unit circle +dft_points = + Enum.map(0..(n_dft - 1), fn k -> + angle = 2 * :math.pi() * k / n_dft + %{re: :math.cos(angle), im: :math.sin(angle), transform: "DFT (full circle)"} + end) + +# Zoom FFT arc: M points on the arc from f1 to f2 +f1 = 45.0 +f2 = 55.0 +fs_plane = 200.0 + +zoom_points = + Enum.map(0..(m_zoom - 1), fn k -> + freq = f1 + k * (f2 - f1) / m_zoom + angle = 2 * :math.pi() * freq / fs_plane + %{re: :math.cos(angle), im: :math.sin(angle), transform: "Zoom FFT (45–55 Hz arc)"} + end) + +contour_data = dft_points ++ zoom_points + +Tucan.scatter(contour_data, "re", "im", point_size: 30) +|> Tucan.color_by("transform") +|> Tucan.Axes.set_x_title("Re(z)") +|> Tucan.Axes.set_y_title("Im(z)") +|> Tucan.set_title("Evaluation points in the z-plane") +|> Tucan.set_width(400) +|> Tucan.set_height(400) +``` + +The DFT points (grey) are spread thinly around the whole circle. The Zoom FFT +points (orange) are clustered on the small arc between 45 Hz and 55 Hz. We have the +same $M$ points covering $\frac{f_2-f_1}{f_s} = 5\%$ of the circle, giving +$20\times$ finer spacing in that region. + +## Evaluating the z-transform off the unit circle + +The CZT can evaluate the z-transform at any points in the complex plane, not +just the unit circle. For example, on a circle of radius $r > 1$ you can +study the pole-zero structure of a system without being constrained to the +unit circle where poles and zeros have the most influence. + +Here we evaluate the z-transform of a simple FIR filter at 200 points on +the unit circle and on a circle of radius 1.5: + +```elixir +# FIR low-pass filter: h = [1, 2, 3, 2, 1] (normalised) +h = Nx.tensor([1.0, 2.0, 3.0, 2.0, 1.0]) |> Nx.divide(9.0) + +m = 200 + +# Unit-circle response (standard DFT) +hz_unit = NxSignal.czt(h, m: m) + +# Outer circle: radius r = 1.5, starting at angle 0 +r = 1.5 +a = Nx.complex(r, 0.0) +w = Nx.complex(:math.cos(-2 * :math.pi() / m), :math.sin(-2 * :math.pi() / m)) + +hz_outer = NxSignal.czt(h, m: m, a: a, w: w) + +angles = Nx.linspace(0, 360, n: m, endpoint: false, type: {:f, 64}) |> Nx.to_flat_list() + +contours_data = + Enum.zip(angles, Nx.to_flat_list(Nx.abs(hz_unit))) + |> Enum.map(fn {a, amp} -> %{angle: a, amplitude: amp, contour: "|z| = 1 (unit circle)"} end) + |> Kernel.++( + Enum.zip(angles, Nx.to_flat_list(Nx.abs(hz_outer))) + |> Enum.map(fn {a, amp} -> %{angle: a, amplitude: amp, contour: "|z| = 1.5"} end) + ) + +Tucan.lineplot(contours_data, "angle", "amplitude", tooltip: true) +|> Tucan.color_by("contour") +|> Tucan.Axes.set_x_title("Angle (degrees)") +|> Tucan.Axes.set_y_title("|H(z)|") +|> Tucan.set_title("FIR frequency response on two contours") +|> Tucan.set_width(680) +|> Tucan.set_height(280) +``` + +The unit-circle response ($|z|=1$) is the standard frequency response. The +outer-circle response ($|z|=1.5$) shows the same low-pass shape but with +reduced dynamic range — consistent with moving away from the unit circle +where the filter's poles and zeros have most influence. + +## Summary + +| Function | What it does | +| ------------------------------------ | ------------------------------------------------------------------- | +| `NxSignal.czt(x)` | Generalised DFT on any spiral contour; defaults to the standard DFT | +| `NxSignal.czt(x, m: m, a: a, w: w)` | $M$-point evaluation at $z_k = A \, W^{-k}$ | +| `NxSignal.zoom_fft(x, f1, f2)` | High-resolution spectrum over normalised $[f_1, f_2]$ | +| `NxSignal.zoom_fft(x, f1, f2, m: m)` | Same with an explicit output bin count | + +Both functions return complex tensors; use `Nx.abs/1` for magnitude and +`Nx.phase/1` for phase. diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex index ec22ed7..a255ba8 100644 --- a/lib/nx_signal.ex +++ b/lib/nx_signal.ex @@ -733,4 +733,155 @@ defmodule NxSignal do |> Tuple.delete_at(idx) |> put_elem(idx, out_len) end + + @doc ~S""" + Chirp Z-Transform. + + Evaluates the z-transform of `x` on a spiral contour in the z-plane, + returning $M$ output points. Uses Bluestein's identity to express the + sum as a convolution, giving $O((N+M)\log(N+M))$ complexity. + + $$ + X[k] = \sum_{n=0}^{N-1} x[n] \, A^{-n} \, W^{nk}, \quad k = 0, 1, \ldots, M-1 + $$ + + When $A = 1$, $W = e^{-j2\pi/M}$, and $M = N$, the result is identical + to the standard DFT. + + See also: `zoom_fft/4` + + ## Options + + * `:m` - number of output points $M$. Defaults to `Nx.size(x)`. + * `:w` - contour ratio $W$ (complex scalar tensor). + Defaults to $e^{-j2\pi/M}$, which gives the DFT. + * `:a` - contour starting point $A$ (complex scalar tensor). + Defaults to $1+0j$. + + ## Examples + + iex> x = Nx.tensor([1.0, 0.0, 0.0, 0.0]) + iex> result = NxSignal.czt(x) + iex> expected = Nx.fft(Nx.as_type(x, {:c, 64})) + iex> Nx.all_close(Nx.abs(result), Nx.abs(expected), atol: 1.0e-5) == Nx.tensor(1, type: :u8) + true + + """ + @doc type: :transforms + deftransform czt(x, opts \\ []) do + opts = Keyword.validate!(opts, [:m, :w, :a]) + + {n} = Nx.shape(x) + m = opts[:m] || n + + fft_len = czt_next_pow2(n + m - 1) + + w = + opts[:w] || + Nx.complex(:math.cos(-2 * :math.pi() / m), :math.sin(-2 * :math.pi() / m)) + + a = opts[:a] || Nx.complex(1.0, 0.0) + + czt_n(x, w, a, m: m, n: n, fft_len: fft_len) + end + + deftransformp czt_next_pow2(n) when n <= 1, do: 1 + + deftransformp czt_next_pow2(n) do + Integer.pow(2, ceil(:math.log2(n))) + end + + defnp czt_n(x, w, a, opts) do + m = opts[:m] + n = opts[:n] + fft_len = opts[:fft_len] + + x = Nx.as_type(x, {:c, 64}) + w = Nx.as_type(w, {:c, 64}) + a = Nx.as_type(a, {:c, 64}) + + log_w = Nx.log(w) + log_a = Nx.log(a) + + # {:c, 64} has 32-bit components; use {:f, 32} so no silent promotion occurs + n_idx = Nx.iota({n}, type: {:f, 32}) + k_idx = Nx.iota({m}, type: {:f, 32}) + + # Pre-multiply: yn[n] = x[n] · a^{-n} · w^{n²/2} + yn = x * Nx.exp(-n_idx * log_a) * Nx.exp(n_idx * n_idx / 2 * log_w) + + # Zero-pad yn to FFT length + yn_padded = Nx.pad(yn, Nx.as_type(0, {:c, 64}), [{0, fft_len - n, 0}]) + + # Build chirp kernel h_padded of length fft_len arranged for circular convolution: + # positions 0..M-1: h[k] = w^{-k²/2} + # positions M..L-N: 0 + # positions L-N+1..L-1: h[-j] = w^{-j²/2} for j = N-1 downto 1 + i = Nx.iota({fft_len}, type: {:f, 32}) + pos_vals = Nx.exp(-i * i / 2 * log_w) + neg_vals = Nx.exp(-(fft_len - i) * (fft_len - i) / 2 * log_w) + + zero_c = Nx.as_type(0, {:c, 64}) + h_padded = Nx.select(i < m, pos_vals, Nx.select(i > fft_len - n, neg_vals, zero_c)) + + # FFT-based convolution + g = Nx.ifft(Nx.fft(yn_padded) * Nx.fft(h_padded)) + + # Post-multiply: X[k] = w^{k²/2} · g[k], k = 0…M-1 + Nx.slice(g, [0], [m]) * Nx.exp(k_idx * k_idx / 2 * log_w) + end + + @doc ~S""" + Zoom FFT: high-resolution DFT over a narrow frequency band. + + Computes $M$ output samples of the DFT concentrated in the normalised + frequency range $[f_1, f_2]$. This is a thin wrapper around `czt/2` that + sets: + + $$ + A = e^{j2\pi f_1}, \quad W = e^{-j2\pi(f_2 - f_1)/M} + $$ + + The output frequencies are $f_k = f_1 + k\,(f_2 - f_1)/M$ for $k = 0, \ldots, M-1$. + + See also: `czt/2` + + ## Arguments + + * `x` - input signal, shape `{n}`. + * `f1` - lower bound of the frequency range (normalised, 0 to 1). + * `f2` - upper bound of the frequency range (normalised, 0 to 1). + + ## Options + + * `:m` - number of output points $M$. Defaults to `Nx.size(x)`. + + ## Examples + + iex> x = Nx.tensor([1.0, 0.0, 0.0, 0.0]) + iex> full = NxSignal.zoom_fft(x, 0.0, 1.0) + iex> dft = Nx.fft(Nx.as_type(x, {:c, 64})) + iex> Nx.all_close(Nx.abs(full), Nx.abs(dft), atol: 1.0e-5) == Nx.tensor(1, type: :u8) + true + + """ + @doc type: :transforms + deftransform zoom_fft(x, f1, f2, opts \\ []) do + opts = Keyword.validate!(opts, [:m]) + + {n} = Nx.shape(x) + m = opts[:m] || n + + # A = e^{+j2πf1}: the CZT evaluates at A^{-n}, so the positive exponent here + # places the first evaluation point at frequency f1. + a = Nx.complex(:math.cos(2 * :math.pi() * f1), :math.sin(2 * :math.pi() * f1)) + + w = + Nx.complex( + :math.cos(-2 * :math.pi() * (f2 - f1) / m), + :math.sin(-2 * :math.pi() * (f2 - f1) / m) + ) + + czt(x, m: m, w: w, a: a) + end end diff --git a/test/nx_signal_test.exs b/test/nx_signal_test.exs index 0c3dd31..dda2843 100644 --- a/test/nx_signal_test.exs +++ b/test/nx_signal_test.exs @@ -1,4 +1,85 @@ defmodule NxSignalTest do use NxSignal.Case doctest NxSignal + + describe "czt/2" do + test "default params equal DFT" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0]) + result = NxSignal.czt(x) + expected = Nx.fft(Nx.as_type(x, {:c, 64})) + + assert_all_close(result, expected, atol: 1.0e-5) + end + + test "unit impulse returns all ones" do + x = Nx.tensor([1.0, 0.0, 0.0, 0.0]) + result = NxSignal.czt(x, m: 6) + ones = Nx.broadcast(Nx.complex(1.0, 0.0), {6}) + + assert_all_close(result, ones, atol: 1.0e-5) + end + + test "M > N is equivalent to zero-padded DFT" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0]) + result = NxSignal.czt(x, m: 8) + expected = Nx.fft(Nx.as_type(x, {:c, 64}), length: 8) + + assert_all_close(result, expected, atol: 1.0e-5) + end + + test "custom a evaluates z-transform off the unit circle" do + # x = [1, 1], so X(z) = 1 + z^{-1}. At z = 2: X(2) = 1 + 0.5 = 1.5 + x = Nx.tensor([1.0, 1.0]) + a = Nx.complex(2.0, 0.0) + w = Nx.complex(-1.0, 0.0) + result = NxSignal.czt(x, m: 1, a: a, w: w) + + assert_all_close(Nx.real(result[0]), Nx.tensor(1.5, type: {:f, 64}), atol: 1.0e-5) + end + + test "returns complex output" do + result = NxSignal.czt(Nx.tensor([1.0, 2.0, 3.0])) + + assert Nx.Type.complex?(Nx.type(result)) + end + end + + describe "zoom_fft/4" do + test "full range [0, 1] equals DFT" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0]) + result = NxSignal.zoom_fft(x, 0.0, 1.0) + expected = Nx.fft(Nx.as_type(x, {:c, 64})) + + assert_all_close(result, expected, atol: 1.0e-5) + end + + test "resolves two closely-spaced tones" do + # Two tones at normalised frequencies 0.49 and 0.51. + # A 16-point DFT cannot separate them; zoom_fft over [0.45, 0.55] can. + n = 128 + t = Nx.iota({n}, type: :f64) + pi2 = 2 * :math.pi() + + x = + Nx.add( + Nx.cos(Nx.multiply(t, pi2 * 0.49)), + Nx.cos(Nx.multiply(t, pi2 * 0.51)) + ) + + result = NxSignal.zoom_fft(x, 0.45, 0.55, m: 64) + magnitudes = Nx.abs(result) + + max_val = Nx.reduce_max(magnitudes) |> Nx.to_number() + median_val = Nx.median(magnitudes) |> Nx.to_number() + + assert max_val > 5 * median_val + end + + test "returns m output points" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]) + result = NxSignal.zoom_fft(x, 0.1, 0.4, m: 16) + + assert Nx.shape(result) == {16} + end + end end diff --git a/test/support/nx_signal_case.ex b/test/support/nx_signal_case.ex index 6b0f18d..cc0b1a1 100644 --- a/test/support/nx_signal_case.ex +++ b/test/support/nx_signal_case.ex @@ -21,7 +21,8 @@ defmodule NxSignal.Case do :windowing, :filters, :waveforms, - :peak_finding + :peak_finding, + :transforms ] def validate_doc_metadata(module) do From fc03e0e5772fdb38b8b48c9eaf777a1a5db25d71 Mon Sep 17 00:00:00 2001 From: thomaspmurphy Date: Mon, 15 Jun 2026 15:33:13 +0100 Subject: [PATCH 2/6] refactor: address review feedback on CZT Preserve input precision via Nx.Type.to_complex/to_real, add axis option with vectorize-based batching, fix doctest assertions, simplify a^-n term. --- lib/nx_signal.ex | 80 ++++++++++++++++++++++++++++------------- test/nx_signal_test.exs | 25 +++++++++++++ 2 files changed, 80 insertions(+), 25 deletions(-) diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex index a255ba8..de16e98 100644 --- a/lib/nx_signal.ex +++ b/lib/nx_signal.ex @@ -760,20 +760,23 @@ defmodule NxSignal do ## Examples - iex> x = Nx.tensor([1.0, 0.0, 0.0, 0.0]) - iex> result = NxSignal.czt(x) - iex> expected = Nx.fft(Nx.as_type(x, {:c, 64})) - iex> Nx.all_close(Nx.abs(result), Nx.abs(expected), atol: 1.0e-5) == Nx.tensor(1, type: :u8) - true + iex> NxSignal.czt(Nx.tensor([1.0, 0.0, 0.0, 0.0])) |> Nx.real() |> Nx.round() + #Nx.Tensor< + f32[4] + [1.0, 1.0, 1.0, 1.0] + > """ @doc type: :transforms deftransform czt(x, opts \\ []) do - opts = Keyword.validate!(opts, [:m, :w, :a]) + opts = Keyword.validate!(opts, [:m, :w, :a, axis: -1]) - {n} = Nx.shape(x) - m = opts[:m] || n + ndim = Nx.rank(x) + axis = opts[:axis] + axis = if axis < 0, do: ndim + axis, else: axis + n = elem(Nx.shape(x), axis) + m = opts[:m] || n fft_len = czt_next_pow2(n + m - 1) w = @@ -782,7 +785,34 @@ defmodule NxSignal do a = opts[:a] || Nx.complex(1.0, 0.0) - czt_n(x, w, a, m: m, n: n, fft_len: fft_len) + complex_type = x |> Nx.type() |> Nx.Type.to_complex() + real_type = Nx.Type.to_real(complex_type) + + other_axes = List.delete(Enum.to_list(0..(ndim - 1)), axis) + perm = other_axes ++ [axis] + + inv_perm = + perm + |> Enum.with_index() + |> Enum.sort_by(fn {v, _} -> v end) + |> Enum.map(fn {_, i} -> i end) + + x_t = Nx.transpose(x, axes: perm) + + result = + if other_axes == [] do + czt_n(x_t, w, a, m: m, n: n, fft_len: fft_len, complex_type: complex_type, real_type: real_type) + else + batch_names = Enum.map(0..(length(other_axes) - 1), fn i -> :"batch_#{i}" end) + x_v = Nx.vectorize(x_t, batch_names) + + result_v = + czt_n(x_v, w, a, m: m, n: n, fft_len: fft_len, complex_type: complex_type, real_type: real_type) + + Nx.devectorize(result_v, keep_names: false) + end + + Nx.transpose(result, axes: inv_perm) end deftransformp czt_next_pow2(n) when n <= 1, do: 1 @@ -795,33 +825,33 @@ defmodule NxSignal do m = opts[:m] n = opts[:n] fft_len = opts[:fft_len] + complex_type = opts[:complex_type] + real_type = opts[:real_type] - x = Nx.as_type(x, {:c, 64}) - w = Nx.as_type(w, {:c, 64}) - a = Nx.as_type(a, {:c, 64}) + x = Nx.as_type(x, complex_type) + w = Nx.as_type(w, complex_type) + a = Nx.as_type(a, complex_type) log_w = Nx.log(w) - log_a = Nx.log(a) - # {:c, 64} has 32-bit components; use {:f, 32} so no silent promotion occurs - n_idx = Nx.iota({n}, type: {:f, 32}) - k_idx = Nx.iota({m}, type: {:f, 32}) + n_idx = Nx.iota({n}, type: real_type) + k_idx = Nx.iota({m}, type: real_type) # Pre-multiply: yn[n] = x[n] · a^{-n} · w^{n²/2} - yn = x * Nx.exp(-n_idx * log_a) * Nx.exp(n_idx * n_idx / 2 * log_w) + yn = x * Nx.pow(a, -n_idx) * Nx.exp(n_idx * n_idx / 2 * log_w) # Zero-pad yn to FFT length - yn_padded = Nx.pad(yn, Nx.as_type(0, {:c, 64}), [{0, fft_len - n, 0}]) + yn_padded = Nx.pad(yn, Nx.as_type(0, complex_type), [{0, fft_len - n, 0}]) # Build chirp kernel h_padded of length fft_len arranged for circular convolution: # positions 0..M-1: h[k] = w^{-k²/2} # positions M..L-N: 0 # positions L-N+1..L-1: h[-j] = w^{-j²/2} for j = N-1 downto 1 - i = Nx.iota({fft_len}, type: {:f, 32}) + i = Nx.iota({fft_len}, type: real_type) pos_vals = Nx.exp(-i * i / 2 * log_w) neg_vals = Nx.exp(-(fft_len - i) * (fft_len - i) / 2 * log_w) - zero_c = Nx.as_type(0, {:c, 64}) + zero_c = Nx.as_type(0, complex_type) h_padded = Nx.select(i < m, pos_vals, Nx.select(i > fft_len - n, neg_vals, zero_c)) # FFT-based convolution @@ -858,11 +888,11 @@ defmodule NxSignal do ## Examples - iex> x = Nx.tensor([1.0, 0.0, 0.0, 0.0]) - iex> full = NxSignal.zoom_fft(x, 0.0, 1.0) - iex> dft = Nx.fft(Nx.as_type(x, {:c, 64})) - iex> Nx.all_close(Nx.abs(full), Nx.abs(dft), atol: 1.0e-5) == Nx.tensor(1, type: :u8) - true + iex> NxSignal.zoom_fft(Nx.tensor([1.0, 0.0, 0.0, 0.0]), 0.0, 1.0) |> Nx.real() |> Nx.round() + #Nx.Tensor< + f32[4] + [1.0, 1.0, 1.0, 1.0] + > """ @doc type: :transforms diff --git a/test/nx_signal_test.exs b/test/nx_signal_test.exs index dda2843..6a3e0da 100644 --- a/test/nx_signal_test.exs +++ b/test/nx_signal_test.exs @@ -42,6 +42,31 @@ defmodule NxSignalTest do assert Nx.Type.complex?(Nx.type(result)) end + + test "preserves f64 precision" do + x = Nx.tensor([1.0, 2.0, 3.0, 4.0], type: {:f, 64}) + result = NxSignal.czt(x) + + assert Nx.type(result) == {:c, 128} + end + + test "2D input with axis: 1 matches row-wise 1D CZT" do + rows = Nx.tensor([[1.0, 2.0, 3.0, 4.0], [4.0, 3.0, 2.0, 1.0]]) + result_2d = NxSignal.czt(rows, axis: 1) + expected = Nx.stack([NxSignal.czt(rows[0]), NxSignal.czt(rows[1])]) + + assert_all_close(result_2d, expected, atol: 1.0e-5) + end + + test "2D input with axis: 0 matches column-wise 1D CZT" do + cols = Nx.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]) + result_2d = NxSignal.czt(cols, axis: 0) + col0 = NxSignal.czt(cols[[.., 0]]) + col1 = NxSignal.czt(cols[[.., 1]]) + expected = Nx.stack([col0, col1], axis: 1) + + assert_all_close(result_2d, expected, atol: 1.0e-5) + end end describe "zoom_fft/4" do From 9ef1ef797f4278fdc2126437f57c8c68b56a1cf4 Mon Sep 17 00:00:00 2001 From: thomaspmurphy Date: Mon, 15 Jun 2026 15:38:32 +0100 Subject: [PATCH 3/6] style: mix format --- lib/nx_signal.ex | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex index de16e98..96ec9a3 100644 --- a/lib/nx_signal.ex +++ b/lib/nx_signal.ex @@ -801,13 +801,25 @@ defmodule NxSignal do result = if other_axes == [] do - czt_n(x_t, w, a, m: m, n: n, fft_len: fft_len, complex_type: complex_type, real_type: real_type) + czt_n(x_t, w, a, + m: m, + n: n, + fft_len: fft_len, + complex_type: complex_type, + real_type: real_type + ) else batch_names = Enum.map(0..(length(other_axes) - 1), fn i -> :"batch_#{i}" end) x_v = Nx.vectorize(x_t, batch_names) result_v = - czt_n(x_v, w, a, m: m, n: n, fft_len: fft_len, complex_type: complex_type, real_type: real_type) + czt_n(x_v, w, a, + m: m, + n: n, + fft_len: fft_len, + complex_type: complex_type, + real_type: real_type + ) Nx.devectorize(result_v, keep_names: false) end From cb1327660eeec53d386a25fbc5abd18df31bd0c2 Mon Sep 17 00:00:00 2001 From: thomaspmurphy Date: Tue, 16 Jun 2026 15:11:00 +0100 Subject: [PATCH 4/6] refactor(czt): address second round of review feedback --- guides/chirp_z.livemd | 32 +++++++++++++++++++------------- lib/nx_signal.ex | 37 +++++++++++++++++++------------------ test/nx_signal_test.exs | 34 ++++++++++++++++++++++------------ 3 files changed, 60 insertions(+), 43 deletions(-) diff --git a/guides/chirp_z.livemd b/guides/chirp_z.livemd index 6be74c3..a820ed0 100644 --- a/guides/chirp_z.livemd +++ b/guides/chirp_z.livemd @@ -79,8 +79,9 @@ The signal itself, 50 ms of the two tones, shows their superposition clearly: n_show = trunc(fs * 0.05) time_data = - Enum.zip(Nx.to_flat_list(t[0..(n_show - 1)]), Nx.to_flat_list(x[0..(n_show - 1)])) - |> Enum.map(fn {ti, xi} -> %{time: ti, amplitude: xi} end) + Enum.zip_with(Nx.to_flat_list(t[0..(n_show - 1)]), Nx.to_flat_list(x[0..(n_show - 1)]), fn ti, xi -> + %{time: ti, amplitude: xi} + end) Tucan.lineplot(time_data, "time", "amplitude") |> Tucan.Axes.set_x_title("Time (s)") @@ -162,10 +163,11 @@ zoom_result = NxSignal.zoom_fft(x2, f1_norm, f2_norm, m: m) zoom_amps = Nx.abs(zoom_result) zoom_data = - Nx.linspace(45.0, 55.0, n: m, endpoint: false, type: {:f, 64}) - |> Nx.to_flat_list() - |> Enum.zip(Nx.to_flat_list(zoom_amps)) - |> Enum.map(fn {f, a} -> %{frequency: f, amplitude: a, method: "Zoom FFT (#{10.0 / m |> Float.round(3)} Hz bins)"} end) + Enum.zip_with( + Nx.linspace(45.0, 55.0, n: m, endpoint: false, type: {:f, 64}) |> Nx.to_flat_list(), + Nx.to_flat_list(zoom_amps), + fn f, a -> %{frequency: f, amplitude: a, method: "Zoom FFT (#{10.0 / m |> Float.round(3)} Hz bins)"} end + ) comparison_data = dft_data ++ zoom_data @@ -272,13 +274,17 @@ hz_outer = NxSignal.czt(h, m: m, a: a, w: w) angles = Nx.linspace(0, 360, n: m, endpoint: false, type: {:f, 64}) |> Nx.to_flat_list() -contours_data = - Enum.zip(angles, Nx.to_flat_list(Nx.abs(hz_unit))) - |> Enum.map(fn {a, amp} -> %{angle: a, amplitude: amp, contour: "|z| = 1 (unit circle)"} end) - |> Kernel.++( - Enum.zip(angles, Nx.to_flat_list(Nx.abs(hz_outer))) - |> Enum.map(fn {a, amp} -> %{angle: a, amplitude: amp, contour: "|z| = 1.5"} end) - ) +unit_data = + Enum.zip_with(angles, Nx.to_flat_list(Nx.abs(hz_unit)), fn a, amp -> + %{angle: a, amplitude: amp, contour: "|z| = 1 (unit circle)"} + end) + +outer_data = + Enum.zip_with(angles, Nx.to_flat_list(Nx.abs(hz_outer)), fn a, amp -> + %{angle: a, amplitude: amp, contour: "|z| = 1.5"} + end) + +contours_data = unit_data ++ outer_data Tucan.lineplot(contours_data, "angle", "amplitude", tooltip: true) |> Tucan.color_by("contour") diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex index 96ec9a3..d7e7db8 100644 --- a/lib/nx_signal.ex +++ b/lib/nx_signal.ex @@ -752,11 +752,11 @@ defmodule NxSignal do ## Options - * `:m` - number of output points $M$. Defaults to `Nx.size(x)`. - * `:w` - contour ratio $W$ (complex scalar tensor). + * `:output_length` - number of output points $M$. Defaults to `Nx.size(x)`. + * `:contour_ratio` - contour ratio $W$ (complex scalar tensor). Defaults to $e^{-j2\pi/M}$, which gives the DFT. - * `:a` - contour starting point $A$ (complex scalar tensor). - Defaults to $1+0j$. + * `:contour_start` - contour starting point $A$ (complex scalar tensor). + Defaults to $1$. ## Examples @@ -769,21 +769,21 @@ defmodule NxSignal do """ @doc type: :transforms deftransform czt(x, opts \\ []) do - opts = Keyword.validate!(opts, [:m, :w, :a, axis: -1]) + opts = Keyword.validate!(opts, [:output_length, :contour_ratio, :contour_start, axis: -1]) ndim = Nx.rank(x) axis = opts[:axis] axis = if axis < 0, do: ndim + axis, else: axis n = elem(Nx.shape(x), axis) - m = opts[:m] || n + m = opts[:output_length] || n fft_len = czt_next_pow2(n + m - 1) w = - opts[:w] || + opts[:contour_ratio] || Nx.complex(:math.cos(-2 * :math.pi() / m), :math.sin(-2 * :math.pi() / m)) - a = opts[:a] || Nx.complex(1.0, 0.0) + a = opts[:contour_start] || 1.0 complex_type = x |> Nx.type() |> Nx.Type.to_complex() real_type = Nx.Type.to_real(complex_type) @@ -809,8 +809,9 @@ defmodule NxSignal do real_type: real_type ) else - batch_names = Enum.map(0..(length(other_axes) - 1), fn i -> :"batch_#{i}" end) - x_v = Nx.vectorize(x_t, batch_names) + batch_shape = x_t |> Nx.shape() |> Tuple.to_list() |> Enum.drop(-1) + batch_axes = Enum.with_index(batch_shape, fn size, i -> {:"batch_#{i}", size} end) + x_v = Nx.revectorize(x_t, x_t.vectorized_axes ++ batch_axes, target_shape: {n}) result_v = czt_n(x_v, w, a, @@ -827,9 +828,9 @@ defmodule NxSignal do Nx.transpose(result, axes: inv_perm) end - deftransformp czt_next_pow2(n) when n <= 1, do: 1 + defp czt_next_pow2(n) when n <= 1, do: 1 - deftransformp czt_next_pow2(n) do + defp czt_next_pow2(n) do Integer.pow(2, ceil(:math.log2(n))) end @@ -850,7 +851,7 @@ defmodule NxSignal do k_idx = Nx.iota({m}, type: real_type) # Pre-multiply: yn[n] = x[n] · a^{-n} · w^{n²/2} - yn = x * Nx.pow(a, -n_idx) * Nx.exp(n_idx * n_idx / 2 * log_w) + yn = x * a ** -n_idx * Nx.exp(n_idx ** 2 / 2 * log_w) # Zero-pad yn to FFT length yn_padded = Nx.pad(yn, Nx.as_type(0, complex_type), [{0, fft_len - n, 0}]) @@ -870,7 +871,7 @@ defmodule NxSignal do g = Nx.ifft(Nx.fft(yn_padded) * Nx.fft(h_padded)) # Post-multiply: X[k] = w^{k²/2} · g[k], k = 0…M-1 - Nx.slice(g, [0], [m]) * Nx.exp(k_idx * k_idx / 2 * log_w) + Nx.slice(g, [0], [m]) * Nx.exp(k_idx ** 2 / 2 * log_w) end @doc ~S""" @@ -896,7 +897,7 @@ defmodule NxSignal do ## Options - * `:m` - number of output points $M$. Defaults to `Nx.size(x)`. + * `:output_length` - number of output points $M$. Defaults to `Nx.size(x)`. ## Examples @@ -909,10 +910,10 @@ defmodule NxSignal do """ @doc type: :transforms deftransform zoom_fft(x, f1, f2, opts \\ []) do - opts = Keyword.validate!(opts, [:m]) + opts = Keyword.validate!(opts, [:output_length]) {n} = Nx.shape(x) - m = opts[:m] || n + m = opts[:output_length] || n # A = e^{+j2πf1}: the CZT evaluates at A^{-n}, so the positive exponent here # places the first evaluation point at frequency f1. @@ -924,6 +925,6 @@ defmodule NxSignal do :math.sin(-2 * :math.pi() * (f2 - f1) / m) ) - czt(x, m: m, w: w, a: a) + czt(x, output_length: m, contour_ratio: w, contour_start: a) end end diff --git a/test/nx_signal_test.exs b/test/nx_signal_test.exs index 6a3e0da..b60c5b2 100644 --- a/test/nx_signal_test.exs +++ b/test/nx_signal_test.exs @@ -13,7 +13,7 @@ defmodule NxSignalTest do test "unit impulse returns all ones" do x = Nx.tensor([1.0, 0.0, 0.0, 0.0]) - result = NxSignal.czt(x, m: 6) + result = NxSignal.czt(x, output_length: 6) ones = Nx.broadcast(Nx.complex(1.0, 0.0), {6}) assert_all_close(result, ones, atol: 1.0e-5) @@ -21,18 +21,18 @@ defmodule NxSignalTest do test "M > N is equivalent to zero-padded DFT" do x = Nx.tensor([1.0, 2.0, 3.0, 4.0]) - result = NxSignal.czt(x, m: 8) + result = NxSignal.czt(x, output_length: 8) expected = Nx.fft(Nx.as_type(x, {:c, 64}), length: 8) assert_all_close(result, expected, atol: 1.0e-5) end - test "custom a evaluates z-transform off the unit circle" do + test "custom contour_start evaluates z-transform off the unit circle" do # x = [1, 1], so X(z) = 1 + z^{-1}. At z = 2: X(2) = 1 + 0.5 = 1.5 x = Nx.tensor([1.0, 1.0]) a = Nx.complex(2.0, 0.0) w = Nx.complex(-1.0, 0.0) - result = NxSignal.czt(x, m: 1, a: a, w: w) + result = NxSignal.czt(x, output_length: 1, contour_start: a, contour_ratio: w) assert_all_close(Nx.real(result[0]), Nx.tensor(1.5, type: {:f, 64}), atol: 1.0e-5) end @@ -56,6 +56,7 @@ defmodule NxSignalTest do expected = Nx.stack([NxSignal.czt(rows[0]), NxSignal.czt(rows[1])]) assert_all_close(result_2d, expected, atol: 1.0e-5) + assert_all_close(NxSignal.czt(rows), result_2d, atol: 1.0e-7) end test "2D input with axis: 0 matches column-wise 1D CZT" do @@ -80,7 +81,9 @@ defmodule NxSignalTest do test "resolves two closely-spaced tones" do # Two tones at normalised frequencies 0.49 and 0.51. - # A 16-point DFT cannot separate them; zoom_fft over [0.45, 0.55] can. + # A 16-point DFT cannot separate them; zoom_fft over [0.45, 0.55] with + # 64 bins places bin spacing at 0.1/64 ≈ 0.00156, so the two tones land + # near bins 26 and 38 respectively. n = 128 t = Nx.iota({n}, type: :f64) pi2 = 2 * :math.pi() @@ -91,20 +94,27 @@ defmodule NxSignalTest do Nx.cos(Nx.multiply(t, pi2 * 0.51)) ) - result = NxSignal.zoom_fft(x, 0.45, 0.55, m: 64) + result = NxSignal.zoom_fft(x, 0.45, 0.55, output_length: 64) magnitudes = Nx.abs(result) - max_val = Nx.reduce_max(magnitudes) |> Nx.to_number() - median_val = Nx.median(magnitudes) |> Nx.to_number() + {_vals, peak_indices} = Nx.top_k(magnitudes, k: 2) + [p1, p2] = peak_indices |> Nx.to_flat_list() |> Enum.sort() - assert max_val > 5 * median_val + # tone at 0.49 → bin (0.49 - 0.45) / (0.1 / 64) ≈ 25.6 + # tone at 0.51 → bin (0.51 - 0.45) / (0.1 / 64) ≈ 38.4 + assert abs(p1 - 26) <= 2 + assert abs(p2 - 38) <= 2 end - test "returns m output points" do + test "output equals equivalent czt call" do x = Nx.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]) - result = NxSignal.zoom_fft(x, 0.1, 0.4, m: 16) + result = NxSignal.zoom_fft(x, 0.1, 0.4, output_length: 16) - assert Nx.shape(result) == {16} + a = Nx.complex(:math.cos(2 * :math.pi() * 0.1), :math.sin(2 * :math.pi() * 0.1)) + w = Nx.complex(:math.cos(-2 * :math.pi() * 0.3 / 16), :math.sin(-2 * :math.pi() * 0.3 / 16)) + expected = NxSignal.czt(x, output_length: 16, contour_ratio: w, contour_start: a) + + assert_all_close(result, expected, atol: 1.0e-5) end end end From 6a66d5d6110bbe089c2f82c99692bf455ce6aba7 Mon Sep 17 00:00:00 2001 From: thomaspmurphy Date: Wed, 17 Jun 2026 16:35:15 +0100 Subject: [PATCH 5/6] refactor(czt): simplify batch vectorization with revectorize auto --- lib/nx_signal.ex | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex index d7e7db8..d340f9b 100644 --- a/lib/nx_signal.ex +++ b/lib/nx_signal.ex @@ -809,9 +809,10 @@ defmodule NxSignal do real_type: real_type ) else + existing_vax = x_t.vectorized_axes batch_shape = x_t |> Nx.shape() |> Tuple.to_list() |> Enum.drop(-1) - batch_axes = Enum.with_index(batch_shape, fn size, i -> {:"batch_#{i}", size} end) - x_v = Nx.revectorize(x_t, x_t.vectorized_axes ++ batch_axes, target_shape: {n}) + + x_v = Nx.revectorize(x_t, [batch: :auto], target_shape: {n}) result_v = czt_n(x_v, w, a, @@ -822,7 +823,7 @@ defmodule NxSignal do real_type: real_type ) - Nx.devectorize(result_v, keep_names: false) + Nx.revectorize(result_v, existing_vax, target_shape: List.to_tuple(batch_shape ++ [m])) end Nx.transpose(result, axes: inv_perm) From a24efabb2090a9b95be73d24353a1c48cfb515b7 Mon Sep 17 00:00:00 2001 From: thomaspmurphy Date: Sun, 21 Jun 2026 15:48:37 +0100 Subject: [PATCH 6/6] fix(czt): show exact complex output in doctest --- lib/nx_signal.ex | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/nx_signal.ex b/lib/nx_signal.ex index d340f9b..473851c 100644 --- a/lib/nx_signal.ex +++ b/lib/nx_signal.ex @@ -760,10 +760,10 @@ defmodule NxSignal do ## Examples - iex> NxSignal.czt(Nx.tensor([1.0, 0.0, 0.0, 0.0])) |> Nx.real() |> Nx.round() + iex> NxSignal.czt(Nx.tensor([1.0, 0.0, 0.0, 0.0])) #Nx.Tensor< - f32[4] - [1.0, 1.0, 1.0, 1.0] + c64[4] + [1.0-1.5893256e-8i, 0.99999994+0.0i, 1.0+1.5893256e-8i, 0.99999994+0.0i] > """