Skip to content

Feature/cosmology jax #269

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Jammy2211 merged 28 commits into from
Jan 30, 2026
Merged

Feature/cosmology jax #269

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
28 commits
Select commit Hold shift + click to select a range
9822561
combining into single module and added method angular_diameter_distan…
Jan 28, 2026
638c0b5
arcsec_tokpc functions added with jax support
Jan 28, 2026
9fabde4
test__arcsec_to_kpc_conversion
Jan 28, 2026
361fdc7
test__angular_diameter_distances
Jan 28, 2026
dea0e8e
test__angular_diameter_distances
Jan 28, 2026
8af5b24
test__angular_diameter_distances -> make numeric
Jan 28, 2026
5ad37c5
cosmic_average_density_from
Jan 29, 2026
2b00bb8
test__critical_surface_mass_densities
Jan 29, 2026
dadf5c4
added test__scaling_factor_between_redshifts_fro, doesnt pass yet
Jan 29, 2026
a00208c
test__scaling_factor_between_redshifts_from
Jan 29, 2026
049d7ed
cosmology model all works
Jan 29, 2026
dc8a0c5
added critical_density method
Jan 30, 2026
3c0b102
remove astropy units in dark calculations for mcr
Jan 30, 2026
39d3c9c
lots of fixes to make tests in test__values_of_quantities_for_real_co…
Jan 30, 2026
3b70800
fix skip
Jan 30, 2026
7548dc9
more units removal
Jan 30, 2026
43bedee
remove another multiupole comps
Jan 30, 2026
37d2e59
remove astropy from mock cosmology
Jan 30, 2026
9151b99
black
Jan 30, 2026
78e771b
minor
Jan 30, 2026
3cf4f81
Update autogalaxy/cosmology/model.py
Jammy2211 Jan 30, 2026
50636a0
Update autogalaxy/cosmology/model.py
Jammy2211 Jan 30, 2026
bedcf9a
Initial plan
Copilot Jan 30, 2026
562d1b9
Update autogalaxy/cosmology/model.py
Jammy2211 Jan 30, 2026
0f13b2f
Update FlatLambdaCDM docstring to reflect JAX-compatible implementation
Copilot Jan 30, 2026
1042a71
Remove trailing whitespace from docstring
Copilot Jan 30, 2026
7f6ff50
Merge pull request #270 from Jammy2211/copilot/sub-pr-269
Jammy2211 Jan 30, 2026
bd4b579
Merge branch 'main' into feature/cosmology_jax
Jammy2211 Jan 30, 2026
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
1 change: 1 addition & 0 deletions autogalaxy/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from autoconf.dictable import from_dict, from_json, output_to_json, to_dict
from autoarray.dataset import preprocess # noqa

# from autoarray.dataset.interferometer.w_tilde import (
# load_curvature_preload_if_compatible,
# )
Expand Down
4 changes: 2 additions & 2 deletions autogalaxy/analysis/analysis/analysis.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from autogalaxy.galaxy.galaxy import Galaxy
from autogalaxy.galaxy.galaxies import Galaxies
from autogalaxy.cosmology.lensing import LensingCosmology
from autogalaxy.cosmology.model import LensingCosmology

logger = logging.getLogger(__name__)

Expand Down Expand Up @@ -36,7 +36,7 @@ def __init__(
The Cosmology assumed for this analysis.
"""

from autogalaxy.cosmology.wrap import Planck15
from autogalaxy.cosmology.model import Planck15

self.cosmology = cosmology or Planck15()
self.preloads = preloads
Expand Down
2 changes: 1 addition & 1 deletion autogalaxy/analysis/analysis/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import autoarray as aa

from autogalaxy.analysis.adapt_images.adapt_images import AdaptImages
from autogalaxy.cosmology.lensing import LensingCosmology
from autogalaxy.cosmology.model import LensingCosmology
from autogalaxy.analysis.analysis.analysis import Analysis
from autogalaxy.analysis.result import ResultDataset

Expand Down
46 changes: 1 addition & 45 deletions autogalaxy/config/priors/cosmology.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,54 +1,10 @@
model.LambdaCDMWrap:
model.FlatLambdaCDM:
H0:
type: Constant
value: 67.66
Om0:
type: Constant
value: 0.30966
Ode0:
type: Constant
value: 0.69034
Tcmb0:
type: Constant
value: 2.7255
Neff:
type: Constant
value: 3.046
m_nu:
type: Constant
value: 0.06
Ob0:
type: Constant
value: 0.04897
model.FlatLambdaCDMWrap:
H0:
type: Constant
value: 67.66
Om0:
type: Constant
value: 0.30966
Tcmb0:
type: Constant
value: 2.7255
Neff:
type: Constant
value: 3.046
m_nu:
type: Constant
value: 0.06
Ob0:
type: Constant
value: 0.04897
model.FlatwCDMWrap:
H0:
type: Constant
value: 67.66
Om0:
type: Constant
value: 0.30966
w0:
type: Constant
value: -1.0
Tcmb0:
type: Constant
value: 2.7255
Expand Down
11 changes: 7 additions & 4 deletions autogalaxy/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,9 +331,12 @@ def multipole_comps_from(
-------
The multipole component parameters.
"""
from astropy import units
# Degrees -> radians conversion
deg_to_rad = xp.asarray(np.pi / 180.0)

multipole_comp_0 = k_m * xp.sin(phi_m * float(m) * units.deg.to(units.rad))
multipole_comp_1 = k_m * xp.cos(phi_m * float(m) * units.deg.to(units.rad))
angle = phi_m * float(m) * deg_to_rad

return (multipole_comp_0, multipole_comp_1)
multipole_comp_0 = k_m * xp.sin(angle)
multipole_comp_1 = k_m * xp.cos(angle)

return multipole_comp_0, multipole_comp_1
4 changes: 1 addition & 3 deletions autogalaxy/cosmology/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1 @@
from .lensing import LensingCosmology
from .wrap import Planck15
from .model import FlatwCDMWrap, FlatLambdaCDMWrap
from .model import FlatLambdaCDM, Planck15, LensingCosmology
Loading
Loading