Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
56 commits
Select commit Hold shift + click to select a range
ba4b4b5
jax wrapper
Jammy2211 Apr 29, 2024
06ed3cb
Merge branch 'main' into feature/jax_wrapper
Jammy2211 Jun 4, 2024
75bbca6
fix test (once conf merged)
rhayes777 Jun 10, 2024
d93a5cc
hack-around to get jax script running
rhayes777 Jun 10, 2024
252ecfd
revert array casting change
rhayes777 Jun 10, 2024
7f0bcec
merge
Jammy2211 Jun 25, 2024
044febb
When using JAX don't import Numpy
CKrawczyk Jun 28, 2024
218d500
gMerge branch 'main' into feature/jax_wrapper
Jammy2211 Jul 19, 2024
ad84a98
merge main
Jammy2211 Jul 19, 2024
2f836de
Merge pull request #180 from Jammy2211/feature/jaxify_geometry_profiles
CKrawczyk Aug 6, 2024
9a477fb
merge
rhayes777 Aug 12, 2024
9444b1e
numpy to jax wrapper conversion
rhayes777 Aug 12, 2024
17ee713
Merge branch 'main' into feature/jax_wrapper
rhayes777 Aug 23, 2024
5e30644
Jax version of MGE Sersic profile
CKrawczyk Sep 18, 2024
3210f4f
Update isothermal and power law profiles to play nice at (0, 0)
CKrawczyk Sep 18, 2024
f8418bd
Merge branch 'main' into feature/jax_wrapper
rhayes777 Sep 30, 2024
32e491c
Small fixes
CKrawczyk Oct 3, 2024
e6540a2
Merge branch 'feature/jax_wrapper' into feature/jaxify_geometry_profiles
CKrawczyk Oct 3, 2024
68d3371
Update mge_numpy.py
CKrawczyk Oct 3, 2024
0e97b8a
Merge pull request #196 from Jammy2211/feature/jaxify_geometry_profiles
CKrawczyk Oct 3, 2024
cec8dee
Merge branch 'main' into feature/jax_wrapper
Jammy2211 Oct 3, 2024
f7600b3
removed preloads
Jammy2211 Oct 3, 2024
9cd4e70
black
Jammy2211 Oct 3, 2024
c322507
restore fit_from
Jammy2211 Oct 3, 2024
033fc09
Update Faddeeva function
CKrawczyk Oct 25, 2024
7da67fe
Merge pull request #197 from Jammy2211/feature/jax_remove_preload
Jammy2211 Oct 25, 2024
ed7c379
Merge branch 'feature/jax_wrapper' into main
Jammy2211 Oct 25, 2024
6c1f19e
Merge pull request #201 from Jammy2211/main
Jammy2211 Oct 25, 2024
b6d157c
Update functions for Tracer to be jax safe
CKrawczyk Oct 25, 2024
686ddad
Revert "Main to JAX"
Jammy2211 Oct 25, 2024
d199bc3
Merge pull request #203 from Jammy2211/revert-201-main
Jammy2211 Oct 25, 2024
860d142
Revert "Revert "Main to JAX""
Jammy2211 Oct 25, 2024
8e4d556
Merge pull request #205 from Jammy2211/revert-203-revert-201-main
Jammy2211 Oct 25, 2024
e591a26
Merge pull request #202 from Jammy2211/feature/jax_tracer
CKrawczyk Oct 25, 2024
8dc5e47
Merge branch 'feature/jax_wrapper' into feature/update_faddeeva_function
CKrawczyk Oct 25, 2024
01b1742
Merge pull request #200 from Jammy2211/feature/update_faddeeva_function
CKrawczyk Oct 25, 2024
8d82399
Use auto diff to get at jacobian of deflection angles
CKrawczyk Oct 25, 2024
758acbb
Merge pull request #206 from Jammy2211/feature/jax_jacobian
CKrawczyk Oct 25, 2024
f5b24a0
Calculate critical curves using JAX
CKrawczyk Nov 4, 2024
da3d13c
Merge pull request #207 from Jammy2211/feature/jax_critical_curves
CKrawczyk Nov 4, 2024
5e8a273
config changes
Jammy2211 Jan 22, 2025
2b5b461
Merge branch 'main' into feature/jax_merge
Jammy2211 Jan 22, 2025
7f16637
merge complete
Jammy2211 Jan 22, 2025
df9f901
merge more changes from main
Jammy2211 Jan 22, 2025
525b7de
more changes
Jammy2211 Jan 22, 2025
a35e209
minor
Jammy2211 Mar 13, 2025
bb79b1e
Merge branch 'main' into feature/jax_merge
Jammy2211 Mar 13, 2025
cb2f01c
Merge pull request #215 from Jammy2211/feature/jax_merge
Jammy2211 Mar 13, 2025
904683b
resolve conflicts
Jammy2211 Mar 13, 2025
9281c4d
updates to fix unit tests
Jammy2211 Apr 2, 2025
58cf1f1
Merge branch 'main' into feature/jax_wrapper
Jammy2211 Apr 2, 2025
dc7f1e1
remove convolver for psf
Jammy2211 Apr 2, 2025
545fc09
black
Jammy2211 Apr 2, 2025
fa34dba
Merge pull request #223 from Jammy2211/feature/remove_convolver
Jammy2211 Apr 3, 2025
82622f7
remove grid relocate radial
Jammy2211 Apr 3, 2025
2ace7ea
pylops docs remove
Jammy2211 Apr 4, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions autogalaxy/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,6 @@
from autoarray.inversion.pixelization.border_relocator import BorderRelocator
from autoarray.mask.mask_1d import Mask1D # noqa
from autoarray.mask.mask_2d import Mask2D # noqa
from autoarray.operators.convolver import Convolver # noqa
from autoarray.operators.convolver import Convolver # noqa
from autoarray.operators.transformer import TransformerDFT # noqa
from autoarray.operators.transformer import TransformerNUFFT # noqa
from autoarray.layout.layout import Layout2D # noqa
Expand All @@ -48,8 +46,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
37 changes: 9 additions & 28 deletions autogalaxy/aggregator/imaging/fit_imaging.py
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` loaded directory `Fit` or sqlite database `Fit` object.
Expand Down Expand Up @@ -53,10 +50,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 @@ -71,25 +64,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 @@ -98,7 +80,6 @@ def _fit_imaging_from(
dataset_model=dataset_model,
adapt_images=adapt_images,
settings_inversion=settings_inversion,
preloads=preloads,
)
)

Expand All @@ -110,7 +91,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 @@ -145,15 +125,17 @@ 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.

Parameters
----------
aggregator
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.
"""
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 @@ -167,7 +149,7 @@ def object_via_gen_from(
----------
fit
A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
an output directory or from an sqlite database..
an output directory or from an sqlite database..
instance
A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
randomly from the PDF).
Expand All @@ -176,5 +158,4 @@ def object_via_gen_from(
fit=fit,
instance=instance,
settings_inversion=self.settings_inversion,
use_preloaded_grid=self.use_preloaded_grid,
)
83 changes: 32 additions & 51 deletions autogalaxy/aggregator/interferometer/fit_interferometer.py
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 @@ -19,7 +17,6 @@ def _fit_interferometer_from(
fit: af.Fit,
instance: Optional[af.ModelInstance] = None,
settings_inversion: aa.SettingsInversion = None,
use_preloaded_grid: bool = True,
) -> List[FitInterferometer]:
"""
Returns a list of `FitInterferometer` objects from a `PyAutoFit` loaded directory `Fit` or sqlite database `Fit` object.
Expand Down Expand Up @@ -54,10 +51,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,50 +93,51 @@ 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 `FitInterferometer` objects from the
results of a model-fit.
Interfaces with an `PyAutoFit` aggregator object to create instances of `FitInterferometer` objects from the
results of a model-fit.

The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
The results of a model-fit can be loaded from hard-disk or stored in a sqlite database, including the following
attributes of the fit:

- The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
- The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
- The settings of inversions used by the fit (`dataset/settings_inversion.json`).
- The interferometer data, noise-map, uv-wavelengths and settings as .fits files (e.g. `dataset/data.fits`).
- The real space mask defining the grid of the interferometer for the FFT (`dataset/real_space_mask.fits`).
- The settings of inversions used by the fit (`dataset/settings_inversion.json`).

The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
can load them all at once and create an `FitInterferometer` object via the `_fit_interferometer_from` method.

The `aggregator` contains the path to each of these files, and they can be loaded individually. This class
can load them all at once and create an `FitInterferometer` object via the `_fit_interferometer_from` method.
This class's methods returns generators which create the instances of the `FitInterferometer` objects. This ensures
that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
`FitInterferometer` instances in the memory at once.

This class's methods returns generators which create the instances of the `FitInterferometer` objects. This ensures
that large sets of results can be efficiently loaded from the hard-disk and do not require storing all
`FitInterferometer` instances in the memory at once.
For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
creates instances of the corresponding 3 `FitInterferometer` objects.

For example, if the `aggregator` contains 3 model-fits, this class can be used to create a generator which
creates instances of the corresponding 3 `FitInterferometer` objects.
If multiple `Imaging` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
`FitImaging` objects.

If multiple `Imaging` objects were fitted simultaneously via analysis summing, the `fit.child_values()` method
is instead used to load lists of the data, noise-map, PSF and mask and combine them into a list of
`FitImaging` objects.
This can be done manually, but this object provides a more concise API.

This can be done manually, but this object provides a more concise API.
Parameters
----------
aggregator
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.

Parameters
----------
aggregator
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.
Parameters
----------
aggregator
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.
"""
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 @@ -169,7 +151,7 @@ def object_via_gen_from(
----------
fit
A `PyAutoFit` `Fit` object which contains the results of a model-fit as an entry which has been loaded from
an output directory or from an sqlite database..
an output directory or from an sqlite database..
instance
A manual instance that overwrites the max log likelihood instance in fit (e.g. for drawing the instance
randomly from the PDF).
Expand All @@ -178,5 +160,4 @@ def object_via_gen_from(
fit=fit,
instance=instance,
settings_inversion=self.settings_inversion,
use_preloaded_grid=self.use_preloaded_grid,
)
32 changes: 16 additions & 16 deletions autogalaxy/analysis/analysis/analysis.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ def profile_log_likelihood_function(
An `info_dict` is also created which stores information on aspects of the model and dataset that dictate
run times, so the profiled times can be interpreted with this context.

The results of this profiling are then output to hard-disk in the `preloads` folder of the model-fit results,
The results of this profiling are then output to hard-disk in the `profiling` folder of the model-fit results,
which they can be inspected to ensure run-times are as expected.

Parameters
Expand Down Expand Up @@ -192,33 +192,33 @@ def profile_log_likelihood_function(

try:
info_dict["image_pixels"] = self.dataset.grids.lp.shape_slim
info_dict[
"sub_total_light_profiles"
] = self.dataset.grids.lp.over_sampler.sub_total
info_dict["sub_total_light_profiles"] = (
self.dataset.grids.lp.over_sampler.sub_total
)
except AttributeError:
pass

if fit.model_obj.has(cls=aa.Pixelization):
info_dict["use_w_tilde"] = fit.inversion.settings.use_w_tilde
try:
info_dict[
"sub_total_pixelization"
] = self.dataset.grids.pixelization.over_sampler.sub_total
info_dict["sub_total_pixelization"] = (
self.dataset.grids.pixelization.over_sampler.sub_total
)
except AttributeError:
pass
info_dict[
"use_positive_only_solver"
] = fit.inversion.settings.use_positive_only_solver
info_dict[
"force_edge_pixels_to_zeros"
] = fit.inversion.settings.force_edge_pixels_to_zeros
info_dict["use_positive_only_solver"] = (
fit.inversion.settings.use_positive_only_solver
)
info_dict["force_edge_pixels_to_zeros"] = (
fit.inversion.settings.force_edge_pixels_to_zeros
)
info_dict["use_w_tilde_numpy"] = fit.inversion.settings.use_w_tilde_numpy
info_dict["source_pixels"] = len(fit.inversion.reconstruction)

if hasattr(fit.inversion, "w_tilde"):
info_dict[
"w_tilde_curvature_preload_size"
] = fit.inversion.w_tilde.curvature_preload.shape[0]
info_dict["w_tilde_curvature_preload_size"] = (
fit.inversion.w_tilde.curvature_preload.shape[0]
)

self.output_profiling_info(
paths=paths, run_time_dict=run_time_dict, info_dict=info_dict
Expand Down
Loading
Loading