Skip to content

Commit

Permalink
Merge pull request #204 from Jammy2211/feature/jax_wrapper
Browse files Browse the repository at this point in the history 
Feature/jax wrapper
  • Loading branch information
@Jammy2211
Jammy2211 authored Oct 25, 2024
2 parents dc1cee8 + d199bc3 commit 1b53f4b
Show file tree
Hide file tree
Showing 54 changed files with 1,176 additions and 1,586 deletions.
5 changes: 2 additions & 3 deletions autogalaxy/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
from autoarray.dataset.interferometer.dataset import Interferometer # noqa
from autoarray.dataset.dataset_model import DatasetModel
from autoarray.dataset.over_sampling import OverSamplingDataset
from autoarray.inversion.inversion.mapper_valued import MapperValued
from autoarray.inversion.pixelization import mesh # noqa
from autoarray.inversion import regularization as reg # noqa
from autoarray.inversion.pixelization import image_mesh
Expand Down Expand Up @@ -53,8 +52,6 @@

from .analysis.adapt_images.adapt_images import AdaptImages
from .analysis.adapt_images.adapt_image_maker import AdaptImageMaker
from .analysis.maker import FitMaker
from .analysis.preloads import Preloads
from . import aggregator as agg
from . import exc
from . import plot
Expand Down Expand Up @@ -111,6 +108,8 @@
from .gui.clicker import Clicker
from .gui.scribbler import Scribbler

from .analysis.clump_model import ClumpModel

from autoconf import conf

