docs: correct false simulator-JAX claims (registration, eager-on-JAX, TransformerNUFFT)

notebooks/README.md

@@ -1,7 +1,7 @@
-The `autogalaxy_workspace/scripts` folder contains the packages for different science use-cases, where each package
+The `scripts` folder contains the packages for different science use-cases, where each package
 only contains Python scripts (e.g. `.py` files) illustrating how to perform a specific task.
 
-For Jupyter Notebook examples illustrating the same use-cases, see the `autogalaxy_workspace/notebooks` folder.
+For Jupyter Notebook examples illustrating the same use-cases, see the `notebooks` folder.
 
 Refer to the `autogalaxy_workspace/README.md` file or the readthedocs documentation for a description of which folder
 you should go into next if you are unsure.

notebooks/guides/using_jax.ipynb

@@ -44,27 +44,28 @@
     "\n",
     "__Writing @jax.jit Yourself__\n",
     "\n",
-    "Pass `use_jax=True` to a simulator constructor and wrap your call in `@jax.jit` when you want to render many\n",
-    "datasets fast \u2014 parameter sweeps, mock-data studies, batch figure generation:\n",
+    "Pass `use_jax=True` to a simulator constructor to run the image calculation through JAX, for parameter sweeps,\n",
+    "mock-data studies or batch figure generation:\n",
     "\n",
     "```python\n",
-    "import jax\n",
-    "\n",
     "simulator = ag.SimulatorImaging(\n",
     "    exposure_time=300.0, psf=psf, background_sky_level=0.1, use_jax=True\n",
     ")\n",
     "\n",
-    "@jax.jit\n",
-    "def simulate(galaxies):\n",
-    "    return simulator.via_galaxies_from(galaxies=galaxies, grid=grid)\n",
+    "dataset = simulator.via_galaxies_from(galaxies=galaxies, grid=grid)\n",
     "```\n",
     "\n",
-    "The simulator handles pytree registration internally, so you write nothing JAX-specific beyond the decorator.\n",
-    "Note that eager `simulator.via_galaxies_from(galaxies, grid)` (no `@jax.jit`) already runs on JAX and is\n",
-    "sufficient for one-off simulations \u2014 the `@jax.jit` wrap only pays off when you call the function many times.\n",
+    "**Wrapping that call in `@jax.jit` does not currently work.** Two things stop it, and it is worth knowing which\n",
+    "is which:\n",
+    "\n",
+    "- **You must register the pytrees yourself first.** Nothing in the library does it for you, and nothing can: JAX\n",
+    "  flattens a jitted function's arguments at trace time, *before* entering the callee, so a simulator that\n",
+    "  registered internally would already be too late. The one-time call is\n",
+    "  `autogalaxy.jax.register_galaxies_classes(galaxies)`.\n",
+    "- **Even with that, the jitted simulator call fails inside autoarray** on array sites that do not yet thread\n",
+    "  `xp` \u2014 see PyAutoLabs/PyAutoArray for the tracked issue. Until it is fixed, use the eager call above.\n",
     "\n",
-    "The per-dataset-type `simulator.py` scripts (`scripts/imaging/simulator.py`,\n",
-    "`scripts/interferometer/simulator.py`) each show this pattern in their `__JAX Variant__` section.\n",
+    "Note the eager call returns a dataset whose `.data.array` is a `numpy.ndarray`, not a `jax.Array`.\n",
     "\n",
     "__Custom Likelihood Functions__\n",
     "\n",
@@ -117,8 +118,10 @@
     "Omit `xp=jnp` and the fit falls back to NumPy internals, raising `TracerArrayConversionError` the moment JAX\n",
     "traces it.\n",
     "\n",
-    "For interferometer data the same shape applies with `ag.FitInterferometer`, with one constraint: use\n",
-    "`TransformerDFT` (the default). `TransformerNUFFT` is not JAX-traceable.\n",
+    "For interferometer data the same shape applies with `ag.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",
+    "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",
     "object, which maps a raw parameter vector to a model instance, calls the analysis, and returns the figure of\n",

notebooks/imaging/simulator.ipynb

@@ -481,14 +481,11 @@
     "\n",
     "__JAX Variant__\n",
     "\n",
-    "For an order-of-magnitude speedup on large or repeated simulations\n",
-    "(parameter sweeps, mock-data studies, batch figure generation), construct\n",
-    "the simulator with `use_jax=True` and wrap your call in `@jax.jit`. The\n",
-    "simulator handles pytree registration internally.\n",
+    "For large or repeated simulations (parameter sweeps, mock-data studies,\n",
+    "batch figure generation), construct the simulator with `use_jax=True` so\n",
+    "the image calculation runs through JAX:\n",
     "\n",
     "```python\n",
-    "import jax\n",
-    "\n",
     "simulator_jax = ag.SimulatorImaging(\n",
     "    exposure_time=300.0,\n",
     "    psf=psf,\n",
@@ -497,20 +494,23 @@
     "    use_jax=True,\n",
     ")\n",
     "\n",
-    "@jax.jit\n",
-    "def simulate(galaxies):\n",
-    "    return simulator_jax.via_galaxies_from(galaxies=galaxies, grid=grid)\n",
-    "\n",
-    "dataset_jax = simulate(galaxies)   # Imaging with jax.Array data\n",
+    "dataset_jax = simulator_jax.via_galaxies_from(galaxies=galaxies, grid=grid)\n",
     "```\n",
     "\n",
-    "The `dataset_jax.data.array` is a `jax.Array`; `aplt.fits_imaging` and the\n",
-    "plotters call `numpy.asarray()` internally, so saving / plotting works\n",
-    "without manual conversion.\n",
+    "The returned `dataset_jax.data.array` is a `numpy.ndarray`. `aplt.fits_imaging`\n",
+    "and the plotters call `numpy.asarray()` internally, so saving / plotting\n",
+    "works either way.\n",
+    "\n",
+    "**Wrapping the call in `@jax.jit` does not currently work.** Two separate\n",
+    "things stop it:\n",
     "\n",
-    "Note: eager `simulator_jax.via_galaxies_from(galaxies, grid)` (no `@jax.jit`)\n",
-    "already runs on JAX and is sufficient for one-off simulations. The\n",
-    "`@jax.jit` wrap is only beneficial when you call the function many times.\n",
+    "- **Pytree registration is yours to do, before the first jitted call** \u2014\n",
+    "  `autogalaxy.jax.register_galaxies_classes(galaxies)`. Nothing in the\n",
+    "  library does it for you, and nothing can: JAX flattens a jitted\n",
+    "  function's arguments at trace time, before entering the callee.\n",
+    "- **Even with that, the jitted call fails inside autoarray** on array sites\n",
+    "  that do not yet thread `xp`. Tracked in PyAutoArray; until it is fixed,\n",
+    "  use the eager call above.\n",
     "\n",
     "See `scripts/guides/data_structures.py` for the broader \"JIT-it-\n",
     "yourself\" pattern."

notebooks/interferometer/features/linear_light_profiles/README.md

@@ -25,4 +25,4 @@ A full guide to result analysis is given at `autogalaxy_workspace/*/guides/resul
 # Imaging Equivalent
 
 For the CCD-imaging version of these scripts, see
-`autogalaxy_workspace/scripts/imaging/features/linear_light_profiles`.
+`scripts/imaging/features/linear_light_profiles`.

notebooks/interferometer/features/multi_gaussian_expansion/README.md

@@ -27,4 +27,4 @@ A full guide to result analysis is given at `autogalaxy_workspace/*/guides/resul
 # Imaging Equivalent
 
 For the CCD-imaging version of these scripts, see
-`autogalaxy_workspace/scripts/imaging/features/multi_gaussian_expansion`.
+`scripts/imaging/features/multi_gaussian_expansion`.

notebooks/interferometer/features/shapelets/README.md

@@ -27,4 +27,4 @@ A full guide to result analysis is given at `autogalaxy_workspace/*/guides/resul
 
 # Imaging Equivalent
 
-For the CCD-imaging version of these scripts, see `autogalaxy_workspace/scripts/imaging/features/shapelets`.
+For the CCD-imaging version of these scripts, see `scripts/imaging/features/shapelets`.

notebooks/interferometer/simulator.ipynb

@@ -388,41 +388,46 @@
     "__JAX Variant__\n",
     "\n",
     "For fast repeated interferometer simulations, construct the simulator\n",
-    "with `use_jax=True` and wrap the call in `@jax.jit`. The simulator\n",
-    "handles pytree registration internally.\n",
+    "with `use_jax=True` so the image calculation and transform run through\n",
+    "JAX:\n",
     "\n",
     "```python\n",
-    "import jax\n",
     "import jax.numpy as jnp\n",
     "\n",
     "simulator_jax = ag.SimulatorInterferometer(\n",
     "    uv_wavelengths=uv_wavelengths,\n",
     "    exposure_time=300.0,\n",
     "    noise_sigma=0.1,\n",
-    "    transformer_class=ag.TransformerDFT,  # NUFFT (pynufft) is not JAX-traceable\n",
     "    use_jax=True,\n",
     ")\n",
     "\n",
