Extend PyAutoLens / derived API for general use case

[Jammy2211]

In PyAutoLens, most calculations are performed using a Tracer object:

https://github.com/Jammy2211/PyAutoLens/blob/main/autolens/lens/ray_tracing.py

This is created within the log_likelihood_function:

https://github.com/Jammy2211/PyAutoLens/blob/main/autolens/analysis/analysis.py

    def tracer_via_instance_from(
        self,
        instance: af.ModelInstance,
        run_time_dict: Optional[Dict] = None,
        tracer_cls=Tracer,
    ) -> Tracer:
        """
        Create a `Tracer` from the galaxies contained in a model instance.

        If PyAutoFit's profiling tools are used with the analysis class, this function may receive a `run_time_dict`
        which times how long each set of the model-fit takes to perform.

        Parameters
        ----------
        instance
            An instance of the model that is fitted to the data by this analysis (whose parameters may have been set
            via a non-linear search).

        Returns
        -------
        Tracer
            An instance of the Tracer class that is used to then fit the dataset.
        """
        if hasattr(instance, "perturb"):
            instance.galaxies.subhalo = instance.perturb

        # TODO : Need to think about how we do this without building it into the model attribute names.
        # TODO : A Subhalo class that extends the Galaxy class maybe?

        if hasattr(instance.galaxies, "subhalo"):
            subhalo_centre = ray_tracing_util.grid_2d_at_redshift_from(
                galaxies=instance.galaxies,
                redshift=instance.galaxies.subhalo.redshift,
                grid=aa.Grid2DIrregular(values=[instance.galaxies.subhalo.mass.centre]),
                cosmology=self.cosmology,
            )

            instance.galaxies.subhalo.mass.centre = tuple(subhalo_centre.in_list[0])

        if hasattr(instance, "clumps"):
            return Tracer.from_galaxies(
                galaxies=instance.galaxies + instance.clumps,
                cosmology=self.cosmology,
                run_time_dict=run_time_dict,
            )

        if hasattr(instance, "cosmology"):
            cosmology = instance.cosmology
        else:
            cosmology = self.cosmology

        return tracer_cls.from_galaxies(
            galaxies=instance.galaxies,
            cosmology=cosmology,
            run_time_dict=run_time_dict,
        )

The majority of derived quantities we want to output are in the Tracer object, meaning this is where we want to put our derived_property decorators.

However, the Tracer is not used in the model composition API:

# Lens:

bulge = af.Model(al.lp.Sersic)

mass = af.Model(al.mp.Isothermal)

shear = af.Model(al.mp.ExternalShear)

lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)

# Source:

source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.Sersic)

# Overall Lens Model:

model = af.Collection(galaxies=af.Collection(lens=lens, source=source))

The question is, how do we implement derived quantities in autolens? There are two options:

1) The autolens API is updated to build models from a Tracer object. This is feasible, but I think would be quite confusing for a user, as they need to understand even more objects to build a model. The Tracer does some things internally with the model galaxies, for example ordering them in ascending redshift order into Plane objects, which may conflict with the model API.

2) Extend the derived PyAutoFit API to support objects no included in model composition (e.g. a custom method in the Analysis object). I'm sure this bring with it its own complexities in terms of source code implementation.

I guess a decent API for 1 would be something like:

# Lens:

bulge = af.Model(al.lp.Sersic)

mass = af.Model(al.mp.Isothermal)

shear = af.Model(al.mp.ExternalShear)

lens = af.Model(al.Galaxy, redshift=0.5, bulge=bulge, mass=mass, shear=shear)

# Source:

source = af.Model(al.Galaxy, redshift=1.0, bulge=al.lp.Sersic)

# Overall Lens Model:

tracer =  af.Model(Tracer, galaxies=af.Collection(lens=lens, source=source))

model = af.Collection(tracer=tracer)

It would remove all the stuff like if hasattr(instance.galaxies, "subhalo") from the tracer_via_instance_from so would be nicer.

[rhayes777]

To me it seems most natural to build the Tracer into the model - I think I actually suggested it a few years ago. In the above example there would be no need to create the final model = af.Collection(tracer) model as the model would be the af.Model(Tracer, ...)

For reporting purposes we could add custom derived quantities. For example:

model = af.Collection(
    lens=lens, 
    source=source, 
    derived_quantities={
        'tracer.quantity': lambda instance: Tracer(instance...).quantity
    }
)

This might not be terribly efficient

[Jammy2211]

The current ways to create a Tracer are:

class Tracer(ABC, ag.OperateImageGalaxies, ag.OperateDeflections):
    def __init__(
        self,
        planes,
        cosmology: ag.cosmo.LensingCosmology,
        run_time_dict: Optional[Dict] = None,
    ):
        """
        Ray-tracer for a lens system with any number of planes.

        The redshift of these planes are specified by the redshits of the galaxies; there is a unique plane redshift \
        for every unique galaxy redshift (galaxies with identical redshifts are put in the same plane).

        To perform multi-plane ray-tracing, a cosmology must be supplied so that deflection-angles can be rescaled \
        according to the lens-geometry of the multi-plane system. All galaxies input to the tracer must therefore \
        have redshifts.

        This tracer has only one grid (see gridStack) which is used for ray-tracing.

        Parameters
        ----------
        galaxies : [Galaxy]
            The list of galaxies in the ray-tracing calculation.
        image_plane_grid : grid_stacks.GridStack
            The image-plane grid which is traced. (includes the grid, sub-grid, blurring-grid, etc.).
        border : masks.GridBorder
            The border of the grid, which is used to relocate demagnified traced pixels to the \
            source-plane borders.
        cosmology : astropy.cosmology
            The cosmology of the ray-tracing calculation.
        """
        self.planes = planes
        self.plane_redshifts = [plane.redshift for plane in planes]
        self.cosmology = cosmology

        self.run_time_dict = run_time_dict

    @classmethod
    def from_galaxies(
        cls,
        galaxies,
        cosmology: ag.cosmo.LensingCosmology = ag.cosmo.Planck15(),
        run_time_dict: Optional[Dict] = None,
    ):
        planes = ag.util.plane.planes_via_galaxies_from(
            galaxies=galaxies, run_time_dict=run_time_dict
        )

        return cls(planes=planes, cosmology=cosmology, run_time_dict=run_time_dict)

I don't think we want to bring Plane objects into model composition API, as this would make it even more cluttered and is not valid as Plane's are determined by the Galaxy redshifts, which evidently could be free parameters which change.

We could make the input to the Tracer object a list of Galaxy objects, and have the planes attribute be a property of the Tracer. We could then use this object for the model API above.

This seems like the best option.

[rhayes777]

Yeah I think that makes sense