conf.instance.register(__file__)
Expand Down
99 changes: 0 additions & 99 deletions autogalaxy/aggregator/agg_util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -68,102 +68,3 @@ def adapt_images_from(
adapt_images_list.append(adapt_images)

return adapt_images_list


def mesh_grids_of_planes_list_from(
fit: af.Fit, total_fits: int, use_preloaded_grid: bool
) -> List[Optional[aa.Grid2D]]:
"""
Returns the image-plane pixelization grid(s) used by the fit.
A subset of image-mesh objects (e.g. `Hilbert`, `KMeans`) create the grid of points that act as the mesh
centres (e.g. the centers of Voronoi cells) in the image-plane. For lensing calculations this may then be
traced to the source-plane to form the pixelization.
This calculation can depend on the library used to determine the image-plane grid and may have a random element
associated with it. This means that performing an analysis on a super computer and then downloading the results
for inspection on your laptop may produce different image-plane grids, changing the result and quantities like
the `log_likelihood`.
By storing this grid as a result in the `files` folder and loading it via the database before creating a fit
this possible mismatch is removed, ensuring results on a local computer are identical to those computer elsewhere.
Parameters
----------
fit
A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry in a sqlite database.
total_fits
The total number of `Analysis` objects summed to create the fit.
use_preloaded_grid
Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
as the fit.
Returns
-------
The list of image-plane mesh centres used when creating the overall fit.
"""

if use_preloaded_grid:
if not fit.children:
return [fit.value(name="preload_mesh_grids_of_planes")]
else:
try:
return fit.child_values(name="preload_mesh_grids_of_planes")
except AttributeError:
return [None] * total_fits
else:
return [None] * total_fits


def preloads_from(
preloads_cls,
use_preloaded_grid: bool,
mesh_grids_of_planes: List,
use_w_tilde: Optional[bool] = False,
) -> aa.Preloads:
"""
Returns a `Preloads` object associated with a fit loaded via the database.
When loading results via the database, the preloads class may have certain attributes associated with it
in order to perform the fit. The main purpose is to use the same image-plane mesh centres for pixelization where
the mesh is computed in the image-plane (see `agg_util.mesh_grids_of_planes_list_from`).
The preloads may also switch off `w_tilde` so fits are computed faster locally as they do not need to recompute
w_tilde.
Parameters
----------
preloads_cls
The `Preloads` object used to create the preloads (this varies across
projects like `autogalaxy` and `autolens`).
use_preloaded_grid
Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
as the fit.
mesh_grids_of_planes
The list of image-plane mesh centres used when creating the overall fit which are associated with the
preloads.
use_w_tilde
Whether to use the w-tilde formalism when recomputing fits.
Returns
-------
The preloads object associated with the fit loaded via the database.
"""
preloads = preloads_cls()

if use_preloaded_grid:
if mesh_grids_of_planes is not None:
preloads = preloads_cls(
image_plane_mesh_grid_pg_list=mesh_grids_of_planes,
use_w_tilde=use_w_tilde,
)

if len(preloads.image_plane_mesh_grid_pg_list) == 2:
if type(preloads.image_plane_mesh_grid_pg_list[1]) != list:
preloads.image_plane_mesh_grid_pg_list[1] = [
preloads.image_plane_mesh_grid_pg_list[1]
]

return preloads
1 change: 1 addition & 0 deletions autogalaxy/aggregator/ellipse/fit_ellipse.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,7 @@ def _fit_ellipse_from(
for dataset, ellipse_list, multipole_lists in zip(
dataset_list, ellipse_list_list, multipole_list_list
):

for ellipse, multipole_list in zip(ellipse_list, multipole_lists):
fit_dataset_list.append(
FitEllipse(
Expand Down
8 changes: 4 additions & 4 deletions autogalaxy/aggregator/galaxies.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -41,14 +41,14 @@ def _galaxies_from(fit: af.Fit, instance: af.ModelInstance) -> List[Galaxy]:
if instance is not None:
galaxies = instance.galaxies

if hasattr(instance, "extra_galaxies"):
galaxies = galaxies + fit.instance.extra_galaxies
if hasattr(instance, "clumps"):
galaxies = galaxies + fit.instance.clumps

else:
galaxies = fit.instance.galaxies

if hasattr(fit.instance, "extra_galaxies"):
galaxies = galaxies + fit.instance.extra_galaxies
if hasattr(fit.instance, "clumps"):
galaxies = galaxies + fit.instance.clumps

if fit.children is not None:
if len(fit.children) > 0:
Expand Down
28 changes: 1 addition & 27 deletions autogalaxy/aggregator/imaging/fit_imaging.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@
import autofit as af
import autoarray as aa

from autogalaxy.analysis.preloads import Preloads

from autogalaxy.aggregator.imaging.imaging import _imaging_from
from autogalaxy.aggregator.galaxies import _galaxies_from
from autogalaxy.aggregator.dataset_model import _dataset_model_from
Expand All @@ -19,7 +17,6 @@ def _fit_imaging_from(
fit: af.Fit,
instance: Optional[af.ModelInstance] = None,
settings_inversion: aa.SettingsInversion = None,
use_preloaded_grid: bool = True,
) -> List[FitImaging]:
"""
Returns a list of `FitImaging` objects from a `PyAutoFit` sqlite database `Fit` object.
Expand Down Expand Up @@ -51,10 +48,6 @@ def _fit_imaging_from(
randomly from the PDF).
settings_inversion
Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
use_preloaded_grid
Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
as the fit.
"""

from autogalaxy.imaging.fit_imaging import FitImaging
Expand All @@ -69,25 +62,14 @@ def _fit_imaging_from(

settings_inversion = settings_inversion or fit.value(name="settings_inversion")

mesh_grids_of_planes_list = agg_util.mesh_grids_of_planes_list_from(
fit=fit, total_fits=len(dataset_list), use_preloaded_grid=use_preloaded_grid
)

fit_dataset_list = []

for dataset, galaxies, dataset_model, adapt_images, mesh_grids_of_planes in zip(
for dataset, galaxies, dataset_model, adapt_images in zip(
dataset_list,
galaxies_list,
dataset_model_list,
adapt_images_list,
mesh_grids_of_planes_list,
):
preloads = agg_util.preloads_from(
preloads_cls=Preloads,
use_preloaded_grid=use_preloaded_grid,
mesh_grids_of_planes=mesh_grids_of_planes,
use_w_tilde=False,
)

fit_dataset_list.append(
FitImaging(
Expand All @@ -96,7 +78,6 @@ def _fit_imaging_from(
dataset_model=dataset_model,
adapt_images=adapt_images,
settings_inversion=settings_inversion,
preloads=preloads,
)
)

Expand All @@ -108,7 +89,6 @@ def __init__(
self,
aggregator: af.Aggregator,
settings_inversion: Optional[aa.SettingsInversion] = None,
use_preloaded_grid: bool = True,
):
"""
Interfaces with an `PyAutoFit` aggregator object to create instances of `FitImaging` objects from the results
Expand Down Expand Up @@ -142,15 +122,10 @@ def __init__(
A `PyAutoFit` aggregator object which can load the results of model-fits.
settings_inversion
Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
use_preloaded_grid
Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
used as the fit.
"""
super().__init__(aggregator=aggregator)

self.settings_inversion = settings_inversion
self.use_preloaded_grid = use_preloaded_grid

def object_via_gen_from(
self, fit, instance: Optional[af.ModelInstance] = None
Expand All @@ -172,5 +147,4 @@ def object_via_gen_from(
fit=fit,
instance=instance,
settings_inversion=self.settings_inversion,
use_preloaded_grid=self.use_preloaded_grid,
)
28 changes: 1 addition & 27 deletions autogalaxy/aggregator/interferometer/fit_interferometer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,6 @@
import autofit as af
import autoarray as aa

from autogalaxy.analysis.preloads import Preloads

from autogalaxy.aggregator import agg_util
from autogalaxy.aggregator.interferometer.interferometer import _interferometer_from
from autogalaxy.aggregator.dataset_model import _dataset_model_from
Expand All @@ -20,7 +18,6 @@ def _fit_interferometer_from(
instance: Optional[af.ModelInstance] = None,
real_space_mask: Optional[aa.Mask2D] = None,
settings_inversion: aa.SettingsInversion = None,
use_preloaded_grid: bool = True,
) -> List[FitInterferometer]:
"""
Returns a list of `FitInterferometer` objects from a `PyAutoFit` sqlite database `Fit` object.
Expand Down Expand Up @@ -53,10 +50,6 @@ def _fit_interferometer_from(
randomly from the PDF).
settings_inversion
Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
use_preloaded_grid
Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This grid
may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is used
as the fit.
"""
from autogalaxy.interferometer.fit_interferometer import FitInterferometer

Expand All @@ -73,25 +66,14 @@ def _fit_interferometer_from(

settings_inversion = settings_inversion or fit.value(name="settings_inversion")

mesh_grids_of_planes_list = agg_util.mesh_grids_of_planes_list_from(
fit=fit, total_fits=len(dataset_list), use_preloaded_grid=use_preloaded_grid
)

fit_dataset_list = []

for dataset, galaxies, dataset_model, adapt_images, mesh_grids_of_planes in zip(
for dataset, galaxies, dataset_model, adapt_images in zip(
dataset_list,
galaxies_list,
dataset_model_list,
adapt_images_list,
mesh_grids_of_planes_list,
):
preloads = agg_util.preloads_from(
preloads_cls=Preloads,
use_preloaded_grid=use_preloaded_grid,
mesh_grids_of_planes=mesh_grids_of_planes,
use_w_tilde=False,
)

fit_dataset_list.append(
FitInterferometer(
Expand All @@ -100,7 +82,6 @@ def _fit_interferometer_from(
dataset_model=dataset_model,
adapt_images=adapt_images,
settings_inversion=settings_inversion,
preloads=preloads,
)
)

Expand All @@ -112,7 +93,6 @@ def __init__(
self,
aggregator: af.Aggregator,
settings_inversion: Optional[aa.SettingsInversion] = None,
use_preloaded_grid: bool = True,
real_space_mask: Optional[aa.Mask2D] = None,
):
"""
Expand Down Expand Up @@ -147,15 +127,10 @@ def __init__(
A `PyAutoFit` aggregator object which can load the results of model-fits.
settings_inversion
Optionally overwrite the `SettingsInversion` of the `Inversion` object that is created from the fit.
use_preloaded_grid
Certain pixelization's construct their mesh in the source-plane from a stochastic KMeans algorithm. This
grid may be output to hard-disk after the model-fit and loaded via the database to ensure the same grid is
used as the fit.
"""
super().__init__(aggregator=aggregator)

self.settings_inversion = settings_inversion
self.use_preloaded_grid = use_preloaded_grid
self.real_space_mask = real_space_mask

def object_via_gen_from(
Expand All @@ -178,5 +153,4 @@ def object_via_gen_from(
fit=fit,
instance=instance,
settings_inversion=self.settings_inversion,
use_preloaded_grid=self.use_preloaded_grid,
)
Loading

0 comments on commit 1b53f4b

Please sign in to comment.