Skip to content

docs(latent): migrate custom-latent examples to the Latent class (Phase 2) #111

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 1 commit into from
Jun 9, 2026
Merged

docs(latent): migrate custom-latent examples to the Latent class (Phase 2) #111

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 31 additions & 26 deletions notebooks/guides/results/latent_variables.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,15 +232,16 @@
"\n",
"The library catalogue is intentionally small. If you want a different derived quantity \u2014 a Sersic effective\n",
"radius converted to kiloparsecs, a colour from two-band fits, a velocity dispersion from a virial estimate \u2014\n",
"you subclass ``ag.AnalysisImaging`` and override ``LATENT_KEYS`` + ``compute_latent_variables``.\n",
"subclass ``ag.LatentGalaxy`` (override ``keys`` and ``variables``) and declare it on your analysis through the\n",
"``Latent`` class attribute. This is the same first-class, swappable mechanism you use to customise a\n",
"``Visualizer`` \u2014 the latent catalogue is a component of the analysis, not a pair of methods you monkey-patch.\n",
"\n",
"The example below adds ``bulge_axis_ratio`` (the axis ratio of the Sersic bulge, computed from its\n",
"``ell_comps`` parameters). The same pattern works for any function of the model instance.\n",
"\n",
"A subclass-level ``LATENT_KEYS`` shadows the library's config-driven ``@property``, so once you define your\n",
"own list the workspace yaml stops controlling it. If you want to keep the library latents AND add your own,\n",
"make ``LATENT_KEYS`` a property that returns ``super().LATENT_KEYS + [\"your.key\"]`` \u2014 the Euclid pipeline\n",
"(``euclid_strong_lens_modeling_pipeline/util.py``) shows this pattern in practice.\n",
"``ell_comps`` parameters). The same pattern works for any function of the model instance. Calling\n",
"``ag.LatentGalaxy.keys(analysis)`` / ``ag.LatentGalaxy.variables(...)`` from your overrides keeps the\n",
"config-driven library latents alongside your custom one (subclass the base ``ag.Latent`` instead if you want a\n",
"catalogue from scratch). The Euclid pipeline (``euclid_strong_lens_modeling_pipeline/util.py``, ``LatentEuclid``)\n",
"uses this exact composition pattern in production.\n",
"'''"
],
"outputs": [],
Expand All @@ -254,30 +255,24 @@
"import numpy as np\n",
"\n",
"\n",
"class AnalysisImagingWithAxisRatio(ag.AnalysisImaging):\n",
"class LatentBulgeAxisRatio(ag.LatentGalaxy):\n",
" \"\"\"\n",
" AnalysisImaging extended with a custom ``bulge_axis_ratio`` latent \u2014 the axis ratio of the Sersic\n",
" bulge derived from its ``ell_comps``. Demonstrates how to add a user-defined latent without modifying\n",
" the library.\n",
" The library galaxy latents plus a custom ``bulge_axis_ratio`` \u2014 the axis\n",
" ratio of the Sersic bulge derived from its ``ell_comps``. Demonstrates adding\n",
" a user-defined latent without modifying the library: subclass\n",
" ``ag.LatentGalaxy`` and compose its ``keys`` / ``variables`` static methods.\n",
" \"\"\"\n",
"\n",
" @property\n",
" def LATENT_KEYS(self):\n",
" return list(super().LATENT_KEYS) + [\"bulge_axis_ratio\"]\n",
" @staticmethod\n",
" def keys(analysis):\n",
" return list(ag.LatentGalaxy.keys(analysis)) + [\"bulge_axis_ratio\"]\n",
"\n",
" def compute_latent_variables(self, parameters, model):\n",
" from autogalaxy.imaging.model.latent import LATENT_FUNCTIONS\n",
" @staticmethod\n",
" def variables(analysis, parameters, model):\n",
" library_values = ag.LatentGalaxy.variables(analysis, parameters, model)\n",
"\n",
" xp = self._xp\n",
" magzero = self.kwargs.get(\"magzero\", None)\n",
" instance = model.instance_from_vector(vector=parameters)\n",
" fit = self.fit_from(instance=instance)\n",
" context = {\"fit\": fit, \"magzero\": magzero, \"xp\": xp}\n",
"\n",
" library_keys = [k for k in super().LATENT_KEYS]\n",
" library_values = tuple(LATENT_FUNCTIONS[k](**context) for k in library_keys)\n",
"\n",
" # Custom latent: axis ratio of the Sersic bulge from its ell_comps.\n",
" xp = analysis._xp\n",
" try:\n",
" ell_y, ell_x = instance.galaxies.galaxy.bulge.ell_comps\n",
" axis_ratio = (1.0 - np.sqrt(ell_y**2 + ell_x**2)) / (\n",
Expand All @@ -286,7 +281,17 @@
" except AttributeError:\n",
" axis_ratio = xp.nan\n",
"\n",
" return library_values + (axis_ratio,)\n"
" return library_values + (axis_ratio,)\n",
"\n",
"\n",
"class AnalysisImagingWithAxisRatio(ag.AnalysisImaging):\n",
" \"\"\"\n",
" ``AnalysisImaging`` that swaps in the custom ``LatentBulgeAxisRatio`` catalogue\n",
" via the ``Latent`` class attribute \u2014 the same one-line mechanism used to\n",
" declare a custom ``Visualizer`` or ``Result``. No library code is modified.\n",
" \"\"\"\n",
"\n",
" Latent = LatentBulgeAxisRatio\n"
],
"outputs": [],
"execution_count": null
Expand Down
54 changes: 11 additions & 43 deletions notebooks/guides/results/workflow/csv_make.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -156,55 +156,23 @@
" n_like_max=300, # samples capped for quick result generation\n",
" )\n",
"\n",
" class AnalysisLatent(ag.AnalysisImaging):\n",
"\n",
" LATENT_KEYS = [\"galaxies.galaxy.bulge.sersic_index_x2\"]\n",
"\n",
" def compute_latent_variables(self, parameters, model):\n",
" \"\"\"\n",
" A latent variable is not a model parameter but can be derived from the model. Its value and errors may be\n",
" of interest and aid in the interpretation of a model-fit.\n",
"\n",
" This code implements a simple example of a latent variable, the magn\n",
"\n",
" By overwriting this method we can manually specify latent variables that are calculated and output to\n",
" a `latent.csv` file, which mirrors the `samples.csv` file.\n",
"\n",
" In the example below, the `latent.csv` file will contain at least two columns with the shear magnitude and\n",
" angle sampled by the non-linear search.\n",
"\n",
" You can add your own custom latent variables here, if you have particular quantities that you\n",
" would like to output to the `latent.csv` file.\n",
" class LatentSersicIndex(ag.Latent):\n",
" \"\"\"Custom catalogue replacing library defaults; subclass ag.Latent (base) and override keys/variables;\n",
" declare via Latent = LatentSersicIndex; note: subclass ag.LatentGalaxy instead to keep library latents.\"\"\"\n",
"\n",
" This function is called at the end of search, following one of two schemes depending on the settings in\n",
" `output.yaml`:\n",
" @staticmethod\n",
" def keys(analysis):\n",
" return [\"galaxies.galaxy.bulge.sersic_index_x2\"]\n",
"\n",
" 1) Call for every search sample, which produces a complete `latent/samples.csv` which mirrors the normal\n",
" `samples.csv` file but takes a long time to compute.\n",
"\n",
" 2) Call only for N random draws from the posterior inferred at the end of the search, which only produces a\n",
" `latent/latent_summary.json` file with the median and 1 and 3 sigma errors of the latent variables but is\n",
" fast to compute.\n",
"\n",
"\n",
" Parameters\n",
" ----------\n",
" parameters : array-like\n",
" The parameter vector of the model sample. This will typically come from the non-linear search.\n",
" Inside this method it is mapped back to a model instance via `model.instance_from_vector`.\n",
" model : Model\n",
" The model object defining how the parameter vector is mapped to an instance. Passed explicitly\n",
" so that this function can be used inside JAX transforms (`vmap`, `jit`) with `functools.partial`.\n",
"\n",
" Returns\n",
" -------\n",
" A dictionary mapping every latent variable name to its value.\n",
"\n",
" \"\"\"\n",
" @staticmethod\n",
" def variables(analysis, parameters, model):\n",
" instance = model.instance_from_vector(vector=parameters)\n",
"\n",
" return (instance.galaxies.galaxy.bulge.sersic_index * 2.0,)\n",
"\n",
" class AnalysisLatent(ag.AnalysisImaging):\n",
" Latent = LatentSersicIndex\n",
"\n",
" analysis = AnalysisLatent(dataset=dataset)\n",
"\n",
" result = search.fit(model=model, analysis=analysis)"
Expand Down
52 changes: 11 additions & 41 deletions notebooks/guides/results/workflow/fits_make.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -167,53 +167,23 @@
" n_like_max=300, # samples capped for quick result generation\n",
" )\n",
"\n",
" class AnalysisLatent(ag.AnalysisImaging):\n",
" LATENT_KEYS = [\"galaxies.galaxy.bulge.sersic_index\"]\n",
"\n",
" def compute_latent_variables(self, parameters, model):\n",
" \"\"\"\n",
" A latent variable is not a model parameter but can be derived from the model. Its value and errors may be\n",
" of interest and aid in the interpretation of a model-fit.\n",
"\n",
" This code implements a simple example of a latent variable, the magn\n",
"\n",
" By overwriting this method we can manually specify latent variables that are calculated and output to\n",
" a `latent.csv` file, which mirrors the `samples.csv` file.\n",
"\n",
" In the example below, the `latent.csv` file will contain at least two columns with the shear magnitude and\n",
" angle sampled by the non-linear search.\n",
"\n",
" You can add your own custom latent variables here, if you have particular quantities that you\n",
" would like to output to the `latent.csv` file.\n",
" class LatentSersicIndex(ag.Latent):\n",
" \"\"\"Custom catalogue replacing library defaults; subclass ag.Latent (base) and override keys/variables;\n",
" declare via Latent = LatentSersicIndex; note: subclass ag.LatentGalaxy instead to keep library latents.\"\"\"\n",
"\n",
" This function is called at the end of search, following one of two schemes depending on the settings in\n",
" `output.yaml`:\n",
" @staticmethod\n",
" def keys(analysis):\n",
" return [\"galaxies.galaxy.bulge.sersic_index\"]\n",
"\n",
" 1) Call for every search sample, which produces a complete `latent/samples.csv` which mirrors the normal\n",
" `samples.csv` file but takes a long time to compute.\n",
"\n",
" 2) Call only for N random draws from the posterior inferred at the end of the search, which only produces a\n",
" `latent/latent_summary.json` file with the median and 1 and 3 sigma errors of the latent variables but is\n",
" fast to compute.\n",
"\n",
" Parameters\n",
" ----------\n",
" parameters : array-like\n",
" The parameter vector of the model sample. This will typically come from the non-linear search.\n",
" Inside this method it is mapped back to a model instance via `model.instance_from_vector`.\n",
" model : Model\n",
" The model object defining how the parameter vector is mapped to an instance. Passed explicitly\n",
" so that this function can be used inside JAX transforms (`vmap`, `jit`) with `functools.partial`.\n",
"\n",
" Returns\n",
" -------\n",
" A dictionary mapping every latent variable name to its value.\n",
"\n",
" \"\"\"\n",
" @staticmethod\n",
" def variables(analysis, parameters, model):\n",
" instance = model.instance_from_vector(vector=parameters)\n",
"\n",
" return (instance.galaxies.galaxy.bulge.sersic_index * 2.0,)\n",
"\n",
" class AnalysisLatent(ag.AnalysisImaging):\n",
" Latent = LatentSersicIndex\n",
"\n",
" analysis = AnalysisLatent(dataset=dataset)\n",
"\n",
" result = search.fit(model=model, analysis=analysis)"
Expand Down
52 changes: 11 additions & 41 deletions notebooks/guides/results/workflow/png_make.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -172,53 +172,23 @@
" n_like_max=300, # samples capped for quick result generation\n",
" )\n",
"\n",
" class AnalysisLatent(ag.AnalysisImaging):\n",
" LATENT_KEYS = [\"galaxies.galaxy.bulge.sersic_index\"]\n",
"\n",
" def compute_latent_variables(self, parameters, model):\n",
" \"\"\"\n",
" A latent variable is not a model parameter but can be derived from the model. Its value and errors may be\n",
" of interest and aid in the interpretation of a model-fit.\n",
"\n",
" This code implements a simple example of a latent variable, the magn\n",
"\n",
" By overwriting this method we can manually specify latent variables that are calculated and output to\n",
" a `latent.csv` file, which mirrors the `samples.csv` file.\n",
"\n",
" In the example below, the `latent.csv` file will contain at least two columns with the shear magnitude and\n",
" angle sampled by the non-linear search.\n",
"\n",
" You can add your own custom latent variables here, if you have particular quantities that you\n",
" would like to output to the `latent.csv` file.\n",
" class LatentSersicIndex(ag.Latent):\n",
" \"\"\"Custom catalogue replacing library defaults; subclass ag.Latent (base) and override keys/variables;\n",
" declare via Latent = LatentSersicIndex; note: subclass ag.LatentGalaxy instead to keep library latents.\"\"\"\n",
"\n",
" This function is called at the end of search, following one of two schemes depending on the settings in\n",
" `output.yaml`:\n",
" @staticmethod\n",
" def keys(analysis):\n",
" return [\"galaxies.galaxy.bulge.sersic_index\"]\n",
"\n",
" 1) Call for every search sample, which produces a complete `latent/samples.csv` which mirrors the normal\n",
" `samples.csv` file but takes a long time to compute.\n",
"\n",
" 2) Call only for N random draws from the posterior inferred at the end of the search, which only produces a\n",
" `latent/latent_summary.json` file with the median and 1 and 3 sigma errors of the latent variables but is\n",
" fast to compute.\n",
"\n",
" Parameters\n",
" ----------\n",
" parameters : array-like\n",
" The parameter vector of the model sample. This will typically come from the non-linear search.\n",
" Inside this method it is mapped back to a model instance via `model.instance_from_vector`.\n",
" model : Model\n",
" The model object defining how the parameter vector is mapped to an instance. Passed explicitly\n",
" so that this function can be used inside JAX transforms (`vmap`, `jit`) with `functools.partial`.\n",
"\n",
" Returns\n",
" -------\n",
" A dictionary mapping every latent variable name to its value.\n",
"\n",
" \"\"\"\n",
" @staticmethod\n",
" def variables(analysis, parameters, model):\n",
" instance = model.instance_from_vector(vector=parameters)\n",
"\n",
" return (instance.galaxies.galaxy.bulge.sersic_index * 2.0,)\n",
"\n",
" class AnalysisLatent(ag.AnalysisImaging):\n",
" Latent = LatentSersicIndex\n",
"\n",
" analysis = AnalysisLatent(dataset=dataset)\n",
"\n",
" result = search.fit(model=model, analysis=analysis)\n"
Expand Down
Loading
Loading