diff --git a/autometacal/python/__init__.py b/autometacal/python/__init__.py index 1e067d1..251a76e 100644 --- a/autometacal/python/__init__.py +++ b/autometacal/python/__init__.py @@ -1,5 +1,7 @@ -import autometacal.python.datasets as datasets +import autometacal.python.datasets as datasets import autometacal.python.tf_ngmix as tf_ngmix +import autometacal.python.metacal_psf as metacal_psf from autometacal.python.metacal import generate_mcal_image, get_metacal_response, get_metacal_response_finitediff from autometacal.python.fitting import fit_multivariate_gaussian, get_ellipticity from autometacal.python.moments import get_moment_ellipticities, gaussian_moments +from autometacal.python.utils import * diff --git a/autometacal/python/datasets/CFIS.py b/autometacal/python/datasets/CFIS.py index c6a4feb..e44c357 100644 --- a/autometacal/python/datasets/CFIS.py +++ b/autometacal/python/datasets/CFIS.py @@ -23,9 +23,9 @@ def __init__(self, *, galaxy_type="parametric", pixel_scale: pixel scale of stamps in arcsec. **kwargs: keyword arguments forwarded to super. """ - v1 = tfds.core.Version("0.0.1") + v1 = tfds.core.Version("0.5.0") super(CFISConfig, self).__init__( - description=("Galaxy stamps"), + description=("CFIS-like Galaxy stamps"), version=v1, **kwargs) @@ -49,7 +49,7 @@ class CFIS(tfds.core.GeneratorBasedBuilder): VERSION = tfds.core.Version('0.0.1') RELEASE_NOTES = { - '0.0.1': 'Initial release.', + '0.5.0': 'Initial release.', } BUILDER_CONFIGS = [CFISConfig(name="parametric_1k", galaxy_type="parametric", data_set_size=1000), @@ -68,12 +68,7 @@ def _info(self) -> tfds.core.DatasetInfo: 'psf': tfds.features.Tensor(shape=[self.builder_config.stamp_size, self.builder_config.stamp_size], dtype=tf.float32), - # 'gal_kimage': tfds.features.Tensor(shape=[2, self.builder_config.kstamp_size, - # self.builder_config.kstamp_size], - # dtype=tf.float32), - # 'psf_kimage': tfds.features.Tensor(shape=[2, self.builder_config.kstamp_size, - # self.builder_config.kstamp_size], - # dtype=tf.float32), + "noise_std": tfds.features.Tensor(shape=[1], dtype=tf.float32), "mag": tfds.features.Tensor(shape=[1], dtype=tf.float32), }), @@ -109,9 +104,6 @@ def _generate_examples(self, size): psf = gs.Kolmogorov(fwhm=self.builder_config.psf_fwhm, flux=1.0) psf = psf.shear(g1=self.builder_config.psf_e1, g2=self.builder_config.psf_e2) - # Prepare borders for kimage - Nk = self.builder_config.kstamp_size - bounds = gs._BoundsI(-Nk//2, Nk//2-1, -Nk//2, Nk//2-1) for i in range(size): # retrieving galaxy and magnitude @@ -134,17 +126,7 @@ def _generate_examples(self, size): scale=self.builder_config.pixel_scale ).array.astype('float32') - # gal_kimage = gal.drawKImage(bounds=bounds, - # scale=2.*np.pi/(self.builder_config.stamp_size*self.builder_config.pixel_scale), - # recenter=False).array.astype('complex64') - - # psf_kimage = psf.drawKImage(bounds=bounds, - # scale=2.*np.pi/(self.builder_config.stamp_size*self.builder_config.pixel_scale), - # recenter=False).array.astype('complex64') - yield '%d'%i, {"obs": gal_stamp, "psf": psf_stamp, - # "gal_kimage": np.stack([gal_kimage.real, gal_kimage.imag]), - # "psf_kimage": np.stack([psf_kimage.real, psf_kimage.imag]), "noise_std": np.array([np.sqrt(sky_level)]).astype('float32'), "mag": np.array([gal_mag]).astype('float32')} \ No newline at end of file diff --git a/autometacal/python/datasets/__init__.py b/autometacal/python/datasets/__init__.py index 12a30af..53c6cbb 100644 --- a/autometacal/python/datasets/__init__.py +++ b/autometacal/python/datasets/__init__.py @@ -1,4 +1,4 @@ -"""gal_gen dataset.""" +"""galaxy datasets""" -from .gal_gen import GalGen -from .CFIS import CFIS \ No newline at end of file +from .galgen import GalGen +from .CFIS import CFIS diff --git a/autometacal/python/datasets/gal_gen.py b/autometacal/python/datasets/gal_gen.py deleted file mode 100644 index 2995511..0000000 --- a/autometacal/python/datasets/gal_gen.py +++ /dev/null @@ -1,131 +0,0 @@ -"""gal_gen dataset.""" -import tensorflow_datasets as tfds -import tensorflow as tf -import numpy as np - -from .galaxies import gs_generate_images, gs_drawKimage - -_DESCRIPTION = "This tfds generates random toy-model galaxy stamps." -_CITATION = "{NEEDED}" -_URL = "https://github.com/CosmoStat/autometacal" - -class GalGenConfig(tfds.core.BuilderConfig): - """BuilderConfig for GalGen.""" - - def __init__(self, *,dataset_size=None, stamp_size=None, pixel_scale=None, flux=None, interp_factor=None, padding_factor=None, **kwargs): - """BuilderConfig for SQUAD. - Args: - pixel_scale: pixel_scale of the image in arcsec/pixel. - stamp_size: image stamp size in pixels. - flux: flux of the profile. - **kwargs: keyword arguments forwarded to super. - """ - v2 = tfds.core.Version("2.0.0") - super(GalGenConfig, self).__init__( - description=("GalGen %d stamps in %d x %d resolution, %.2f arcsec/pixel, with flux of %.2f ." % - (dataset_size,stamp_size, stamp_size, pixel_scale, flux)), - version=v2, - release_notes={ - "2.0.0": "New split API (https://tensorflow.org/datasets/splits)", - }, - **kwargs) - self.dataset_size = dataset_size - self.stamp_size = stamp_size - self.pixel_scale = pixel_scale - self.flux = flux - self.interp_factor = 2 - self.padding_factor = 1 - self.kstamp_size = self.interp_factor*self.padding_factor*self.stamp_size - - - -class GalGen(tfds.core.GeneratorBasedBuilder): - """Random galaxy image generator.""" - - MANUAL_DOWNLOAD_INSTRUCTIONS = """\ - Nothing to download. DataSet is generated at first call. - """ - - BUILDER_CONFIGS = [ - GalGenConfig(name='small_stamp_100k', dataset_size=100000, stamp_size=50, pixel_scale=.2, flux=1e5), - GalGenConfig(name='large_stamp_100k', dataset_size=100000, stamp_size=100, pixel_scale=.2, flux=1e5), - GalGenConfig(name='small_stamp_100', dataset_size=100, stamp_size=50, pixel_scale=.2, flux=1e5), - GalGenConfig(name='large_stamp_100', dataset_size=100, stamp_size=100, pixel_scale=.2, flux=1e5) - ] - - VERSION = tfds.core.Version('0.1.0') - RELEASE_NOTES = {'0.1.0': "Basic functionalities, that work."} - - def _info(self): - return tfds.core.DatasetInfo( - builder=self, - # Description and homepage used for documentation - description=_DESCRIPTION, - homepage=_URL, - features=tfds.features.FeaturesDict({'label': tfds.features.Tensor(shape=[2], dtype=tf.float32), - 'gal_model': tfds.features.Tensor(shape=[self.builder_config.stamp_size, - self.builder_config.stamp_size], - dtype=tf.float32), - 'obs_image': tfds.features.Tensor(shape=[self.builder_config.stamp_size, - self.builder_config.stamp_size], - dtype=tf.float32), - 'psf_image': tfds.features.Tensor(shape=[self.builder_config.stamp_size, - self.builder_config.stamp_size], - dtype=tf.float32), - 'obs_kimage': tfds.features.Tensor(shape=[2,self.builder_config.kstamp_size, - self.builder_config.kstamp_size], - dtype=tf.float32), - 'psf_kimage': tfds.features.Tensor(shape=[2,self.builder_config.kstamp_size, - self.builder_config.kstamp_size], - dtype=tf.float32), - 'psf_deconv': tfds.features.Tensor(shape=[2,self.builder_config.kstamp_size, - self.builder_config.kstamp_size], - dtype=tf.float32) - }), - supervised_keys=("obs_image","label"), - citation=_CITATION) - - def _split_generators(self,dl): - """Returns generators according to split.""" - return {tfds.Split.TRAIN: self._generate_examples(self.builder_config.dataset_size, - self.builder_config.stamp_size, - self.builder_config.pixel_scale, - self.builder_config.flux, - self.builder_config.interp_factor, - self.builder_config.padding_factor)} - - def _generate_examples(self, dataset_size, stamp_size, pixel_scale, flux, interp_factor, padding_factor): - """Yields examples.""" - np.random.seed(31415) - - for i in range(dataset_size): - - #generate example - label, model, obs_img, psf_img, obs_kimg, psf_kimg, psf_deconv = gs_generate_images(stamp_size = stamp_size, - pixel_scale = pixel_scale, - flux = flux, - interp_factor = interp_factor, - padding_factor = padding_factor) - - - #store complex arrays in 2,N,N - obs_kimg = decomplexify(obs_kimg.numpy()) - psf_kimg = decomplexify(psf_kimg.numpy()) - psf_deconv = decomplexify(psf_deconv.numpy()) - - - yield '%d'%i, {'gal_model': model.numpy(), #noiseless PSFless galaxy model - 'obs_image': obs_img.numpy(), #observed image - 'psf_image': psf_img.numpy(), #psf image - 'obs_kimage': obs_kimg, #obs k image - 'psf_kimage': psf_kimg, #psf k image - 'psf_deconv': psf_deconv, #psf deconv kernel - 'label': label.numpy() } - - -def decomplexify(arr): - arr=np.array([arr.real,arr.imag]) - return arr - -def recomplexify(arl): - return arl[0]+1j*arl[1] diff --git a/autometacal/python/datasets/galaxies.py b/autometacal/python/datasets/galaxies.py deleted file mode 100644 index 63c781f..0000000 --- a/autometacal/python/datasets/galaxies.py +++ /dev/null @@ -1,257 +0,0 @@ -import galsim -import tensorflow as tf -import numpy as np -from scipy.stats import truncnorm, norm - -""" -Simple wrappers for toy galaxy generation -Functions with gs_ prefix depend on galsim -""" - - -def gs_generate_images(**kwargs): - - """ Random Galaxy Generator - Generates noiseless galaxy images with a simple light profile. The resulting image is before the convolution with a PSF. - Galaxy shapes follow a bivariate normal distribution centered in zero. - - Args: - g_range: galaxy shapes go from -g_range and + g_range in each g1, g2 - g_scatter: galaxy shapes scatter - flux: galaxy flux (counts) - pixel_scale: intended pixel scale in arcsec^2/pixel - stamp_size: size in pixels of the NxN resulting image - method: galsim drawing method - mean_radius: mean half-light radii of generated galaxies - scatter_radius: scatter of half-light radii of generated galaxies - mean_snr: average snr of generated stamps (snr as per galsim definition) - scatter_snr: scatter in snr of generated stamps - interp_factor: interpolation factor for drawing k space images - padding_factor: padding factor for drawing k space images - Returns: - g1, g2: galaxy shape parameters - gal: tensor with a 2-d array representing an observed image of a galaxy (with convolved psf) - psf: tensor with a 2-d array representing the model of the psf - gal_k: k space image of gal - psf_k: k space image of psf - """ - - defaults = {'g_range' : 0.8, #elipticity - 'g_scatter' : 0.25, # - 'mean_radius': 1.0, #size - 'scatter_radius': 0.1, # - 'psf_beta': 5, #psf - 'psf_fwhm': 0.7, # - 'mean_snr': 200, #snr - 'scatter_snr': 20, # - 'flux' : 1e5, #flux - 'pixel_scale' : 0.2, # - 'stamp_size' : 50, # - 'method' : "no_pixel", # - 'interp_factor': 2, #kimage interpolation - 'padding_factor': 1 #kimage padding - } - - defaults.update(kwargs) - - #ellipticity range - a, b = (-defaults['g_range'] - 0) / defaults['g_scatter'], (defaults['g_range'] - 0) / defaults['g_scatter'] - g1=g2=1 - - #select ellipticities, ensure g<=1 - while g1**2+g2**2>1: - g1 = truncnorm.rvs(a, b, loc=0, scale=defaults['g_scatter']) - g2 = truncnorm.rvs(a, b, loc=0, scale=defaults['g_scatter']) - - re = norm.rvs(defaults['mean_radius'], defaults['scatter_radius']) - - #very simple galaxy model - gal = galsim.Exponential(flux=defaults['flux'] , - half_light_radius=re) - - #apply shear - gal = gal.shear(g1=g1,g2=g2) - - - #create constant psf - psf = galsim.Moffat(beta=defaults['psf_beta'], - fwhm=defaults['psf_fwhm']) - - #draw galaxy before convolution - model_Image = gal.drawImage(nx=defaults['stamp_size'], - ny=defaults['stamp_size'], - scale=defaults['pixel_scale'], - method=defaults['method']) - - #convolve galaxy and psf - gal = galsim.Convolve([gal,psf]) - - #draw psf image - psf_Image = psf.drawImage(nx=defaults['stamp_size'], - ny=defaults['stamp_size'], - scale=defaults['pixel_scale'], - method=defaults['method']) - - - #draw final observed image - obs_Image = gal.drawImage(nx=defaults['stamp_size'], - ny=defaults['stamp_size'], - scale=defaults['pixel_scale'], - method=defaults['method']) - - #add noise to image with a given SNR - noise = galsim.GaussianNoise() - snr = norm.rvs(defaults['mean_snr'],defaults['scatter_snr'],) - obs_Image.addNoiseSNR(noise,snr=snr) - - #draw kimage of galaxy - obs_k = gs_drawKimage(obs_Image.array, - defaults['pixel_scale'], - interp_factor=defaults['interp_factor'], - padding_factor=defaults['padding_factor']) - - # draw kimage of the psf - psf_k = gs_drawKimage(psf_Image.array, - defaults['pixel_scale'], - interp_factor=defaults['interp_factor'], - padding_factor=defaults['padding_factor']) - - #draw psf deconvolution kernel - psf_deconv = gs_Deconvolve(psf_Image.array, - defaults['pixel_scale'], - interp_factor=defaults['interp_factor'], - padding_factor=defaults['padding_factor']) - - #output everything to tf tensors - # tfds names: - g = tf.convert_to_tensor([g1,g2],dtype=tf.float32) # label - model = tf.convert_to_tensor(model_Image.array,dtype=tf.float32) # model_image - obs = tf.convert_to_tensor(obs_Image.array,dtype=tf.float32) # obs_image - psf = tf.convert_to_tensor(psf_Image.array,dtype=tf.float32) # psf_image - - return g, model, obs, psf, obs_k, psf_k, psf_deconv - - -def gs_drawKimage(image, - pixel_scale=0.2, - interp_factor=2, - padding_factor=1): - """ - Args: - image: numpy array - pixel_scale: telescope pixel scale in arcsec/pixel - interp_factor: interpolation factor for superresolution - padding_factor: padding added by fraction - - returns: - tensor with k-image of object - """ - - #prepare borders - N = len(image) - Nk = N*interp_factor*padding_factor - bounds = galsim._BoundsI(-Nk//2, Nk//2-1, -Nk//2, Nk//2-1) - - #interpolated galsim object from input image - img_galsim = galsim.InterpolatedImage(galsim.Image(image,scale=pixel_scale)) - - #draw galsim output image - result = img_galsim.drawKImage(bounds=bounds, - scale=2.*np.pi/(N*padding_factor*pixel_scale), - recenter=False) - - return tf.convert_to_tensor(result.array,dtype=tf.complex64) - - -def gs_Deconvolve(psf_img, - pixel_scale=0.2, - interp_factor=2, - padding_factor=1): - """ - Returns a deconvolution kernel of a psf image. - - Args: - psf_img: numpy array representing the psf model - pixel_scale: the pixel scale of the image, in arcsec/pixel - interp_factor: the interpolation factor for super-resolution - padding_factor: a factor to add side pads to the image - Returns: - A complex tensorflow tensor that is a deconvolution kernel. - - """ - - N = len(psf_img) - - psf_galsim=galsim.InterpolatedImage(galsim.Image(psf_img,scale=pixel_scale)) - ipsf=galsim.Deconvolve(psf_galsim) - Nk = N*interp_factor*padding_factor - bounds = galsim._BoundsI(-Nk//2, - Nk//2-1, - -Nk//2, - Nk//2-1) - imipsf = ipsf.drawKImage(bounds=bounds, - scale=2.*np.pi/(N*padding_factor* pixel_scale), - recenter=False) - return tf.convert_to_tensor(imipsf.array,dtype=tf.complex64) - -def gs_noise_generator(stamp_size=50,variance=5,pixel_scale=.2,interp_factor=2,padding_factor=1): - """ Generate a noise k space image using GalSim. - - Args: - stamp_size: in pixels - variance: noise variance - pixel_scale: in arcsec/pixel - interp_factor: interpolation factor for k space - padding_factor: padding wrap factor - Returns: - A complex64 numpy array. - - """ - noise = galsim.GaussianNoise().withVariance(variance) - noise_image = galsim.Image(stamp_size,stamp_size, scale=pixel_scale) - noise.applyTo(noise_image) - noise_image = galsim.InterpolatedImage(noise_image) - Nk = stamp_size*padding_factor*interp_factor - from galsim.bounds import _BoundsI - - bounds = _BoundsI(-Nk//2, Nk//2-1, -Nk//2, Nk//2-1) - imnos = noise_image.drawKImage(bounds=bounds, - scale=2.*np.pi/(stamp_size*padding_factor*pixel_scale), - recenter=False) - return imnos.array.astype('complex64') - -def make_data(Ngals=1, - snr = 200, - scale = 0.2, - stamp_size = 51, - psf_fwhm = 0.9, - gal_hlr = 0.7, - gal_g1 = [0], - gal_g2 = [0], - flux=1.e5): - """Simple exponetial profile toy model galaxy""" - - gal_list = [] - psf_list = [] - - for n in range(Ngals): - psf = galsim.Moffat(beta=2.5, - fwhm=psf_fwhm) - - obj0 = galsim.Exponential(half_light_radius=gal_hlr,flux=flux).shear(g1=gal_g1[n],g2=gal_g2[n]) - obj = galsim.Convolve(psf, obj0) - - psf_image = psf.drawImage(nx=stamp_size, ny=stamp_size, scale=scale).array - gal_image = obj.drawImage(nx=stamp_size, ny=stamp_size, scale=scale) - noise = galsim.GaussianNoise() - gal_image.addNoiseSNR(noise,snr=snr) - - gal_image = tf.convert_to_tensor(gal_image.array) - psf_image = tf.convert_to_tensor(psf_image) - gal_list.append(gal_image) - psf_list.append(psf_image) - - gal_image_stack = tf.stack(gal_list) - psf_image_stack = tf.stack(psf_list) - - return gal_image_stack, psf_image_stack \ No newline at end of file diff --git a/autometacal/python/datasets/galgen.py b/autometacal/python/datasets/galgen.py new file mode 100644 index 0000000..f91b917 --- /dev/null +++ b/autometacal/python/datasets/galgen.py @@ -0,0 +1,181 @@ +"""gal_gen dataset.""" +import tensorflow_datasets as tfds +import tensorflow as tf +import numpy as np +from scipy.stats import truncnorm, norm +import galsim + +_DESCRIPTION = "This tfds generates random toy-model galaxy stamps." +_CITATION = "{NEEDED}" +_URL = "https://github.com/CosmoStat/autometacal" + +class GalGenConfig(tfds.core.BuilderConfig): + """BuilderConfig for GalGen.""" + + def __init__(self, *,dataset_size=None, stamp_size=None, pixel_scale=None, **kwargs): + """BuilderConfig for GalGen. + Args: + pixel_scale: pixel_scale of the image in arcsec/pixel. + stamp_size: image stamp size in pixels. + flux: flux of the profile. + **kwargs: keyword arguments forwarded to super. + """ + v2 = tfds.core.Version("2.0.0") + super(GalGenConfig, self).__init__( + description=("GalGen %d stamps in %d x %d resolution, %.2f arcsec/pixel, with flux of 1." % + (dataset_size,stamp_size, stamp_size, pixel_scale)), + version=v2, + release_notes={ + "2.0.0": "New split API (https://tensorflow.org/datasets/splits)", + }, + **kwargs) + self.dataset_size = dataset_size + self.stamp_size = stamp_size + +class GalGen(tfds.core.GeneratorBasedBuilder): + """Random galaxy image generator.""" + + MANUAL_DOWNLOAD_INSTRUCTIONS = """\ + Nothing to download. DataSet is generated at first call. + """ + + BUILDER_CONFIGS = [ + GalGenConfig(name='small_stamp_100k', dataset_size=100000, stamp_size=51, pixel_scale=.2), + GalGenConfig(name='large_stamp_100k', dataset_size=100000, stamp_size=101, pixel_scale=.2), + GalGenConfig(name='small_stamp_1k', dataset_size=1000, stamp_size=51, pixel_scale=.2), + GalGenConfig(name='large_stamp_1k', dataset_size=1000, stamp_size=101, pixel_scale=.2) + ] + + VERSION = tfds.core.Version('0.5.0') + RELEASE_NOTES = {'0.5.0': "Updated functionalities, simpler."} + + def _info(self): + return tfds.core.DatasetInfo( + builder=self, + # Description and homepage used for documentation + description=_DESCRIPTION, + homepage=_URL, + features=tfds.features.FeaturesDict({'label': tfds.features.Tensor(shape=[2], dtype=tf.float32), + 'gal_model': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + 'obs_image': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + 'psf_image': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + }), + supervised_keys=("obs_image","label"), + citation=_CITATION) + + def _split_generators(self,dl): + """Returns generators according to split.""" + return {tfds.Split.TRAIN: self._generate_examples(self.builder_config.dataset_size, + self.builder_config.stamp_size)} + + def _generate_examples(self, dataset_size, stamp_size): + """Yields examples.""" + np.random.seed(31415) + + for i in range(dataset_size): + + #generate example + label, model, obs_img, psf_img, = gs_generate_images(stamp_size = stamp_size, + ) + yield '%d'%i, {'gal_model': model, #noiseless PSFless galaxy model + 'obs_image': obs_img, #observed image + 'psf_image': psf_img, #psf image + 'label': label} + + +def gs_generate_images(**kwargs): + + """ Random Galaxy Generator + Generates noiseless galaxy images with a simple light profile. The resulting image is before the convolution with a PSF. + Galaxy shapes follow a bivariate normal distribution centered in zero. + + Args: + g_range: galaxy shapes go from -g_range and + g_range in each g1, g2 + g_scatter: galaxy shapes scatter + flux: galaxy flux (counts) + pixel_scale: intended pixel scale in arcsec^2/pixel + stamp_size: size in pixels of the NxN resulting image + method: galsim drawing method + mean_radius: mean half-light radii of generated galaxies + scatter_radius: scatter of half-light radii of generated galaxies + mean_snr: average snr of generated stamps (snr as per galsim definition) + scatter_snr: scatter in snr of generated stamps + Returns: + g1, g2: galaxy shape parameters + mode: galaxy model without psf or noise + gal: tensor with a 2-d array representing an observed image of a galaxy (with convolved psf) + psf: tensor with a 2-d array representing the model of the psf + """ + + defaults = {'g_range' : 0.7, #elipticity + 'g_scatter' : 0.25, # + 'mean_hlr': .9, #size + 'scatter_hlr': 0.1, # + 'psf_beta': 5.0, #psf + 'psf_fwhm': 0.7, # + 'mean_snr': 200, #snr + 'scatter_snr': 20, # + 'pixel_scale' : 0.2, # + 'stamp_size' : 51, # + 'method' : "no_pixel", # + 'interp_factor': 2, #kimage interpolation + 'padding_factor': 1 #kimage padding + } + + defaults.update(kwargs) + + #ellipticity range + a, b = (-defaults['g_range'] - 0) / defaults['g_scatter'], (defaults['g_range'] - 0) / defaults['g_scatter'] + + g1 = truncnorm.rvs(a, b, loc=0, scale=defaults['g_scatter']) + g2 = truncnorm.rvs(a, b, loc=0, scale=defaults['g_scatter']) + + hlr = norm.rvs(defaults['mean_hlr'], defaults['scatter_hlr']) + + #very simple galaxy model + gal = galsim.Exponential(half_light_radius=hlr) + + #apply 'shear' + gal = gal.shear(g1=g1,g2=g2) + + + #create constant psf + psf = galsim.Moffat(beta=defaults['psf_beta'], + fwhm=defaults['psf_fwhm']) + + #draw galaxy before convolution + model_image = gal.drawImage(nx=defaults['stamp_size'], + ny=defaults['stamp_size'], + scale=defaults['pixel_scale'], + method=defaults['method']) + + #convolve galaxy and psf + gal = galsim.Convolve([gal,psf]) + + #draw psf image + psf_image = psf.drawImage(nx=defaults['stamp_size'], + ny=defaults['stamp_size'], + scale=defaults['pixel_scale'], + method=defaults['method']) + + + #draw final observed image + obs_image = gal.drawImage(nx=defaults['stamp_size'], + ny=defaults['stamp_size'], + scale=defaults['pixel_scale'], + method=defaults['method']) + + #add noise to image with a given SNR + noise = galsim.GaussianNoise() + snr = norm.rvs(defaults['mean_snr'],defaults['scatter_snr'],) + obs_image.addNoiseSNR(noise,snr=snr) + + #output everything to tf tensors + return [g1,g2], model_image.array.astype('float32'), obs_image.array.astype('float32'), psf_image.array.astype('float32') + diff --git a/autometacal/python/datasets/simple.py b/autometacal/python/datasets/simple.py new file mode 100644 index 0000000..5e6547b --- /dev/null +++ b/autometacal/python/datasets/simple.py @@ -0,0 +1,135 @@ +"""gal_gen dataset.""" +import tensorflow_datasets as tfds +import tensorflow as tf +import numpy as np + +import galsim + +_DESCRIPTION = "This tfds generates random toy-model galaxy stamps." +_CITATION = "{NEEDED}" +_URL = "https://github.com/CosmoStat/autometacal" + +class SimpleConfig(tfds.core.BuilderConfig): + """BuilderConfig for GalGen.""" + + def __init__(self, *,dataset_size=1000, stamp_size=45, + shear_g1=0.0, shear_g2=0.0,**kwargs): + """BuilderConfig for Simple. + Args: + pixel_scale: pixel_scale of the image in arcsec/pixel. + stamp_size: image stamp size in pixels. + flux: flux of the profile. + g1, g2: shear applied to all 1k stamps + **kwargs: keyword arguments forwarded to super. + """ + v2 = tfds.core.Version("1.0.0") + super(SimpleConfig, self).__init__( + description=("Simple %d stamps in %d x %d resolution, 0.263 arcsec/pixel, with flux of 1." % + (dataset_size,stamp_size, stamp_size)), + version=v2, + **kwargs) + self.dataset_size = dataset_size + self.stamp_size = stamp_size + self.g1 = shear_g1 + self.g2 = shear_g2 + +class Simple(tfds.core.GeneratorBasedBuilder): + """Random galaxy image generator.""" + + MANUAL_DOWNLOAD_INSTRUCTIONS = """\ + Nothing to download. DataSet is generated at first call. + """ + + BUILDER_CONFIGS = [ + SimpleConfig(name='small_1k_noshear', dataset_size=1000, stamp_size=51), + SimpleConfig(name='small_1k_g1p', dataset_size=1000, stamp_size=51,shear_g1=.01), + SimpleConfig(name='small_1k_g1m', dataset_size=1000, stamp_size=51,shear_g1=-.01), + SimpleConfig(name='small_1k_g2p', dataset_size=1000, stamp_size=51,shear_g2=.01), + SimpleConfig(name='small_1k_g2m', dataset_size=1000, stamp_size=51,shear_g2=-.01), + SimpleConfig(name='large_1k_noshear', dataset_size=1000, stamp_size=101), + SimpleConfig(name='large_1k_g1p', dataset_size=1000, stamp_size=101,shear_g1=.01), + SimpleConfig(name='large_1k_g1m', dataset_size=1000, stamp_size=101,shear_g1=-.01), + SimpleConfig(name='large_1k_g2p', dataset_size=1000, stamp_size=101,shear_g2=.01), + SimpleConfig(name='large_1k_g2m', dataset_size=1000, stamp_size=101,shear_g2=-.01), + SimpleConfig(name='small_100k_noshear', dataset_size=100000, stamp_size=51), + SimpleConfig(name='small_100k_g1p', dataset_size=100000, stamp_size=51,shear_g1=.01), + SimpleConfig(name='small_100k_g1m', dataset_size=100000, stamp_size=51,shear_g1=-.01), + SimpleConfig(name='small_100k_g2p', dataset_size=100000, stamp_size=51,shear_g2=.01), + SimpleConfig(name='small_100k_g2m', dataset_size=100000, stamp_size=51,shear_g2=-.01), + SimpleConfig(name='large_100k_noshear', dataset_size=100000, stamp_size=101), + SimpleConfig(name='large_100k_g1p', dataset_size=100000, stamp_size=101,shear_g1=.01), + SimpleConfig(name='large_100k_g1m', dataset_size=100000, stamp_size=101,shear_g1=-.01), + SimpleConfig(name='large_100k_g2p', dataset_size=100000, stamp_size=101,shear_g2=.01), + SimpleConfig(name='large_100k_g2m', dataset_size=100000, stamp_size=101,shear_g2=-.01), + ] + + VERSION = tfds.core.Version('0.5.0') + RELEASE_NOTES = {'0.5.0': "Updated functionalities, simpler."} + + def _info(self): + return tfds.core.DatasetInfo( + builder=self, + # Description and homepage used for documentation + description=_DESCRIPTION, + homepage=_URL, + features=tfds.features.FeaturesDict({'label': tfds.features.Tensor(shape=[2], dtype=tf.float32), + 'gal_model': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + 'obs': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + 'psf_image': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + 'noise': tfds.features.Tensor(shape=[self.builder_config.stamp_size, + self.builder_config.stamp_size], + dtype=tf.float32), + }), + supervised_keys=("obs","label"), + citation=_CITATION) + + def _split_generators(self,dl): + """Returns generators according to split.""" + return {tfds.Split.TRAIN: self._generate_examples(self.builder_config.dataset_size, + self.builder_config.stamp_size, + self.builder_config.g1, + self.builder_config.g2)} + + def _generate_examples(self, dataset_size, stamp_size,g1,g2): + """Yields examples.""" + rng = np.random.RandomState(31415) + + for i in range(dataset_size): + + #generate example + model, obs_img, psf_img, noise_img = make_data(rng,stamp_size = stamp_size,g1=g1,g2=g2) + yield '%d'%i, {'gal_model': model, #noiseless PSFless galaxy model + 'obs': obs_img, #observed image + 'psf_image': psf_img, #psf image + 'noise' : noise_img, + 'label': np.array([g1,g2],dtype='float32')} + +def make_data(rng,stamp_size = 51,g1=0.,g2=0.): + scale = 0.263 + psf_fwhm = 0.9 + gal_hlr = 0.5 + noise = 1e-6 + + """Simple exponetial profile toy model galaxy""" + + psf = galsim.Moffat(beta=2.5,fwhm=psf_fwhm) + + obj = galsim.Exponential(half_light_radius=gal_hlr).shear(g1=g1,g2=g2) + obs = galsim.Convolve(psf, obj) + + psf_image = psf.drawImage(nx=stamp_size, ny=stamp_size, scale=scale).array #psf image + gal_image = obj.drawImage(nx=stamp_size, ny=stamp_size, scale=scale).array #model image + + obs_image = obs.drawImage(nx=stamp_size, ny=stamp_size, scale=scale).array #observed image, prenoise + + noise_image = rng.normal(scale=noise, size=obs_image.shape) + obs_image += noise_image + + return gal_image, obs_image, psf_image, noise_image.astype('float32') + diff --git a/notebooks/New_Datasets_Test.ipynb b/notebooks/New_Datasets_Test.ipynb new file mode 100644 index 0000000..7eb1eb2 --- /dev/null +++ b/notebooks/New_Datasets_Test.ipynb @@ -0,0 +1,459 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "db161b2d-6041-44f1-9186-266d08bc218a", + "metadata": {}, + "source": [ + "# New Datasets Test" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "abad247e-8940-4ebe-b352-e0d4665acedc", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Populating the interactive namespace from numpy and matplotlib\n" + ] + } + ], + "source": [ + "%pylab inline" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "c8b9fa3e-049e-4d59-818a-49e84e6fd4fb", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + ":219: RuntimeWarning: scipy._lib.messagestream.MessageStream size changed, may indicate binary incompatibility. Expected 56 from C header, got 64 from PyObject\n" + ] + } + ], + "source": [ + "import autometacal as amc\n", + "import tensorflow as tf\n", + "import tensorflow_datasets as tfds" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "e3d83b2c-cd8a-4a93-bf4c-29b8d139cae5", + "metadata": {}, + "outputs": [], + "source": [ + "from scipy.stats import norm" + ] + }, + { + "cell_type": "markdown", + "id": "f759386c-c8d7-4b26-8141-e11e7c85cfdb", + "metadata": {}, + "source": [ + "## Gal Gen" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "6d29e9dc-d862-4d83-9695-cf43391239d3", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[1mDownloading and preparing dataset Unknown size (download: Unknown size, generated: Unknown size, total: Unknown size) to /local/home/az264973/tensorflow_datasets/gal_gen/small_stamp_1k/2.0.0...\u001b[0m\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Generating splits...: 0%| | 0/1 [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "style.use('default')\n", + "figure(figsize=(10,10))\n", + "for i in range(16):\n", + " subplot(4,4,i+1)\n", + " imshow(gal_images[i])" + ] + }, + { + "cell_type": "markdown", + "id": "2dedaecc-a46d-4c7b-8b34-c761467fdd89", + "metadata": {}, + "source": [ + "## CFIS" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "9900ef4e-76c4-42ab-926a-334ea26eb401", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[1mDownloading and preparing dataset Unknown size (download: Unknown size, generated: Unknown size, total: Unknown size) to /local/home/az264973/tensorflow_datasets/cfis/parametric_1k/0.5.0...\u001b[0m\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Generating splits...: 0%| | 0/1 [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "style.use('default')\n", + "figure(figsize=(10,10))\n", + "for i in range(16):\n", + " subplot(4,4,i+1)\n", + " imshow(gal_images[i])" + ] + }, + { + "cell_type": "markdown", + "id": "867fc339-811c-4b81-8b0a-ebb9744e9a6e", + "metadata": {}, + "source": [ + "## Simple" + ] + }, + { + "cell_type": "code", + "execution_count": 10, + "id": "286ab869-5a32-4cf0-81d2-9a6f5edc7f90", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\u001b[1mDownloading and preparing dataset Unknown size (download: Unknown size, generated: Unknown size, total: Unknown size) to /local/home/az264973/tensorflow_datasets/simple/small_1k_g1m/1.0.0...\u001b[0m\n" + ] + }, + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Generating splits...: 0%| | 0/1 [00:00" + ] + }, + "metadata": {}, + "output_type": "display_data" + } + ], + "source": [ + "style.use('default')\n", + "figure(figsize=(10,10))\n", + "for i in range(16):\n", + " subplot(4,4,i+1)\n", + " imshow(gal_images[i])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e7c35c9b-f654-40c5-a681-0601860914b2", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "917671e8-49a5-47bf-8e20-cf71ad5a3871", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ba7f3272-b363-4876-a3b1-281faa0d9005", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.8" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/autometacal_vs_ngmix_test.ipynb b/notebooks/autometacal_vs_ngmix_test.ipynb index c769d8d..ce38ed4 100644 --- a/notebooks/autometacal_vs_ngmix_test.ipynb +++ b/notebooks/autometacal_vs_ngmix_test.ipynb @@ -110,9 +110,6 @@ "execution_count": 4, "id": "e6588dba-e224-4784-b5fa-40e16f9546bd", "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [] }, "outputs": [], @@ -307,207 +304,6 @@ "\n" ] }, - { - "cell_type": "raw", - "id": "c60e7c2c-1f7a-49e6-88c1-000721937764", - "metadata": {}, - "source": [ - "# We now create the autometacal function which returns (e, R)\n", - "@tf.function\n", - "def get_autometacal_shape(im, psf, rpsf):\n", - " method = lambda x: amc.get_moment_ellipticities(x, scale=0.263, fwhm=weight_fwhm)\n", - " return amc.get_metacal_response(im, psf, rpsf, method)\n", - "\n", - "\n", - "@tf.function\n", - "def get_finitediff_shape(im, psf, rpsf):\n", - " method = lambda x: amc.get_moment_ellipticities(x, scale=0.263, fwhm=weight_fwhm)\n", - " return amc.get_metacal_response_finitediff(im, psf, rpsf,0.01, method) " - ] - }, - { - "cell_type": "raw", - "id": "4143957a-bd8e-4185-b4c9-9b5c3623cfb3", - "metadata": {}, - "source": [ - "\n", - "def get_metacal_response_finitediff(gal_image,psf_image,reconv_psf_image,step,method):\n", - " \"\"\"\n", - " Gets shear response as a finite difference operation, \n", - " instead of automatic differentiation.\n", - " \"\"\"\n", - " \n", - " img0s = generate_mcal_image(\n", - " gal_image,\n", - " psf_image,\n", - " reconv_psf_image,\n", - " [[0,0]]\n", - " ) \n", - " img1p = generate_mcal_image(\n", - " gal_image,\n", - " psf_image,\n", - " reconv_psf_image,\n", - " [[step,0]]\n", - " ) \n", - " img1m = generate_mcal_image(\n", - " gal_image,\n", - " psf_image,\n", - " reconv_psf_image,\n", - " [[-step,0]]\n", - " ) \n", - " img2p = generate_mcal_image(\n", - " gal_image,\n", - " psf_image,\n", - " reconv_psf_image,\n", - " [[0,step]]\n", - " ) \n", - " img2m = generate_mcal_image(\n", - " gal_image,\n", - " psf_image,\n", - " reconv_psf_image,\n", - " [[0,-step]]\n", - " ) \n", - " \n", - " g0s = method(img0s)\n", - " g1p = method(img1p)\n", - " g1m = method(img1m)\n", - " g2p = method(img2p)\n", - " g2m = method(img2m)\n", - " \n", - " R11 = (g1p[:,0]-g1m[:,0])/(2*step)\n", - " R21 = (g1p[:,1]-g1m[:,1])/(2*step) \n", - " R12 = (g2p[:,0]-g2m[:,0])/(2*step)\n", - " R22 = (g2p[:,1]-g2m[:,1])/(2*step)\n", - " \n", - " #N. B.:The matrix is correct. \n", - " #The transposition will swap R12 with R21 across a batch correctly.\n", - " R = tf.transpose(tf.convert_to_tensor(\n", - " [[R11,R21],\n", - " [R12,R22]],dtype=tf.float32)\n", - " )\n", - " \n", - " ellip_dict = {\n", - " 'noshear':g0s,\n", - " '1p':g1p,\n", - " '1m':g1m,\n", - " '2p':g2p,\n", - " '2m':g2m, \n", - " } \n", - "\n", - " return ellip_dict, R" - ] - }, - { - "cell_type": "raw", - "id": "8ff2370e-707d-4404-a856-61fa38215f90", - "metadata": {}, - "source": [ - "def generate_mcal_image(gal_images,\n", - " psf_images,\n", - " reconvolution_psf_image,\n", - " g):\n", - " \"\"\" Generate a metacalibrated image given input and target PSFs.\n", - " \n", - " Args: \n", - " gal_images: tf.Tensor or np.array\n", - " (batch_size, N, N ) image of galaxies\n", - " psf_images: tf.Tensor or np.array\n", - " (batch_size, N, N ) image of psf model\n", - " reconvolution_psf_image: tf.Tensor\n", - " (N, N ) tensor of reconvolution psf model\n", - " g: tf.Tensor or np.array\n", - " [batch_size, 2] input shear\n", - " Returns:\n", - " img: tf.Tensor\n", - " tf tensor containing image of galaxy after deconvolution by psf_deconv, \n", - " shearing by g, and reconvolution with reconvolution_psf_image.\n", - " \n", - " \"\"\"\n", - " #cast stuff as float32 tensors\n", - " gal_images = tf.convert_to_tensor(gal_images, dtype=tf.float32) \n", - " psf_images = tf.convert_to_tensor(psf_images, dtype=tf.float32) \n", - " reconvolution_psf_image = tf.convert_to_tensor(reconvolution_psf_image, dtype=tf.float32) \n", - " g = tf.convert_to_tensor(g, dtype=tf.float32) \n", - " \n", - " #Get batch info\n", - " batch_size, nx, ny = gal_images.get_shape().as_list() \n", - " \n", - " #add pads in real space\n", - " padfactor = 3 #total width of image after padding\n", - " fact = (padfactor - 1)//2 #how many image sizes to one direction\n", - " paddings = tf.constant([[0, 0,], [nx*fact, nx*fact], [ny*fact, ny*fact]])\n", - " \n", - " padded_gal_images = tf.pad(gal_images,paddings)\n", - " padded_psf_images = tf.pad(psf_images,paddings)\n", - " padded_reconvolution_psf_image = tf.pad(reconvolution_psf_image,paddings)\n", - " \n", - " #Convert galaxy images to k space\n", - " im_shift = tf.signal.ifftshift(padded_gal_images,axes=[1,2]) # The ifftshift is to remove the phase for centered objects\n", - " im_complex = tf.cast(im_shift, tf.complex64)\n", - " im_fft = tf.signal.fft2d(im_complex)\n", - " imk = tf.signal.fftshift(im_fft, axes=[1,2])#the fftshift is to put the 0 frequency at the center of the k image\n", - " \n", - " #Convert psf images to k space \n", - " psf_complex = tf.cast(padded_psf_images, tf.complex64)\n", - " psf_fft = tf.signal.fft2d(psf_complex)\n", - " psf_fft_abs = tf.abs(psf_fft)\n", - " psf_fft_abs_complex = tf.cast(psf_fft_abs,tf.complex64)\n", - " kpsf = tf.signal.fftshift(psf_fft_abs_complex,axes=[1,2])\n", - "\n", - " #Convert reconvolution psf image to k space \n", - " rpsf_complex = tf.cast(padded_reconvolution_psf_image, tf.complex64)\n", - " rpsf_fft = tf.signal.fft2d(rpsf_complex)\n", - " rpsf_fft_abs = tf.abs(rpsf_fft)\n", - " psf_fft_abs_complex = tf.cast(rpsf_fft_abs,tf.complex64)\n", - " krpsf = tf.signal.fftshift(psf_fft_abs_complex,axes=[1,2])\n", - "\n", - " # Compute Fourier mask for high frequencies\n", - " # careful, this is not exactly the correct formula for fftfreq\n", - " kx, ky = tf.meshgrid(tf.linspace(-0.5,0.5,padfactor*nx),\n", - " tf.linspace(-0.5,0.5,padfactor*ny))\n", - " mask = tf.cast(tf.math.sqrt(kx**2 + ky**2) <= 0.5, dtype='complex64')\n", - " mask = tf.expand_dims(mask, axis=0)\n", - "\n", - " # Deconvolve image from input PSF\n", - " im_deconv = imk * ( (1./(kpsf+1e-10))*mask)\n", - "\n", - " # Apply shear\n", - " im_sheared = gf.shear(tf.expand_dims(im_deconv,-1), g[...,0], g[...,1])[...,0]\n", - "\n", - " # Reconvolve with target PSF\n", - " im_reconv = tf.signal.ifft2d(tf.signal.ifftshift(im_sheared * krpsf * mask))\n", - "\n", - " # Compute inverse Fourier transform\n", - " img = tf.math.real(tf.signal.fftshift(im_reconv))\n", - "\n", - " return img[:,fact*nx:-fact*nx,fact*ny:-fact*ny]\n", - "\n", - "def get_metacal_response(gal_images,\n", - " psf_images,\n", - " reconvolution_psf_image,\n", - " method):\n", - " \"\"\"\n", - " Convenience function to compute the shear response\n", - " \"\"\" \n", - " gal_images = tf.convert_to_tensor(gal_images, dtype=tf.float32)\n", - " psf_images = tf.convert_to_tensor(psf_images, dtype=tf.float32)\n", - " reconvolution_psf_image = tf.convert_to_tensor(reconvolution_psf_image, dtype=tf.float32)\n", - " batch_size, _ , _ = gal_images.get_shape().as_list()\n", - " g = tf.zeros([batch_size, 2])\n", - " with tf.GradientTape() as tape:\n", - " tape.watch(g)\n", - " # Measure ellipticity under metacal\n", - " e = method(generate_mcal_image(gal_images,\n", - " psf_images,\n", - " reconvolution_psf_image,\n", - " g))\n", - " \n", - " # Compute response matrix\n", - "\n", - " R = tape.batch_jacobian(e, g)\n", - " return e, R\n" - ] - }, { "cell_type": "code", "execution_count": 6, @@ -1232,7 +1028,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.5" + "version": "3.8.8" } }, "nbformat": 4, diff --git a/tests/test_metacal_ngmix.py b/tests/test_metacal_ngmix.py index ca6ec80..34dc76e 100644 --- a/tests/test_metacal_ngmix.py +++ b/tests/test_metacal_ngmix.py @@ -264,7 +264,7 @@ def get_finitediff_shape(im, psf, rpsf): st = make_struct(res=sres, obs=obsdict[stype], shear_type=stype) dlist.append(st) - # Same thing with autometacal + # Same thing with autometacal im = obs.image.reshape(1,45,45).astype('float32') psf = obs.psf.image.reshape(1,45,45).astype('float32') rpsf = obsdict['noshear'].psf.image.reshape(1,45,45).astype( diff --git a/tests/test_tf_ngmix.py b/tests/test_tf_ngmix.py index 6338886..b50244e 100644 --- a/tests/test_tf_ngmix.py +++ b/tests/test_tf_ngmix.py @@ -1,45 +1,130 @@ # This module tests our implementation against ngmix from numpy.testing import assert_allclose + import autometacal import ngmix import numpy as np import tensorflow as tf +import galsim -scale = .2 stamp_size=51 Ngals = 1000 + +def make_data(rng, noise, shear): + """ + simulate an exponential object with moffat psf + + Parameters + ---------- + rng: np.random.RandomState + The random number generator + noise: float + Noise for the image + shear: (g1, g2) + The shear in each component + + Returns + ------- + ngmix.Observation + """ + + psf_noise = 1.0e-6 + + scale = 0.263 + stamp_size = 51 + psf_fwhm = 0.9 + gal_hlr = 0.5 + # We keep things centered for now + dy, dx = 0.,0. #rng.uniform(low=-scale/2, high=scale/2, size=2) + + psf = galsim.Moffat(beta=2.5, fwhm=psf_fwhm).shear( + g1=0.02, + g2=-0.01, + ) + + obj0 = galsim.Exponential(half_light_radius=gal_hlr).shear( + g1=shear[0], + g2=shear[1], + ) + + obj = galsim.Convolve(psf, obj0) + + psf_im = psf.drawImage(nx=stamp_size, ny=stamp_size, scale=scale).array + im = obj.drawImage(nx=stamp_size, ny=stamp_size, scale=scale).array + + psf_im += rng.normal(scale=psf_noise, size=psf_im.shape) + im += rng.normal(scale=noise, size=im.shape) + + cen = (np.array(im.shape)-1.0)/2.0 + psf_cen = (np.array(psf_im.shape)-1.0)/2.0 + + jacobian = ngmix.DiagonalJacobian( + row=cen[0] + dy/scale, + col=cen[1] + dx/scale, + scale=scale, + ) + psf_jacobian = ngmix.DiagonalJacobian( + row=psf_cen[0], + col=psf_cen[1], + scale=scale, + ) + + wt = im*0 + 1.0/noise**2 + psf_wt = psf_im*0 + 1.0/psf_noise**2 + + psf_obs = ngmix.Observation( + psf_im, + weight=psf_wt, + jacobian=psf_jacobian, + ) + + obs = ngmix.Observation( + im, + weight=wt, + jacobian=jacobian, + psf=psf_obs, + ) + + return obs + + def test_tf_ngmix(): """ This test generates a simple galaxy and measure moments with ngmix, vs. tf_ngmix. """ + scale = 0.263 + imlist = [] + results_ngmix=[] + rng =np.random.RandomState(31415) + fitter = ngmix.gaussmom.GaussMom(fwhm=1.2) + for i in range(Ngals): + obs = make_data(rng, + 1e-6, + [np.random.uniform(-.7,.7), + np.random.uniform(-.7,.7)]) + + e = fitter._measure_moments(obs)['e'] + results_ngmix.append(e) + imlist.append(obs.image.reshape(stamp_size,stamp_size).astype('float32')) - gals, _ = autometacal.datasets.galaxies.make_data(Ngals=Ngals, snr=100, - gal_g1=np.random.uniform(-.7,.7,Ngals), - gal_g2=np.random.uniform(-.7,.7,Ngals), - scale=scale) weight_fwhm = scale*stamp_size/2 # <- this sets everything for the window function - results_ngmix=[] - + + # ngmix version fitter = ngmix.gaussmom.GaussMom(fwhm=weight_fwhm) - for gal in gals: - obs = ngmix.Observation(gal.numpy(),jacobian=ngmix.DiagonalJacobian(row=stamp_size//2, - col=stamp_size//2, - scale=scale)) - e = fitter._measure_moments(obs)['e'] - results_ngmix.append(e) + results_ngmix = np.array(results_ngmix) - + gals = tf.stack(imlist) # our version: @tf.function def get_ellipticity(im): - return autometacal.get_moment_ellipticities(im, scale=0.2, fwhm=weight_fwhm) + return autometacal.get_moment_ellipticities(im, scale=scale, fwhm=1.2) result_tf_ngmix = get_ellipticity(gals) assert_allclose(results_ngmix,result_tf_ngmix,rtol=1e-6,atol=1e-6) if __name__=='__main__': - test_tf_ngmix() \ No newline at end of file + test_tf_ngmix()