Skip to content

Commit 9c26997

Browse files
b-remybeckermr
b-remy
and
beckermr
authored
add to_galsim for Image class, demo1.py, & demo2.py (#163)
* dev hsm * ShapeData running * demo1.py runnning * adding demo1.py demo2.py * update tests * apply back & ruff * fix test_hsm.py erros * import jax_galsim examples/demo*.py * add to-from image.py * add FindAdaptiveMom to Image & remove hsm.py * clean branch and add from-to-galsim test * enable image to-from-galsim test * remove demo2.py * add again hsm allowed failures * fix: update test submodule * test: add symlink for good measure * fix: more fixes for test suite * fix: update submodule for more test fixes * fix: convert args and kwargs if needed * add demo2.py --------- Co-authored-by: Matthew R. Becker <beckermr@users.noreply.github.com> Co-authored-by: beckermr <becker.mr@gmail.com>
1 parent 89b9d56 commit 9c26997

8 files changed

Lines changed: 360 additions & 3 deletions

File tree

‎examples/demo1.py‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
# Copyright (c) 2012-2026 by the GalSim developers team on GitHub
2+
# https://github.com/GalSim-developers
3+
#
4+
# This file is part of GalSim: The modular galaxy image simulation toolkit.
5+
# https://github.com/GalSim-developers/GalSim
6+
#
7+
# GalSim is free software: redistribution and use in source and binary forms,
8+
# with or without modification, are permitted provided that the following
9+
# conditions are met:
10+
#
11+
# 1. Redistributions of source code must retain the above copyright notice, this
12+
# list of conditions, and the disclaimer given in the accompanying LICENSE
13+
# file.
14+
# 2. Redistributions in binary form must reproduce the above copyright notice,
15+
# this list of conditions, and the disclaimer given in the documentation
16+
# and/or other materials provided with the distribution.
17+
#
18+
"""
19+
Demo #1
20+
21+
This is the first script in our tutorial about using JAX-GalSim in python scripts: examples/demo*.py.
22+
(This file is designed to be viewed in a window 100 characters wide.)
23+
24+
Each of these demo*.py files are designed to be equivalent to the corresponding demo*.yaml file
25+
(or demo*.json -- found in the json directory). If you are new to python, you should probably
26+
look at those files first as they will probably have a quicker learning curve for you. Then you
27+
can look through these python scripts, which show how to do the same thing. Of course, experienced
28+
pythonistas may prefer to start with these scripts and then look at the corresponding YAML files.
29+
30+
To run this script, simply write:
31+
32+
python demo1.py
33+
34+
35+
This first script is about as simple as it gets. We draw an image of a single galaxy convolved
36+
with a PSF and write it to disk. We use a circular Gaussian profile for both the PSF and the
37+
galaxy, and add a constant level of Gaussian noise to the image.
38+
39+
In each demo, we list the new features introduced in that demo file. These will differ somewhat
40+
between the .py and .yaml (or .json) versions, since the two methods implement things in different
41+
ways. (demo*.py are python scripts, while demo*.yaml and demo*.json are configuration files.)
42+
43+
New features introduced in this demo:
44+
45+
- obj = jax_galsim.Gaussian(flux, sigma)
46+
- obj = jax_galsim.Convolve([list of objects])
47+
- image = obj.drawImage(scale)
48+
- image.added_flux (Only present after a drawImage command.)
49+
- noise = jax_galsim.GaussianNoise(sigma)
50+
- image.addNoise(noise)
51+
- image.write(file_name)
52+
- image.FindAdaptiveMom()
53+
"""
54+
55+
import logging
56+
import math
57+
import os
58+
import sys
59+
60+
import jax_galsim
61+
62+
63+
def main(argv):
64+
"""
65+
About as simple as it gets:
66+
- Use a circular Gaussian profile for the galaxy.
67+
- Convolve it by a circular Gaussian PSF.
68+
- Add Gaussian noise to the image.
69+
"""
70+
# In non-script code, use getLogger(__name__) at module scope instead.
71+
logging.basicConfig(format="%(message)s", level=logging.INFO, stream=sys.stdout)
72+
logger = logging.getLogger("demo1")
73+
74+
gal_flux = 1.0e5 # total counts on the image
75+
gal_sigma = 2.0 # arcsec
76+
psf_sigma = 1.0 # arcsec
77+
pixel_scale = 0.2 # arcsec / pixel
78+
noise = 30.0 # standard deviation of the counts in each pixel
79+
80+
logger.info("Starting demo script 1 using:")
81+
logger.info(
82+
" - circular Gaussian galaxy (flux = %.1e, sigma = %.1f),",
83+
gal_flux,
84+
gal_sigma,
85+
)
86+
logger.info(" - circular Gaussian PSF (sigma = %.1f),", psf_sigma)
87+
logger.info(" - pixel scale = %.2f,", pixel_scale)
88+
logger.info(" - Gaussian noise (sigma = %.2f).", noise)
89+
90+
# Define the galaxy profile
91+
gal = jax_galsim.Gaussian(flux=gal_flux, sigma=gal_sigma)
92+
logger.debug("Made galaxy profile")
93+
94+
# Define the PSF profile
95+
psf = jax_galsim.Gaussian(flux=1.0, sigma=psf_sigma) # PSF flux should always = 1
96+
logger.debug("Made PSF profile")
97+
98+
# Final profile is the convolution of these
99+
# Can include any number of things in the list, all of which are convolved
100+
# together to make the final flux profile.
101+
final = jax_galsim.Convolve([gal, psf])
102+
logger.debug("Convolved components into final profile")
103+
104+
# Draw the image with a particular pixel scale, given in arcsec/pixel.
105+
# The returned image has a member, added_flux, which is gives the total flux actually added to
106+
# the image. One could use this value to check if the image is large enough for some desired
107+
# accuracy level. Here, we just ignore it.
108+
image = final.drawImage(scale=pixel_scale)
109+
logger.debug(
110+
"Made image of the profile: flux = %f, added_flux = %f",
111+
gal_flux,
112+
image.added_flux,
113+
)
114+
115+
# Add Gaussian noise to the image with specified sigma
116+
image.addNoise(jax_galsim.GaussianNoise(sigma=noise))
117+
logger.debug("Added Gaussian noise")
118+
119+
# Write the image to a file
120+
if not os.path.isdir("output"):
121+
os.mkdir("output")
122+
file_name = os.path.join("output", "demo1.fits")
123+
# Note: if the file already exists, this will overwrite it.
124+
image.write(file_name)
125+
logger.info(
126+
"Wrote image to %r" % file_name
127+
) # using %r adds quotes around filename for us
128+
129+
results = image.FindAdaptiveMom()
130+
131+
logger.info("HSM reports that the image has observed shape and size:")
132+
logger.info(
133+
" e1 = %.3f, e2 = %.3f, sigma = %.3f (pixels)",
134+
results.observed_shape.e1,
135+
results.observed_shape.e2,
136+
results.moments_sigma,
137+
)
138+
logger.info(
139+
"Expected values in the limit that pixel response and noise are negligible:"
140+
)
141+
logger.info(
142+
" e1 = %.3f, e2 = %.3f, sigma = %.3f",
143+
0.0,
144+
0.0,
145+
math.sqrt(gal_sigma**2 + psf_sigma**2) / pixel_scale,
146+
)
147+
148+
149+
if __name__ == "__main__":
150+
main(sys.argv)

‎examples/demo2.py‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
# Copyright (c) 2012-2026 by the GalSim developers team on GitHub
2+
# https://github.com/GalSim-developers
3+
#
4+
# This file is part of GalSim: The modular galaxy image simulation toolkit.
5+
# https://github.com/GalSim-developers/GalSim
6+
#
7+
# GalSim is free software: redistribution and use in source and binary forms,
8+
# with or without modification, are permitted provided that the following
9+
# conditions are met:
10+
#
11+
# 1. Redistributions of source code must retain the above copyright notice, this
12+
# list of conditions, and the disclaimer given in the accompanying LICENSE
13+
# file.
14+
# 2. Redistributions in binary form must reproduce the above copyright notice,
15+
# this list of conditions, and the disclaimer given in the documentation
16+
# and/or other materials provided with the distribution.
17+
#
18+
"""
19+
Demo #2
20+
21+
The second script in our tutorial about using JAX-GalSim in python scripts: examples/demo*.py.
22+
(This file is designed to be viewed in a window 100 characters wide.)
23+
24+
This script is a bit more sophisticated, but still pretty basic. We're still only making
25+
a single image, but now the galaxy has an exponential radial profile and is sheared.
26+
The PSF is a circular Moffat profile. The noise is drawn from a Poisson distribution
27+
using the flux from both the object and a background sky level to determine the
28+
variance in each pixel.
29+
30+
New features introduced in this demo:
31+
32+
- obj = jax_galsim.Exponential(flux, scale_radius)
33+
- obj = jax_galsim.Moffat(beta, flux, half_light_radius)
34+
- obj = obj.shear(g1, g2) -- with explanation of other ways to specify shear
35+
- rng = jax_galsim.BaseDeviate(seed)
36+
- noise = jax_galsim.PoissonNoise(rng, sky_level)
37+
- galsim.hsm.EstimateShear(image, image_epsf)
38+
"""
39+
40+
import logging
41+
import os
42+
import sys
43+
44+
import galsim
45+
46+
import jax_galsim
47+
48+
49+
def main(argv):
50+
"""
51+
A little bit more sophisticated, but still pretty basic:
52+
- Use a sheared, exponential profile for the galaxy.
53+
- Convolve it by a circular Moffat PSF.
54+
- Add Poisson noise to the image.
55+
"""
56+
# In non-script code, use getLogger(__name__) at module scope instead.
57+
logging.basicConfig(format="%(message)s", level=logging.INFO, stream=sys.stdout)
58+
logger = logging.getLogger("demo2")
59+
60+
gal_flux = 1.0e5 # counts
61+
gal_r0 = 2.7 # arcsec
62+
g1 = 0.1 #
63+
g2 = 0.2 #
64+
psf_beta = 5 #
65+
psf_re = 1.0 # arcsec
66+
pixel_scale = 0.2 # arcsec / pixel
67+
sky_level = 2.5e3 # counts / arcsec^2
68+
69+
# This time use a particular seed, so the image is deterministic.
70+
# This is the same seed that is used in demo2.yaml, which means the images
71+
# produced by the two methods will be precisely identical.
72+
random_seed = 1534225
73+
74+
# The first thing the config layer does with the random seed is to scramble
75+
# it a bit. Specifically, it makes a random number generator (BaseDeviate)
76+
# using that seed and asks for a raw value. This becomes the seed that
77+
# actually gets used.
78+
# The reason for this extra step is that eventually (cf. demo4) the config
79+
# layer will want to increment these seed values when building multiple
80+
# objects or images. If the user is likewise incrementing seed values for
81+
# multiple runs of a given config file, these can interfere leading to
82+
# surprising (and typically bad) results.
83+
random_seed = jax_galsim.BaseDeviate(random_seed).raw()
84+
85+
logger.info("Starting demo script 2 using:")
86+
logger.info(
87+
" - sheared (%.2f,%.2f) exponential galaxy (flux = %.1e, scale radius = %.2f),",
88+
g1,
89+
g2,
90+
gal_flux,
91+
gal_r0,
92+
)
93+
logger.info(" - circular Moffat PSF (beta = %.1f, re = %.2f),", psf_beta, psf_re)
94+
logger.info(" - pixel scale = %.2f,", pixel_scale)
95+
logger.info(" - Poisson noise (sky level = %.1e).", sky_level)
96+
97+
# Initialize the (pseudo-)random number generator that we will be using below.
98+
# For a technical reason that will be explained later (demo9.py), we add 1 to the
99+
# given random seed here.
100+
rng = jax_galsim.BaseDeviate(random_seed + 1)
101+
102+
# Define the galaxy profile.
103+
gal = jax_galsim.Exponential(flux=gal_flux, scale_radius=gal_r0)
104+
105+
# Shear the galaxy by some value.
106+
# There are quite a few ways you can use to specify a shape.
107+
# q, beta Axis ratio and position angle: q = b/a, 0 < q < 1
108+
# e, beta Ellipticity and position angle: |e| = (1-q^2)/(1+q^2)
109+
# g, beta ("Reduced") Shear and position angle: |g| = (1-q)/(1+q)
110+
# eta, beta Conformal shear and position angle: eta = ln(1/q)
111+
# e1,e2 Ellipticity components: e1 = e cos(2 beta), e2 = e sin(2 beta)
112+
# g1,g2 ("Reduced") shear components: g1 = g cos(2 beta), g2 = g sin(2 beta)
113+
# eta1,eta2 Conformal shear components: eta1 = eta cos(2 beta), eta2 = eta sin(2 beta)
114+
gal = gal.shear(g1=g1, g2=g2)
115+
logger.debug("Made galaxy profile")
116+
117+
# Define the PSF profile.
118+
psf = jax_galsim.Moffat(beta=psf_beta, flux=1.0, half_light_radius=psf_re)
119+
logger.debug("Made PSF profile")
120+
121+
# Final profile is the convolution of these.
122+
final = jax_galsim.Convolve([gal, psf])
123+
logger.debug("Convolved components into final profile")
124+
125+
# Draw the image with a particular pixel scale.
126+
image = final.drawImage(scale=pixel_scale)
127+
# The "effective PSF" is the PSF as drawn on an image, which includes the convolution
128+
# by the pixel response. We label it epsf here.
129+
image_epsf = psf.drawImage(scale=pixel_scale)
130+
logger.debug("Made image of the profile")
131+
132+
# To get Poisson noise on the image, we will use a class called PoissonNoise.
133+
# However, we want the noise to correspond to what you would get with a significant
134+
# flux from tke sky. This is done by telling PoissonNoise to add noise from a
135+
# sky level in addition to the counts currently in the image.
136+
#
137+
# One wrinkle here is that the PoissonNoise class needs the sky level in each pixel,
138+
# while we have a sky_level in counts per arcsec^2. So we need to convert:
139+
sky_level_pixel = sky_level * pixel_scale**2
140+
noise = jax_galsim.PoissonNoise(rng, sky_level=sky_level_pixel)
141+
image.addNoise(noise)
142+
logger.debug("Added Poisson noise")
143+
144+
# Write the image to a file.
145+
if not os.path.isdir("output"):
146+
os.mkdir("output")
147+
file_name = os.path.join("output", "demo2.fits")
148+
file_name_epsf = os.path.join("output", "demo2_epsf.fits")
149+
image.write(file_name)
150+
image_epsf.write(file_name_epsf)
151+
logger.info("Wrote image to %r", file_name)
152+
logger.info("Wrote effective PSF image to %r", file_name_epsf)
153+
154+
results = galsim.hsm.EstimateShear(image.to_galsim(), image_epsf.to_galsim())
155+
156+
logger.info("HSM reports that the image has observed shape and size:")
157+
logger.info(
158+
" e1 = %.3f, e2 = %.3f, sigma = %.3f (pixels)",
159+
results.observed_shape.e1,
160+
results.observed_shape.e2,
161+
results.moments_sigma,
162+
)
163+
logger.info(
164+
"When carrying out Regaussianization PSF correction, HSM reports distortions"
165+
)
166+
logger.info(" e1, e2 = %.3f, %.3f", results.corrected_e1, results.corrected_e2)
167+
logger.info(
168+
"Expected values in the limit that noise and non-Gaussianity are negligible:"
169+
)
170+
exp_shear = galsim.Shear(g1=g1, g2=g2)
171+
logger.info(" g1, g2 = %.3f, %.3f", exp_shear.e1, exp_shear.e2)
172+
173+
174+
if __name__ == "__main__":
175+
main(sys.argv)

‎jax_galsim/image.py‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1094,15 +1094,45 @@ def tree_unflatten(cls, aux_data, children):
10941094
@classmethod
10951095
def from_galsim(cls, galsim_image):
10961096
"""Create a `Image` from a `galsim.Image` instance."""
1097+
wcs = (
1098+
BaseWCS.from_galsim(galsim_image.wcs)
1099+
if galsim_image.wcs is not None
1100+
else None
1101+
)
10971102
im = cls(
10981103
array=galsim_image.array,
1099-
wcs=BaseWCS.from_galsim(galsim_image.wcs),
1104+
wcs=wcs,
11001105
bounds=Bounds.from_galsim(galsim_image.bounds),
11011106
)
11021107
if hasattr(galsim_image, "header"):
11031108
im.header = galsim_image.header
11041109
return im
11051110

1111+
def to_galsim(self):
1112+
"""Create a galsim `Image` from a `jax_galsim.Image` object."""
1113+
wcs = self.wcs.to_galsim() if self.wcs is not None else None
1114+
return _galsim.Image(
1115+
np.asarray(self.array), bounds=self.bounds.to_galsim(), wcs=wcs
1116+
)
1117+
1118+
@implements(
1119+
_galsim.Image.FindAdaptiveMom,
1120+
lax_description=(
1121+
"This method converts the current `jax_galsim.Image` to a native "
1122+
"`galsim.Image` and delegates the computation to "
1123+
"`galsim.hsm.FindAdaptiveMom`. The returned object is GalSim's "
1124+
"`ShapeData`."
1125+
),
1126+
)
1127+
def FindAdaptiveMom(self, *args, **kwargs):
1128+
args_ = [arg.to_galsim() if hasattr(arg, "to_galsim") else arg for arg in args]
1129+
kwargs_ = {
1130+
key: val.to_galsim() if hasattr(val, "to_galsim") else val
1131+
for key, val in kwargs.items()
1132+
}
1133+
gs_image = self.to_galsim()
1134+
return gs_image.FindAdaptiveMom(*args_, **kwargs_)
1135+
11061136

11071137
@implements(
11081138
_galsim._Image,

‎tests/GalSim‎

‎tests/SBProfile_comparison_images‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
GalSim/tests/SBProfile_comparison_images

‎tests/fits_file‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
GalSim/tests/fits_files

‎tests/galsim_tests_config.yaml‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,6 @@ allowed_failures:
8686
- "'Image' object has no attribute 'bin'"
8787
- "module 'jax_galsim' has no attribute 'InterpolatedKImage'"
8888
- "module 'jax_galsim' has no attribute 'CorrelatedNoise'"
89-
- "'Image' object has no attribute 'FindAdaptiveMom'"
9089
- "CelestialCoord.precess is too slow" # cannot get jax to warmup but once it does it passes
9190
- "ValueError not raised by from_xyz"
9291
- "ValueError not raised by greatCirclePoint"

‎tests/jax/test_api.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -604,6 +604,7 @@ def _reg_sfun(g1):
604604
def test_api_image(obj):
605605
_run_object_checks(obj, obj.__class__, "docs-methods")
606606
_run_object_checks(obj, obj.__class__, "pickle-eval-repr-img")
607+
_run_object_checks(obj, obj.__class__, "to-from-galsim")
607608

608609
# JAX tracing should be an identity
609610
assert obj.__class__.tree_unflatten(*((obj.tree_flatten())[::-1])) == obj

0 commit comments

Comments
 (0)