-    "@jax.jit\n",
-    "def simulate(galaxies):\n",
-    "    galaxy_obj = ag.Galaxies(galaxies=galaxies)\n",
-    "    image = galaxy_obj.image_2d_from(grid=real_space_grid, xp=jnp)\n",
-    "    return simulator_jax.via_image_from(image=image)\n",
-    "\n",
-    "dataset_jax = simulate(galaxies)   # Interferometer with jax.Array visibilities\n",
+    "image = ag.Galaxies(galaxies=galaxies).image_2d_from(\n",
+    "    grid=real_space_grid, xp=jnp\n",
+    ")\n",
+    "dataset_jax = simulator_jax.via_image_from(image=image)\n",
     "```\n",
     "\n",
-    "Two notes:\n",
-    "\n",
-    "- Use `TransformerDFT` (the default) under JAX. `TransformerNUFFT`\n",
-    "  (pynufft) is faster on large UV sets but is not JAX-traceable; the\n",
-    "  `nufftax` replacement is a research path (see\n",
-    "  `autolens_workspace_test/scripts/interferometer/nufft.py`).\n",
-    "- Eager `simulator_jax.via_image_from(image)` already runs on JAX without\n",
-    "  the `@jax.jit` wrap; the JIT only matters for repeated calls.\n",
-    "\n",
-    "See `scripts/guides/data_structures.py` for the broader \"JIT-it-\n",
-    "yourself\" pattern."
+    "**Wrapping the call in `@jax.jit` does not currently work.** Two separate\n",
+    "things stop it:\n",
+    "\n",
+    "- **Pytree registration is yours to do, before the first jitted call** \u2014\n",
+    "  `autogalaxy.jax.register_galaxies_classes(galaxies)`. Nothing in the\n",
+    "  library does it for you, and nothing can: JAX flattens a jitted\n",
+    "  function's arguments at trace time, before entering the callee.\n",
+    "- **Even with that, the jitted call fails inside autoarray** on array\n",
+    "  sites that do not yet thread `xp`. Tracked in PyAutoArray; until it is\n",
+    "  fixed, use the eager call above.\n",
+    "\n",
+    "On transformers: `TransformerDFT` (the `SimulatorInterferometer` default)\n",
+    "and the nufftax-backed `TransformerNUFFT` (the `Interferometer` default,\n",
+    "so what a fit uses) are both JAX-traceable. Only the legacy pynufft-backed\n",
+    "`TransformerNUFFTPyNUFFT` is not \u2014 see\n",
+    "`autolens_workspace_test/scripts/interferometer/nufft.py` for the parity\n",
+    "work.\n",
+    "\n",
+    "See `scripts/guides/using_jax.py` for the full picture and\n",
+    "`scripts/guides/data_structures.py` for the broader \"JIT-it-yourself\"\n",
+    "pattern."
    ]
   }
  ],

scripts/README.md

@@ -1,7 +1,7 @@
-The `autogalaxy_workspace/scripts` folder contains the packages for different science use-cases, where each package
+The `scripts` folder contains the packages for different science use-cases, where each package
 only contains Python scripts (e.g. `.py` files) illustrating how to perform a specific task.
 
-For Jupyter Notebook examples illustrating the same use-cases, see the `autogalaxy_workspace/notebooks` folder.
+For Jupyter Notebook examples illustrating the same use-cases, see the `notebooks` folder.
 
 Refer to the `autogalaxy_workspace/README.md` file or the readthedocs documentation for a description of which folder
 you should go into next if you are unsure.

scripts/guides/using_jax.py

@@ -39,27 +39,28 @@
 
 __Writing @jax.jit Yourself__
 
-Pass `use_jax=True` to a simulator constructor and wrap your call in `@jax.jit` when you want to render many
-datasets fast — parameter sweeps, mock-data studies, batch figure generation:
+Pass `use_jax=True` to a simulator constructor to run the image calculation through JAX, for parameter sweeps,
+mock-data studies or batch figure generation:
 
 ```python
-import jax
-
 simulator = ag.SimulatorImaging(
     exposure_time=300.0, psf=psf, background_sky_level=0.1, use_jax=True
 )
 
-@jax.jit
-def simulate(galaxies):
-    return simulator.via_galaxies_from(galaxies=galaxies, grid=grid)
+dataset = simulator.via_galaxies_from(galaxies=galaxies, grid=grid)
 ```
 
-The simulator handles pytree registration internally, so you write nothing JAX-specific beyond the decorator.
-Note that eager `simulator.via_galaxies_from(galaxies, grid)` (no `@jax.jit`) already runs on JAX and is
-sufficient for one-off simulations — the `@jax.jit` wrap only pays off when you call the function many times.
+**Wrapping that call in `@jax.jit` does not currently work.** Two things stop it, and it is worth knowing which
+is which:
+
+- **You must register the pytrees yourself first.** Nothing in the library does it for you, and nothing can: JAX
+  flattens a jitted function's arguments at trace time, *before* entering the callee, so a simulator that
+  registered internally would already be too late. The one-time call is
+  `autogalaxy.jax.register_galaxies_classes(galaxies)`.
+- **Even with that, the jitted simulator call fails inside autoarray** on array sites that do not yet thread
+  `xp` — see PyAutoLabs/PyAutoArray for the tracked issue. Until it is fixed, use the eager call above.
 
