feat: implement Chirp Z-Transform (CZT) and zoom_fft by thomaspmurphy · Pull Request #32 · elixir-nx/nx_signal

thomaspmurphy · 2026-06-15T13:53:06Z

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.

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.

polvalente

Thanks for the contribution! I haven't read the guide yet, but there are some minor adjustments for the implementation!

Preserve input precision via Nx.Type.to_complex/to_real, add axis option with vectorize-based batching, fix doctest assertions, simplify a^-n term.

polvalente · 2026-06-15T14:59:17Z

+      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]
+      >


Suggested change

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]

>

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]

>

Assert on the regular output here. Nx.round doesn't have precision, so anything between 0.5 and 1.5 would round to 1

Hmm the output type is c64[4], not f32[4]. It also has floating-point noise from Bluestein's (~6×10⁻⁸), so the suggested expected output would fail on both type and values. Possible alternative that avoids Nx.round entirely?

iex> NxSignal.czt(Nx.tensor([1.0, 2.0, 3.0, 4.0]), output_length: 1) #Nx.Tensor< c64[1] [10.0+0.0i] >

This evaluates the z-transform at a single point (z=1, sum of all samples) which gives an exact complex result. What do you think?

You can assert on the actual result type and value. It's fine if the doctest is not "nice and round"

feat: implement Chirp Z-Transform (CZT) and zoom_fft

3cfe70c

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.