feat: integration tests for model composition stability

[Jammy2211]

Overview

Add integration test scripts in autofit_workspace_test and autolens_workspace_test that compose models at realistic science-run complexity and assert structural properties of the composed af.Model graph. These tests protect against silent PyAutoFit model composition regressions — changes that alter prior count, prior identity, parameter ordering, path structure, or the output identifier hash without any visible error.

The identifier stability check is particularly valuable: if a PyAutoFit refactor changes how a model is composed, the identifier changes, which means existing users' output folders no longer match and backwards compatibility is silently broken.

Plan

• Add autofit_workspace_test/scripts/model_composition/model_composition.py — exercises the core af.Model/af.Collection API with multi-component models, linked priors, nested collections, serialization round-trips, and hardcoded identifier regression anchors
• Add autolens_workspace_test/scripts/model_composition/multi_galaxy_mge.py — composes a realistic multi-galaxy lens model (MGE light + Isothermal mass + shear) and asserts prior count, prior identity, path structure, and identifier stability
• No library source code or unit test changes — purely workspace-test integration scripts

Detailed implementation plan

Affected Repositories

• autofit_workspace_test
• autolens_workspace_test

Work Classification

Workspace work (integration test scripts only)

Branch Survey

Repository	Current Branch	Dirty?
autofit_workspace_test	main	clean
autolens_workspace_test	main	modified datasets (unrelated)

Suggested branch: feature/model-composition-integration
Worktree root: ~/Code/PyAutoLabs-wt/model-composition-integration/

Implementation Steps

Script 1: autofit_workspace_test/scripts/model_composition/model_composition.py

Exercises the PyAutoFit model composition API at realistic complexity, mirroring concepts from cookbooks/model_internal.py. Uses simple classes (Gaussian, Exponential) — no autolens/autogalaxy dependency.

Sections and assertions:

1. Basic model composition — af.Model, af.Collection with multiple components. Assert prior_count, total_free_parameters, unique_prior_paths structure and ordering, prior types match defaults.

2. Prior linking / shared priors — assign same prior across params and across components. Assert prior_count drops, identity checks (is), instance_from_vector gives same value to linked params, unique_prior_tuples deduplicates correctly.

3. Nested collections — multi-level af.Collection nesting. Assert prior_count is sum of children, unique_prior_paths have correct depth, cross-collection linking works.

4. Prior types and customisation — swap priors to Uniform, Gaussian, TruncatedGaussian. Assert isinstance checks, bounds/sigma/mean values, prior_count unchanged by type swap.

5. Fixed parameters — set params to constants. Assert prior_count drops, instances use fixed values.

6. Parameter vector ordering — prior_tuples_ordered_by_id. Assert creation-order sorting, instance_from_vector maps correctly, vector_from_unit_vector([0.5, ...]) returns medians.

7. Serialization round-trip — model.dict() → from_dict(). Assert prior_count, shared-prior identity, prior types, path structure all preserved.

8. Identifier stability — compose fixed model + search, compute identifier. Assert deterministic, changes with prior bounds, changes with structure, hardcode known-good value as regression anchor.

9. Model subsetting — with_paths, without_paths. Assert correct prior_count and filtered paths.

10. Assertions (model constraints) — add_assertion. Assert valid vectors pass, invalid raise FitException.

Script 2: autolens_workspace_test/scripts/model_composition/multi_galaxy_mge.py

Composes a realistic lens model as used in science runs.

1. Lens galaxy — MGE bulge (20 Gaussians, gaussian_per_basis=2) + Isothermal mass + ExternalShear. Assert per-component prior counts (6 + 5 + 2).

2. Source galaxy — MGE bulge (20 Gaussians, gaussian_per_basis=1). Assert prior count = 4.

3. Full model — af.Collection(galaxies=af.Collection(lens=lens, source=source)). Assert total prior_count = 17, unique_prior_paths nesting structure.

4. Prior identity within MGE — Gaussians within a basis share centres/ell_comps (is), across bases share centres but not ell_comps (is not).

5. Identifier stability — hardcode known-good identifier as regression anchor.

6. Serialization round-trip — assert prior_count and paths preserved.

7. Model info — assert key component names present in model.info.

Key Files

• autofit_workspace_test/scripts/model_composition/model_composition.py — new
• autofit_workspace_test/scripts/model_composition/__init__.py — new (empty)
• autolens_workspace_test/scripts/model_composition/multi_galaxy_mge.py — new
• autolens_workspace_test/scripts/model_composition/__init__.py — new (empty)

Original Prompt

Click to expand starting prompt

There are unit tests again model composition, but we dont have anything which tests against the full level of complex of full complex models we typically use in science runs.

For example, in PyAutoLens we compose models made of multiple galaxies and made with a multi gaussina expansion (e.g. @autolens_workspace/scripts/imaging/modeling.py).

We also have a script which gives a run through of the autofit model composition aPI at the lower level @autofit_workspace/scripts/cookbooks/model_internal.py

I would like some integrations tests in autofit_workspace_test/scripts and autolens_workspace_test/scripts which basically compose models of realistic complexity, make asertions against their number of parameters, aspects of their composition and their random unique_id which defines the folder results are output into. We can basically assume the source code is ok, but this will mean I can maintain confidence that upcoming large refactors of autofit do not impact the actual model composition.

Another benefit specific of the unique_id is if this changes, it means a source code change has change how a model is composed and maybe break backwards compatibiltiy for autolens users. So its a good sanity check on if my dev work has an unexpected impact.

[Jammy2211]

Shipped — 2026-04-12

PRs

• feat: model composition integration tests autofit_workspace_test#6
• feat: multi-galaxy MGE model composition integration tests autolens_workspace_test#25

Summary

Added two integration test scripts that guard against silent PyAutoFit model composition regressions:

1. autofit_workspace_test — exercises the core af.Model/af.Collection API with multi-component models, linked priors, nested collections, serialization round-trips, and a hardcoded identifier regression anchor. Uses simple Gaussian/Exponential classes with no autogalaxy/autolens dependency.

2. autolens_workspace_test — composes a realistic 17-parameter multi-galaxy lens model (MGE light + Isothermal mass + ExternalShear) and asserts prior count, prior identity within/across MGE bases, path structure, identifier stability, and kwarg variants.

Session Notes

• The model identifier (md5 hash) does not capture prior identity/linking — only distribution parameters. Two models with the same prior distributions but different sharing structure produce the same identifier. This is by design (Prior.id is excluded from the hash traversal in identifier.py:137).
• Serialization round-trip (dict → from_dict) preserves prior identity (shared priors remain the same object) and prior count, but creates new prior IDs. This means the ordering of prior_tuples_ordered_by_id and unique_prior_paths can change after round-trip. Tests compare these as sorted sets rather than ordered lists.