Skip to content

refactor: reduce magic in transform decorator #271

@Jammy2211

Description

@Jammy2211

Overview

The @aa.grid_dec.transform decorator reduces repeated code across ~58 profile methods but relies on implicit is_transformed kwargs and requires manual back-rotation for deflection vectors. This refactor makes the decorator less "magical" while keeping its DRY benefits, split into two phases.

Plan

  • Phase 1: Move is_transformed from kwargs onto a grid property, add decorator stacking validation, and update all profile methods to use the new pattern
  • Phase 2: Add optional automatic back-rotation for vector outputs to the transform decorator, eliminating manual rotated_grid_from_reference_frame_from calls in deflection methods. Validate with direct unit tests first before rolling out across all profiles.
  • Update the add_a_profile.py workspace guide to reflect the simplified API
Detailed implementation plan

Affected Repositories

  • PyAutoArray (primary)
  • PyAutoGalaxy
  • autolens_workspace

Work Classification

Library (both phases), then workspace guide update

Branch Survey

Repository Current Branch Dirty?
PyAutoArray main clean
PyAutoGalaxy main clean
autolens_workspace main clean

Suggested branch: feature/transform-decorator
Worktree root: ~/Code/PyAutoLabs-wt/transform-decorator/ (created later by /start_library)

Phase 1: Explicit transform state + validation

  1. Add is_transformed property to grid classes (autoarray/structures/grids/)

    • Add boolean attribute to Grid2D, Grid2DIrregular, and base grid classes
    • Ensure it's preserved through grid operations (slicing, copying)
  2. Update the transform decorator (autoarray/structures/decorators/transform.py)

    • Read grid.is_transformed instead of kwargs.get("is_transformed")
    • Set is_transformed = True on the transformed grid copy
    • Keep backward compat with kwargs during transition
  3. Add decorator stacking validation (autoarray/structures/decorators/abstract.py)

    • Tag decorators with priority; raise ValueError on wrong order
  4. Update profile methods in PyAutoGalaxy (autogalaxy/profiles/)

    • Remove is_transformed kwarg propagation where it only served this purpose
    • Verify all existing tests still pass

Phase 2: Automatic back-rotation for deflection vectors

  1. Write direct unit tests first against a few known profiles (e.g. Isothermal, NFW)

    • Compute deflections with current manual back-rotation
    • Compute deflections with new automatic back-rotation
    • Assert results match to numerical precision
  2. Enhance the transform decorator (autoarray/structures/decorators/transform.py)

    • Add parameter to signal that output is a vector requiring back-rotation
    • Implement inverse rotation using rotated_grid_from_reference_frame_from logic
    • The @aa.grid_dec.to_vector_yx decorator or a new combined decorator could handle this
  3. Update deflection methods (~15-20 methods across autogalaxy/profiles/mass/)

    • Remove manual rotated_grid_from_reference_frame_from calls
    • Use new decorator parameter for automatic back-rotation
  4. Update add_a_profile.py guide in autolens_workspace

    • Simplify the "Do I Rotate Back?" section
    • Show the new pattern

Key Files

  • autoarray/structures/decorators/transform.py — decorator implementation
  • autoarray/structures/decorators/to_vector_yx.py — vector output decorator
  • autoarray/structures/grids/ — grid classes for is_transformed property
  • autogalaxy/profiles/geometry_profiles.py — transform logic
  • autogalaxy/profiles/mass/total/isothermal.py — example profile to update
  • autolens_workspace/scripts/guides/advanced/add_a_profile.py — user guide

Original Prompt

Click to expand starting prompt

The transform decorator is used on nearly every mass and light profile in @PyAutoGalaxy/autogalaxy/profiles, for example: Isothermal(PowerLaw) with @aa.grid_dec.transform on deflections_yx_2d_from.

The decorator shifts the input grid to the profile centre and rotates by elliptical components. It has "magic" via is_transformed kwargs to prevent double-transforms. Users have reported not understanding where transforms are performed. The trade-off is DRY code vs readability/transparency.

Decision: keep the decorator but reduce the magic — move is_transformed to a grid property, add optional automatic back-rotation for vector outputs, and add decorator stacking validation.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type
    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions