Overview
The decorator system at autoarray/structures/decorators/ handles input→output type mapping for all profile methods (e.g. Grid2D in → Array2D out, Grid2DIrregular → ArrayIrregular). It works correctly but is verbose (~648 lines), contains dead code (Grid1D path), duplicates logic across three Maker classes, and has near-zero dedicated test coverage. This task cleans it up in phases, starting with tests and targeted cleanup, with an optional follow-up to consolidate to a registry-based architecture.
Plan
- Write comprehensive decorator unit tests covering all grid types × all decorators × edge cases (ndarray passthrough, JAX passthrough, list returns) — safety net before any refactoring
- Remove the dead Grid1D code path from all makers — Grid1D is never passed to decorated profile methods
- Extract shared list-return handling into a single helper used by all makers
- Move over-sampling metadata extraction into Grid2D's construction rather than try/except in GridMaker
- Rename
grid_dec → decorators in the public API (aa.grid_dec → aa.decorators) with grep-and-replace across PyAutoGalaxy/PyAutoLens
- Phase 3 (follow-up): If tests are solid and cleanup goes well, consolidate the three Maker classes into one generic wrapper + flat lookup table (~530 lines → ~150 lines)
Detailed implementation plan
Affected Repositories
- PyAutoArray (primary)
- PyAutoGalaxy (consumer — decorator usage, rename)
- PyAutoLens (consumer — decorator usage, rename)
Work Classification
Library
Branch Survey
| Repository |
Current Branch |
Dirty? |
| PyAutoArray |
main |
clean |
| PyAutoGalaxy |
main |
clean |
| PyAutoLens |
main |
clean |
Suggested branch: feature/data-typing-simplify
Worktree root: ~/Code/PyAutoLabs-wt/data-typing-simplify/
Implementation Steps
Phase 1 — Tests (zero-risk)
- Expand
test_autoarray/test_decorators.py with tests for:
@to_array with Grid2D input → asserts Array2D output with correct mask@to_array with Grid2DIrregular input → asserts ArrayIrregular output@to_array with raw np.ndarray → asserts np.ndarray passthrough@to_vector_yx with Grid2D → asserts VectorYX2D@to_vector_yx with Grid2DIrregular → asserts VectorYX2DIrregular@to_grid with Grid2D → asserts Grid2D with over-sampling metadata preserved@to_grid with Grid2DIrregular → asserts Grid2DIrregular@transform — coordinate frame transformation applied and reversed- List return handling for all decorators
- JAX array passthrough (if jax available)
Phase 2 — Targeted cleanup (low risk)
- Remove
Grid1D isinstance branch and via_grid_1d() from all three makers (to_array.py, to_grid.py, to_vector_yx.py) and the 1D projection logic from abstract.py
- Extract list-return handling from each maker into a shared helper in
abstract.py
- Move over-sampling metadata try/except from
GridMaker.via_grid_2d() into cleaner getattr calls or into Grid2D's own construction
- Rename
grid_dec → decorators in autoarray/__init__.py and update all aa.grid_dec. references in PyAutoGalaxy and PyAutoLens
Phase 3 — Registry consolidation (follow-up, moderate risk)
- Replace the three Maker classes with a single generic wrapper + lookup table:
{(Grid2D, "array"): Array2D, (Grid2DIrregular, "array"): ArrayIrregular, ...}
- Keep external API (
to_array, to_vector_yx, to_grid) as thin wrappers so downstream code doesn't change
- Run full test suites across PyAutoArray, PyAutoGalaxy, PyAutoLens
Key Files
autoarray/structures/decorators/abstract.py — base maker class, isinstance routingautoarray/structures/decorators/to_array.py — ArrayMakerautoarray/structures/decorators/to_grid.py — GridMaker (over-sampling complexity)autoarray/structures/decorators/to_vector_yx.py — VectorYXMakerautoarray/structures/decorators/transform.py — coordinate frame decoratortest_autoarray/test_decorators.py — currently near-empty (11 lines)autoarray/__init__.py — exports grid_dec
Original Prompt
Click to expand starting prompt
The different data structures in PyAutoArray are quite confusing, in some ways:
- Array2D
- ArrayIrregular
- Grid2D
- GridIrregular
And so on.
They serve an important purpose, unifying the API and abstractions in a way which ensures a user can understand
that data and grids can be paired to something uniform or not. Furthermore, the slam / native API is important
in making both accessible, and streamlining how masked data vectors are stored.
First, we should assess if this code can be simplified at all or if its acceptable. I dont have a better idea.
You would probably benefit from reading @autolens_workspace/scripts/guides/data_structures.py to see how a user
interfaces with these objects.
However, things get more complex, as these objects are used to define mappings at the profile level of
@PyAutoGalaxy/autogalaxy/profiles. For example for this function:
class Isothermal(PowerLaw):
def init(self, centre, ell_comps, einstein_radius):
super().init(centre=centre, ell_comps=ell_comps, einstein_radius=einstein_radius, slope=2.0)
@aa.grid_dec.to_vector_yx
@aa.grid_dec.transform
def deflections_yx_2d_from(self, grid, xp=np, **kwargs):
...
The decorator @aa.grid_dec.to_vector_yx is used to understand that for this function, the result
that comes out must be a VectorYX object. It handles more typing, for example if a Grid2DIrregular comes in
a VectorYXIrregular comes out, but a Grid2D produces a VectorYX object.
There is also behaviour where if a numpy ndarray comes in, a numpy ndarray comes out without type casting,
with the same behaviour for a JAX array.
The problem is really just how complex things got in this decorator, which is all handled at
@PyAutoArray/autoarray/structures/decorators. Its complicated, messy and hard to trace.
So, can you give me your assessment of whether there's a quite large, sweeping restructure that can simplify this
code but retain the desired functionality and behaviour? Think hard, this could require a good chunk of planning
and no doubt extensive testing after!
Overview
The decorator system at
autoarray/structures/decorators/handles input→output type mapping for all profile methods (e.g. Grid2D in → Array2D out, Grid2DIrregular → ArrayIrregular). It works correctly but is verbose (~648 lines), contains dead code (Grid1D path), duplicates logic across three Maker classes, and has near-zero dedicated test coverage. This task cleans it up in phases, starting with tests and targeted cleanup, with an optional follow-up to consolidate to a registry-based architecture.Plan
grid_dec→decoratorsin the public API (aa.grid_dec→aa.decorators) with grep-and-replace across PyAutoGalaxy/PyAutoLensDetailed implementation plan
Affected Repositories
Work Classification
Library
Branch Survey
Suggested branch:
feature/data-typing-simplifyWorktree root:
~/Code/PyAutoLabs-wt/data-typing-simplify/Implementation Steps
Phase 1 — Tests (zero-risk)
test_autoarray/test_decorators.pywith tests for:@to_arraywith Grid2D input → asserts Array2D output with correct mask@to_arraywith Grid2DIrregular input → asserts ArrayIrregular output@to_arraywith raw np.ndarray → asserts np.ndarray passthrough@to_vector_yxwith Grid2D → asserts VectorYX2D@to_vector_yxwith Grid2DIrregular → asserts VectorYX2DIrregular@to_gridwith Grid2D → asserts Grid2D with over-sampling metadata preserved@to_gridwith Grid2DIrregular → asserts Grid2DIrregular@transform— coordinate frame transformation applied and reversedPhase 2 — Targeted cleanup (low risk)
Grid1Disinstance branch andvia_grid_1d()from all three makers (to_array.py,to_grid.py,to_vector_yx.py) and the 1D projection logic fromabstract.pyabstract.pyGridMaker.via_grid_2d()into cleanergetattrcalls or into Grid2D's own constructiongrid_dec→decoratorsinautoarray/__init__.pyand update allaa.grid_dec.references in PyAutoGalaxy and PyAutoLensPhase 3 — Registry consolidation (follow-up, moderate risk)
to_array,to_vector_yx,to_grid) as thin wrappers so downstream code doesn't changeKey Files
autoarray/structures/decorators/abstract.py— base maker class, isinstance routingautoarray/structures/decorators/to_array.py— ArrayMakerautoarray/structures/decorators/to_grid.py— GridMaker (over-sampling complexity)autoarray/structures/decorators/to_vector_yx.py— VectorYXMakerautoarray/structures/decorators/transform.py— coordinate frame decoratortest_autoarray/test_decorators.py— currently near-empty (11 lines)autoarray/__init__.py— exportsgrid_decOriginal Prompt
Click to expand starting prompt
The different data structures in PyAutoArray are quite confusing, in some ways:
And so on.
They serve an important purpose, unifying the API and abstractions in a way which ensures a user can understand
that data and grids can be paired to something uniform or not. Furthermore, the slam / native API is important
in making both accessible, and streamlining how masked data vectors are stored.
First, we should assess if this code can be simplified at all or if its acceptable. I dont have a better idea.
You would probably benefit from reading @autolens_workspace/scripts/guides/data_structures.py to see how a user
interfaces with these objects.
However, things get more complex, as these objects are used to define mappings at the profile level of
@PyAutoGalaxy/autogalaxy/profiles. For example for this function:
class Isothermal(PowerLaw):
def init(self, centre, ell_comps, einstein_radius):
super().init(centre=centre, ell_comps=ell_comps, einstein_radius=einstein_radius, slope=2.0)
The decorator @aa.grid_dec.to_vector_yx is used to understand that for this function, the result
that comes out must be a VectorYX object. It handles more typing, for example if a Grid2DIrregular comes in
a VectorYXIrregular comes out, but a Grid2D produces a VectorYX object.
There is also behaviour where if a numpy ndarray comes in, a numpy ndarray comes out without type casting,
with the same behaviour for a JAX array.
The problem is really just how complex things got in this decorator, which is all handled at
@PyAutoArray/autoarray/structures/decorators. Its complicated, messy and hard to trace.
So, can you give me your assessment of whether there's a quite large, sweeping restructure that can simplify this
code but retain the desired functionality and behaviour? Think hard, this could require a good chunk of planning
and no doubt extensive testing after!