feat: scripts/imaging/ — model_fit + visualization + JIT visualizer scripts

config/build/env_vars.yaml

@@ -21,3 +21,7 @@ overrides:
   # JAX likelihood functions test JIT compilation — need JAX enabled and full-size datasets
   - pattern: "jax_likelihood_functions/"
     unset: [PYAUTO_WORKSPACE_SMALL_DATASETS, PYAUTO_DISABLE_JAX]
+  # imaging/visualization.py asserts subplot PNG / FITS files land on disk —
+  # PYAUTO_FAST_PLOTS skips savefig and would break those assertions.
+  - pattern: "imaging/visualization"
+    unset: [PYAUTO_FAST_PLOTS]

scripts/imaging/config/visualize/plots.yaml

@@ -0,0 +1,31 @@
+subplot_format: [png]
+fits_are_zoomed: false
+
+dataset:
+  subplot_dataset: true
+  fits_dataset: true
+
+fit:
+  subplot_fit: true
+  subplot_fit_log10: false
+  subplot_of_galaxies: false
+  subplot_galaxy_images: false
+  fits_fit: false
+  fits_galaxy_images: false
+  fits_model_galaxy_images: false
+
+fit_imaging: {}
+
+galaxies:
+  subplot_galaxies: false
+  subplot_galaxy_images: false
+  fits_galaxy_images: false
+
+inversion:
+  subplot_inversion: false
+  subplot_mappings: false
+  csv_reconstruction: false
+
+adapt:
+  subplot_adapt_images: true
+  fits_adapt_images: true

scripts/imaging/config_source/visualize/plots.yaml

@@ -0,0 +1,31 @@
+subplot_format: [png]
+fits_are_zoomed: false
+
+dataset:
+  subplot_dataset: false
+  fits_dataset: false
+
+fit:
+  subplot_fit: true
+  subplot_fit_log10: false
+  subplot_of_galaxies: false
+  subplot_galaxy_images: false
+  fits_fit: false
+  fits_galaxy_images: false
+  fits_model_galaxy_images: false
+
+fit_imaging: {}
+
+galaxies:
+  subplot_galaxies: true
+  subplot_galaxy_images: false
+  fits_galaxy_images: false
+
+inversion:
+  subplot_inversion: true
+  subplot_mappings: false
+  csv_reconstruction: false
+
+adapt:
+  subplot_adapt_images: false
+  fits_adapt_images: false

scripts/imaging/model_fit.py

@@ -0,0 +1,110 @@
+"""
+Modeling: Single Galaxy Sersic Fit
+==================================
+
+End-to-end imaging model-fit on the autogalaxy single-galaxy ``jax_test`` dataset.
+Exercises ``AnalysisImaging`` -> ``FitImaging`` with a Nautilus search.
+
+Galaxy: a single ``Sersic`` bulge — no lens / mass / source split (this is autogalaxy,
+not autolens).
+"""
+
+import os
+from os import path
+
+import autofit as af
+import autogalaxy as ag
+import autogalaxy.plot as aplt
+
+
+"""
+__Dataset__
+
+Reuse the ``jax_test`` dataset already used by ``scripts/jax_likelihood_functions/imaging``.
+"""
+dataset_path = path.join("dataset", "imaging", "jax_test")
+
+if not path.exists(path.join(dataset_path, "data.fits")):
+    import subprocess
+    import sys
+
+    subprocess.run(
+        [sys.executable, "scripts/jax_likelihood_functions/imaging/simulator.py"],
+        check=True,
+    )
+
+dataset = ag.Imaging.from_fits(
+    data_path=path.join(dataset_path, "data.fits"),
+    psf_path=path.join(dataset_path, "psf.fits"),
+    noise_map_path=path.join(dataset_path, "noise_map.fits"),
+    pixel_scales=0.2,
+)
+
+aplt.plot_array(array=dataset.data)
+
+
+"""
+__Mask__
+"""
+mask_radius = 3.0
+
+mask = ag.Mask2D.circular(
+    shape_native=dataset.shape_native,
+    pixel_scales=dataset.pixel_scales,
+    radius=mask_radius,
+)
+
+dataset = dataset.apply_mask(mask=mask)
+
+aplt.plot_array(array=dataset.data)
+
+
+"""
+__Model__
+
+Single galaxy with a parametric ``Sersic`` bulge.
+"""
+bulge = af.Model(ag.lp.Sersic)
+bulge.centre.centre_0 = 0.0
+bulge.centre.centre_1 = 0.0
+
+galaxy = af.Model(ag.Galaxy, redshift=0.5, bulge=bulge)
+
+model = af.Collection(galaxies=af.Collection(galaxy=galaxy))
+
+
+"""
+__Search__
+"""
+search = af.Nautilus(
+    path_prefix=path.join("build", "model_fit", "imaging"),
+    n_live=50,
+    n_like_max=300,
+    number_of_cores=2,
+)
+
+
+"""
+__Analysis__
+"""
+analysis = ag.AnalysisImaging(dataset=dataset)
+
+
+"""
+__Model-Fit__
+"""
+result = search.fit(model=model, analysis=analysis)
+
+
+"""
+__Result__
+"""
+print(result.max_log_likelihood_instance)
+
+aplt.subplot_galaxies(
+    galaxies=result.max_log_likelihood_galaxies, grid=result.grids.lp
+)
+
+aplt.subplot_fit_imaging(fit=result.max_log_likelihood_fit)
+
+aplt.corner_cornerpy(samples=result.samples)