diff --git a/markdown/interferometer/simulator.md b/markdown/interferometer/simulator.md index 5feaa17cc..fa0197ec9 100644 --- a/markdown/interferometer/simulator.md +++ b/markdown/interferometer/simulator.md @@ -407,7 +407,7 @@ simulator_jax = al.SimulatorInterferometer( uv_wavelengths=uv_wavelengths, exposure_time=300.0, noise_sigma=0.1, - transformer_class=al.TransformerDFT, # NUFFT (pynufft) is not JAX-traceable + transformer_class=al.TransformerDFT, # exact; NUFFT is faster on large UV sets use_jax=True, ) @@ -421,7 +421,7 @@ dataset_jax = simulate(tracer) # Interferometer with jax.Array visibilities Two notes specific to interferometer: -- Use `TransformerDFT` (the default) under JAX. `TransformerNUFFT` (pynufft) +- Use `TransformerDFT` (the default) under JAX. `TransformerNUFFT` (nufftax) is faster on large UV sets but is not JAX-traceable. The `nufftax` research path is tracking a JAX-native NUFFT replacement; see `autolens_workspace_test/scripts/interferometer/nufft.py` for the diff --git a/markdown/interferometer/start_here.md b/markdown/interferometer/start_here.md index c876bb9c0..87e6e8519 100644 --- a/markdown/interferometer/start_here.md +++ b/markdown/interferometer/start_here.md @@ -36,10 +36,10 @@ __JAX__ PyAutoLens runs interferometer model-fits on JAX by default (JAX installs with `autolens` itself) — `al.AnalysisInterferometer(dataset=dataset)` -below auto-enables `use_jax=True`. Use `TransformerDFT` (the default in -this script) under JAX — `TransformerNUFFT` (pynufft) is faster on large -UV sets but is not JAX-traceable; the `nufftax` replacement (see the -`__NUFFT (nufftax)__` section below) is a research path tracking that. +below auto-enables `use_jax=True`. Both `TransformerDFT` (the default in +this script) and `TransformerNUFFT` are JAX-traceable, so either works +under JAX; `TransformerNUFFT` is nufftax-backed and is much faster on +large UV sets (see the `__NUFFT (nufftax)__` section below). For the broader JAX principles (when you write `@jax.jit` yourself, the return-type contract), see the top-level `autolens_workspace/start_here.py` @@ -62,8 +62,9 @@ or required switching to a pixelized source reconstruction. Pixelized sources ar complex, irregular source morphologies (see `features/pixelization`), but they are no longer a performance requirement for large datasets. -If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed -transformer (`TransformerNUFFTPyNUFFT`) is also available as a non-JAX fallback. +If `nufftax` is not installed, install it via `pip install nufftax`. Note that `nufftax` +requires JAX; where JAX is unavailable (notably Intel macOS, for which JAX ships no wheels) +use `TransformerDFT`, which is exact and pure-numpy but scales as O(N_vis x N_pix). __Number of Visibilities__ diff --git a/markdown/start_here.md b/markdown/start_here.md index fe3301c17..4db3fd051 100644 --- a/markdown/start_here.md +++ b/markdown/start_here.md @@ -594,7 +594,7 @@ Modeling of interferometer data from submillimeter (e.g. ALMA) and radio (e.g. L Visibilities data is fitted directly in the uv-plane, circumventing issues that arise when fitting a dirty image such as correlated noise. This uses the non-uniform fast fourier transform algorithm -[PyNUFFT](https://github.com/jyhmiinlin/pynufft) to efficiently map the galaxy model images to the uv-plane. +[nufftax](https://github.com/GragasLab/nufftax) to efficiently map the galaxy model images to the uv-plane. Checkout the`autolens_workspace/*/interferometer` package to get started. diff --git a/notebooks/guides/using_jax.ipynb b/notebooks/guides/using_jax.ipynb index 71db34089..14924be1f 100644 --- a/notebooks/guides/using_jax.ipynb +++ b/notebooks/guides/using_jax.ipynb @@ -156,8 +156,8 @@ "`Tracer` is an argument.\n", "\n", "For interferometer data the same shape applies with `al.FitInterferometer`. Both `TransformerDFT` and the\n", - "nufftax-backed `TransformerNUFFT` are JAX-traceable, so either works; only the legacy pynufft-backed\n", - "`TransformerNUFFTPyNUFFT` is not. Note the defaults differ by class: `Interferometer` (what a fit uses) defaults\n", + "nufftax-backed `TransformerNUFFT` are JAX-traceable, so either works. Note the defaults differ by class:\n", + "`Interferometer` (what a fit uses) defaults\n", "to `TransformerNUFFT`, while `SimulatorInterferometer` defaults to `TransformerDFT`.\n", "\n", "**Via `Fitness` \u2014 the production path.** A non-linear search does not call your function; it calls a `Fitness`\n", diff --git a/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb b/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb index 3b0dd8b16..7270ef258 100644 --- a/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb +++ b/notebooks/interferometer/features/linear_light_profiles/modeling.ipynb @@ -81,8 +81,9 @@ "NUFFT calls, so a model with N linear light profiles costs only N forward-NUFFTs per iteration on the GPU \u2014\n", "fast enough that linear inversions in the visibility plane are now routinely practical.\n", "\n", - "If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed transformer\n", - "(`TransformerNUFFTPyNUFFT`) is available as a non-JAX fallback but is not recommended for linear light profiles.\n", + "If `nufftax` is not installed, install it via `pip install nufftax`. Note that `nufftax` requires JAX; where\n", + "JAX is unavailable (notably Intel macOS, for which JAX ships no wheels) `TransformerDFT` is the only option,\n", + "though it is not recommended for linear light profiles at realistic visibility counts.\n", "\n", "__Positive Only Solver__\n", "\n", diff --git a/notebooks/interferometer/simulator.ipynb b/notebooks/interferometer/simulator.ipynb index 0626edbf8..95bb95db2 100644 --- a/notebooks/interferometer/simulator.ipynb +++ b/notebooks/interferometer/simulator.ipynb @@ -602,9 +602,9 @@ "Two notes specific to interferometer:\n", "\n", "- `TransformerNUFFT` is backed by the JAX-native `nufftax` library (see `__Many Visibilities__` above), so it\n", - " supports `jax.jit` and scales to large UV sets. The legacy pynufft-backed `TransformerNUFFTPyNUFFT` is not\n", - " JAX-traceable; see `autolens_workspace_test/scripts/interferometer/nufft.py` for the parity work. For small\n", - " visibility counts `TransformerDFT` (the simulator default) is also JAX-traceable.\n", + " supports `jax.jit` and scales to large UV sets; see\n", + " `autolens_workspace_test/scripts/interferometer/nufft.py` for the accuracy check against the exact DFT. For\n", + " small visibility counts `TransformerDFT` (the simulator default) is also JAX-traceable.\n", "- The eager call above works and is the supported route today. Note it returns a dataset whose visibilities are\n", " NumPy-backed, not `jax.Array`.\n", "\n", diff --git a/notebooks/interferometer/start_here.ipynb b/notebooks/interferometer/start_here.ipynb index e22a89359..01c7c8b03 100644 --- a/notebooks/interferometer/start_here.ipynb +++ b/notebooks/interferometer/start_here.ipynb @@ -38,10 +38,10 @@ "\n", "PyAutoLens runs interferometer model-fits on JAX by default (JAX installs\n", "with `autolens` itself) \u2014 `al.AnalysisInterferometer(dataset=dataset)`\n", - "below auto-enables `use_jax=True`. Use `TransformerDFT` (the default in\n", - "this script) under JAX \u2014 `TransformerNUFFT` (pynufft) is faster on large\n", - "UV sets but is not JAX-traceable; the `nufftax` replacement (see the\n", - "`__NUFFT (nufftax)__` section below) is a research path tracking that.\n", + "below auto-enables `use_jax=True`. Both `TransformerDFT` (the default in\n", + "this script) and `TransformerNUFFT` are JAX-traceable, so either works\n", + "under JAX; `TransformerNUFFT` is nufftax-backed and is much faster on\n", + "large UV sets (see the `__NUFFT (nufftax)__` section below).\n", "\n", "For the broader JAX principles (when you write `@jax.jit` yourself, the\n", "return-type contract), see the top-level `autolens_workspace/start_here.py`\n", @@ -64,8 +64,9 @@ "complex, irregular source morphologies (see `features/pixelization`), but they are no longer a\n", "performance requirement for large datasets.\n", "\n", - "If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed\n", - "transformer (`TransformerNUFFTPyNUFFT`) is also available as a non-JAX fallback.\n", + "If `nufftax` is not installed, install it via `pip install nufftax`. Note that `nufftax`\n", + "requires JAX; where JAX is unavailable (notably Intel macOS, for which JAX ships no wheels)\n", + "use `TransformerDFT`, which is exact and pure-numpy but scales as O(N_vis x N_pix).\n", "\n", "__Number of Visibilities__\n", "\n", diff --git a/scripts/guides/using_jax.py b/scripts/guides/using_jax.py index 0e9b58f02..c8991752b 100644 --- a/scripts/guides/using_jax.py +++ b/scripts/guides/using_jax.py @@ -151,8 +151,7 @@ def log_likelihood(instance): `Tracer` is an argument. For interferometer data the same shape applies with `al.FitInterferometer`. Both `TransformerDFT` and the -nufftax-backed `TransformerNUFFT` are JAX-traceable, so either works; only the legacy pynufft-backed -`TransformerNUFFTPyNUFFT` is not. Note the defaults differ by class: `Interferometer` (what a fit uses) defaults +nufftax-backed `TransformerNUFFT` are JAX-traceable, so either works. Note the defaults differ by class: `Interferometer` (what a fit uses) defaults to `TransformerNUFFT`, while `SimulatorInterferometer` defaults to `TransformerDFT`. **Via `Fitness` — the production path.** A non-linear search does not call your function; it calls a `Fitness` diff --git a/scripts/interferometer/features/linear_light_profiles/modeling.py b/scripts/interferometer/features/linear_light_profiles/modeling.py index c47b70fb0..bc2d79819 100644 --- a/scripts/interferometer/features/linear_light_profiles/modeling.py +++ b/scripts/interferometer/features/linear_light_profiles/modeling.py @@ -76,8 +76,9 @@ NUFFT calls, so a model with N linear light profiles costs only N forward-NUFFTs per iteration on the GPU — fast enough that linear inversions in the visibility plane are now routinely practical. -If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed transformer -(`TransformerNUFFTPyNUFFT`) is available as a non-JAX fallback but is not recommended for linear light profiles. +If `nufftax` is not installed, install it via `pip install nufftax`. Note that `nufftax` requires JAX; where +JAX is unavailable (notably Intel macOS, for which JAX ships no wheels) `TransformerDFT` is the only option, +though it is not recommended for linear light profiles at realistic visibility counts. __Positive Only Solver__ diff --git a/scripts/interferometer/simulator.py b/scripts/interferometer/simulator.py index 91e4663f8..cae8cddb3 100644 --- a/scripts/interferometer/simulator.py +++ b/scripts/interferometer/simulator.py @@ -362,9 +362,9 @@ Two notes specific to interferometer: - `TransformerNUFFT` is backed by the JAX-native `nufftax` library (see `__Many Visibilities__` above), so it - supports `jax.jit` and scales to large UV sets. The legacy pynufft-backed `TransformerNUFFTPyNUFFT` is not - JAX-traceable; see `autolens_workspace_test/scripts/interferometer/nufft.py` for the parity work. For small - visibility counts `TransformerDFT` (the simulator default) is also JAX-traceable. + supports `jax.jit` and scales to large UV sets; see + `autolens_workspace_test/scripts/interferometer/nufft.py` for the accuracy check against the exact DFT. For + small visibility counts `TransformerDFT` (the simulator default) is also JAX-traceable. - The eager call above works and is the supported route today. Note it returns a dataset whose visibilities are NumPy-backed, not `jax.Array`. diff --git a/scripts/interferometer/start_here.py b/scripts/interferometer/start_here.py index aa70c0f11..70d5ed46c 100644 --- a/scripts/interferometer/start_here.py +++ b/scripts/interferometer/start_here.py @@ -33,10 +33,10 @@ PyAutoLens runs interferometer model-fits on JAX by default (JAX installs with `autolens` itself) — `al.AnalysisInterferometer(dataset=dataset)` -below auto-enables `use_jax=True`. Use `TransformerDFT` (the default in -this script) under JAX — `TransformerNUFFT` (pynufft) is faster on large -UV sets but is not JAX-traceable; the `nufftax` replacement (see the -`__NUFFT (nufftax)__` section below) is a research path tracking that. +below auto-enables `use_jax=True`. Both `TransformerDFT` (the default in +this script) and `TransformerNUFFT` are JAX-traceable, so either works +under JAX; `TransformerNUFFT` is nufftax-backed and is much faster on +large UV sets (see the `__NUFFT (nufftax)__` section below). For the broader JAX principles (when you write `@jax.jit` yourself, the return-type contract), see the top-level `autolens_workspace/start_here.py` @@ -59,8 +59,9 @@ complex, irregular source morphologies (see `features/pixelization`), but they are no longer a performance requirement for large datasets. -If `nufftax` is not installed, install it via `pip install nufftax`. A legacy pynufft-backed -transformer (`TransformerNUFFTPyNUFFT`) is also available as a non-JAX fallback. +If `nufftax` is not installed, install it via `pip install nufftax`. Note that `nufftax` +requires JAX; where JAX is unavailable (notably Intel macOS, for which JAX ships no wheels) +use `TransformerDFT`, which is exact and pure-numpy but scales as O(N_vis x N_pix). __Number of Visibilities__ diff --git a/start_here.ipynb b/start_here.ipynb index f5b0250e0..8d3a5e292 100644 --- a/start_here.ipynb +++ b/start_here.ipynb @@ -619,7 +619,7 @@ "\n", "Visibilities data is fitted directly in the uv-plane, circumventing issues that arise when fitting a dirty image\n", "such as correlated noise. This uses the non-uniform fast fourier transform algorithm\n", - "[PyNUFFT](https://github.com/jyhmiinlin/pynufft) to efficiently map the galaxy model images to the uv-plane.\n", + "[nufftax](https://github.com/GragasLab/nufftax) to efficiently map the galaxy model images to the uv-plane.\n", "\n", "Checkout the`autolens_workspace/*/interferometer` package to get started.\n", "\n", diff --git a/start_here.py b/start_here.py index 886adb9bf..84822efc7 100644 --- a/start_here.py +++ b/start_here.py @@ -482,7 +482,7 @@ Visibilities data is fitted directly in the uv-plane, circumventing issues that arise when fitting a dirty image such as correlated noise. This uses the non-uniform fast fourier transform algorithm -[PyNUFFT](https://github.com/jyhmiinlin/pynufft) to efficiently map the galaxy model images to the uv-plane. +[nufftax](https://github.com/GragasLab/nufftax) to efficiently map the galaxy model images to the uv-plane. Checkout the`autolens_workspace/*/interferometer` package to get started.