-The per-dataset-type `simulator.py` scripts (`scripts/imaging/simulator.py`,
-`scripts/interferometer/simulator.py`) each show this pattern in their `__JAX Variant__` section.
+Note the eager call returns a dataset whose `.data.array` is a `numpy.ndarray`, not a `jax.Array`.
 
 __Custom Likelihood Functions__
 
@@ -112,8 +113,10 @@ def log_likelihood(instance):
 Omit `xp=jnp` and the fit falls back to NumPy internals, raising `TracerArrayConversionError` the moment JAX
 traces it.
 
-For interferometer data the same shape applies with `ag.FitInterferometer`, with one constraint: use
-`TransformerDFT` (the default). `TransformerNUFFT` is not JAX-traceable.
+For interferometer data the same shape applies with `ag.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
+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`
 object, which maps a raw parameter vector to a model instance, calls the analysis, and returns the figure of

scripts/imaging/simulator.py

@@ -251,14 +251,11 @@
 
 __JAX Variant__
 
-For an order-of-magnitude speedup on large or repeated simulations
-(parameter sweeps, mock-data studies, batch figure generation), construct
-the simulator with `use_jax=True` and wrap your call in `@jax.jit`. The
-simulator handles pytree registration internally.
+For large or repeated simulations (parameter sweeps, mock-data studies,
+batch figure generation), construct the simulator with `use_jax=True` so
+the image calculation runs through JAX:
 
 ```python
-import jax
-
 simulator_jax = ag.SimulatorImaging(
     exposure_time=300.0,
     psf=psf,
@@ -267,20 +264,23 @@
     use_jax=True,
 )
 
-@jax.jit
-def simulate(galaxies):
-    return simulator_jax.via_galaxies_from(galaxies=galaxies, grid=grid)
-
-dataset_jax = simulate(galaxies)   # Imaging with jax.Array data
+dataset_jax = simulator_jax.via_galaxies_from(galaxies=galaxies, grid=grid)
 ```
 
-The `dataset_jax.data.array` is a `jax.Array`; `aplt.fits_imaging` and the
-plotters call `numpy.asarray()` internally, so saving / plotting works
-without manual conversion.
+The returned `dataset_jax.data.array` is a `numpy.ndarray`. `aplt.fits_imaging`
+and the plotters call `numpy.asarray()` internally, so saving / plotting
+works either way.
+
+**Wrapping the call in `@jax.jit` does not currently work.** Two separate
+things stop it:
 
-Note: eager `simulator_jax.via_galaxies_from(galaxies, grid)` (no `@jax.jit`)
-already runs on JAX and is sufficient for one-off simulations. The
-`@jax.jit` wrap is only beneficial when you call the function many times.
+- **Pytree registration is yours to do, before the first jitted call** —
+  `autogalaxy.jax.register_galaxies_classes(galaxies)`. Nothing in the
+  library does it for you, and nothing can: JAX flattens a jitted
+  function's arguments at trace time, before entering the callee.
+- **Even with that, the jitted call fails inside autoarray** on array sites
+  that do not yet thread `xp`. Tracked in PyAutoArray; until it is fixed,
+  use the eager call above.
 
 See `scripts/guides/data_structures.py` for the broader "JIT-it-
 yourself" pattern.

scripts/interferometer/features/linear_light_profiles/README.md

@@ -25,4 +25,4 @@ A full guide to result analysis is given at `autogalaxy_workspace/*/guides/resul
 # Imaging Equivalent
 
 For the CCD-imaging version of these scripts, see
-`autogalaxy_workspace/scripts/imaging/features/linear_light_profiles`.
+`scripts/imaging/features/linear_light_profiles`.

scripts/interferometer/features/multi_gaussian_expansion/README.md

@@ -27,4 +27,4 @@ A full guide to result analysis is given at `autogalaxy_workspace/*/guides/resul
 # Imaging Equivalent
 
 For the CCD-imaging version of these scripts, see
-`autogalaxy_workspace/scripts/imaging/features/multi_gaussian_expansion`.
+`scripts/imaging/features/multi_gaussian_expansion`.

scripts/interferometer/features/shapelets/README.md

@@ -27,4 +27,4 @@ A full guide to result analysis is given at `autogalaxy_workspace/*/guides/resul
 
 # Imaging Equivalent
 
-For the CCD-imaging version of these scripts, see `autogalaxy_workspace/scripts/imaging/features/shapelets`.
+For the CCD-imaging version of these scripts, see `scripts/imaging/features/shapelets`.