New data generators

autometacal/python/__init__.py

@@ -2,4 +2,4 @@
 import autometacal.python.tf_ngmix as tf_ngmix
 from autometacal.python.metacal import generate_mcal_image, get_metacal_response
 from autometacal.python.fitting import fit_multivariate_gaussian, get_ellipticity
-from autometacal.python.moments import get_moment_ellipticities, gaussian_moments
+from autometacal.python.moments import gaussian_moments, get_moment_ellipticities

autometacal/python/datasets/__init__.py

@@ -1,3 +1,4 @@
 """gal_gen dataset."""
 
 from .gal_gen import GalGen
+from .galaxies import *

autometacal/python/datasets/gal_gen.py

@@ -2,8 +2,8 @@
 import tensorflow_datasets as tfds
 import tensorflow as tf
 import numpy as np
-
-from .galaxies import gs_generate_images, gs_drawKimage
+from scipy.stats import truncnorm
+from .galaxies import simple
 
 _DESCRIPTION = "This tfds generates random toy-model galaxy stamps."
 _CITATION = "{NEEDED}"
@@ -12,7 +12,13 @@
 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):
+  def __init__(self,
+               *,
+               dataset_size=None,
+               stamp_size=None,
+               pixel_scale=None,
+               flux=None,
+               **kwargs):
     """BuilderConfig for SQUAD.
     Args:
       pixel_scale: pixel_scale of the image in arcsec/pixel.
@@ -33,24 +39,28 @@ def __init__(self, *,dataset_size=None, stamp_size=None, pixel_scale=None, flux=
     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 = """\ 
+  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)
+      GalGenConfig(name="simple_100",
+                   dataset_size=100,
+                   stamp_size=51,
+                   pixel_scale=.2,
+                   flux=1.e5),
+      GalGenConfig(name="simple_1k",
+                   dataset_size=1000,
+                   stamp_size=51,
+                   pixel_scale=.2,
+                   flux=1.e5),
    ]
 
   VERSION = tfds.core.Version('0.1.0')
@@ -62,70 +72,57 @@ def _info(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"),
+      features=tfds.features.FeaturesDict(
+        {'label': tfds.features.Tensor(shape=[2], dtype=tf.float32),
+          'gal_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=("gal_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):
+    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
+      )
+    }
+
+  def _generate_examples(self, dataset_size, stamp_size, pixel_scale, flux):
     """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() }
+    a, b = -.7/.3, .7/.3 #(max_ellip/ellip_sigma)
 
 
-def decomplexify(arr):
-  arr=np.array([arr.real,arr.imag])
-  return arr
+    for i in range(dataset_size):
+      g1=g2=1
+      while g1**2+g2**2>1:
+        g1 = truncnorm.rvs(a, b, loc=0, scale=.3)
+        g2 = truncnorm.rvs(a, b, loc=0, scale=.3)
+
+      gals, psfs = simple(
+        snr = 100,
+        scale = pixel_scale,
+        stamp_size = stamp_size,
+        psf_fwhm = 0.9,
+        gal_hlr = 0.7,
+        gal_g1 = g1,
+        gal_g2 = g2,
+        flux=flux
+      )
 
-def recomplexify(arl):
-  return arl[0]+1j*arl[1]
+
+      #get example
+      yield '%d'%i, {'gal_image': gal.numpy(), #galaxy image
+                     'psf_image': psf.numpy(), #psf image
+                     'label': [g1,g2]} #ellipticity