From 954ab426559b761cf3816ed27e9b4377251de288 Mon Sep 17 00:00:00 2001
From: Bretonniere Hubert <hbretonniere@tycho57.obspm.fr>
Date: Tue, 11 Aug 2020 14:54:02 +0200
Subject: [PATCH 001/149] first (buggued) version of candels

---
 galaxy2galaxy/data_generators/candels.py      | 446 ++++++++++++++++++
 .../check_candels_emulation.py                |  18 +
 2 files changed, 464 insertions(+)
 create mode 100644 galaxy2galaxy/data_generators/candels.py
 create mode 100644 galaxy2galaxy/data_generators/check_candels_emulation.py

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
new file mode 100644
index 0000000..799e3ad
--- /dev/null
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -0,0 +1,446 @@
+""" HSC Datasets """
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from . import hsc_utils
+from . import astroimage_utils
+
+from tensor2tensor.data_generators import generator_utils
+from tensor2tensor.data_generators import image_utils
+from tensor2tensor.data_generators import problem
+from tensor2tensor.layers import modalities, common_layers
+from tensor2tensor.utils import metrics
+
+from galaxy2galaxy.utils import registry
+
+from scipy.ndimage import gaussian_filter
+
+import tensorflow as tf
+import numpy as np
+import fits2hdf.pyhdfits as fits
+from astropy.table import Table
+from astropy.visualization import make_lupton_rgb
+import h5py
+import glob
+import os
+
+
+# HSC default pixel scale
+_HSC_PIXEL_SCALE=0.168 #arcsec
+# Path to sql files for HSC samples
+_HSC_SAMPLE_SQL_DIR=os.path.join(os.path.dirname(os.path.realpath(__file__)), 'data')
+
+def _resize_image(im, size):
+  centh = im.shape[0]/2
+  centw = im.shape[1]/2
+  lh, rh = int(centh-size/2), int(centh+size/2)
+  lw, rw = int(centw-size/2), int(centw+size/2)
+  cropped = im[lh:rh, lw:rw, :]
+  assert cropped.shape[0]==size and cropped.shape[1]==size, f"Wrong size! Still {cropped.shape}"
+  return cropped
+
+
+@registry.register_problem
+class Img2imgCandelsAll(astroimage_utils.AstroImageProblem):
+  """Base class for image problems created from HSC Public Data Release.
+  """
+
+#   @property
+#   def dataset_splits(self):
+#     """Splits of data to produce and number of output shards for each.
+#        Note that each shard will be produced in parallel.
+#        We are going to split the GalSim data into shards of 1000 galaxies each,
+#        with 80 shards for training, 2 shards for validation.
+#     """
+#     return [{
+#         "split": problem.DatasetSplit.TRAIN,
+#         "shards": 1,
+#     }, {
+#         "split": problem.DatasetSplit.EVAL,
+#         "shards": 0,
+#     }]
+
+  # START: Subclass interface
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.img_len = 128
+    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850w']
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "targets": None}
+    p.add_hparam("psf", None)
+
+  @property
+  def num_bands(self):
+    """Number of bands."""
+    p = self.get_hparams()
+    return len(p.filters)
+
+  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
+    """ Generator yielding individual postage stamps.
+    """
+    p = self.get_hparams()
+    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='egs   ')] = 'EGS'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='UDS   ')] = 'UDS'
+    
+    cube_psf = np.zeros((167, 167, len(p.filters)))
+    
+    for i, filter in enumerate(p.filters):
+        cube_psf[:, :, i] = fits.open(tmp_dir + 'psf_' + filter +'.fits')[0].data
+
+#     # Step 1: Maybe regenerate the dataset
+#     self.maybe_build_dataset(tmp_dir)
+#     import timeit
+
+    # Step 2: Extract postage stamps, resize them to requested size
+    ''' Genrate GDS and GDN'''
+    for field in ['GDS', 'GDN'] :
+        print(f"\n generating{field}\n")
+
+        n=0
+        index = 0
+        im = np.zeros((128, 128, 5))
+        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
+        for i in sub_cat['RB_ID']:
+                if i == index :
+                    continue
+                index = i
+                print(field, i)
+#             try:
+                gal = glob.glob(os.path.join(tmp_dir, field, p.filters[0])+'/galaxy_'+str(index)+'_*')[0]
+                im[:,:, 0] = fits.open(gal)[0].data
+                for i, filter in enumerate(p.filters[1:4]):
+                    try :
+                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
+                        im[:,:, i+1] = fits.open(tmp_file)[0].data
+                    except Exception:
+                        print('Galaxy not seen in every filter')
+                        continue
+                im = _resize_image(im, p.img_len)
+
+                n+=1
+#                 attributes = {k: all_cat[k][0] for k in p.attributes}
+                flag = [1, 1, 1, 1, 0]
+
+
+    #           imCp = np.fft.fft2(psf, s=stamp_size, stamp_size)
+                ps = np.random.normal(0, 1e-4, (p.img_len, p.img_len))
+                im_psf = cube_psf
+
+#                 serialized_output = {"image/encoded": [im.tostring()],
+#                 "image/format": ["raw"],
+#                 "psf/encoded": [im_psf.tostring()],
+#                 "psf/format": ["raw"],
+#                 "ps/encoded": [ps.tostring()],
+#                 "ps/format": ["raw"],
+#                 "flag/filters":flag}
+                
+                serialized_output = {"image/encoded": [im.tostring()],
+                "image/format": ["raw"],
+                "psf/encoded": [im_psf.tostring()],
+                "psf/format": ["raw"]}
+        
+        
+                if n > 20:
+                    break
+                yield serialized_output
+#             except Exception:
+#                 print('No image corresponding to the Index')
+#                 n+=1
+  def example_reading_spec(self):
+    """Define how data is serialized to file and read back.
+
+    Returns:
+      data_fields: A dictionary mapping data names to its feature type.
+      data_items_to_decoders: A dictionary mapping data names to TF Example
+         decoders, to be used when reading back TF examples from disk.
+    """
+    p = self.get_hparams()
+
+    data_fields = {
+        "image/encoded": tf.FixedLenFeature((), tf.string),
+        "image/format": tf.FixedLenFeature((), tf.string),
+
+        "psf/encoded": tf.FixedLenFeature((), tf.string),
+        "psf/format": tf.FixedLenFeature((), tf.string),
+
+#         "ps/encoded": tf.FixedLenFeature((), tf.string),
+#         "ps/format": tf.FixedLenFeature((), tf.string),
+    }
+
+
+    # Adds additional attributes to be decoded as specified in the configuration
+    if hasattr(p, 'attributes'):
+        for k in p.attributes:
+            data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
+
+    data_items_to_decoders = {
+        "inputs": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="image/encoded",
+                format_key="image/format",
+                channels=self.num_bands,
+                shape=[p.img_len, p.img_len, self.num_bands],
+                dtype=tf.float32),
+
+        "psf": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="psf/encoded",
+                format_key="psf/format",
+                channels=self.num_bands,
+                shape=[p.img_len, p.img_len, self.num_bands],
+                dtype=tf.float32),
+
+#         "ps": tf.contrib.slim.tfexample_decoder.Image(
+#                 image_key="ps/encoded",
+#                 format_key="ps/format",
+#                 channels=self.num_bands,
+#                 shape=[p.img_len, p.img_len // 2 + 1],
+#                 dtype=tf.float32),
+    }
+    if hasattr(p, 'attributes'):
+        for k in p.attributes:
+            data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
+
+@registry.register_problem
+class Img2imgCandelsGoods(astroimage_utils.AstroImageProblem):
+  """Base class for image problems created from HSC Public Data Release.
+  """
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each.
+       Note that each shard will be produced in parallel.
+       We are going to split the GalSim data into shards of 1000 galaxies each,
+       with 80 shards for training, 2 shards for validation.
+    """
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 1,
+    }]
+
+  # START: Subclass interface
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.img_len = 128
+    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w']#, 'acs_f850w']
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "targets": None}
+    p.add_hparam("psf", None)
+
+  @property
+  def num_bands(self):
+    """Number of bands."""
+    p = self.get_hparams()
+    return len(p.filters)
+
+  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
+    """ Generator yielding individual postage stamps.
+    """
+    p = self.get_hparams()
+    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
+    
+    cube_psf = np.zeros((167, 167, len(p.filters)))
+    
+    for i, filter in enumerate(p.filters):
+        cube_psf[:, :, i] = fits.open(tmp_dir + 'psf_' + filter +'.fits')[0].data
+
+#     # Step 1: Maybe regenerate the dataset
+#     self.maybe_build_dataset(tmp_dir)
+#     import timeit
+
+    # Step 2: Extract postage stamps, resize them to requested size
+    ''' Genrate GDS and GDN'''
+    for field in ['GDS', 'GDN'] :
+        print(f"\n generating{field}\n")
+        n=0
+        index = 0
+        im = np.zeros((128, 128, len(p.filters)))
+        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
+        for i in sub_cat['RB_ID']:
+                if i == index :
+                    continue
+                index = i
+                print(field, i)
+#             try:
+                for i, filter in enumerate(p.filters):
+#                     try :
+                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
+                        im[:,:, i] = fits.open(tmp_file)[0].data
+#                     except Exception:
+#                         print('Galaxy not seen in every filter')
+#                         continue
+                im = _resize_image(im, p.img_len)
+
+                n+=1
+                if hasattr(p, 'attributes'):
+                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
+                    print(attributes)
+                else:
+                    attributes=None
+
+
+    #           imCp = np.fft.fft2(psf, s=stamp_size, stamp_size)
+                ps = np.random.normal(0, 1e-4, (p.img_len, p.img_len))
+                im_psf = cube_psf
+                im_psf = _resize_image(im_psf, p.img_len)
+
+                print(f"\n shapes : {im.shape}, {im_psf.shape}, !!!!! \n")
+
+                serialized_output = {"image/encoded": [im.tostring()],
+                "image/format": ["raw"],
+                "psf/encoded": [im_psf.tostring()],
+                "psf/format": ["raw"],
+                "ps/encoded": [ps.tostring()],
+                "ps/format": ["raw"]}
+                
+                if attributes is not None:
+                    for k in attributes:
+                        serialized_output['attrs/'+k] = [attributes[k]]
+                if n > 5:
+                    break
+                yield serialized_output
+#             except Exception:
+#                 print('No image corresponding to the Index')
+#                 n+=1
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    """ Preprocess the examples, can be used for further augmentation or
+    image standardization.
+    """
+    p = self.get_hparams()
+    image = example["inputs"]
+
+    # Clip to 1 the values of the image
+    # image = tf.clip_by_value(image, -1, 1)
+
+    # Aggregate the conditions
+    if hasattr(p, 'attributes'):
+      example['attributes'] = tf.stack([example[k] for k in p.attributes])
+
+    example["inputs"] = image
+    example["targets"] = image
+    return example
+
+  def example_reading_spec(self):
+    """Define how data is serialized to file and read back.
+
+    Returns:
+      data_fields: A dictionary mapping data names to its feature type.
+      data_items_to_decoders: A dictionary mapping data names to TF Example
+         decoders, to be used when reading back TF examples from disk.
+    """
+    p = self.get_hparams()
+
+    data_fields = {
+        "image/encoded": tf.FixedLenFeature((), tf.string),
+        "image/format": tf.FixedLenFeature((), tf.string),
+
+        "psf/encoded": tf.FixedLenFeature((), tf.string),
+        "psf/format": tf.FixedLenFeature((), tf.string),
+
+        "ps/encoded": tf.FixedLenFeature((), tf.string),
+        "ps/format": tf.FixedLenFeature((), tf.string),
+    }
+
+
+    # Adds additional attributes to be decoded as specified in the configuration
+    if hasattr(p, 'attributes'):
+        for k in p.attributes:
+            data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
+
+    data_items_to_decoders = {
+        "inputs": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="image/encoded",
+                format_key="image/format",
+                channels=self.num_bands,
+                shape=[p.img_len, p.img_len, self.num_bands],
+                dtype=tf.float32),
+
+        "psf": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="psf/encoded",
+                format_key="psf/format",
+                channels=self.num_bands,
+                shape=[p.img_len, p.img_len, self.num_bands],
+                dtype=tf.float32),
+
+        "ps": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="ps/encoded",
+                format_key="ps/format",
+                channels=self.num_bands,
+                shape=[p.img_len, p.img_len // 2 + 1],
+                dtype=tf.float32),
+    }
+    if hasattr(p, 'attributes'):
+        for k in p.attributes:
+            data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
+
+    return data_fields, data_items_to_decoders
+#     END: Subclass interface
+    
+    
+  @property
+  def is_generate_per_split(self):
+    return False
+
+
+@registry.register_problem
+class Attrs2imgCandelsAll64Euclid(Img2imgCandelsAll):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 128
+    p.example_per_shard = 1
+    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850w']
+
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = []#['mag', 're', 'q']
+    
+    
+@registry.register_problem
+class Attrs2imgCandelsGoods64Euclid(Img2imgCandelsGoods):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 128
+    p.example_per_shard = 1
+    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w']#, 'acs_f850l']
+
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+#     p.attributes = ['mag', 're', 'q']
+    
diff --git a/galaxy2galaxy/data_generators/check_candels_emulation.py b/galaxy2galaxy/data_generators/check_candels_emulation.py
new file mode 100644
index 0000000..0882721
--- /dev/null
+++ b/galaxy2galaxy/data_generators/check_candels_emulation.py
@@ -0,0 +1,18 @@
+from galaxy2galaxy import problems
+import tensorflow as tf
+import matplotlib.pyplot as plt
+import numpy as np
+Modes = tf.estimator.ModeKeys
+
+candels = problems.problem('attrs2img_candels_goods64_euclid')
+dset = candels.dataset(Modes.TRAIN, data_dir='/data57/hbretonniere/datasets/test_candels')
+batch = dset.batch(1)
+
+print("data set :", dset, "\n \n")
+
+print("batch : ", batch, "\n \n")
+
+
+it = batch.make_one_shot_iterator().get_next()
+sess = tf.Session()
+batch_1 = sess.run(it)
\ No newline at end of file

From c95b783371da3f14faa2236db417f07f6712529f Mon Sep 17 00:00:00 2001
From: Bretonniere Hubert <hbretonniere@tycho57.obspm.fr>
Date: Tue, 11 Aug 2020 14:57:43 +0200
Subject: [PATCH 002/149] add candels to the possible problems

---
 galaxy2galaxy/data_generators/all_problems.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/data_generators/all_problems.py b/galaxy2galaxy/data_generators/all_problems.py
index fd8064c..38a4eda 100644
--- a/galaxy2galaxy/data_generators/all_problems.py
+++ b/galaxy2galaxy/data_generators/all_problems.py
@@ -14,6 +14,7 @@
 try:
   import galsim
   MODULES += ["galaxy2galaxy.data_generators.cosmos"]
+  MODULES += ["galaxy2galaxy.data_generators.candels"]
 except:
   print("Could not import GalSim, excluding some data generators")
 

From a519d357be29875e40cf9d7142e521bcb0a9a5c6 Mon Sep 17 00:00:00 2001
From: Hubert Bretonniere <hubert.bretonniere@ias.u-psud.fr>
Date: Fri, 23 Oct 2020 11:06:50 +0200
Subject: [PATCH 003/149] new release

---
 galaxy2galaxy/data_generators/all_problems.py |   2 +-
 galaxy2galaxy/data_generators/candels.py      | 507 +++++++++++++-----
 .../check_candels_emulation.py                |   2 +-
 galaxy2galaxy/data_generators/cosmos.py       | 131 ++++-
 galaxy2galaxy/models/autoencoders.py          |  12 +-
 galaxy2galaxy/models/autoregressive.py        |   2 +-
 6 files changed, 495 insertions(+), 161 deletions(-)

diff --git a/galaxy2galaxy/data_generators/all_problems.py b/galaxy2galaxy/data_generators/all_problems.py
index 38a4eda..aca90af 100644
--- a/galaxy2galaxy/data_generators/all_problems.py
+++ b/galaxy2galaxy/data_generators/all_problems.py
@@ -14,7 +14,7 @@
 try:
   import galsim
   MODULES += ["galaxy2galaxy.data_generators.cosmos"]
-  MODULES += ["galaxy2galaxy.data_generators.candels"]
+#   MODULES += ["galaxy2galaxy.data_generators.candels"]
 except:
   print("Could not import GalSim, excluding some data generators")
 
diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 799e3ad..c603f1c 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -25,6 +25,8 @@
 import h5py
 import glob
 import os
+from skimage.transform import resize
+
 
 
 # HSC default pixel scale
@@ -41,118 +43,123 @@ def _resize_image(im, size):
   assert cropped.shape[0]==size and cropped.shape[1]==size, f"Wrong size! Still {cropped.shape}"
   return cropped
 
-
 @registry.register_problem
-class Img2imgCandelsAll(astroimage_utils.AstroImageProblem):
-  """Base class for image problems created from HSC Public Data Release.
+class Img2imgCandelsOnefilter(astroimage_utils.AstroImageProblem):
+  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
   """
 
-#   @property
-#   def dataset_splits(self):
-#     """Splits of data to produce and number of output shards for each.
-#        Note that each shard will be produced in parallel.
-#        We are going to split the GalSim data into shards of 1000 galaxies each,
-#        with 80 shards for training, 2 shards for validation.
-#     """
-#     return [{
-#         "split": problem.DatasetSplit.TRAIN,
-#         "shards": 1,
-#     }, {
-#         "split": problem.DatasetSplit.EVAL,
-#         "shards": 0,
-#     }]
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each.
+      .
+    """
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 0,
+    }]
 
   # START: Subclass interface
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.img_len = 128
-    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850w']
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": None,
                     "targets": None}
     p.add_hparam("psf", None)
-
+    
   @property
   def num_bands(self):
     """Number of bands."""
-    p = self.get_hparams()
-    return len(p.filters)
+    return 1
 
   def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
-    """ Generator yielding individual postage stamps.
+    """ 
+    Generator yielding individual postage stamps.
     """
+    
     p = self.get_hparams()
+    
+    '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='egs   ')] = 'EGS'
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='UDS   ')] = 'UDS'
     
-    cube_psf = np.zeros((167, 167, len(p.filters)))
-    
-    for i, filter in enumerate(p.filters):
-        cube_psf[:, :, i] = fits.open(tmp_dir + 'psf_' + filter +'.fits')[0].data
-
-#     # Step 1: Maybe regenerate the dataset
-#     self.maybe_build_dataset(tmp_dir)
-#     import timeit
-
+    ''' Load the psfs for each filter and resize'''
+    psf = np.expand_dims(fits.open(tmp_dir + '/psfs/psf_acs_f606w.fits')[0].data, -1)
+    psf = _resize_image(psf, p.img_len)
     # Step 2: Extract postage stamps, resize them to requested size
-    ''' Genrate GDS and GDN'''
-    for field in ['GDS', 'GDN'] :
+    ''' Loop on the fields'''
+    for n_field, field in enumerate(['GDS', 'GDN', 'EGS', 'UDS', 'COSMOS']):
+        
         print(f"\n generating{field}\n")
-
-        n=0
+        n_gal_creat = 0
         index = 0
-        im = np.zeros((128, 128, 5))
+        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
         sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-        for i in sub_cat['RB_ID']:
-                if i == index :
+
+        ''' Loop on all the galaxies of the field '''
+        for gal in sub_cat['RB_ID']:
+                print(f"\n generating{gal}\n")
+
+                if gal == index :     # To take care of the redudency inside the cat
                     continue
-                index = i
-                print(field, i)
-#             try:
-                gal = glob.glob(os.path.join(tmp_dir, field, p.filters[0])+'/galaxy_'+str(index)+'_*')[0]
-                im[:,:, 0] = fits.open(gal)[0].data
-                for i, filter in enumerate(p.filters[1:4]):
-                    try :
-                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
-                        im[:,:, i+1] = fits.open(tmp_file)[0].data
-                    except Exception:
-                        print('Galaxy not seen in every filter')
-                        continue
-                im = _resize_image(im, p.img_len)
+                index = gal
 
-                n+=1
-#                 attributes = {k: all_cat[k][0] for k in p.attributes}
-                flag = [1, 1, 1, 1, 0]
+#             try:
 
+                tmp_file = glob.glob(os.path.join(tmp_dir, field, 'acs_f606w')+ '/galaxy_'+str(index)+'_*')[0]
+                im = np.zeros((128, 128, 1))
+                im[:, :, 0] = fits.open(tmp_file)[0].data
 
-    #           imCp = np.fft.fft2(psf, s=stamp_size, stamp_size)
-                ps = np.random.normal(0, 1e-4, (p.img_len, p.img_len))
-                im_psf = cube_psf
+                        
+                ''' Resize the image to the wanted size'''
+                im = _resize_image(im, p.img_len)
+                
+                ''' Load the wanted physical parameters of the galaxy '''
+                if hasattr(p, 'attributes'):
+                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
 
-#                 serialized_output = {"image/encoded": [im.tostring()],
-#                 "image/format": ["raw"],
-#                 "psf/encoded": [im_psf.tostring()],
-#                 "psf/format": ["raw"],
-#                 "ps/encoded": [ps.tostring()],
-#                 "ps/format": ["raw"],
-#                 "flag/filters":flag}
+                else:
+                    attributes=None
                 
-                serialized_output = {"image/encoded": [im.tostring()],
+                ''' Add a flag corresponding to the field '''
+                field_info = np.asarray(n_field)
+
+                ''' Create the output to match T2T format '''
+                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
                 "image/format": ["raw"],
-                "psf/encoded": [im_psf.tostring()],
-                "psf/format": ["raw"]}
-        
-        
-                if n > 20:
+                "psf/encoded": [psf.astype('float32').tostring()],
+                "psf/format": ["raw"],
+                "sigma_noise/encoded": [str(4.8e-3)],
+                "sigma_noise/format": ["raw"],
+                "field/encoded": [nb_field.astype('float32').tostring()],
+                "field/format": ["raw"]}
+
+                if attributes is not None:
+                    for k in attributes:
+                        serialized_output['attrs/'+k] = [attributes[k]]
+#                     serialized_output['attrs/field'] = [attributes['field']]
+                
+                ''' Increment the number of galaxy created on the shard '''
+                n_gal_creat += 1
+                
+                if n_gal_creat > p.example_per_shard:
                     break
                 yield serialized_output
 #             except Exception:
 #                 print('No image corresponding to the Index')
 #                 n+=1
+
+  @property
+  def is_generate_per_split(self):
+    return False
+
   def example_reading_spec(self):
     """Define how data is serialized to file and read back.
 
@@ -170,16 +177,17 @@ def example_reading_spec(self):
         "psf/encoded": tf.FixedLenFeature((), tf.string),
         "psf/format": tf.FixedLenFeature((), tf.string),
 
-#         "ps/encoded": tf.FixedLenFeature((), tf.string),
-#         "ps/format": tf.FixedLenFeature((), tf.string),
-    }
-
+        "sigma_noise/encoded": tf.FixedLenFeature((), tf.string),
+        "sigma_noise/format": tf.FixedLenFeature((), tf.string),
+        
+        "field/encoded": tf.FixedLenFeature((), tf.string),
+        "field/format": tf.FixedLenFeature((), tf.string)
+        }
 
     # Adds additional attributes to be decoded as specified in the configuration
     if hasattr(p, 'attributes'):
         for k in p.attributes:
             data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
-
     data_items_to_decoders = {
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
@@ -195,42 +203,51 @@ def example_reading_spec(self):
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
-#         "ps": tf.contrib.slim.tfexample_decoder.Image(
-#                 image_key="ps/encoded",
-#                 format_key="ps/format",
-#                 channels=self.num_bands,
-#                 shape=[p.img_len, p.img_len // 2 + 1],
-#                 dtype=tf.float32),
+        "sigma": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="sigma_noise/encoded",
+                format_key="sigma_noise/format",
+                channels=self.num_bands,
+                shape=[self.num_bands],
+                dtype=tf.float32),
+        
+        "field": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="field/encoded",
+                format_key="field/format",
+                shape=[1],
+                dtype=tf.float32),
     }
     if hasattr(p, 'attributes'):
         for k in p.attributes:
             data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
+#         data_items_to_decoders['field'] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/field')
+
+    return data_fields, data_items_to_decoders
+#     END: Subclass interface
 
 @registry.register_problem
-class Img2imgCandelsGoods(astroimage_utils.AstroImageProblem):
-  """Base class for image problems created from HSC Public Data Release.
+class Img2imgCandelsAll(astroimage_utils.AstroImageProblem):
+  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
   """
 
   @property
   def dataset_splits(self):
     """Splits of data to produce and number of output shards for each.
-       Note that each shard will be produced in parallel.
-       We are going to split the GalSim data into shards of 1000 galaxies each,
-       with 80 shards for training, 2 shards for validation.
+      .
     """
     return [{
         "split": problem.DatasetSplit.TRAIN,
         "shards": 1,
     }, {
         "split": problem.DatasetSplit.EVAL,
-        "shards": 1,
+        "shards": 0,
     }]
 
   # START: Subclass interface
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.img_len = 128
-    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w']#, 'acs_f850w']
+    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1.5e-3]
+    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f814w', 'acs_f850l']
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": None,
@@ -244,71 +261,240 @@ def num_bands(self):
     return len(p.filters)
 
   def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
-    """ Generator yielding individual postage stamps.
+    """ 
+    Generator yielding individual postage stamps.
     """
+    
     p = self.get_hparams()
+    
+    '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='egs   ')] = 'EGS'
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='UDS   ')] = 'UDS'
     
+    ''' Load the psfs for each filter and resize'''
     cube_psf = np.zeros((167, 167, len(p.filters)))
-    
     for i, filter in enumerate(p.filters):
-        cube_psf[:, :, i] = fits.open(tmp_dir + 'psf_' + filter +'.fits')[0].data
-
-#     # Step 1: Maybe regenerate the dataset
-#     self.maybe_build_dataset(tmp_dir)
-#     import timeit
+        cube_psf[:, :, i] = fits.open(tmp_dir + '/psfs/psf_' + filter +'.fits')[0].data
+                
 
     # Step 2: Extract postage stamps, resize them to requested size
-    ''' Genrate GDS and GDN'''
-    for field in ['GDS', 'GDN'] :
-        print(f"\n generating{field}\n")
-        n=0
+    ''' Loop on the fields'''
+    for n_field, field in enumerate(['GDS', 'GDN', 'EGS', 'UDS', 'COSMOS']):
+        
+        print(f"\n generating{field}, {n_field}\n")
+        n_gal_creat = 0
         index = 0
         im = np.zeros((128, 128, len(p.filters)))
+        sigmas = np.zeros(len(p.filters)) + 10
+        im_psf = np.zeros((167, 167, len(p.filters)))
+        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
         sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-        for i in sub_cat['RB_ID']:
-                if i == index :
+        
+        if field in ['GDS', 'GDN']:
+            filters_field = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850l']
+            bands = [0, 1, 2, 3, 4, 5, 7]
+        else :
+            filters_field = ['f125w', 'wfc3_f160w', 'acs_f606w', 'acs_f814w']
+            bands = [1, 2, 4, 6]
+
+        sigmas[bands] = p.sigmas[bands]
+        im_psf[bands] = cube_psf[bands]
+        ''' Loop on all the galaxies of the field '''
+        for gal in sub_cat['RB_ID']:
+                if gal == index :     # To take care of the redudency inside the cat
                     continue
-                index = i
-                print(field, i)
+                index = gal
+
 #             try:
-                for i, filter in enumerate(p.filters):
-#                     try :
+                ''' Loop on the filters '''
+                for n_filter, filter in enumerate(filters_field):
+                    try :
+                        ''' Open the image corresponding to the index of the current galaxy'''
                         tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
-                        im[:,:, i] = fits.open(tmp_file)[0].data
-#                     except Exception:
-#                         print('Galaxy not seen in every filter')
-#                         continue
+                        im[:, :, bands[n_filter]] = fits.open(tmp_file)[0].data
+                    except Exception:
+                        print('Galaxy not seen in every filter')
+                        continue
+                        
+                ''' Resize the image to the wanted size'''
                 im = _resize_image(im, p.img_len)
-
-                n+=1
+                print(self.num_bands)
+                
+                ''' Load the wanted physical parameters of the galaxy '''
                 if hasattr(p, 'attributes'):
                     attributes = {k: float(all_cat[k][index]) for k in p.attributes}
-                    print(attributes)
+
                 else:
                     attributes=None
+                
+#                 ''' Create the power spectrum '''
+#                 ps = np.zeros((p.img_len, p.img_len//2+1, len(p.filters)))
+#                 for i in range(len(p.filters)):
+#                     ps[:, :, n_filter] = np.random.normal(0, p.sigmas[n_filter], (p.img_len, p.img_len // 2 + 1))
+                
+                ''' Add a flag corresponding to the field '''
+                field_info = np.asarray(n_field)
+
+                ''' Create the output to match T2T format '''
+                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
+                "image/format": ["raw"],
+                "psf/encoded": [im_psf.astype('float32').tostring()],
+                "psf/format": ["raw"],
+                "sigma_noise/encoded": [np.asarray(sigmas).astype('float32').tostring()],
+                "sigma_noise/format": ["raw"],
+                "field/encoded": [field_info.astype('float32').tostring()],
+                "field/format": ["raw"]}
+                
+                if attributes is not None:
+                    for k in attributes:
+                        serialized_output['attrs/'+k] = [attributes[k]]
+#                     serialized_output['attrs/field'] = [attributes['field']]
+                
+                ''' Increment the number of galaxy created on the shard '''
+                n_gal_creat += 1
+                
+                if n_gal_creat > p.example_per_shard:
+                    print('out')
+                    break
+                yield serialized_output
+#             except Exception:
+#                 print('No image corresponding to the Index')
+#                 n+=1
+
 
+@registry.register_problem
+class Img2imgCandelsGoods(astroimage_utils.AstroImageProblem):
+  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
+  """
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each.
+      .
+    """
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 0,
+    }]
+
+  # START: Subclass interface
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.img_len = 128
+    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1.5e-3]
+    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850l']
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "targets": None}
+    p.add_hparam("psf", None)
+
+  @property
+  def num_bands(self):
+    """Number of bands."""
+    p = self.get_hparams()
+    return len(p.filters)
 
-    #           imCp = np.fft.fft2(psf, s=stamp_size, stamp_size)
-                ps = np.random.normal(0, 1e-4, (p.img_len, p.img_len))
-                im_psf = cube_psf
-                im_psf = _resize_image(im_psf, p.img_len)
+  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
+    """ 
+    Generator yielding individual postage stamps.
+    """
+    
+    p = self.get_hparams()
+    
+    '''Load the catalogue containing every fields and every filter'''
+    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
+    
+    ''' Load the psfs for each filter and resize'''
+    cube_psf = np.zeros((167, 167, len(p.filters)))
+    for i, filter in enumerate(p.filters):
+        cube_psf[:, :, i] = fits.open(tmp_dir + '/psfs/psf_' + filter +'.fits')[0].data
+    
+    scaling = p.pixel_scale/0.06
+    im_psf = resize(cube_psf, (np.ceil(128/scaling)+1, np.ceil(128/scaling)+1,len(p.filters)))
+    im_psf = _resize_image(cube_psf, p.img_len)
+    
+    im_tmp = np.zeros((128, 128, len(p.filters)))
+    sigmas = np.asarray(p.sigmas)
 
-                print(f"\n shapes : {im.shape}, {im_psf.shape}, !!!!! \n")
+    # Step 2: Extract postage stamps, resize them to requested size
+    ''' Loop on the two fields'''
+    for n_field, field in enumerate(['GDS']):#, 'GDN']):
+        
+        print(f"\n generating{field}, {n_field}\n")
+        n_gal_creat = 0
+        index = 0
+        
+        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
+        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
+        
+        ''' Loop on all the galaxies of the field '''
+        for gal in sub_cat['RB_ID']:
+                if gal == index :     # To take care of the redudency inside the cat
+                    continue
+                index = gal
+
+#             try:
+                ''' Loop on the filters '''
+                for n_filter, filter in enumerate(p.filters):
+                    try :
+                        ''' Open the image corresponding to the index of the current galaxy'''
+                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
+                        if np.max(fits.open(tmp_file)[0].data) == 0.:
+                            sigmas[n_filter] = 10
+                        im_tmp[:, :, n_filter] = fits.open(tmp_file)[0].data
 
-                serialized_output = {"image/encoded": [im.tostring()],
+                    except Exception:
+                        print('Galaxy not seen in every filter')
+                        continue
+                        
+                ''' Resize the image to the wanted size'''
+                im = resize(im_tmp, (np.ceil(128/scaling)+1, np.ceil(128/scaling)+1, len(p.filters)))
+                
+                im = _resize_image(im, p.img_len)
+                
+                ''' Load the wanted physical parameters of the galaxy '''
+                if hasattr(p, 'attributes'):
+                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
+
+                else:
+                    attributes=None
+                
+#                 ''' Create the power spectrum '''
+#                 ps = np.zeros((p.img_len, p.img_len//2+1, len(p.filters)))
+#                 for i in range(len(p.filters)):
+#                     ps[:, :, n_filter] = np.random.normal(0, p.sigmas[n_filter], (p.img_len, p.img_len // 2 + 1))
+                
+                ''' Add a flag corresponding to the field '''
+                field_info = np.asarray(n_field)
+
+                ''' Create the output to match T2T format '''
+                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
                 "image/format": ["raw"],
-                "psf/encoded": [im_psf.tostring()],
+                "psf/encoded": [im_psf.astype('float32').tostring()],
                 "psf/format": ["raw"],
-                "ps/encoded": [ps.tostring()],
-                "ps/format": ["raw"]}
+                "sigma_noise/encoded": [sigmas.astype('float32').tostring()],
+                "sigma_noise/format": ["raw"],
+                "field/encoded": [field_info.astype('float32').tostring()],
+                "field/format": ["raw"]}
                 
                 if attributes is not None:
                     for k in attributes:
                         serialized_output['attrs/'+k] = [attributes[k]]
-                if n > 5:
+                
+                ''' Increment the number of galaxy created on the shard '''
+                n_gal_creat += 1
+                
+                if n_gal_creat > p.example_per_shard:
+                    print('out')
                     break
                 yield serialized_output
 #             except Exception:
@@ -329,6 +515,7 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
     if hasattr(p, 'attributes'):
       example['attributes'] = tf.stack([example[k] for k in p.attributes])
 
+
     example["inputs"] = image
     example["targets"] = image
     return example
@@ -350,41 +537,49 @@ def example_reading_spec(self):
         "psf/encoded": tf.FixedLenFeature((), tf.string),
         "psf/format": tf.FixedLenFeature((), tf.string),
 
-        "ps/encoded": tf.FixedLenFeature((), tf.string),
-        "ps/format": tf.FixedLenFeature((), tf.string),
-    }
+        "sigma_noise/encoded": tf.FixedLenFeature((), tf.string),
+        "sigma_noise/format": tf.FixedLenFeature((), tf.string),
 
+        "field/encoded": tf.FixedLenFeature((), tf.string),
+        "field/format": tf.FixedLenFeature((), tf.string),
+    }
 
     # Adds additional attributes to be decoded as specified in the configuration
     if hasattr(p, 'attributes'):
         for k in p.attributes:
             data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
-
     data_items_to_decoders = {
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
-                channels=self.num_bands,
+#                 channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "psf": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="psf/encoded",
                 format_key="psf/format",
-                channels=self.num_bands,
+#                 channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
-        "ps": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="ps/encoded",
-                format_key="ps/format",
-                channels=self.num_bands,
-                shape=[p.img_len, p.img_len // 2 + 1],
+        "sigma": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="sigma_noise/encoded",
+                format_key="sigma_noise/format",
+#                 channels=self.num_bands,
+                shape=[self.num_bands],
+                dtype=tf.float32),
+        
+        "field": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="field/encoded",
+                format_key="field/format",
+                shape=[1],
                 dtype=tf.float32),
     }
     if hasattr(p, 'attributes'):
         for k in p.attributes:
             data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
+#         data_items_to_decoders['field'] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/field')
 
     return data_fields, data_items_to_decoders
 #     END: Subclass interface
@@ -393,10 +588,32 @@ def example_reading_spec(self):
   @property
   def is_generate_per_split(self):
     return False
+  
+@registry.register_problem
+class Attrs2imgCandelsOnefilter64Euclid(Img2imgCandelsOnefilter):
+  """
+  """
 
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 64
+    p.example_per_shard = 10
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag', 're', 'q']
+    
 
 @registry.register_problem
-class Attrs2imgCandelsAll64Euclid(Img2imgCandelsAll):
+class Attrs2imgCandelsGoods64Euclid(Img2imgCandelsGoods):
   """
   """
 
@@ -407,9 +624,10 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = 0.1
-    p.img_len = 128
-    p.example_per_shard = 1
-    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850w']
+    p.img_len = 64
+    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1.5e-3]
+    p.example_per_shard = 5
+    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850l']
 
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
@@ -417,11 +635,10 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "attributes": None,
                     "targets": None}
-    p.attributes = []#['mag', 're', 'q']
-    
+    p.attributes = ['mag', 're', 'q']
     
 @registry.register_problem
-class Attrs2imgCandelsGoods64Euclid(Img2imgCandelsGoods):
+class Attrs2imgCandelsAll64Euclid(Img2imgCandelsAll):
   """
   """
 
@@ -432,9 +649,10 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = 0.1
-    p.img_len = 128
-    p.example_per_shard = 1
-    p.filters = ['wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w']#, 'acs_f850l']
+    p.img_len = 64
+    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1e-4, 1.5e-3]
+    p.example_per_shard = 5
+    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f814w', 'acs_f850l']
 
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
@@ -442,5 +660,4 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "attributes": None,
                     "targets": None}
-#     p.attributes = ['mag', 're', 'q']
-    
+    p.attributes = ['mag', 're', 'q']
\ No newline at end of file
diff --git a/galaxy2galaxy/data_generators/check_candels_emulation.py b/galaxy2galaxy/data_generators/check_candels_emulation.py
index 0882721..7a3091e 100644
--- a/galaxy2galaxy/data_generators/check_candels_emulation.py
+++ b/galaxy2galaxy/data_generators/check_candels_emulation.py
@@ -5,7 +5,7 @@
 Modes = tf.estimator.ModeKeys
 
 candels = problems.problem('attrs2img_candels_goods64_euclid')
-dset = candels.dataset(Modes.TRAIN, data_dir='/data57/hbretonniere/datasets/test_candels')
+dset = candels.dataset(Modes.TRAIN, data_dir='/data57/hbretonniere/datasets/candels_for_g2g/test_candels')
 batch = dset.batch(1)
 
 print("data set :", dset, "\n \n")
diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 030e7be..f419e3c 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -49,7 +49,7 @@ def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = 0.03
     p.img_len = 64
-    p.example_per_shard = 1000
+
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": None,
@@ -64,9 +64,10 @@ def num_bands(self):
 
   def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     """
-    Generates and yields postage stamps obtained with GalSim.
+    Generates and yields poattrs2img_cosmos64_euclid_pixscale_1_wmorpho_2stage stamps obtained with GalSim.
     """
     p = self.get_hparams()
+
     try:
         # try to use default galsim path to the data
         catalog = galsim.COSMOSCatalog()
@@ -87,6 +88,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
     bparams = cat_param['bulgefit']
     sparams = cat_param['sersicfit']
+    # Parameters for a 2 component fit
     cat_param = append_fields(cat_param, 'bulge_q', bparams[:,11])
     cat_param = append_fields(cat_param, 'bulge_beta', bparams[:,15])
     cat_param = append_fields(cat_param, 'disk_q', bparams[:,3])
@@ -96,8 +98,37 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     cat_param = append_fields(cat_param, 'disk_hlr', cat_param['hlr'][:,2])
     cat_param = append_fields(cat_param, 'disk_flux_log10', np.where(cat_param['use_bulgefit'] ==1, np.log10(cat_param['flux'][:,2]), np.log10(cat_param['flux'][:,0])))
 
+    # Parameters for a single component fit
+    cat_param = append_fields(cat_param, 'sersic_flux_log10', np.log10(sparams[:,0]))
+    cat_param = append_fields(cat_param, 'sersic_q', sparams[:,3])
+    cat_param = append_fields(cat_param, 'sersic_hlr', sparams[:,1])
+    cat_param = append_fields(cat_param, 'sersic_n', sparams[:,2])
+    cat_param = append_fields(cat_param, 'sersic_beta', sparams[:,7])
+
+    passed = 0
+    
+    late = 0
+    irr = 0
     for ind in index:
       # Draw a galaxy using GalSim, any kind of operation can be done here
+#       if cat_param['Mph3'][ind] not in [1,2,3] :
+#             passed += 1
+#             continue
+
+#       if cat_param['Mph3'][ind] == 2:
+#         if late >= 85 :
+#             if ind % 200 == 0:
+#                 print("done all the late of the id")
+#             continue
+#         late += 1
+      
+#       if cat_param['Mph3'][ind] == 3:
+#         if irr >= 85 :
+#             if ind % 200 == 0:
+#                 print("done all the irr of the id")
+#             continue
+#         irr += 1
+        
       gal = catalog.makeGalaxy(ind, noise_pad_size=p.img_len * p.pixel_scale)
 
       # We apply the orginal psf if a different PSF is not requested
@@ -106,12 +137,12 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
       else:
           psf = p.psf
 
-      # Apply random rotation if requested
-      if hasattr(p, "rotation") and p.rotation:
-        rotation_angle = galsim.Angle(-np.random.rand()* 2 * np.pi,
+      # Apply rotation so that the galaxy is at 0 PA
+#       if hasattr(p, "rotation") and p.rotation:
+      rotation_angle = galsim.Angle(-cat_param[ind]['sersic_beta'],
                                       galsim.radians)
-        gal = gal.rotate(rotation_angle)
-        psf = psf.rotate(rotation_angle)
+      gal = gal.rotate(rotation_angle)
+      psf = psf.rotate(rotation_angle)
 
       # We save the corresponding attributes for this galaxy
       if hasattr(p, 'attributes'):
@@ -216,6 +247,7 @@ def hparams(self, defaults, model_hparams):
                   "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": None,
                     "targets": None}
+    
 
 @registry.register_problem
 class Img2imgCosmos128(Img2imgCosmos):
@@ -233,6 +265,23 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "targets": None}
 
+
+@registry.register_problem
+class Img2imgCosmos64(Img2imgCosmos):
+  """ Smaller version of the Img2imgCosmos problem, at half the pixel
+  resolution
+  """
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 64
+    p.example_per_shard = 1000
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "targets": None}
+
 @registry.register_problem
 class Attrs2imgCosmos128(Img2imgCosmos128):
   """ Smaller version of the Img2imgCosmos problem, at half the pixel
@@ -256,6 +305,74 @@ def hparams(self, defaults, model_hparams):
                     "targets": None}
     p.attributes = ['mag_auto', 'flux_radius', 'zphot']
 
+@registry.register_problem
+class Attrs2imgCosmos128Euclid(Img2imgCosmos128):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 128
+    p.example_per_shard = 1000
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag_auto', 'flux_radius', 'sersic_n', 'sersic_q']
+    
+@registry.register_problem
+class Attrs2imgCosmos64Euclid(Img2imgCosmos128):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 64
+    p.example_per_shard = 1000
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+#     p.attributes = ['mag_auto', 'flux_radius', 'sersic_n', 'sersic_q']
+    p.attributes = ['flux_radius', 'sersic_n', 'sersic_q']
+
+@registry.register_problem
+class Attrs2imgCosmos64EuclidWithMorpho(Img2imgCosmos128):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 64
+    p.example_per_shard = 1000
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag_auto', 'flux_radius', 'sersic_n', 'sersic_q', 'Mph3']
+    
+
 @registry.register_problem
 class Attrs2imgCosmos32(Attrs2imgCosmos):
   """ Lower resolution equivalent of conditional generation problem.
diff --git a/galaxy2galaxy/models/autoencoders.py b/galaxy2galaxy/models/autoencoders.py
index c5be069..fba00f8 100644
--- a/galaxy2galaxy/models/autoencoders.py
+++ b/galaxy2galaxy/models/autoencoders.py
@@ -247,16 +247,16 @@ def continuous_autoencoder_residual_128():
   hparams = continuous_autoencoder_basic()
   hparams.optimizer = "Adafactor"
   hparams.clip_grad_norm = 1.0
-  hparams.learning_rate_constant = 0.25
+  hparams.learning_rate_constant = 0.35
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup * rsqrt_decay"
-  hparams.num_hidden_layers = 7
-  hparams.hidden_size = 32
-  hparams.max_hidden_size = 512
-  hparams.batch_size = 32
+  hparams.num_hidden_layers = 6
+  hparams.hidden_size = 64
+  hparams.max_hidden_size = 1024
+  hparams.batch_size = 64
   hparams.bottleneck_bits = 32
 
-  hparams.bottleneck_warmup_steps = 5000
+  hparams.bottleneck_warmup_steps = 2000
 
   hparams.add_hparam("autoregressive_decode_steps", 0)
   hparams.add_hparam("num_residual_layers", 2)
diff --git a/galaxy2galaxy/models/autoregressive.py b/galaxy2galaxy/models/autoregressive.py
index 9b83348..34ae1a3 100644
--- a/galaxy2galaxy/models/autoregressive.py
+++ b/galaxy2galaxy/models/autoregressive.py
@@ -14,7 +14,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-from pixel_cnn_pp.model import model_spec
+# from pixel_cnn_pp.model import model_spec
 
 import tensorflow as tf
 import tensorflow_probability as tfp

From bc84c2ec37ed1a6fc5023a2e3a758fcd0e125306 Mon Sep 17 00:00:00 2001
From: Hubert Bretonniere <hubert.bretonniere@ias.u-psud.fr>
Date: Thu, 29 Oct 2020 10:42:16 +0100
Subject: [PATCH 004/149] working RNG

---
 galaxy2galaxy/bin/concatenate_models.py |  3 ++-
 galaxy2galaxy/models/latent_flow.py     | 20 ++++++++++++--------
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/galaxy2galaxy/bin/concatenate_models.py b/galaxy2galaxy/bin/concatenate_models.py
index 3cd5dc4..3454ef4 100644
--- a/galaxy2galaxy/bin/concatenate_models.py
+++ b/galaxy2galaxy/bin/concatenate_models.py
@@ -24,7 +24,8 @@ def main(argv):
   def generative_model_fn():
     code = hub.Module(FLAGS.flow_module)
     decoder = hub.Module(FLAGS.decoder_module)
-    inputs = {k: tf.placeholder(tf.float32, shape=[None]) for k in code.get_input_info_dict().keys()}
+    input_info = code.get_input_info_dict()
+    inputs = {k: tf.placeholder(tf.float32, shape=input_info[k].get_shape()) for k in input_info.keys()}
     hub.add_signature(inputs=inputs, outputs=decoder(code(inputs)))
     hub.attach_message("stamp_size", decoder.get_attached_message("stamp_size", tf.train.Int64List))
     hub.attach_message("pixel_size", decoder.get_attached_message("pixel_size", tf.train.FloatList))
diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index d409157..76ff3d3 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -25,7 +25,6 @@
 
 class LatentFlow(t2t_model.T2TModel):
   """ Base class for latent flows
-
   This assumes that an already exported tensorflow hub autoencoder is provided
   in hparams.
   """
@@ -66,18 +65,23 @@ def get_flow(inputs, is_training=True):
       y = tf.concat([tf.expand_dims(inputs[k], axis=1) for k in hparamsp.attributes] ,axis=1)
       y = tf.layers.batch_normalization(y, name="y_norm", training=is_training)
       flow = self.normalizing_flow(y, latent_size)
-      code = tf.reshape(flow.sample(tf.shape(y)[0]), code_shape)
-      return code, flow
+      return flow
 
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       # Export the latent flow alone
       def flow_module_spec():
         inputs = {k: tf.placeholder(tf.float32, shape=[None]) for k in hparamsp.attributes}
-        code, _ = get_flow(inputs, is_training=False)
-        hub.add_signature(inputs=inputs, outputs=code)
+        random_normal = tf.placeholder(tf.float32, shape=[None, latent_size])
+        print(f'\n \n \n !!!!!!!!!!! {random_normal} !!!!!!!! \n \n \n')
+        flow = get_flow(inputs, is_training=False)
+        samples = flow._bijector.forward(random_normal)
+        samples = tf.reshape(samples, code_shape)
+        hub.add_signature(inputs={**inputs, 'random_normal': random_normal},
+                          outputs=samples)
       flow_spec = hub.create_module_spec(flow_module_spec)
       flow = hub.Module(flow_spec, name='flow_module')
       hub.register_module_for_export(flow, "code_sampler")
+      cond['random_normal'] = tf.random_normal(shape=[tf.shape(cond[hparamsp.attributes[0]])[0] , latent_size])
       samples = flow(cond)
       return samples, {'loglikelihood': 0}
 
@@ -88,13 +92,13 @@ def flow_module_spec():
       code = encoder(x)
 
     with tf.variable_scope("flow_module"):
-      samples, flow = get_flow(cond)
+      flow = get_flow(cond)
       loglikelihood = flow.log_prob(tf.layers.flatten(code))
 
     # This is the loglikelihood of a batch of images
     tf.summary.scalar('loglikelihood', tf.reduce_mean(loglikelihood))
     loss = - tf.reduce_mean(loglikelihood)
-    return samples, {'training': loss}
+    return code, {'training': loss}
 
 @registry.register_model
 class LatentMAF(LatentFlow):
@@ -179,4 +183,4 @@ def latent_flow_larger():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
-  return hparams
+  return hparams
\ No newline at end of file

From f1a21d2b125fade9c64d38524df37ce66dd9d3fc Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 18 Feb 2021 12:00:47 +0000
Subject: [PATCH 005/149] Added multiresolution support

---
 galaxy2galaxy/data_generators/candels.py | 399 ++++++++++++++++++++++-
 1 file changed, 397 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c603f1c..a26a6ee 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -26,7 +26,11 @@
 import glob
 import os
 from skimage.transform import resize
-
+from scipy.ndimage import binary_dilation  # type: ignore
+from astropy.table import Table
+from scipy.ndimage import rotate
+from scipy.spatial import KDTree
+import sep
 
 
 # HSC default pixel scale
@@ -660,4 +664,395 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "attributes": None,
                     "targets": None}
-    p.attributes = ['mag', 're', 'q']
\ No newline at end of file
+    p.attributes = ['mag', 're', 'q']
+
+
+
+@registry.register_problem
+class Img2imgCandelsGoodsMultires(astroimage_utils.AstroImageProblem):
+  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
+  """
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each.
+      .
+    """
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 1,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 0,
+    }]
+
+  # START: Subclass interface
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.img_len = 128
+    p.sigmas_dic = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.filters_dic = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.resolutions = ["high","low"]
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "targets": None}
+    p.add_hparam("psf", None)
+
+  @property
+  def num_bands(self):
+    """Number of bands."""
+    p = self.get_hparams()
+    return np.sum([len(p.filters[res]) for res in p.resolutions])
+
+  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
+    """ 
+    Generator yielding individual postage stamps.
+    """
+    
+    p = self.get_hparams()
+    band_num = np.sum([len(p.filters[res]) for res in p.resolutions])
+    scalings = {}
+    for res in p.resolutions:
+        scalings[res] = p.pixel_scale[res]/0.06
+    target_pixel_scale = np.min([p.pixel_scale[res] for res in p.resolutions])
+    
+    '''Load the catalogue containing every fields and every filter'''
+    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
+    
+    ''' Load the psfs for each filter and resize'''
+    cube_psf = np.zeros((167, 167, band_num))
+    k = 0
+    for res in p.resolutions:
+        cube_psf_tmp = np.zeros((167, 167, len(p.filters[res])))
+        for i, filt in enumerate(p.filters[res]):
+            cube_psf_tmp[:, :, i] = fits.open(tmp_dir + '/psfs/psf_' + filt +'.fits')[0].data
+        cube_psf_tmp = resize(cube_psf_tmp, (np.ceil(128/scalings[res])+1, np.ceil(128/scalings[res])+1,len(p.filters[res])))
+        cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
+
+        k += len(p.filters[res])
+    
+    im_psf = _resize_image(cube_psf, p.img_len)
+    
+    sigmas = np.asarray(p.sigmas)
+
+    # Step 2: Extract postage stamps, resize them to requested size
+    ''' Loop on the two fields'''
+    for n_field, field in enumerate(['GDS','GDN']):
+        
+        print(f"\n generating{field}, {n_field}\n")
+        n_gal_creat = 0
+        index = 0
+        
+        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
+        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
+        
+        ''' Loop on all the galaxies of the field '''
+        for gal in sub_cat['RB_ID']:
+            if gal == index :     # To take care of the redudency inside the cat
+                continue
+            index = gal
+
+            try:
+                ''' Loop on the filters '''
+                im = np.zeros((128, 128, band_num))
+
+                k = 0
+                for res in p.resolutions:
+                    im_tmp = np.zeros((128, 128, len(p.filters[res])))
+                    for n_filter, filt in enumerate(p.filters[res]):
+                        # try :
+                        ''' Open the image corresponding to the index of the current galaxy'''
+
+                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
+                        if np.max(fits.open(tmp_file)[0].data) == 0.:
+                            sigmas[n_filter] = 10
+                        im_import = fits.open(tmp_file)[0].data
+                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)
+
+                        # except Exception:
+                        #     print('Galaxy not seen in every filter')
+                        #     continue
+                            
+                    ''' Resize the image to the low resolution'''
+                    new_size = (np.ceil(128/scalings[res])+1
+                    im_tmp = resize(im_tmp, new_size, new_size, len(p.filters[res])))
+
+                    ''' Resize the image to the highest resolution to get consistent array sizes'''
+                    im_tmp = resize(im_tmp, (np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, len(p.filters[res])))
+
+                    im[:,:,k:k+len(p.filters[res])] = im_tmp
+                
+                im = _resize_image(im, p.img_len)
+                
+
+                ''' Load the wanted physical parameters of the galaxy '''
+                if hasattr(p, 'attributes'):
+                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
+
+                else:
+                    attributes=None
+                
+#                 ''' Create the power spectrum '''
+#                 ps = np.zeros((p.img_len, p.img_len//2+1, len(p.filters)))
+#                 for i in range(len(p.filters)):
+#                     ps[:, :, n_filter] = np.random.normal(0, p.sigmas[n_filter], (p.img_len, p.img_len // 2 + 1))
+                
+                ''' Add a flag corresponding to the field '''
+                field_info = np.asarray(n_field)
+
+                ''' Create the output to match T2T format '''
+                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
+                "image/format": ["raw"],
+                "psf/encoded": [im_psf.astype('float32').tostring()],
+                "psf/format": ["raw"],
+                "sigma_noise/encoded": [sigmas.astype('float32').tostring()],
+                "sigma_noise/format": ["raw"],
+                "field/encoded": [field_info.astype('float32').tostring()],
+                "field/format": ["raw"]}
+                
+                if attributes is not None:
+                    for k in attributes:
+                        serialized_output['attrs/'+k] = [attributes[k]]
+                
+                ''' Increment the number of galaxy created on the shard '''
+                n_gal_creat += 1
+                
+                if n_gal_creat > p.example_per_shard:
+                    print('out')
+                    break
+                yield serialized_output
+            except ValueError:
+                print(sys.exc_info()[0], sys.exc_info()[1])
+                continue
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    """ Preprocess the examples, can be used for further augmentation or
+    image standardization.
+    """
+    p = self.get_hparams()
+    image = example["inputs"]
+
+    # Clip to 1 the values of the image
+    # image = tf.clip_by_value(image, -1, 1)
+
+    # Aggregate the conditions
+    if hasattr(p, 'attributes'):
+      example['attributes'] = tf.stack([example[k] for k in p.attributes])
+
+
+    example["inputs"] = image
+    example["targets"] = image
+    return example
+
+  def example_reading_spec(self):
+    """Define how data is serialized to file and read back.
+
+    Returns:
+      data_fields: A dictionary mapping data names to its feature type.
+      data_items_to_decoders: A dictionary mapping data names to TF Example
+         decoders, to be used when reading back TF examples from disk.
+    """
+    p = self.get_hparams()
+
+    data_fields = {
+        "image/encoded": tf.FixedLenFeature((), tf.string),
+        "image/format": tf.FixedLenFeature((), tf.string),
+
+        "psf/encoded": tf.FixedLenFeature((), tf.string),
+        "psf/format": tf.FixedLenFeature((), tf.string),
+
+        "sigma_noise/encoded": tf.FixedLenFeature((), tf.string),
+        "sigma_noise/format": tf.FixedLenFeature((), tf.string),
+
+        "field/encoded": tf.FixedLenFeature((), tf.string),
+        "field/format": tf.FixedLenFeature((), tf.string),
+    }
+
+    # Adds additional attributes to be decoded as specified in the configuration
+    if hasattr(p, 'attributes'):
+        for k in p.attributes:
+            data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
+    data_items_to_decoders = {
+        "inputs": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="image/encoded",
+                format_key="image/format",
+#                 channels=self.num_bands,
+                shape=[p.img_len, p.img_len, self.num_bands],
+                dtype=tf.float32),
+
+        "psf": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="psf/encoded",
+                format_key="psf/format",
+#                 channels=self.num_bands,
+                shape=[p.img_len, p.img_len, self.num_bands],
+                dtype=tf.float32),
+
+        "sigma": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="sigma_noise/encoded",
+                format_key="sigma_noise/format",
+#                 channels=self.num_bands,
+                shape=[self.num_bands],
+                dtype=tf.float32),
+        
+        "field": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="field/encoded",
+                format_key="field/format",
+                shape=[1],
+                dtype=tf.float32),
+    }
+    if hasattr(p, 'attributes'):
+        for k in p.attributes:
+            data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
+#         data_items_to_decoders['field'] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/field')
+
+    return data_fields, data_items_to_decoders
+#     END: Subclass interface
+    
+    
+  @property
+  def is_generate_per_split(self):
+    return False
+
+
+
+@registry.register_problem
+class Attrs2imgCandelsGoodsEuclid64(Img2imgCandelsGoodsEuclid):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
+    p.img_len = 64
+    p.sigmas_dic = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.filters_dic = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.resolutions = ["high","low"]
+    p.example_per_shard = 5
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag', 're', 'q']
+
+
+
+
+
+def find_central(sex_cat):
+
+    n_detect = len(sex_cat)
+    
+    ''' Match the pred and true cat'''
+    pred_pos = np.zeros((n_detect, 2))
+    pred_pos[:, 0] = sex_cat['x']
+    pred_pos[:, 1] = sex_cat['y']
+
+    true_pos = np.zeros((1, 2))
+    true_pos[:, 0] = 64
+    true_pos[:, 1] = 64
+
+    _, match_index = KDTree(pred_pos).query(true_pos)
+    
+    return match_index
+
+
+import re 
+
+def sort_nicely( l ): 
+    """ Sort the given list in the way that humans expect. 
+    """ 
+    convert = lambda text: int(text) if text.isdigit() else text 
+    alphanum_key = lambda key: [ convert(c) for c in re.split('([0-9]+)', key) ] 
+    l.sort( key=alphanum_key ) 
+    return l
+
+
+def mask_out_pixels(img, segmap, segval,
+                    n_iter: int = 5, shuffle: bool = False,
+                    noise_factor: int = 1):
+    """
+    Replace central galaxy neighbours with background noise
+
+    Basic recipe to replace the detected sources around the central galaxy
+    with either randomly selected pixels from the background, or a random
+    realisation of the background noise.
+
+    """
+    masked_img = img.copy()
+    # Create binary masks of all segmented sources
+    sources = binary_dilation(segmap, iterations=n_iter)
+
+    background_mask = np.logical_not(sources)
+    # Create binary mask of the central galaxy
+    central_source = binary_dilation(np.where(segmap == segval, 1, 0),
+                                     iterations=n_iter)
+    print("Flux fraction : ",np.sum(img[np.where(segmap == segval)])/np.sum(img[central_source]))
+    # Compute the binary mask of all sources BUT the central galaxy
+    sources_except_central = np.logical_xor(sources, central_source)
+
+    if shuffle:
+        # Select random pixels from the noise in the image
+        n_pixels_to_fill_in = sources_except_central.sum()
+        random_background_pixels = np.random.choice(
+            img[background_mask],
+            size=n_pixels_to_fill_in
+        )
+        # Fill in the voids with these pixels
+        masked_img[sources_except_central] = random_background_pixels
+    else:
+        # Create a realisation of the background for the std value
+        background_std = np.std(img[background_mask])
+        # background_std = 0.007220430274502116
+        random_background = np.random.normal(scale=background_std, size=img.shape)
+        masked_img[sources_except_central] = random_background[sources_except_central]
+        masked_img[np.where(masked_img==0.0)] = random_background[np.where(masked_img==0.0)]
+        
+    return masked_img.astype(img.dtype), sources, background_mask, central_source, sources_except_central
+
+def clean_rotate_stamp(img, eps=5, sigma_sex=2):
+
+    ''' Sex for clean'''
+    img = img.byteswap().newbyteorder()
+    bkg = sep.Background(img)
+   
+    cat,sex_seg = sep.extract(img-bkg,sigma_sex,err=bkg.globalrms,segmentation_map=True)
+    
+    if len(cat) == 0:
+        raise ValueError('No galaxy detected in the field')
+    
+    middle_pos = [cat[find_central(cat)[0]]['x'],cat[find_central(cat)[0]]['y']]
+    
+    distance = np.sqrt((middle_pos[0]-64)**2 + (middle_pos[1]-64)**2)
+    if distance > 10 :
+        raise ValueError('No galaxy detected in the center')
+
+    middle = np.max(sex_seg[int(round(middle_pos[0]))-eps:int(round(middle_pos[0]))+eps, int(round(middle_pos[1]))-eps:int(round(middle_pos[1]))+eps])
+    if middle == 0:
+        raise ValueError('No galaxy detected in the center')
+
+    cleaned, _, _, central, _ = mask_out_pixels(img, sex_seg, middle,n_iter=5)
+    
+    if np.any(np.logical_and(np.not_equal(sex_seg[central],0),np.not_equal(sex_seg[central],middle))):
+       raise ValueError('Blending suspected')
+
+
+    '''Rotate'''
+    PA = cat[find_central(cat)[0]][4]
+    img_rotate = rotate(cleaned, PA, reshape=False)
+    
+    '''Add noise'''
+    background_mask = np.where(sex_seg == 0, 1, 0)
+    background_std = np.std(img * background_mask)
+    random_background = np.random.normal(scale=background_std, size=img_rotate.shape)
+    rotated = np.where(img_rotate == 0, random_background, img_rotate)
+
+    return rotated
\ No newline at end of file

From d4d2e24168f29907376d14ddb950de8ab13abd5b Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 18 Feb 2021 18:23:56 +0000
Subject: [PATCH 006/149] Correct typo

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index a26a6ee..0909623 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -777,8 +777,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                         #     continue
                             
                     ''' Resize the image to the low resolution'''
-                    new_size = (np.ceil(128/scalings[res])+1
-                    im_tmp = resize(im_tmp, new_size, new_size, len(p.filters[res])))
+                    new_size = np.ceil(128/scalings[res])+1
+                    im_tmp = resize(im_tmp, new_size, new_size, len(p.filters[res]))
 
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
                     im_tmp = resize(im_tmp, (np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, len(p.filters[res])))

From 647304a5e43611ae806b7f2ab275cff21a6b136e Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 18 Feb 2021 19:57:40 +0000
Subject: [PATCH 007/149] Added candels to problems

---
 galaxy2galaxy/data_generators/all_problems.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/all_problems.py b/galaxy2galaxy/data_generators/all_problems.py
index aca90af..38a4eda 100644
--- a/galaxy2galaxy/data_generators/all_problems.py
+++ b/galaxy2galaxy/data_generators/all_problems.py
@@ -14,7 +14,7 @@
 try:
   import galsim
   MODULES += ["galaxy2galaxy.data_generators.cosmos"]
-#   MODULES += ["galaxy2galaxy.data_generators.candels"]
+  MODULES += ["galaxy2galaxy.data_generators.candels"]
 except:
   print("Could not import GalSim, excluding some data generators")
 

From d404c27a1c2aed362e41fd793f31e7dd89918785 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 18 Feb 2021 20:19:38 +0000
Subject: [PATCH 008/149] Correct typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 0909623..2668602 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -919,7 +919,7 @@ def is_generate_per_split(self):
 
 
 @registry.register_problem
-class Attrs2imgCandelsGoodsEuclid64(Img2imgCandelsGoodsEuclid):
+class Attrs2imgCandelsGoodsEuclid64(Img2imgCandelsGoodsMultires):
   """
   """
 

From c5c9ec7598657ba5f84ab08df613766d46264c2b Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 10:25:13 +0000
Subject: [PATCH 009/149] Change parameters for testing

---
 galaxy2galaxy/data_generators/candels.py | 30 +++++++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 2668602..0db08fc 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -740,7 +740,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
     # Step 2: Extract postage stamps, resize them to requested size
     ''' Loop on the two fields'''
-    for n_field, field in enumerate(['GDS','GDN']):
+    for n_field, field in enumerate(['GDS']):
         
         print(f"\n generating{field}, {n_field}\n")
         n_gal_creat = 0
@@ -945,6 +945,34 @@ def hparams(self, defaults, model_hparams):
 
 
 
+@registry.register_problem
+class Attrs2imgCandelsGoodsEuclid64Test(Img2imgCandelsGoodsMultires):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
+    p.img_len = 64
+    p.sigmas_dic = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.filters_dic = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.resolutions = ["high","low"]
+    p.example_per_shard = 5
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag', 're', 'q']
+
+
+
+
 
 
 def find_central(sex_cat):

From 640b0578165a9ddc3db1abf05919c57b27d23d71 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 10:44:37 +0000
Subject: [PATCH 010/149] Correct variable names

---
 galaxy2galaxy/data_generators/candels.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 0db08fc..469d97d 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -690,8 +690,8 @@ def dataset_splits(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.img_len = 128
-    p.sigmas_dic = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
-    p.filters_dic = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
@@ -931,8 +931,8 @@ def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
     p.img_len = 64
-    p.sigmas_dic = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
-    p.filters_dic = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.example_per_shard = 5
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
@@ -958,8 +958,8 @@ def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
     p.img_len = 64
-    p.sigmas_dic = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
-    p.filters_dic = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.sigmas = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.example_per_shard = 5
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,

From 685667d8ce01632feda54af38264937f5731099b Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 10:47:38 +0000
Subject: [PATCH 011/149] changed tmp_dir to data_dir

---
 galaxy2galaxy/data_generators/candels.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 469d97d..acd21a2 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -718,7 +718,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     target_pixel_scale = np.min([p.pixel_scale[res] for res in p.resolutions])
     
     '''Load the catalogue containing every fields and every filter'''
-    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
+    all_cat = Table.read(os.path.join(data_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
     
@@ -728,7 +728,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     for res in p.resolutions:
         cube_psf_tmp = np.zeros((167, 167, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
-            cube_psf_tmp[:, :, i] = fits.open(tmp_dir + '/psfs/psf_' + filt +'.fits')[0].data
+            cube_psf_tmp[:, :, i] = fits.open(data_dir + '/psfs/psf_' + filt +'.fits')[0].data
         cube_psf_tmp = resize(cube_psf_tmp, (np.ceil(128/scalings[res])+1, np.ceil(128/scalings[res])+1,len(p.filters[res])))
         cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
 
@@ -766,7 +766,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                         # try :
                         ''' Open the image corresponding to the index of the current galaxy'''
 
-                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
+                        tmp_file = glob.glob(os.path.join(data_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
                         if np.max(fits.open(tmp_file)[0].data) == 0.:
                             sigmas[n_filter] = 10
                         im_import = fits.open(tmp_file)[0].data

From d43e955a5f003fa808a7b4fecaad54cdbaa5765e Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 11:47:01 +0000
Subject: [PATCH 012/149] Correct PSF size

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index acd21a2..aef1c18 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -729,7 +729,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         cube_psf_tmp = np.zeros((167, 167, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
             cube_psf_tmp[:, :, i] = fits.open(data_dir + '/psfs/psf_' + filt +'.fits')[0].data
-        cube_psf_tmp = resize(cube_psf_tmp, (np.ceil(128/scalings[res])+1, np.ceil(128/scalings[res])+1,len(p.filters[res])))
+        cube_psf_tmp = resize(cube_psf_tmp, 167,167,len(p.filters[res]))
         cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
 
         k += len(p.filters[res])

From 3cf0f0a8aed785489443515235afa8b560b69624 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 11:50:25 +0000
Subject: [PATCH 013/149] Correct PSF size

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index aef1c18..c656f8a 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -729,7 +729,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         cube_psf_tmp = np.zeros((167, 167, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
             cube_psf_tmp[:, :, i] = fits.open(data_dir + '/psfs/psf_' + filt +'.fits')[0].data
-        cube_psf_tmp = resize(cube_psf_tmp, 167,167,len(p.filters[res]))
+        cube_psf_tmp = resize(cube_psf_tmp, (167,167,len(p.filters[res])))
         cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
 
         k += len(p.filters[res])

From 3e564c2a475d399c0d4b5311a16cb9fcee4dcb93 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 11:53:12 +0000
Subject: [PATCH 014/149] Correct image size

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c656f8a..99571b6 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -778,7 +778,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                             
                     ''' Resize the image to the low resolution'''
                     new_size = np.ceil(128/scalings[res])+1
-                    im_tmp = resize(im_tmp, new_size, new_size, len(p.filters[res]))
+                    im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
 
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
                     im_tmp = resize(im_tmp, (np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, len(p.filters[res])))

From 885155363e2bdb31725aae9d70fa037e01ce7702 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 11:56:40 +0000
Subject: [PATCH 015/149] Missing import

---
 galaxy2galaxy/data_generators/candels.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 99571b6..c823f3b 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -25,6 +25,7 @@
 import h5py
 import glob
 import os
+import sys
 from skimage.transform import resize
 from scipy.ndimage import binary_dilation  # type: ignore
 from astropy.table import Table

From 9df2cc2034f39d93c9dc3017e526802f2258a5d6 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 15:05:16 +0000
Subject: [PATCH 016/149] Fix issue with sigmas

---
 galaxy2galaxy/data_generators/candels.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c823f3b..a118c92 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -737,7 +737,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     
     im_psf = _resize_image(cube_psf, p.img_len)
     
-    sigmas = np.asarray(p.sigmas)
+    sigmas = p.sigmas
 
     # Step 2: Extract postage stamps, resize them to requested size
     ''' Loop on the two fields'''
@@ -769,7 +769,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
                         tmp_file = glob.glob(os.path.join(data_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
                         if np.max(fits.open(tmp_file)[0].data) == 0.:
-                            sigmas[n_filter] = 10
+                            sigmas[res][n_filter] = 10
                         im_import = fits.open(tmp_file)[0].data
                         im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)
 
@@ -804,6 +804,11 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 ''' Add a flag corresponding to the field '''
                 field_info = np.asarray(n_field)
 
+                sigmas_array = []
+                for res in resolutions:
+                    sigmas_array += sigmas[res]
+                sigmas_array = np.array(sigmas)
+
                 ''' Create the output to match T2T format '''
                 serialized_output = {"image/encoded": [im.astype('float32').tostring()],
                 "image/format": ["raw"],

From c10e2c252e0271ca76ec032e1fb2cbed79f79348 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 16:03:29 +0000
Subject: [PATCH 017/149] Hopefully fixing shape problem

---
 galaxy2galaxy/data_generators/candels.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index a118c92..c5f01bf 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -717,6 +717,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     for res in p.resolutions:
         scalings[res] = p.pixel_scale[res]/0.06
     target_pixel_scale = np.min([p.pixel_scale[res] for res in p.resolutions])
+    target_scaling = target_pixel_scale/0.06
     
     '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(data_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
@@ -782,7 +783,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
 
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = resize(im_tmp, (np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, np.ceil(new_size/(target_pixel_scale/p.pixel_scale[res]))+1, len(p.filters[res])))
+                    im_tmp = resize(im_tmp, (np.ceil(128/target_scaling)+1, np.ceil(18/target_scaling)+1, len(p.filters[res])))
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp
                 

From 91de7acabbcc0deef2c342da4127400cba5ba7c7 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 17:06:24 +0000
Subject: [PATCH 018/149] Really fixing size issue

---
 galaxy2galaxy/data_generators/candels.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c5f01bf..6f4f8e6 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -756,10 +756,12 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             if gal == index :     # To take care of the redudency inside the cat
                 continue
             index = gal
+            print(index)
 
             try:
                 ''' Loop on the filters '''
-                im = np.zeros((128, 128, band_num))
+                target_size = np.ceil(128/scalings[res])+1
+                im = np.zeros((target_size, target_size, band_num))
 
                 k = 0
                 for res in p.resolutions:
@@ -783,7 +785,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
 
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = resize(im_tmp, (np.ceil(128/target_scaling)+1, np.ceil(18/target_scaling)+1, len(p.filters[res])))
+                    im_tmp = resize(im_tmp, (target_size, target_size, len(p.filters[res])))
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp
                 
@@ -1030,7 +1032,6 @@ def mask_out_pixels(img, segmap, segval,
     # Create binary mask of the central galaxy
     central_source = binary_dilation(np.where(segmap == segval, 1, 0),
                                      iterations=n_iter)
-    print("Flux fraction : ",np.sum(img[np.where(segmap == segval)])/np.sum(img[central_source]))
     # Compute the binary mask of all sources BUT the central galaxy
     sources_except_central = np.logical_xor(sources, central_source)
 

From dc7ed7a922d9ba064b8bfd259589ec28155a3e4d Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 17:22:50 +0000
Subject: [PATCH 019/149] Integer error

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 6f4f8e6..7c4fb75 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -760,7 +760,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
             try:
                 ''' Loop on the filters '''
-                target_size = np.ceil(128/scalings[res])+1
+                target_size = int(np.ceil(128/scalings[res])+1)
                 im = np.zeros((target_size, target_size, band_num))
 
                 k = 0

From 6025147b723efe95ee89c6d3f9ea52d4a171b6af Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 17:30:07 +0000
Subject: [PATCH 020/149] Typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 7c4fb75..745b16d 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -808,7 +808,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 field_info = np.asarray(n_field)
 
                 sigmas_array = []
-                for res in resolutions:
+                for res in p.resolutions:
                     sigmas_array += sigmas[res]
                 sigmas_array = np.array(sigmas)
 

From 0ecd0a40b15879e923d6527f927eec44b4a4a369 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 17:35:38 +0000
Subject: [PATCH 021/149] Typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 745b16d..5ea1306 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -817,7 +817,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 "image/format": ["raw"],
                 "psf/encoded": [im_psf.astype('float32').tostring()],
                 "psf/format": ["raw"],
-                "sigma_noise/encoded": [sigmas.astype('float32').tostring()],
+                "sigma_noise/encoded": [sigmas_array.astype('float32').tostring()],
                 "sigma_noise/format": ["raw"],
                 "field/encoded": [field_info.astype('float32').tostring()],
                 "field/format": ["raw"]}

From 3ca91fedaa8eef103fc422cb307d329b647cc477 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 17:40:45 +0000
Subject: [PATCH 022/149] Typo

---
 galaxy2galaxy/data_generators/candels.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 5ea1306..2da9562 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -767,6 +767,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 for res in p.resolutions:
                     im_tmp = np.zeros((128, 128, len(p.filters[res])))
                     for n_filter, filt in enumerate(p.filters[res]):
+                        print(n_filter)
                         # try :
                         ''' Open the image corresponding to the index of the current galaxy'''
 
@@ -774,7 +775,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                         if np.max(fits.open(tmp_file)[0].data) == 0.:
                             sigmas[res][n_filter] = 10
                         im_import = fits.open(tmp_file)[0].data
-                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)
+                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1)
 
                         # except Exception:
                         #     print('Galaxy not seen in every filter')
@@ -811,6 +812,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 for res in p.resolutions:
                     sigmas_array += sigmas[res]
                 sigmas_array = np.array(sigmas)
+                print(sigmas_array)
 
                 ''' Create the output to match T2T format '''
                 serialized_output = {"image/encoded": [im.astype('float32').tostring()],

From 68b8de9f6b9a1b3b60a7d16847010ac9e52f9f31 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 17:46:40 +0000
Subject: [PATCH 023/149] Typo

---
 galaxy2galaxy/data_generators/candels.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 2da9562..fc436aa 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -811,8 +811,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 sigmas_array = []
                 for res in p.resolutions:
                     sigmas_array += sigmas[res]
-                sigmas_array = np.array(sigmas)
-                print(sigmas_array)
+                sigmas_array = np.array(sigmas_array)
 
                 ''' Create the output to match T2T format '''
                 serialized_output = {"image/encoded": [im.astype('float32').tostring()],

From 985121575a833404bc50faf87aaf579e9cf9a2d3 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 19 Feb 2021 18:37:56 +0000
Subject: [PATCH 024/149] Fix filter drawing

---
 galaxy2galaxy/data_generators/candels.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index fc436aa..f9bdb18 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -789,6 +789,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     im_tmp = resize(im_tmp, (target_size, target_size, len(p.filters[res])))
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp
+                    k += len(p.filters[res])
                 
                 im = _resize_image(im, p.img_len)
                 

From 0bc24156a24628e522e25a2ce56806336a75578d Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Sat, 20 Feb 2021 13:11:21 +0000
Subject: [PATCH 025/149] Troubleshooting galaxy number

---
 galaxy2galaxy/data_generators/candels.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index f9bdb18..f976dbc 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -750,6 +750,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         
         ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
         sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
+        print("Expected number of galaxies : ",np.unique(subcat['RB_ID']))
         
         ''' Loop on all the galaxies of the field '''
         for gal in sub_cat['RB_ID']:

From 63f438044a9ed58ad71490356d11fec73e04cdef Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Sat, 20 Feb 2021 15:51:52 +0000
Subject: [PATCH 026/149] Put a higher number of galaxies per shard

---
 galaxy2galaxy/data_generators/candels.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index f976dbc..ed998f9 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -750,8 +750,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         
         ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
         sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-        print("Expected number of galaxies : ",np.unique(subcat['RB_ID']))
-        
+
         ''' Loop on all the galaxies of the field '''
         for gal in sub_cat['RB_ID']:
             if gal == index :     # To take care of the redudency inside the cat
@@ -946,7 +945,7 @@ def hparams(self, defaults, model_hparams):
     p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
     p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
-    p.example_per_shard = 5
+    p.example_per_shard = 1000
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
@@ -973,7 +972,7 @@ def hparams(self, defaults, model_hparams):
     p.sigmas = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
     p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
-    p.example_per_shard = 5
+    p.example_per_shard = 1000
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}

From cd5a12d714c72694ade0d6600da861dd39b9eef3 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Sat, 20 Feb 2021 22:44:22 +0000
Subject: [PATCH 027/149] Fix the inconsistant size

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index ed998f9..68ee778 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -760,7 +760,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
             try:
                 ''' Loop on the filters '''
-                target_size = int(np.ceil(128/scalings[res])+1)
+                target_size = int(np.ceil(128/scalings[p.resolutions[0]])+1)
                 im = np.zeros((target_size, target_size, band_num))
 
                 k = 0

From 9a047ce882b8ce536fd1308697ba6c821a47c027 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 09:53:14 +0000
Subject: [PATCH 028/149] Get the right base resolution ?

---
 galaxy2galaxy/data_generators/candels.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 68ee778..2e2b4d7 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -715,9 +715,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     band_num = np.sum([len(p.filters[res]) for res in p.resolutions])
     scalings = {}
     for res in p.resolutions:
-        scalings[res] = p.pixel_scale[res]/0.06
+        scalings[res] = p.pixel_scale[res]/p.base_pixel_scale[res]
     target_pixel_scale = np.min([p.pixel_scale[res] for res in p.resolutions])
-    target_scaling = target_pixel_scale/0.06
+    target_scaling = target_pixel_scale/p.base_pixel_scale["high"]
     
     '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(data_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
@@ -941,6 +941,7 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
+    p.base_pixel_scale = {'high' : 0.03,'low' : 0.13}
     p.img_len = 64
     p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
     p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
@@ -968,6 +969,7 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
+    p.base_pixel_scale = {'high' : 0.03,'low' : 0.13}
     p.img_len = 64
     p.sigmas = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
     p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}

From 7241b4bbca609bbe702fe8f249b75c7b5169c695 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 09:55:21 +0000
Subject: [PATCH 029/149] Small changes

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 2e2b4d7..fa15da5 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -753,7 +753,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
         ''' Loop on all the galaxies of the field '''
         for gal in sub_cat['RB_ID']:
-            if gal == index :     # To take care of the redudency inside the cat
+            if gal == index or galaxy == 15431 :     # To take care of the redudency inside the cat
                 continue
             index = gal
             print(index)
@@ -775,7 +775,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                         if np.max(fits.open(tmp_file)[0].data) == 0.:
                             sigmas[res][n_filter] = 10
                         im_import = fits.open(tmp_file)[0].data
-                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1)
+                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)
 
                         # except Exception:
                         #     print('Galaxy not seen in every filter')

From ab94a463a9159f42a0a69603885e8cbc5f284b50 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 10:57:20 +0000
Subject: [PATCH 030/149] CHanged parameters and added some prints

---
 galaxy2galaxy/data_generators/candels.py | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index fa15da5..33dfa2a 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -716,8 +716,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     scalings = {}
     for res in p.resolutions:
         scalings[res] = p.pixel_scale[res]/p.base_pixel_scale[res]
-    target_pixel_scale = np.min([p.pixel_scale[res] for res in p.resolutions])
-    target_scaling = target_pixel_scale/p.base_pixel_scale["high"]
+    target_pixel_scale = p.pixel_scale[p.resolutions[0]]
+    target_scaling = target_pixel_scale/p.base_pixel_scale[p.resolutions[0]]
+    print("scalings and all ",scalings,target_pixel_scale,target_scaling)
     
     '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(data_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
@@ -753,21 +754,22 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
         ''' Loop on all the galaxies of the field '''
         for gal in sub_cat['RB_ID']:
-            if gal == index or galaxy == 15431 :     # To take care of the redudency inside the cat
+            if gal == index or gal == 15431 :     # To take care of the redudency inside the cat
                 continue
             index = gal
             print(index)
 
             try:
                 ''' Loop on the filters '''
-                target_size = int(np.ceil(128/scalings[p.resolutions[0]])+1)
+                target_size = int(np.ceil(128/target_scaling)+1)
+                print("target ",target_size)
                 im = np.zeros((target_size, target_size, band_num))
 
                 k = 0
                 for res in p.resolutions:
                     im_tmp = np.zeros((128, 128, len(p.filters[res])))
                     for n_filter, filt in enumerate(p.filters[res]):
-                        print(n_filter)
+                        print(filt)
                         # try :
                         ''' Open the image corresponding to the index of the current galaxy'''
 
@@ -969,7 +971,7 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
-    p.base_pixel_scale = {'high' : 0.03,'low' : 0.13}
+    p.base_pixel_scale = {'high' : 0.05,'low' : 0.13}
     p.img_len = 64
     p.sigmas = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
     p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}

From 0807b8608c4bc7c1d8c0a016b9a21dbb5db54301 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 11:29:51 +0000
Subject: [PATCH 031/149] Get the right resolution

---
 galaxy2galaxy/data_generators/candels.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 33dfa2a..165c079 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -26,7 +26,7 @@
 import glob
 import os
 import sys
-from skimage.transform import resize
+from skimage.transform import resize,rescale
 from scipy.ndimage import binary_dilation  # type: ignore
 from astropy.table import Table
 from scipy.ndimage import rotate
@@ -786,9 +786,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     ''' Resize the image to the low resolution'''
                     new_size = np.ceil(128/scalings[res])+1
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
-
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = resize(im_tmp, (target_size, target_size, len(p.filters[res])))
+                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/p.target_pixel_scale)
+                    im_tmp = _resize_image(im_tmp, (target_size, target_size, len(p.filters[res])))
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp
                     k += len(p.filters[res])

From ddab0d2233f1585504714948f02da652203d1641 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 11:55:03 +0000
Subject: [PATCH 032/149] Fix typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 165c079..fdba825 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -787,7 +787,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     new_size = np.ceil(128/scalings[res])+1
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/p.target_pixel_scale)
+                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale)
                     im_tmp = _resize_image(im_tmp, (target_size, target_size, len(p.filters[res])))
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp

From 09af55dfffd795f1375f78869f90ee862ae67c5d Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 12:00:53 +0000
Subject: [PATCH 033/149] Fix typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index fdba825..3e824d2 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -788,7 +788,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
                     im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale)
-                    im_tmp = _resize_image(im_tmp, (target_size, target_size, len(p.filters[res])))
+                    im_tmp = _resize_image(im_tmp,target_size)
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp
                     k += len(p.filters[res])

From 83fea88b2ff2160708d64fa7823ebeb7f2722bf3 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 14:29:32 +0000
Subject: [PATCH 034/149] Fix rescaling problem

---
 galaxy2galaxy/data_generators/candels.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 3e824d2..fb4c822 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -762,7 +762,6 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             try:
                 ''' Loop on the filters '''
                 target_size = int(np.ceil(128/target_scaling)+1)
-                print("target ",target_size)
                 im = np.zeros((target_size, target_size, band_num))
 
                 k = 0
@@ -787,7 +786,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     new_size = np.ceil(128/scalings[res])+1
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale)
+                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel='True')
                     im_tmp = _resize_image(im_tmp,target_size)
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp

From 60dbc55a57e16dc44a563e012e4728cf212d1b92 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 14:31:49 +0000
Subject: [PATCH 035/149] FIx typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index fb4c822..cc3a952 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -786,7 +786,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     new_size = np.ceil(128/scalings[res])+1
                     im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
                     ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel='True')
+                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel=True)
                     im_tmp = _resize_image(im_tmp,target_size)
 
                     im[:,:,k:k+len(p.filters[res])] = im_tmp

From 1b16371c19dbea955de070fc173569ce8c9c5b2f Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Feb 2021 19:59:58 +0000
Subject: [PATCH 036/149] Fix resolutions and correct noise level and new
 blending definition

---
 galaxy2galaxy/data_generators/candels.py | 27 +++++++++++++++++-------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index cc3a952..d41da3a 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -776,7 +776,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                         if np.max(fits.open(tmp_file)[0].data) == 0.:
                             sigmas[res][n_filter] = 10
                         im_import = fits.open(tmp_file)[0].data
-                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)
+                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5,noise_level=p.sigmas[res][n_filter])
 
                         # except Exception:
                         #     print('Galaxy not seen in every filter')
@@ -942,7 +942,7 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
-    p.base_pixel_scale = {'high' : 0.03,'low' : 0.13}
+    p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
     p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
     p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
@@ -970,9 +970,9 @@ def eval_metrics(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
-    p.base_pixel_scale = {'high' : 0.05,'low' : 0.13}
+    p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
-    p.sigmas = {"high" : [3.4e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.sigmas = {"high" : [0.017380157011262], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
     p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.example_per_shard = 1000
@@ -1059,7 +1059,7 @@ def mask_out_pixels(img, segmap, segval,
         
     return masked_img.astype(img.dtype), sources, background_mask, central_source, sources_except_central
 
-def clean_rotate_stamp(img, eps=5, sigma_sex=2):
+def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
 
     ''' Sex for clean'''
     img = img.byteswap().newbyteorder()
@@ -1082,8 +1082,16 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2):
 
     cleaned, _, _, central, _ = mask_out_pixels(img, sex_seg, middle,n_iter=5)
     
-    if np.any(np.logical_and(np.not_equal(sex_seg[central],0),np.not_equal(sex_seg[central],middle))):
-       raise ValueError('Blending suspected')
+    blended_pixels = np.logical_and(np.not_equal(sex_seg,0),np.not_equal(sex_seg,middle))*central
+    blend_flux = np.sum(img[np.nonzero(blended_pixels)])
+    if np.any(blended_pixels):
+        loc = np.argwhere(blended_pixels==True)[0]
+        blended_galaxy = sex_seg[loc[0],loc[1]]
+        blended_galaxy_flux = np.sum(img[np.where(sex_seg==blended_galaxy)])
+    else:
+        blended_galaxy_flux = np.inf
+    if blend_flux/blended_galaxy_flux > 0.5:
+        raise ValueError('Blending suspected')
 
 
     '''Rotate'''
@@ -1092,7 +1100,10 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2):
     
     '''Add noise'''
     background_mask = np.where(sex_seg == 0, 1, 0)
-    background_std = np.std(img * background_mask)
+    if noise_level == None:
+        background_std = np.std(img * background_mask)
+    else:
+        background_std = noise_level
     random_background = np.random.normal(scale=background_std, size=img_rotate.shape)
     rotated = np.where(img_rotate == 0, random_background, img_rotate)
 

From d1d694ec1545f45d9ce0150400ec4ee695bb5d1c Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 23 Feb 2021 12:28:15 +0000
Subject: [PATCH 037/149] Correct noise

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index d41da3a..1fca8de 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -944,7 +944,7 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
     p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
-    p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
+    p.sigmas = {"high" : [1e-4], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
     p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.example_per_shard = 1000
@@ -972,7 +972,7 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
     p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
-    p.sigmas = {"high" : [0.017380157011262], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
+    p.sigmas = {"high" : [0.0034001764449330513], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
     p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.example_per_shard = 1000

From 70e052714ef9aeaf6f45d8072c9dafa6934651fb Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 23 Feb 2021 14:31:44 +0000
Subject: [PATCH 038/149] Add missing noise

---
 galaxy2galaxy/data_generators/candels.py | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 1fca8de..367643b 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -1020,7 +1020,7 @@ def sort_nicely( l ):
 
 def mask_out_pixels(img, segmap, segval,
                     n_iter: int = 5, shuffle: bool = False,
-                    noise_factor: int = 1):
+                    noise_factor: int = 1, noise_level=None):
     """
     Replace central galaxy neighbours with background noise
 
@@ -1051,8 +1051,10 @@ def mask_out_pixels(img, segmap, segval,
         masked_img[sources_except_central] = random_background_pixels
     else:
         # Create a realisation of the background for the std value
-        background_std = np.std(img[background_mask])
-        # background_std = 0.007220430274502116
+        if noise_level == None:
+            background_std = np.std(img[background_mask])
+        else:
+            background_std = noise_level
         random_background = np.random.normal(scale=background_std, size=img.shape)
         masked_img[sources_except_central] = random_background[sources_except_central]
         masked_img[np.where(masked_img==0.0)] = random_background[np.where(masked_img==0.0)]
@@ -1080,7 +1082,7 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
     if middle == 0:
         raise ValueError('No galaxy detected in the center')
 
-    cleaned, _, _, central, _ = mask_out_pixels(img, sex_seg, middle,n_iter=5)
+    cleaned, _, _, central, _ = mask_out_pixels(img, sex_seg, middle,n_iter=5,noise_level=noise_level)
     
     blended_pixels = np.logical_and(np.not_equal(sex_seg,0),np.not_equal(sex_seg,middle))*central
     blend_flux = np.sum(img[np.nonzero(blended_pixels)])

From faaead3f2acda7424082ab4d7ebd42641f6d572c Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 23 Feb 2021 15:50:52 +0000
Subject: [PATCH 039/149] Commented possibly problematic lines

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 367643b..b8a69a1 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -773,8 +773,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                         ''' Open the image corresponding to the index of the current galaxy'''
 
                         tmp_file = glob.glob(os.path.join(data_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
-                        if np.max(fits.open(tmp_file)[0].data) == 0.:
-                            sigmas[res][n_filter] = 10
+                        # if np.max(fits.open(tmp_file)[0].data) == 0.:
+                        #     sigmas[res][n_filter] = 10
                         im_import = fits.open(tmp_file)[0].data
                         im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5,noise_level=p.sigmas[res][n_filter])
 

From 34fee53a537ebef9c65a567a6c22f06b464eaa64 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Feb 2021 12:26:49 +0000
Subject: [PATCH 040/149] Handled multi channel PSF convolution and changed
 resolutions to 0.06 for candels

---
 galaxy2galaxy/data_generators/candels.py   |  4 ++--
 galaxy2galaxy/models/autoencoders_utils.py | 22 ++++++++++++----------
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index b8a69a1..c4df272 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -718,6 +718,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         scalings[res] = p.pixel_scale[res]/p.base_pixel_scale[res]
     target_pixel_scale = p.pixel_scale[p.resolutions[0]]
     target_scaling = target_pixel_scale/p.base_pixel_scale[p.resolutions[0]]
+    target_size = p.img_len
     print("scalings and all ",scalings,target_pixel_scale,target_scaling)
     
     '''Load the catalogue containing every fields and every filter'''
@@ -761,7 +762,6 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
             try:
                 ''' Loop on the filters '''
-                target_size = int(np.ceil(128/target_scaling)+1)
                 im = np.zeros((target_size, target_size, band_num))
 
                 k = 0
@@ -969,7 +969,7 @@ def eval_metrics(self):
 
   def hparams(self, defaults, model_hparams):
     p = defaults
-    p.pixel_scale = {'high' : 0.1, 'low' : 0.3}
+    p.pixel_scale = {'high' : 0.06, 'low' : 0.06}
     p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
     p.sigmas = {"high" : [0.0034001764449330513], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 71e2177..d76532c 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -214,17 +214,19 @@ def make_model_spec():
                                activation=output_activation)
 
   # Apply channel-wise convolution with the PSF if requested
-  # TODO: Handle multiple bands
   if hparams.apply_psf and 'psf' in features:
-    if self.num_channels > 1:
-      raise NotImplementedError
-    rec_padded = tf.pad(reconstr[:,:,:,0], [[0,0],
-                                            [0, int(hparams.psf_convolution_pad_factor*shape[1])],
-                                            [0, int(hparams.psf_convolution_pad_factor*shape[2])]])
-    psf_padded = tf.pad(features['psf'][...,0], [[0,0],
-                                            [0, int(hparams.psf_convolution_pad_factor*shape[1])],
-                                            [0, int(hparams.psf_convolution_pad_factor*shape[2])]])
-    reconstr = tf.expand_dims(tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded)), tf.complex64)),axis=-1)
+    rec_padded = tf.pad(reconstr, [[0,0],
+                                  [0, int(hparams.psf_convolution_pad_factor*shape[1])],
+                                  [0, int(hparams.psf_convolution_pad_factor*shape[2])],
+                                  [0,0]])
+    rec_padded_t = tf.transpose(rec_padded, perm=[0, 3,  1, 2])
+    psf_padded = tf.pad(features['psf'], [[0,0],
+                                          [0, int(hparams.psf_convolution_pad_factor*shape[1])],
+                                          [0, int(hparams.psf_convolution_pad_factor*shape[2])],
+                                          [0,0]])
+    psf_padded_t = tf.transpose(psf_padded, perm=[0, 3,  1, 2])
+    reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded)), tf.complex64))
+    reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
     reconstr = reconstr[:, :shape[1], :shape[2], :]
 
   # Losses.

From 5f070aa385129cf8224c50c8de94cf3b09da79d3 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Feb 2021 14:50:33 +0000
Subject: [PATCH 041/149] Added power spectrum computation to CANDELS (probably
 not good)

---
 galaxy2galaxy/data_generators/candels.py | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c4df272..c279502 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -803,10 +803,14 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     attributes=None
                 
 #                 ''' Create the power spectrum '''
-#                 ps = np.zeros((p.img_len, p.img_len//2+1, len(p.filters)))
-#                 for i in range(len(p.filters)):
-#                     ps[:, :, n_filter] = np.random.normal(0, p.sigmas[n_filter], (p.img_len, p.img_len // 2 + 1))
-                
+                k = 0
+                noise_im = np.zeros((p.img_len, p.img_len, len(p.filters)))
+                for res in p.resolutions:
+                    for n_filter in range(len(p.filters[res])):
+                        noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
+                    k+=1
+                ps = np.abs(np.fft.rfft2(ps))
+
                 ''' Add a flag corresponding to the field '''
                 field_info = np.asarray(n_field)
 
@@ -820,6 +824,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 "image/format": ["raw"],
                 "psf/encoded": [im_psf.astype('float32').tostring()],
                 "psf/format": ["raw"],
+                "ps/encoded": [ps.astype('float32').tostring()],
+                "ps/format": ["raw"],
                 "sigma_noise/encoded": [sigmas_array.astype('float32').tostring()],
                 "sigma_noise/format": ["raw"],
                 "field/encoded": [field_info.astype('float32').tostring()],

From 7ab7e0f583e806d2b6cb3656aaa99b6e2207bbbe Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Feb 2021 16:24:05 +0000
Subject: [PATCH 042/149] Ensure power spectrum is properly exported

---
 galaxy2galaxy/data_generators/candels.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c279502..e86fdf0 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -882,6 +882,9 @@ def example_reading_spec(self):
         "psf/encoded": tf.FixedLenFeature((), tf.string),
         "psf/format": tf.FixedLenFeature((), tf.string),
 
+        "ps/encoded": tf.FixedLenFeature((), tf.string),
+        "ps/format": tf.FixedLenFeature((), tf.string),
+
         "sigma_noise/encoded": tf.FixedLenFeature((), tf.string),
         "sigma_noise/format": tf.FixedLenFeature((), tf.string),
 
@@ -908,7 +911,14 @@ def example_reading_spec(self):
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
-        "sigma": tf.contrib.slim.tfexample_decoder.Image(
+        "ps": tf.contrib.slim.tfexample_decoder.Image(
+                image_key="ps/encoded",
+                format_key="ps/format",
+#                 channels=self.num_bands,
+                shape=[p.img_len, p.img_len//2+1, self.num_bands],
+                dtype=tf.float32),
+
+        "sigma_noise": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="sigma_noise/encoded",
                 format_key="sigma_noise/format",
 #                 channels=self.num_bands,

From 152f59308827d9e2f32af6adb8084b6328bb5791 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Feb 2021 17:14:38 +0000
Subject: [PATCH 043/149] Fix dimension problem

---
 galaxy2galaxy/models/autoencoders_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index d76532c..139492e 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -37,8 +37,8 @@ def loglikelihood_fn(xin, yin, features, hparams):
   size = xin.get_shape().as_list()[1]
   if hparams.likelihood_type == 'Fourier':
     # Compute FFT normalization factor
-    x = tf.spectral.rfft2d(xin[...,0]) / tf.complex(tf.sqrt(tf.exp(features['ps'])),0.) / size**2 * (2*np.pi)**2
-    y = tf.spectral.rfft2d(yin[...,0]) / tf.complex(tf.sqrt(tf.exp(features['ps'])),0.) / size**2 * (2*np.pi)**2
+    x = tf.spectral.rfft2d(xin[...,0]) / tf.complex(tf.sqrt(tf.exp(features['ps'][...,0])),0.) / size**2 * (2*np.pi)**2
+    y = tf.spectral.rfft2d(yin[...,0]) / tf.complex(tf.sqrt(tf.exp(features['ps'][...,0])),0.) / size**2 * (2*np.pi)**2
 
     pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2]) #/ size**2
     return -pz

From 5126f9518a22c1cb33aebba7a5873bb3b72a98e7 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 25 Feb 2021 10:22:39 +0000
Subject: [PATCH 044/149] Fix dimension issues

---
 galaxy2galaxy/models/autoencoders_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 139492e..69629da 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -118,7 +118,7 @@ def make_model_spec():
       input_layer = tf.placeholder(tf.float32, shape=b_shape)
       x = self.unbottleneck(input_layer, res_size)
       x = self.decoder(x, None)
-      reconstr = tf.layers.dense(x, self.num_channels, name="autoencoder_final",
+      reconstr = tf.layers.dense(x, input_shape[-1], name="autoencoder_final",
                                  activation=output_activation)
       hub.add_signature(inputs=input_layer, outputs=reconstr)
       hub.attach_message("stamp_size", tf.train.Int64List(value=[hparams.problem_hparams.img_len]))
@@ -210,7 +210,7 @@ def make_model_spec():
     res = x[:, :shape[1], :shape[2], :]
 
   with tf.variable_scope('decoder_module'):
-    reconstr = tf.layers.dense(res, self.num_channels, name="autoencoder_final",
+    reconstr = tf.layers.dense(res, shape[-1], name="autoencoder_final",
                                activation=output_activation)
 
   # Apply channel-wise convolution with the PSF if requested
@@ -225,7 +225,7 @@ def make_model_spec():
                                           [0, int(hparams.psf_convolution_pad_factor*shape[2])],
                                           [0,0]])
     psf_padded_t = tf.transpose(psf_padded, perm=[0, 3,  1, 2])
-    reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded)), tf.complex64))
+    reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded_t)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded_t)), tf.complex64))
     reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
     reconstr = reconstr[:, :shape[1], :shape[2], :]
 

From 4a3f24696067104ba2645450f1c165e898140c25 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 25 Feb 2021 10:43:14 +0000
Subject: [PATCH 045/149] Fix power spectrum error

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index e86fdf0..7b094ba 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -804,7 +804,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 
 #                 ''' Create the power spectrum '''
                 k = 0
-                noise_im = np.zeros((p.img_len, p.img_len, len(p.filters)))
+                noise_im = np.zeros((p.img_len, p.img_len, band_num))
                 for res in p.resolutions:
                     for n_filter in range(len(p.filters[res])):
                         noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))

From 2a7302438c5cd3ccc5d1149c4ad8c773654f615f Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 25 Feb 2021 14:12:49 +0000
Subject: [PATCH 046/149] Fix typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 7b094ba..e5a82de 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -809,7 +809,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     for n_filter in range(len(p.filters[res])):
                         noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
                     k+=1
-                ps = np.abs(np.fft.rfft2(ps))
+                ps = np.abs(np.fft.rfft2(noise_im))
 
                 ''' Add a flag corresponding to the field '''
                 field_info = np.asarray(n_field)

From 6e836b60b44ce318313b9faabee5ccea2f57b6ec Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 25 Feb 2021 17:59:28 +0000
Subject: [PATCH 047/149] Add missing transposition for power spectrum

---
 galaxy2galaxy/data_generators/candels.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index e5a82de..0665fae 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -809,7 +809,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     for n_filter in range(len(p.filters[res])):
                         noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
                     k+=1
+                noise_im = np.transpose(noise_im,[2,0,1])
                 ps = np.abs(np.fft.rfft2(noise_im))
+                ps = np.transpose(ps,[1,2,0])
 
                 ''' Add a flag corresponding to the field '''
                 field_info = np.asarray(n_field)

From a1dd1e8a86b23ef711996aabd242fbf21dad4f7a Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 26 Feb 2021 16:22:00 +0000
Subject: [PATCH 048/149] Made loss multiband

---
 galaxy2galaxy/models/autoencoders_utils.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 69629da..7456b5a 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -37,10 +37,12 @@ def loglikelihood_fn(xin, yin, features, hparams):
   size = xin.get_shape().as_list()[1]
   if hparams.likelihood_type == 'Fourier':
     # Compute FFT normalization factor
-    x = tf.spectral.rfft2d(xin[...,0]) / tf.complex(tf.sqrt(tf.exp(features['ps'][...,0])),0.) / size**2 * (2*np.pi)**2
-    y = tf.spectral.rfft2d(yin[...,0]) / tf.complex(tf.sqrt(tf.exp(features['ps'][...,0])),0.) / size**2 * (2*np.pi)**2
+    x = tf.transpose(xin,[0,3,1,2])
+    y = tf.transpose(yin,[0,3,1,2])
+    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(features['ps'])),0.) / size**2 * (2*np.pi)**2
+    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(features['ps'])),0.) / size**2 * (2*np.pi)**2
 
-    pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2]) #/ size**2
+    pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2, -3]) #/ size**2
     return -pz
   elif hparams.likelihood_type == 'Pixel':
     # TODO: include per example noise std

From e60cf53ff24deb6a7d96d6dc10f251665fc24cb4 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Sun, 28 Feb 2021 15:39:15 +0000
Subject: [PATCH 049/149] Fix issue with shards

---
 galaxy2galaxy/data_generators/candels.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 0665fae..ccffbc8 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -681,10 +681,10 @@ def dataset_splits(self):
     """
     return [{
         "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,
+        "shards": 10,
     }, {
         "split": problem.DatasetSplit.EVAL,
-        "shards": 0,
+        "shards": 1,
     }]
 
   # START: Subclass interface
@@ -993,7 +993,7 @@ def hparams(self, defaults, model_hparams):
     p.sigmas = {"high" : [0.0034001764449330513], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
     p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
-    p.example_per_shard = 1000
+    p.example_per_shard = 100
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}

From ba46fb635708e3f998642cf010991179988f1703 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 1 Mar 2021 11:01:50 +0000
Subject: [PATCH 050/149] Clean up to try avoiding the error

---
 galaxy2galaxy/data_generators/candels.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index ccffbc8..476f2c8 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -868,12 +868,14 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
     return example
 
   def example_reading_spec(self):
-    """Define how data is serialized to file and read back.
+    """
+    Define how data is serialized to file and read back.
 
     Returns:
       data_fields: A dictionary mapping data names to its feature type.
       data_items_to_decoders: A dictionary mapping data names to TF Example
-         decoders, to be used when reading back TF examples from disk.
+                              decoders, to be used when reading back TF examples 
+                              from disk.
     """
     p = self.get_hparams()
 
@@ -902,28 +904,28 @@ def example_reading_spec(self):
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
-#                 channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "psf": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="psf/encoded",
                 format_key="psf/format",
-#                 channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "ps": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="ps/encoded",
                 format_key="ps/format",
-#                 channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len//2+1, self.num_bands],
                 dtype=tf.float32),
 
         "sigma_noise": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="sigma_noise/encoded",
                 format_key="sigma_noise/format",
-#                 channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[self.num_bands],
                 dtype=tf.float32),
         
@@ -936,7 +938,6 @@ def example_reading_spec(self):
     if hasattr(p, 'attributes'):
         for k in p.attributes:
             data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
-#         data_items_to_decoders['field'] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/field')
 
     return data_fields, data_items_to_decoders
 #     END: Subclass interface

From 3a4646472e860312014a38e57e321ec9e330d98d Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 1 Mar 2021 15:32:17 +0000
Subject: [PATCH 051/149] Fix exporter

---
 galaxy2galaxy/bin/g2g-exporter | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/bin/g2g-exporter b/galaxy2galaxy/bin/g2g-exporter
index 0183f78..3e3f7eb 100755
--- a/galaxy2galaxy/bin/g2g-exporter
+++ b/galaxy2galaxy/bin/g2g-exporter
@@ -6,6 +6,11 @@ from __future__ import print_function
 
 import os
 
+import tensorflow as tf
+tf.disable_v2_behavior()
+
+import tensorflow_hub as hub
+
 from tensor2tensor.bin import t2t_trainer
 from tensor2tensor.utils import decoding
 from tensor2tensor.utils import t2t_model
@@ -15,8 +20,6 @@ from tensor2tensor.utils import usr_dir
 from galaxy2galaxy import models
 from galaxy2galaxy import problems
 
-import tensorflow as tf
-import tensorflow_hub as hub
 
 FLAGS = tf.flags.FLAGS
 
@@ -97,6 +100,8 @@ def main(_):
   run_config = t2t_trainer.create_run_config(hparams)
 
   estimator = create_estimator(run_config, hparams)
+  hparams.img_len = problem.get_hparams().img_len
+  hparams.attributes = problem.get_hparams().attributes
 
   # Use tf hub to export any module that has been registered
   exporter = hub.LatestModuleExporter("tf_hub",
@@ -110,4 +115,4 @@ def main(_):
 
 if __name__ == "__main__":
   tf.logging.set_verbosity(tf.logging.INFO)
-  tf.app.run()
+  tf.app.run()
\ No newline at end of file

From e27cbfe0e0756f5d207af4d9322a3621f6b6659c Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 1 Mar 2021 15:34:17 +0000
Subject: [PATCH 052/149] Fix exporter

---
 galaxy2galaxy/bin/g2g-exporter | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/bin/g2g-exporter b/galaxy2galaxy/bin/g2g-exporter
index 3e3f7eb..9eb0f44 100755
--- a/galaxy2galaxy/bin/g2g-exporter
+++ b/galaxy2galaxy/bin/g2g-exporter
@@ -101,6 +101,7 @@ def main(_):
 
   estimator = create_estimator(run_config, hparams)
   hparams.img_len = problem.get_hparams().img_len
+  hparams.resolutions = problem.get_hparams().resolutions
   hparams.attributes = problem.get_hparams().attributes
 
   # Use tf hub to export any module that has been registered

From b51366b19cd60b49fe2a0fc5aabaf52769477842 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 1 Mar 2021 15:42:06 +0000
Subject: [PATCH 053/149] Fix exporter

---
 galaxy2galaxy/bin/g2g-exporter | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/bin/g2g-exporter b/galaxy2galaxy/bin/g2g-exporter
index 9eb0f44..6b82be4 100755
--- a/galaxy2galaxy/bin/g2g-exporter
+++ b/galaxy2galaxy/bin/g2g-exporter
@@ -16,6 +16,7 @@ from tensor2tensor.utils import decoding
 from tensor2tensor.utils import t2t_model
 from tensor2tensor.utils import trainer_lib
 from tensor2tensor.utils import usr_dir
+from tensor2tensor.utils import registry
 
 from galaxy2galaxy import models
 from galaxy2galaxy import problems
@@ -100,9 +101,8 @@ def main(_):
   run_config = t2t_trainer.create_run_config(hparams)
 
   estimator = create_estimator(run_config, hparams)
-  hparams.img_len = problem.get_hparams().img_len
-  hparams.resolutions = problem.get_hparams().resolutions
-  hparams.attributes = problem.get_hparams().attributes
+  problem = registry.problem(FLAGS.problem)
+  problem.hparams(hparams,None)
 
   # Use tf hub to export any module that has been registered
   exporter = hub.LatestModuleExporter("tf_hub",

From 2af4a42324aee71f7994dc2a015673a5e8ca40b5 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 1 Mar 2021 15:48:00 +0000
Subject: [PATCH 054/149] Fix hub pixel scale error

---
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 7456b5a..f66ebe7 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -124,7 +124,7 @@ def make_model_spec():
                                  activation=output_activation)
       hub.add_signature(inputs=input_layer, outputs=reconstr)
       hub.attach_message("stamp_size", tf.train.Int64List(value=[hparams.problem_hparams.img_len]))
-      hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale]))
+      hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale[res] for res in hparams.problem_hparams.resolutions]))
     spec = hub.create_module_spec(make_model_spec, drop_collections=['checkpoints'])
     decoder = hub.Module(spec, name="decoder_module")
     hub.register_module_for_export(decoder, "decoder")

From 8ef6cbdf26e8f42e33ba106940f03f4370ce28da Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 3 Mar 2021 22:34:54 +0000
Subject: [PATCH 055/149] Made use of task id

---
 galaxy2galaxy/data_generators/candels.py | 201 +++++++++++------------
 1 file changed, 99 insertions(+), 102 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 476f2c8..21a54f1 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -744,109 +744,106 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
     # Step 2: Extract postage stamps, resize them to requested size
     ''' Loop on the two fields'''
-    for n_field, field in enumerate(['GDS']):
-        
-        print(f"\n generating{field}, {n_field}\n")
-        n_gal_creat = 0
-        index = 0
-        
-        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-
-        ''' Loop on all the galaxies of the field '''
-        for gal in sub_cat['RB_ID']:
-            if gal == index or gal == 15431 :     # To take care of the redudency inside the cat
-                continue
-            index = gal
-            print(index)
-
-            try:
-                ''' Loop on the filters '''
-                im = np.zeros((target_size, target_size, band_num))
-
-                k = 0
-                for res in p.resolutions:
-                    im_tmp = np.zeros((128, 128, len(p.filters[res])))
-                    for n_filter, filt in enumerate(p.filters[res]):
-                        print(filt)
-                        # try :
-                        ''' Open the image corresponding to the index of the current galaxy'''
-
-                        tmp_file = glob.glob(os.path.join(data_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
-                        # if np.max(fits.open(tmp_file)[0].data) == 0.:
-                        #     sigmas[res][n_filter] = 10
-                        im_import = fits.open(tmp_file)[0].data
-                        im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5,noise_level=p.sigmas[res][n_filter])
-
-                        # except Exception:
-                        #     print('Galaxy not seen in every filter')
-                        #     continue
-                            
-                    ''' Resize the image to the low resolution'''
-                    new_size = np.ceil(128/scalings[res])+1
-                    im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
-                    ''' Resize the image to the highest resolution to get consistent array sizes'''
-                    im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel=True)
-                    im_tmp = _resize_image(im_tmp,target_size)
-
-                    im[:,:,k:k+len(p.filters[res])] = im_tmp
-                    k += len(p.filters[res])
-                
-                im = _resize_image(im, p.img_len)
-                
-
-                ''' Load the wanted physical parameters of the galaxy '''
-                if hasattr(p, 'attributes'):
-                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
-
-                else:
-                    attributes=None
-                
+    n_gal_creat = 0
+    index = 0
+    
+    ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
+    sub_cat = all_cat[np.where(np.isin(all_cat["FIELD_1"],["GDS","GDN"]))[0]]
+
+    ''' Loop on all the galaxies of the field '''
+    for m,gal in enumerate(sub_cat['RB_ID']):
+        if gal == index or gal == 15431 or m < task_id*p.example_per_shard:     # To take care of the redudency inside the cat
+            continue
+        index = gal
+        print(index)
+
+        try:
+            ''' Loop on the filters '''
+            im = np.zeros((target_size, target_size, band_num))
+
+            k = 0
+            for res in p.resolutions:
+                im_tmp = np.zeros((128, 128, len(p.filters[res])))
+                for n_filter, filt in enumerate(p.filters[res]):
+                    print(filt)
+                    # try :
+                    ''' Open the image corresponding to the index of the current galaxy'''
+
+                    tmp_file = glob.glob(os.path.join(data_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
+                    # if np.max(fits.open(tmp_file)[0].data) == 0.:
+                    #     sigmas[res][n_filter] = 10
+                    im_import = fits.open(tmp_file)[0].data
+                    im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5,noise_level=p.sigmas[res][n_filter])
+
+                    # except Exception:
+                    #     print('Galaxy not seen in every filter')
+                    #     continue
+                        
+                ''' Resize the image to the low resolution'''
+                new_size = np.ceil(128/scalings[res])+1
+                im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
+                ''' Resize the image to the highest resolution to get consistent array sizes'''
+                im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel=True)
+                im_tmp = _resize_image(im_tmp,target_size)
+
+                im[:,:,k:k+len(p.filters[res])] = im_tmp
+                k += len(p.filters[res])
+            
+            im = _resize_image(im, p.img_len)
+            
+
+            ''' Load the wanted physical parameters of the galaxy '''
+            if hasattr(p, 'attributes'):
+                attributes = {k: float(all_cat[k][index]) for k in p.attributes}
+
+            else:
+                attributes=None
+            
 #                 ''' Create the power spectrum '''
-                k = 0
-                noise_im = np.zeros((p.img_len, p.img_len, band_num))
-                for res in p.resolutions:
-                    for n_filter in range(len(p.filters[res])):
-                        noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
-                    k+=1
-                noise_im = np.transpose(noise_im,[2,0,1])
-                ps = np.abs(np.fft.rfft2(noise_im))
-                ps = np.transpose(ps,[1,2,0])
-
-                ''' Add a flag corresponding to the field '''
-                field_info = np.asarray(n_field)
-
-                sigmas_array = []
-                for res in p.resolutions:
-                    sigmas_array += sigmas[res]
-                sigmas_array = np.array(sigmas_array)
-
-                ''' Create the output to match T2T format '''
-                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
-                "image/format": ["raw"],
-                "psf/encoded": [im_psf.astype('float32').tostring()],
-                "psf/format": ["raw"],
-                "ps/encoded": [ps.astype('float32').tostring()],
-                "ps/format": ["raw"],
-                "sigma_noise/encoded": [sigmas_array.astype('float32').tostring()],
-                "sigma_noise/format": ["raw"],
-                "field/encoded": [field_info.astype('float32').tostring()],
-                "field/format": ["raw"]}
-                
-                if attributes is not None:
-                    for k in attributes:
-                        serialized_output['attrs/'+k] = [attributes[k]]
-                
-                ''' Increment the number of galaxy created on the shard '''
-                n_gal_creat += 1
-                
-                if n_gal_creat > p.example_per_shard:
-                    print('out')
-                    break
-                yield serialized_output
-            except ValueError:
-                print(sys.exc_info()[0], sys.exc_info()[1])
-                continue
+            k = 0
+            noise_im = np.zeros((p.img_len, p.img_len, band_num))
+            for res in p.resolutions:
+                for n_filter in range(len(p.filters[res])):
+                    noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
+                k+=1
+            noise_im = np.transpose(noise_im,[2,0,1])
+            ps = np.abs(np.fft.rfft2(noise_im))
+            ps = np.transpose(ps,[1,2,0])
+
+            ''' Add a flag corresponding to the field '''
+            field_info = np.asarray(n_field)
+
+            sigmas_array = []
+            for res in p.resolutions:
+                sigmas_array += sigmas[res]
+            sigmas_array = np.array(sigmas_array)
+
+            ''' Create the output to match T2T format '''
+            serialized_output = {"image/encoded": [im.astype('float32').tostring()],
+            "image/format": ["raw"],
+            "psf/encoded": [im_psf.astype('float32').tostring()],
+            "psf/format": ["raw"],
+            "ps/encoded": [ps.astype('float32').tostring()],
+            "ps/format": ["raw"],
+            "sigma_noise/encoded": [sigmas_array.astype('float32').tostring()],
+            "sigma_noise/format": ["raw"],
+            "field/encoded": [field_info.astype('float32').tostring()],
+            "field/format": ["raw"]}
+            
+            if attributes is not None:
+                for k in attributes:
+                    serialized_output['attrs/'+k] = [attributes[k]]
+            
+            ''' Increment the number of galaxy created on the shard '''
+            n_gal_creat += 1
+            
+            if n_gal_creat > p.example_per_shard or m >= task_id*p.example_per_shard:
+                print('out')
+                break
+            yield serialized_output
+        except ValueError:
+            print(sys.exc_info()[0], sys.exc_info()[1])
+            continue
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
     """ Preprocess the examples, can be used for further augmentation or

From aebc0bc91cdebbacc27a5b0cf2ad2b80d8ce7a73 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 10:06:40 +0000
Subject: [PATCH 056/149] Fixed fields and changed params

---
 galaxy2galaxy/data_generators/candels.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 21a54f1..2c909b5 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -752,7 +752,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
     ''' Loop on all the galaxies of the field '''
     for m,gal in enumerate(sub_cat['RB_ID']):
-        if gal == index or gal == 15431 or m < task_id*p.example_per_shard:     # To take care of the redudency inside the cat
+        if gal == index or gal == 15431 or m < task_id*2*p.example_per_shard:     # To take care of the redudency inside the cat
             continue
         index = gal
         print(index)
@@ -769,7 +769,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     # try :
                     ''' Open the image corresponding to the index of the current galaxy'''
 
-                    tmp_file = glob.glob(os.path.join(data_dir, field, filt)+'/galaxy_'+str(index)+'_*')[0]
+                    tmp_file = glob.glob(os.path.join(data_dir, sub_cat["FIELD_1"][m], filt)+'/galaxy_'+str(index)+'_*')[0]
                     # if np.max(fits.open(tmp_file)[0].data) == 0.:
                     #     sigmas[res][n_filter] = 10
                     im_import = fits.open(tmp_file)[0].data
@@ -837,7 +837,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ''' Increment the number of galaxy created on the shard '''
             n_gal_creat += 1
             
-            if n_gal_creat > p.example_per_shard or m >= task_id*p.example_per_shard:
+            if n_gal_creat > p.example_per_shard or m >= task_id*2*p.example_per_shard:
                 print('out')
                 break
             yield serialized_output
@@ -988,10 +988,10 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = {'high' : 0.06, 'low' : 0.06}
     p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
-    p.sigmas = {"high" : [0.0034001764449330513], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
-    p.filters = {"high" : ['acs_f775w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.sigmas = {"high" : [0.0034001764449330513], "low" : [0.004017507500562]}
+    p.filters = {"high" : ['acs_f775w'], "low" : ['wfc3_f160w']}
     p.resolutions = ["high","low"]
-    p.example_per_shard = 100
+    p.example_per_shard = 1000
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}

From f1a6e71fe63431cf93c8d1702638b3fa649e3907 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 11:53:06 +0000
Subject: [PATCH 057/149] Changing the main filter

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 2c909b5..c2712d1 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -988,8 +988,8 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = {'high' : 0.06, 'low' : 0.06}
     p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
-    p.sigmas = {"high" : [0.0034001764449330513], "low" : [0.004017507500562]}
-    p.filters = {"high" : ['acs_f775w'], "low" : ['wfc3_f160w']}
+    p.sigmas = {"high" : [0.004094741966557142], "low" : [0.004017507500562]}
+    p.filters = {"high" : ['acs_f606w'], "low" : ['wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.example_per_shard = 1000
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,

From 4dceda810ff703d77958629559263101287c7199 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 12:03:39 +0000
Subject: [PATCH 058/149] Fix empty catalog

---
 galaxy2galaxy/data_generators/candels.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c2712d1..e1331d9 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -719,7 +719,6 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     target_pixel_scale = p.pixel_scale[p.resolutions[0]]
     target_scaling = target_pixel_scale/p.base_pixel_scale[p.resolutions[0]]
     target_size = p.img_len
-    print("scalings and all ",scalings,target_pixel_scale,target_scaling)
     
     '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(data_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
@@ -748,7 +747,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     index = 0
     
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-    sub_cat = all_cat[np.where(np.isin(all_cat["FIELD_1"],["GDS","GDN"]))[0]]
+    sub_cat = all_cat[np.where(np.isin(all_cat["FIELD_1"],["GDS","GDN"]))]
 
     ''' Loop on all the galaxies of the field '''
     for m,gal in enumerate(sub_cat['RB_ID']):

From 4874891c7a94304568b9efdc20ad6f59dd83b163 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 15:53:02 +0000
Subject: [PATCH 059/149] CHeck task id

---
 galaxy2galaxy/data_generators/candels.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index e1331d9..19d5ad3 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -710,6 +710,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     """ 
     Generator yielding individual postage stamps.
     """
+    print(task_id)
     
     p = self.get_hparams()
     band_num = np.sum([len(p.filters[res]) for res in p.resolutions])
@@ -747,7 +748,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     index = 0
     
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-    sub_cat = all_cat[np.where(np.isin(all_cat["FIELD_1"],["GDS","GDN"]))]
+    sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN"]))]
 
     ''' Loop on all the galaxies of the field '''
     for m,gal in enumerate(sub_cat['RB_ID']):
@@ -810,7 +811,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ps = np.transpose(ps,[1,2,0])
 
             ''' Add a flag corresponding to the field '''
-            field_info = np.asarray(n_field)
+            field_info = np.asarray(all_cat["FIELD_1"][m])
 
             sigmas_array = []
             for res in p.resolutions:
@@ -826,7 +827,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             "ps/format": ["raw"],
             "sigma_noise/encoded": [sigmas_array.astype('float32').tostring()],
             "sigma_noise/format": ["raw"],
-            "field/encoded": [field_info.astype('float32').tostring()],
+            "field/encoded": [field_info.tostring()],
             "field/format": ["raw"]}
             
             if attributes is not None:

From 58f7bc50661f251cae92e155f1a23ec75de35fbb Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 15:54:38 +0000
Subject: [PATCH 060/149] Fix typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 19d5ad3..779932d 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -837,7 +837,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ''' Increment the number of galaxy created on the shard '''
             n_gal_creat += 1
             
-            if n_gal_creat > p.example_per_shard or m >= task_id*2*p.example_per_shard:
+            if n_gal_creat > p.example_per_shard or m >= (task_id+1)*2*p.example_per_shard:
                 print('out')
                 break
             yield serialized_output

From 4dff64cb56b776e4c11d6a0e59827286c44f3171 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 16:04:11 +0000
Subject: [PATCH 061/149] Modify code for task_id

---
 galaxy2galaxy/data_generators/candels.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 779932d..138667f 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -746,13 +746,18 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     ''' Loop on the two fields'''
     n_gal_creat = 0
     index = 0
+   
     
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN"]))]
+    assert(task_id > -1)
+    indexes = range(task_id*p.example_per_shard,
+                  min((task_id+1)*p.example_per_shard, len(sub_cat)))
+    sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''
     for m,gal in enumerate(sub_cat['RB_ID']):
-        if gal == index or gal == 15431 or m < task_id*2*p.example_per_shard:     # To take care of the redudency inside the cat
+        if gal == index or gal == 15431:     # To take care of the redudency inside the cat
             continue
         index = gal
         print(index)
@@ -837,7 +842,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ''' Increment the number of galaxy created on the shard '''
             n_gal_creat += 1
             
-            if n_gal_creat > p.example_per_shard or m >= (task_id+1)*2*p.example_per_shard:
+            if n_gal_creat > p.example_per_shard:
                 print('out')
                 break
             yield serialized_output

From 08ff2c94372a6de5f2c171c400ff56934794ddf5 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 16:29:26 +0000
Subject: [PATCH 062/149] Add multiprocessing to candels

---
 galaxy2galaxy/data_generators/candels.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 138667f..6add2f3 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -687,6 +687,11 @@ def dataset_splits(self):
         "shards": 1,
     }]
 
+  @property
+  def multiprocess_generate(self):
+    """Whether to generate the data in multiple parallel processes."""
+    return True
+
   # START: Subclass interface
   def hparams(self, defaults, model_hparams):
     p = defaults

From d9a8fd947d774541fdf78e5ca4bd7c4eec9c9831 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 16:44:19 +0000
Subject: [PATCH 063/149] Fix indexes

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 6add2f3..7a1b5c3 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -756,8 +756,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN"]))]
     assert(task_id > -1)
-    indexes = range(task_id*p.example_per_shard,
-                  min((task_id+1)*p.example_per_shard, len(sub_cat)))
+    indexes = list(range(task_id*p.example_per_shard,
+                  min((task_id+1)*p.example_per_shard, len(sub_cat))))
     sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''

From 691524ed27878e525070856be44179a3a3fbc123 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 4 Mar 2021 17:04:04 +0000
Subject: [PATCH 064/149] Fix problem with field num

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 7a1b5c3..549637d 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -821,7 +821,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ps = np.transpose(ps,[1,2,0])
 
             ''' Add a flag corresponding to the field '''
-            field_info = np.asarray(all_cat["FIELD_1"][m])
+            field_info = np.asarray(1 if all_cat["FIELD_1"][m] == "GDS" else 0)
 
             sigmas_array = []
             for res in p.resolutions:
@@ -837,7 +837,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             "ps/format": ["raw"],
             "sigma_noise/encoded": [sigmas_array.astype('float32').tostring()],
             "sigma_noise/format": ["raw"],
-            "field/encoded": [field_info.tostring()],
+            "field/encoded": [field_info.astype('float32').tostring()],
             "field/format": ["raw"]}
             
             if attributes is not None:

From c9733e3451da23d8ccd709b78436b41b6b546efa Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 5 Mar 2021 11:48:45 +0000
Subject: [PATCH 065/149] Fix weird summary error

---
 galaxy2galaxy/data_generators/candels.py   | 8 ++++----
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 549637d..14e1566 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -696,8 +696,8 @@ def multiprocess_generate(self):
   def hparams(self, defaults, model_hparams):
     p = defaults
     p.img_len = 128
-    p.sigmas = {"high" : [1e-4], "low" : [6.7e-3, 5.4e-3, 4.0e-3]}
-    p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
+    p.sigmas = {"high" : [1e-4], "low" : [4.0e-3]}
+    p.filters = {"high" : ['acs_f814w'], "low" : ['wfc3_f160w']}
     p.resolutions = ["high","low"]
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
@@ -756,8 +756,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN"]))]
     assert(task_id > -1)
-    indexes = list(range(task_id*p.example_per_shard,
-                  min((task_id+1)*p.example_per_shard, len(sub_cat))))
+    indexes = list(range(task_id*4*p.example_per_shard,
+                  min((task_id*4+1)*p.example_per_shard, len(sub_cat))))
     sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index f66ebe7..0e05657 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -53,7 +53,7 @@ def loglikelihood_fn(xin, yin, features, hparams):
 
 def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
   """Helper for image summaries that are safe on TPU."""
-  if len(image_logits.get_shape()) != 4:
+  if len(image_logits.get_shape()) != 4 or not image_logits.get_shape()[4] in [1,3,4]:
     tf.logging.info("Not generating image summary, maybe not an image.")
     return
   return tf.summary.image(name, pack_images(image_logits, rows, cols),

From 71e44bcfa6de02e0893353db280a651d52572d05 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Sun, 7 Mar 2021 14:15:52 +0000
Subject: [PATCH 066/149] Add all fields and correct typo

---
 galaxy2galaxy/data_generators/candels.py   | 5 ++++-
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 14e1566..74d241b 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -729,7 +729,10 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     '''Load the catalogue containing every fields and every filter'''
     all_cat = Table.read(os.path.join(data_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='egs   ')] = 'EGS'
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='UDS   ')] = 'UDS'
+    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='COSMOS   ')] = 'COSMOS'
     
     ''' Load the psfs for each filter and resize'''
     cube_psf = np.zeros((167, 167, band_num))
@@ -754,7 +757,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
    
     
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-    sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN"]))]
+    sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     assert(task_id > -1)
     indexes = list(range(task_id*4*p.example_per_shard,
                   min((task_id*4+1)*p.example_per_shard, len(sub_cat))))
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 0e05657..a0ff56a 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -53,7 +53,7 @@ def loglikelihood_fn(xin, yin, features, hparams):
 
 def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
   """Helper for image summaries that are safe on TPU."""
-  if len(image_logits.get_shape()) != 4 or not image_logits.get_shape()[4] in [1,3,4]:
+  if len(image_logits.get_shape()) != 4 or not image_logits.get_shape()[3] in [1,3,4]:
     tf.logging.info("Not generating image summary, maybe not an image.")
     return
   return tf.summary.image(name, pack_images(image_logits, rows, cols),

From 718c1989a27e2a6ce19239fe47829f99f6fa7be6 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 8 Mar 2021 11:16:45 +0000
Subject: [PATCH 067/149] Change shard number

---
 galaxy2galaxy/data_generators/candels.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 74d241b..162b5de 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -681,10 +681,10 @@ def dataset_splits(self):
     """
     return [{
         "split": problem.DatasetSplit.TRAIN,
-        "shards": 10,
+        "shards": 20,
     }, {
         "split": problem.DatasetSplit.EVAL,
-        "shards": 1,
+        "shards": 2,
     }]
 
   @property
@@ -851,10 +851,10 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             n_gal_creat += 1
             
             if n_gal_creat > p.example_per_shard:
-                print('out')
+                print('out ',n_gal_creat)
                 break
             yield serialized_output
-        except ValueError:
+        except Exception:
             print(sys.exc_info()[0], sys.exc_info()[1])
             continue
 
@@ -976,7 +976,7 @@ def hparams(self, defaults, model_hparams):
     p.sigmas = {"high" : [1e-4], "low" : [0.003954237367399534, 0.003849901319445, 0.004017507500562]}
     p.filters = {"high" : ['acs_f814w'], "low" : ['f105w', 'f125w', 'wfc3_f160w']}
     p.resolutions = ["high","low"]
-    p.example_per_shard = 1000
+    p.example_per_shard = 2000
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}

From 52ee908d000433620d0cb9f6ac94d135cf5f7c36 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 8 Mar 2021 11:24:18 +0000
Subject: [PATCH 068/149] Change task_id thing

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 162b5de..dc63198 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -759,8 +759,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     assert(task_id > -1)
-    indexes = list(range(task_id*4*p.example_per_shard,
-                  min((task_id*4+1)*p.example_per_shard, len(sub_cat))))
+    indexes = list(range(task_id*3*p.example_per_shard,
+                  min((task_id*3+1)*p.example_per_shard, len(sub_cat))))
     sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''

From 0089b6104e61db5bbf5491767f51629cfcd8a2d7 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 8 Mar 2021 11:43:49 +0000
Subject: [PATCH 069/149] Fix task_id

---
 galaxy2galaxy/data_generators/candels.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index dc63198..143df47 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -759,8 +759,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     assert(task_id > -1)
-    indexes = list(range(task_id*3*p.example_per_shard,
-                  min((task_id*3+1)*p.example_per_shard, len(sub_cat))))
+    n_shards = self.dataset_splits()[0]["shards"] + self.dataset_splits()[1]["shards"]
+    indexes = list(range(task_id*len(sub_cat)/n_shards*p.example_per_shard,
+                  min((task_id+1)*len(sub_cat)/n_shards*p.example_per_shard, len(sub_cat))))
     sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''

From 967a05241ff5e3b80f8993ed4acf4abe14461e7e Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 8 Mar 2021 11:53:08 +0000
Subject: [PATCH 070/149] typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 143df47..6439f04 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -759,7 +759,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     assert(task_id > -1)
-    n_shards = self.dataset_splits()[0]["shards"] + self.dataset_splits()[1]["shards"]
+    n_shards = self.dataset_splits[0]["shards"] + self.dataset_splits[1]["shards"]
     indexes = list(range(task_id*len(sub_cat)/n_shards*p.example_per_shard,
                   min((task_id+1)*len(sub_cat)/n_shards*p.example_per_shard, len(sub_cat))))
     sub_cat = sub_cat[indexes]

From c1ddaaa56560cff0d081330e8ef02ad6d6524ffd Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 8 Mar 2021 11:59:18 +0000
Subject: [PATCH 071/149] force int

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 6439f04..84fc507 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -760,8 +760,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     assert(task_id > -1)
     n_shards = self.dataset_splits[0]["shards"] + self.dataset_splits[1]["shards"]
-    indexes = list(range(task_id*len(sub_cat)/n_shards*p.example_per_shard,
-                  min((task_id+1)*len(sub_cat)/n_shards*p.example_per_shard, len(sub_cat))))
+    indexes = list(range(task_id*len(sub_cat)//n_shards*p.example_per_shard,
+                  min((task_id+1)*len(sub_cat)//n_shards*p.example_per_shard, len(sub_cat))))
     sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''

From 2a19c361069673735c9fd7400903a0b069813eb4 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 8 Mar 2021 14:40:05 +0000
Subject: [PATCH 072/149] Corrected mistake

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 84fc507..c9fd6e3 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -760,8 +760,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     assert(task_id > -1)
     n_shards = self.dataset_splits[0]["shards"] + self.dataset_splits[1]["shards"]
-    indexes = list(range(task_id*len(sub_cat)//n_shards*p.example_per_shard,
-                  min((task_id+1)*len(sub_cat)//n_shards*p.example_per_shard, len(sub_cat))))
+    indexes = list(range(task_id*len(sub_cat)//n_shards,
+                  min((task_id+1)*len(sub_cat)//n_shards, len(sub_cat))))
     sub_cat = sub_cat[indexes]
 
     ''' Loop on all the galaxies of the field '''

From bda47e5f5d01aeed5d0a15b07b75815395fc67fc Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 9 Mar 2021 11:39:28 +0000
Subject: [PATCH 073/149] Change noise treatment

---
 galaxy2galaxy/data_generators/candels.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c9fd6e3..acd3eed 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -787,7 +787,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     # if np.max(fits.open(tmp_file)[0].data) == 0.:
                     #     sigmas[res][n_filter] = 10
                     im_import = fits.open(tmp_file)[0].data
-                    im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5,noise_level=p.sigmas[res][n_filter])
+                    im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
 
                     # except Exception:
                     #     print('Galaxy not seen in every filter')
@@ -1063,7 +1063,7 @@ def mask_out_pixels(img, segmap, segval,
     # Create binary masks of all segmented sources
     sources = binary_dilation(segmap, iterations=n_iter)
 
-    background_mask = np.logical_not(sources)
+    background_mask = np.logical_and(np.logical_not(sources),np.array(img,dtype=bool))
     # Create binary mask of the central galaxy
     central_source = binary_dilation(np.where(segmap == segval, 1, 0),
                                      iterations=n_iter)
@@ -1117,21 +1117,19 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
     blended_pixels = np.logical_and(np.not_equal(sex_seg,0),np.not_equal(sex_seg,middle))*central
     blend_flux = np.sum(img[np.nonzero(blended_pixels)])
     if np.any(blended_pixels):
-        loc = np.argwhere(blended_pixels==True)[0]
-        blended_galaxy = sex_seg[loc[0],loc[1]]
-        blended_galaxy_flux = np.sum(img[np.where(sex_seg==blended_galaxy)])
-    else:
-        blended_galaxy_flux = np.inf
-    if blend_flux/blended_galaxy_flux > 0.5:
-        raise ValueError('Blending suspected')
-
+        loc = np.argwhere(blended_pixels==True)
+        blended_galaxies = np.unique(sex_seg[loc])
+        for blended_galaxy in blended_galaxies:
+            blended_galaxy_flux = np.sum(img[np.where(sex_seg==blended_galaxy)])
+            if blend_flux/blended_galaxy_flux > 0.5:
+              raise ValueError('Blending suspected')
 
     '''Rotate'''
     PA = cat[find_central(cat)[0]][4]
     img_rotate = rotate(cleaned, PA, reshape=False)
     
     '''Add noise'''
-    background_mask = np.where(sex_seg == 0, 1, 0)
+    background_mask = np.logical_and(np.logical_not(sex_seg==0),np.array(img,dtype=bool))
     if noise_level == None:
         background_std = np.std(img * background_mask)
     else:

From 99e9b08c95bbe4ee501d53b9dcaa8a673cfcc2fb Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 10 Mar 2021 12:26:37 +0000
Subject: [PATCH 074/149] Fix wrong shape with single band

---
 galaxy2galaxy/models/autoencoders_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index a0ff56a..91ed0d1 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -39,8 +39,8 @@ def loglikelihood_fn(xin, yin, features, hparams):
     # Compute FFT normalization factor
     x = tf.transpose(xin,[0,3,1,2])
     y = tf.transpose(yin,[0,3,1,2])
-    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(features['ps'])),0.) / size**2 * (2*np.pi)**2
-    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(features['ps'])),0.) / size**2 * (2*np.pi)**2
+    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.spectral.rfft2d(x).get_shape()))),0.) / size**2 * (2*np.pi)**2
+    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.spectral.rfft2d(y).get_shape()))),0.) / size**2 * (2*np.pi)**2
 
     pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2, -3]) #/ size**2
     return -pz

From 5d4844c2df24518702b70cc99b2856f4d6c61c8b Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 10 Mar 2021 12:33:50 +0000
Subject: [PATCH 075/149] Fix wrong shape with single band

---
 galaxy2galaxy/models/autoencoders_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 91ed0d1..daaa555 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -39,8 +39,8 @@ def loglikelihood_fn(xin, yin, features, hparams):
     # Compute FFT normalization factor
     x = tf.transpose(xin,[0,3,1,2])
     y = tf.transpose(yin,[0,3,1,2])
-    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.spectral.rfft2d(x).get_shape()))),0.) / size**2 * (2*np.pi)**2
-    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.spectral.rfft2d(y).get_shape()))),0.) / size**2 * (2*np.pi)**2
+    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]).get_shape().as_list()))),0.) / size**2 * (2*np.pi)**2
+    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]).get_shape().as_list()))),0.) / size**2 * (2*np.pi)**2
 
     pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2, -3]) #/ size**2
     return -pz

From 930d2e6971e589517e1378ccf83193741f5844a2 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 10 Mar 2021 12:45:02 +0000
Subject: [PATCH 076/149] Fix wrong shape with single band

---
 galaxy2galaxy/models/autoencoders_utils.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index daaa555..b9e5b91 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -39,8 +39,9 @@ def loglikelihood_fn(xin, yin, features, hparams):
     # Compute FFT normalization factor
     x = tf.transpose(xin,[0,3,1,2])
     y = tf.transpose(yin,[0,3,1,2])
-    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]).get_shape().as_list()))),0.) / size**2 * (2*np.pi)**2
-    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(tf.reshape(features['ps'],tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]).get_shape().as_list()))),0.) / size**2 * (2*np.pi)**2
+    ps = tf.reshape(features['ps'],tf.shape(tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1])))
+    x = tf.transpose(tf.spectral.rfft2d(x),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(ps)),0.) / size**2 * (2*np.pi)**2
+    y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(ps)),0.) / size**2 * (2*np.pi)**2
 
     pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2, -3]) #/ size**2
     return -pz

From 2decc43466e7c468b2c3af7120465563f95b02df Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 10 Mar 2021 14:50:27 +0000
Subject: [PATCH 077/149] Fix export error

---
 galaxy2galaxy/models/autoencoders_utils.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index b9e5b91..ef23e61 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -125,7 +125,10 @@ def make_model_spec():
                                  activation=output_activation)
       hub.add_signature(inputs=input_layer, outputs=reconstr)
       hub.attach_message("stamp_size", tf.train.Int64List(value=[hparams.problem_hparams.img_len]))
-      hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale[res] for res in hparams.problem_hparams.resolutions]))
+      if hparams.problem_hparams.resolutions:
+        hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale[res] for res in hparams.problem_hparams.resolutions]))
+      else:
+        hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale]))
     spec = hub.create_module_spec(make_model_spec, drop_collections=['checkpoints'])
     decoder = hub.Module(spec, name="decoder_module")
     hub.register_module_for_export(decoder, "decoder")

From 4dca50cf1404e89adeae8f40b6ab5a3468d90eca Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 10 Mar 2021 16:07:12 +0000
Subject: [PATCH 078/149] Fix export error

---
 galaxy2galaxy/models/autoencoders_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index ef23e61..1ec2600 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -125,9 +125,9 @@ def make_model_spec():
                                  activation=output_activation)
       hub.add_signature(inputs=input_layer, outputs=reconstr)
       hub.attach_message("stamp_size", tf.train.Int64List(value=[hparams.problem_hparams.img_len]))
-      if hparams.problem_hparams.resolutions:
+      try:
         hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale[res] for res in hparams.problem_hparams.resolutions]))
-      else:
+      except AttributeError:
         hub.attach_message("pixel_size", tf.train.FloatList(value=[hparams.problem_hparams.pixel_scale]))
     spec = hub.create_module_spec(make_model_spec, drop_collections=['checkpoints'])
     decoder = hub.Module(spec, name="decoder_module")

From a9335d355a604aa23cbe432e361f4c023db662ba Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 10 Mar 2021 18:34:01 +0000
Subject: [PATCH 079/149] New attrs2img

---
 galaxy2galaxy/data_generators/cosmos.py | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index f419e3c..8e5ffdb 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -282,6 +282,29 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "targets": None}
 
+@registry.register_problem
+class Attrs2imgCosmos64(Img2imgCosmos64):
+  """ Smaller version of the Img2imgCosmos problem, at half the pixel
+  resolution
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 64
+    p.example_per_shard = 1000
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag_auto', 'flux_radius', 'zphot']
+
 @registry.register_problem
 class Attrs2imgCosmos128(Img2imgCosmos128):
   """ Smaller version of the Img2imgCosmos problem, at half the pixel

From 067589ca2e6212fe76e6a628c2f8689478d308ec Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 10:04:09 +0000
Subject: [PATCH 080/149] Added multiband cosmos

---
 galaxy2galaxy/data_generators/cosmos.py       | 185 ++++++++++++++++++
 galaxy2galaxy/data_generators/galsim_utils.py |  43 ++--
 2 files changed, 210 insertions(+), 18 deletions(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 8e5ffdb..6f5ec27 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -414,3 +414,188 @@ def hparams(self, defaults, model_hparams):
                     "targets": None}
     p.attributes = ['mag_auto', 'flux_radius', 'zphot', 'bulge_q', 'bulge_beta' ,
                     'disk_q', 'disk_beta', 'bulge_hlr', 'disk_hlr']
+
+
+
+@registry.register_problem
+class Img2imgCosmosMultiband(galsim_utils.GalsimProblem):
+  """
+  Img2img problem on GalSim's COSMOS 25.2 sample, at native pixel resolution,
+  on 64px postage stamps.
+  """
+
+  @property
+  def dataset_splits(self):
+    """Splits of data to produce and number of output shards for each.
+       Note that each shard will be produced in parallel.
+       We are going to split the GalSim data into shards of 1000 galaxies each,
+       with 80 shards for training, 2 shards for validation.
+    """
+    return [{
+        "split": problem.DatasetSplit.TRAIN,
+        "shards": 80,
+    }, {
+        "split": problem.DatasetSplit.EVAL,
+        "shards": 2,
+    }]
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.03
+    p.img_len = 64
+    p.flux_ratio = [1.0]
+
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "targets": None}
+    p.add_hparam("psf", None)
+    p.add_hparam("rotation", False)
+
+  @property
+  def num_bands(self):
+    """Number of bands."""
+    return 1
+
+  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
+    """
+    Generates and yields poattrs2img_cosmos64_euclid_pixscale_1_wmorpho_2stage stamps obtained with GalSim.
+    """
+    p = self.get_hparams()
+
+    try:
+        # try to use default galsim path to the data
+        catalog = galsim.COSMOSCatalog()
+    except:
+        # If that fails, tries to use the specified tmp_dir
+        catalog = galsim.COSMOSCatalog(dir=tmp_dir+'/COSMOS_25.2_training_sample')
+
+    # Create a list of galaxy indices for this task, remember, there is a task
+    # per shard, each shard is 1000 galaxies.
+    assert(task_id > -1)
+    index = range(task_id*p.example_per_shard,
+                  min((task_id+1)*p.example_per_shard, catalog.getNObjects()))
+
+    # Extracts additional information about the galaxies
+    cat_param = catalog.param_cat[catalog.orig_index]
+    from numpy.lib.recfunctions import append_fields
+    import numpy as np
+
+    bparams = cat_param['bulgefit']
+    sparams = cat_param['sersicfit']
+    # Parameters for a 2 component fit
+    cat_param = append_fields(cat_param, 'bulge_q', bparams[:,11])
+    cat_param = append_fields(cat_param, 'bulge_beta', bparams[:,15])
+    cat_param = append_fields(cat_param, 'disk_q', bparams[:,3])
+    cat_param = append_fields(cat_param, 'disk_beta', bparams[:,7])
+    cat_param = append_fields(cat_param, 'bulge_hlr', cat_param['hlr'][:,1])
+    cat_param = append_fields(cat_param, 'bulge_flux_log10', np.where(cat_param['use_bulgefit'] ==1, np.log10(cat_param['flux'][:,1]), np.zeros(len(cat_param) )))
+    cat_param = append_fields(cat_param, 'disk_hlr', cat_param['hlr'][:,2])
+    cat_param = append_fields(cat_param, 'disk_flux_log10', np.where(cat_param['use_bulgefit'] ==1, np.log10(cat_param['flux'][:,2]), np.log10(cat_param['flux'][:,0])))
+
+    # Parameters for a single component fit
+    cat_param = append_fields(cat_param, 'sersic_flux_log10', np.log10(sparams[:,0]))
+    cat_param = append_fields(cat_param, 'sersic_q', sparams[:,3])
+    cat_param = append_fields(cat_param, 'sersic_hlr', sparams[:,1])
+    cat_param = append_fields(cat_param, 'sersic_n', sparams[:,2])
+    cat_param = append_fields(cat_param, 'sersic_beta', sparams[:,7])
+
+    passed = 0
+    
+    late = 0
+    irr = 0
+    for ind in index:
+      # Draw a galaxy using GalSim, any kind of operation can be done here
+#       if cat_param['Mph3'][ind] not in [1,2,3] :
+#             passed += 1
+#             continue
+
+#       if cat_param['Mph3'][ind] == 2:
+#         if late >= 85 :
+#             if ind % 200 == 0:
+#                 print("done all the late of the id")
+#             continue
+#         late += 1
+      
+#       if cat_param['Mph3'][ind] == 3:
+#         if irr >= 85 :
+#             if ind % 200 == 0:
+#                 print("done all the irr of the id")
+#             continue
+#         irr += 1
+        
+      gal = catalog.makeGalaxy(ind, noise_pad_size=p.img_len * p.pixel_scale)
+
+      # We apply the orginal psf if a different PSF is not requested
+      if ~hasattr(p, "psf") or p.psf is None:
+          psf = gal.original_psf
+      else:
+          psf = p.psf
+
+      # Apply rotation so that the galaxy is at 0 PA
+#       if hasattr(p, "rotation") and p.rotation:
+      rotation_angle = galsim.Angle(-cat_param[ind]['sersic_beta'],
+                                      galsim.radians)
+      gal = gal.rotate(rotation_angle)
+      psf = psf.rotate(rotation_angle)
+
+      # We save the corresponding attributes for this galaxy
+      if hasattr(p, 'attributes'):
+        params = cat_param[ind]
+        attributes = {k: params[k] for k in p.attributes}
+      else:
+        attributes = None
+
+      # Utility function encodes the postage stamp for serialized features
+      yield galsim_utils.draw_and_encode_stamp(gal, psf,
+                                               stamp_size=p.img_len,
+                                               pixel_scale=p.pixel_scale,
+                                               attributes=attributes,
+                                               flux_r=p.flux_ratio)
+
+  def preprocess_example(self, example, unused_mode, unused_hparams):
+    """ Preprocess the examples, can be used for further augmentation or
+    image standardization.
+    """
+    p = self.get_hparams()
+    image = example["inputs"]
+
+    # Clip to 1 the values of the image
+    # image = tf.clip_by_value(image, -1, 1)
+
+    # Aggregate the conditions
+    if hasattr(p, 'attributes'):
+      example['attributes'] = tf.stack([example[k] for k in p.attributes])
+
+    example["inputs"] = image
+    example["targets"] = image
+    return example
+
+
+@registry.register_problem
+class Attrs2imgCosmosMultiband64(Img2imgCosmos64):
+  """ Smaller version of the Img2imgCosmos problem, at half the pixel
+  resolution
+  """
+  @property
+  def num_bands(self):
+    """Number of bands."""
+    return 2
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = 0.1
+    p.img_len = 64
+    p.example_per_shard = 1000
+    p.flux_ratio = [1,0.5]
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag_auto', 'flux_radius', 'zphot']
\ No newline at end of file
diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index a1577bc..c782dbc 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -142,7 +142,7 @@ def _float_feature(value):
 def _bytes_feature(value):
     return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
 
-def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, attributes=None):
+def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, flux_r = [1.0], attributes=None):
     """
     Draws the galaxy, psf and noise power spectrum on a postage stamp and
     encodes it to be exported in a TFRecord.
@@ -150,27 +150,34 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, attributes=None):
     # Apply the PSF
     gal = galsim.Convolve(gal, psf)
 
+    im_multi = np.zeros((stamp_size,stamp_size,num_bands))
+    psf_multi = np.zeros((stamp_size,stamp_size,num_bands))
     # Draw the Fourier domain image of the galaxy
-    imC = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
-                         np.pi / (pixel_scale * stamp_size))
+    for i in self.num_bands:
+        imC = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
+                             np.pi / (pixel_scale * stamp_size))
 
-    imCp = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
-                          np.pi / (pixel_scale * stamp_size))
+        imCp = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
+                              np.pi / (pixel_scale * stamp_size))
+        gal_temp = gal/flux_r[i]
+        gal_temp.drawKImage(image=imC)
+        psf.drawKImage(image=imCp)
 
-    gal.drawKImage(image=imC)
-    psf.drawKImage(image=imCp)
+        # Keep track of the pixels with 0 value
+        mask = ~(np.fft.fftshift(imC.array)[:, :(stamp_size) // 2 + 1] == 0)
 
-    # Keep track of the pixels with 0 value
-    mask = ~(np.fft.fftshift(imC.array)[:, :(stamp_size) // 2 + 1] == 0)
+        # Inverse Fourier transform of the image
+        # TODO: figure out why we need 2 fftshifts....
+        im = np.fft.fftshift(np.fft.ifft2(
+            np.fft.fftshift(imC.array))).real.astype('float32')
 
-    # Inverse Fourier transform of the image
-    # TODO: figure out why we need 2 fftshifts....
-    im = np.fft.fftshift(np.fft.ifft2(
-        np.fft.fftshift(imC.array))).real.astype('float32')
+        # Transform the psf array into proper format for Theano
+        im_psf = np.fft.fftshift(np.fft.ifft2(
+                np.fft.fftshift(imCp.array))).real.astype('float32')
+
+        im_multi[:,:,i] = im
+        psf_multi[:,:,i] = im_psf
 
-    # Transform the psf array into proper format for Theano
-    im_psf = np.fft.fftshift(np.fft.ifft2(
-            np.fft.fftshift(imCp.array))).real.astype('float32')
 
     # Compute noise power spectrum
     ps = gal.noise._get_update_rootps((stamp_size, stamp_size),
@@ -193,9 +200,9 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, attributes=None):
     # Apply mask to power spectrum so that it is very large outside maxk
     ps = np.where(mask, np.log(ps**2), 10).astype('float32')
 
-    serialized_output = {"image/encoded": [im.tostring()],
+    serialized_output = {"image/encoded": [im_multi.tostring()],
             "image/format": ["raw"],
-            "psf/encoded": [im_psf.tostring()],
+            "psf/encoded": [psf_multi.tostring()],
             "psf/format": ["raw"],
             "ps/encoded": [ps.tostring()],
             "ps/format": ["raw"]}

From 2831ea19eb3021416557032ec7332b3ac6acd190 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 10:33:13 +0000
Subject: [PATCH 081/149] Fix typo

---
 galaxy2galaxy/data_generators/galsim_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index c782dbc..5ff7339 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -150,8 +150,8 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, flux_r = [1.0], att
     # Apply the PSF
     gal = galsim.Convolve(gal, psf)
 
-    im_multi = np.zeros((stamp_size,stamp_size,num_bands))
-    psf_multi = np.zeros((stamp_size,stamp_size,num_bands))
+    im_multi = np.zeros((stamp_size,stamp_size,self.num_bands))
+    psf_multi = np.zeros((stamp_size,stamp_size,self.num_bands))
     # Draw the Fourier domain image of the galaxy
     for i in self.num_bands:
         imC = galsim.ImageCF(stamp_size, stamp_size, scale=2. *

From 1e5c7dd0f44e2aafa258bc834ad530bb72bbb1b6 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 12:40:04 +0000
Subject: [PATCH 082/149] Fix typo

---
 galaxy2galaxy/data_generators/galsim_utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 5ff7339..f37094b 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -142,7 +142,7 @@ def _float_feature(value):
 def _bytes_feature(value):
     return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))
 
-def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, flux_r = [1.0], attributes=None):
+def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux_r = [1.0], attributes=None):
     """
     Draws the galaxy, psf and noise power spectrum on a postage stamp and
     encodes it to be exported in a TFRecord.
@@ -150,10 +150,10 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, flux_r = [1.0], att
     # Apply the PSF
     gal = galsim.Convolve(gal, psf)
 
-    im_multi = np.zeros((stamp_size,stamp_size,self.num_bands))
-    psf_multi = np.zeros((stamp_size,stamp_size,self.num_bands))
+    im_multi = np.zeros((stamp_size,stamp_size,num_bands))
+    psf_multi = np.zeros((stamp_size,stamp_size,num_bands))
     # Draw the Fourier domain image of the galaxy
-    for i in self.num_bands:
+    for i in num_bands:
         imC = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
                              np.pi / (pixel_scale * stamp_size))
 

From f0be20fd7097c5a5f44c272cf422481c7af7be45 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 13:10:35 +0000
Subject: [PATCH 083/149] Fix typo

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index f37094b..a21e82d 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -153,7 +153,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
     im_multi = np.zeros((stamp_size,stamp_size,num_bands))
     psf_multi = np.zeros((stamp_size,stamp_size,num_bands))
     # Draw the Fourier domain image of the galaxy
-    for i in num_bands:
+    for i in range(num_bands):
         imC = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
                              np.pi / (pixel_scale * stamp_size))
 

From 4081c81861d8b781f9e47d230d6b3c75d3efcc9f Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 14:31:15 +0000
Subject: [PATCH 084/149] Test fix error

---
 galaxy2galaxy/data_generators/galsim_utils.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index a21e82d..5c3dc0d 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -92,21 +92,21 @@ def example_reading_spec(self):
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
-                channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "psf": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="psf/encoded",
                 format_key="psf/format",
-                channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "ps": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="ps/encoded",
                 format_key="ps/format",
-                channels=self.num_bands,
+                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len // 2 + 1],
                 dtype=tf.float32),
     }

From d798e87676e0a70ba694b9547b3ec280d4d10e36 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 15:54:21 +0000
Subject: [PATCH 085/149] Test fix error

---
 galaxy2galaxy/data_generators/galsim_utils.py | 43 ++++++++++---------
 1 file changed, 22 insertions(+), 21 deletions(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 5c3dc0d..546e410 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -152,6 +152,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
 
     im_multi = np.zeros((stamp_size,stamp_size,num_bands))
     psf_multi = np.zeros((stamp_size,stamp_size,num_bands))
+    ps_multi = np.zeros((stamp_size,stamp_size//2+1,num_bands))
     # Draw the Fourier domain image of the galaxy
     for i in range(num_bands):
         imC = galsim.ImageCF(stamp_size, stamp_size, scale=2. *
@@ -178,27 +179,27 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         im_multi[:,:,i] = im
         psf_multi[:,:,i] = im_psf
 
-
-    # Compute noise power spectrum
-    ps = gal.noise._get_update_rootps((stamp_size, stamp_size),
-                                      wcs=galsim.PixelScale(pixel_scale))
-
-    # The following comes from correlatednoise.py
-    rt2 = np.sqrt(2.)
-    shape = (stamp_size, stamp_size)
-    ps[0, 0] = rt2 * ps[0, 0]
-    # Then make the changes necessary for even sized arrays
-    if shape[1] % 2 == 0:  # x dimension even
-        ps[0, shape[1] // 2] = rt2 * ps[0, shape[1] // 2]
-    if shape[0] % 2 == 0:  # y dimension even
-        ps[shape[0] // 2, 0] = rt2 * ps[shape[0] // 2, 0]
-        # Both dimensions even
-        if shape[1] % 2 == 0:
-            ps[shape[0] // 2, shape[1] // 2] = rt2 * \
-                ps[shape[0] // 2, shape[1] // 2]
-
-    # Apply mask to power spectrum so that it is very large outside maxk
-    ps = np.where(mask, np.log(ps**2), 10).astype('float32')
+        # Compute noise power spectrum
+        ps = gal_temp.noise._get_update_rootps((stamp_size, stamp_size),
+                                          wcs=galsim.PixelScale(pixel_scale))
+
+        # The following comes from correlatednoise.py
+        rt2 = np.sqrt(2.)
+        shape = (stamp_size, stamp_size)
+        ps[0, 0] = rt2 * ps[0, 0]
+        # Then make the changes necessary for even sized arrays
+        if shape[1] % 2 == 0:  # x dimension even
+            ps[0, shape[1] // 2] = rt2 * ps[0, shape[1] // 2]
+        if shape[0] % 2 == 0:  # y dimension even
+            ps[shape[0] // 2, 0] = rt2 * ps[shape[0] // 2, 0]
+            # Both dimensions even
+            if shape[1] % 2 == 0:
+                ps[shape[0] // 2, shape[1] // 2] = rt2 * \
+                    ps[shape[0] // 2, shape[1] // 2]
+
+        # Apply mask to power spectrum so that it is very large outside maxk
+        ps = np.where(mask, np.log(ps**2), 10).astype('float32')
+        ps_multi[:,i] = ps 
 
     serialized_output = {"image/encoded": [im_multi.tostring()],
             "image/format": ["raw"],

From 63731f2c17f518d01f21590860bf2f19f6100c52 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 16:10:29 +0000
Subject: [PATCH 086/149] Test fix error

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 546e410..118e2ba 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -199,7 +199,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
 
         # Apply mask to power spectrum so that it is very large outside maxk
         ps = np.where(mask, np.log(ps**2), 10).astype('float32')
-        ps_multi[:,i] = ps 
+        ps_multi[:,:,i] = ps 
 
     serialized_output = {"image/encoded": [im_multi.tostring()],
             "image/format": ["raw"],

From 34483ffff185436517ce37fdc4fa5b840dd5a78d Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 19:07:44 +0000
Subject: [PATCH 087/149] Test fix error

---
 galaxy2galaxy/data_generators/cosmos.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 6f5ec27..a0ad203 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -551,7 +551,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                                                stamp_size=p.img_len,
                                                pixel_scale=p.pixel_scale,
                                                attributes=attributes,
-                                               flux_r=p.flux_ratio)
+                                               flux_r=p.flux_ratio,
+                                               num_bands=self.num_bands)
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
     """ Preprocess the examples, can be used for further augmentation or

From 020777015d02e2e6e8c938c3f2ac9398a3a91ce0 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 21:56:43 +0000
Subject: [PATCH 088/149] Test fix error

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 118e2ba..be16bd4 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -205,7 +205,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
             "image/format": ["raw"],
             "psf/encoded": [psf_multi.tostring()],
             "psf/format": ["raw"],
-            "ps/encoded": [ps.tostring()],
+            "ps/encoded": [ps_multi.tostring()],
             "ps/format": ["raw"]}
 
     # Adding the parameters provided

From 6103013818f1ea814b8decba7f73253d8010fdb7 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Thu, 11 Mar 2021 23:27:05 +0000
Subject: [PATCH 089/149] Test fix error

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index be16bd4..07c0fdd 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -107,7 +107,7 @@ def example_reading_spec(self):
                 image_key="ps/encoded",
                 format_key="ps/format",
                 # channels=self.num_bands,
-                shape=[p.img_len, p.img_len // 2 + 1],
+                shape=[p.img_len, p.img_len // 2 + 1,self.num_bands],
                 dtype=tf.float32),
     }
 

From 321394dc4db67994fb90fc2e0733817c37499119 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Fri, 12 Mar 2021 13:01:09 +0000
Subject: [PATCH 090/149] Added some prints

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 ++
 galaxy2galaxy/models/autoencoders_utils.py    | 1 +
 2 files changed, 3 insertions(+)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 07c0fdd..1dc0670 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -201,6 +201,8 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         ps = np.where(mask, np.log(ps**2), 10).astype('float32')
         ps_multi[:,:,i] = ps 
 
+    print(ps_multi)
+
     serialized_output = {"image/encoded": [im_multi.tostring()],
             "image/format": ["raw"],
             "psf/encoded": [psf_multi.tostring()],
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 1ec2600..0afa408 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -44,6 +44,7 @@ def loglikelihood_fn(xin, yin, features, hparams):
     y = tf.transpose(tf.spectral.rfft2d(y),[0,2,3,1]) / tf.complex(tf.sqrt(tf.exp(ps)),0.) / size**2 * (2*np.pi)**2
 
     pz = 0.5 * tf.reduce_sum(tf.abs(x - y)**2, axis=[-1, -2, -3]) #/ size**2
+    tf.print(-pz)
     return -pz
   elif hparams.likelihood_type == 'Pixel':
     # TODO: include per example noise std

From b3fcc7c90f04356d6eef0fdbc8dfa8d7301a3d9b Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Sun, 14 Mar 2021 13:30:47 +0000
Subject: [PATCH 091/149] Changed print

---
 galaxy2galaxy/data_generators/galsim_utils.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 1dc0670..dadbbdf 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -201,7 +201,8 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         ps = np.where(mask, np.log(ps**2), 10).astype('float32')
         ps_multi[:,:,i] = ps 
 
-    print(ps_multi)
+    print(ps_multi.shape)
+    print(num_bands)
 
     serialized_output = {"image/encoded": [im_multi.tostring()],
             "image/format": ["raw"],

From 44b43d25ad4535781f824f8589c382a81f15378e Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 15 Mar 2021 17:10:24 +0000
Subject: [PATCH 092/149] Squeezed to see if it fixes error

---
 galaxy2galaxy/data_generators/galsim_utils.py | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index dadbbdf..7cb94ee 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -92,21 +92,21 @@ def example_reading_spec(self):
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
-                # channels=self.num_bands,
+                channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "psf": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="psf/encoded",
                 format_key="psf/format",
-                # channels=self.num_bands,
+                channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "ps": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="ps/encoded",
                 format_key="ps/format",
-                # channels=self.num_bands,
+                channels=self.num_bands,
                 shape=[p.img_len, p.img_len // 2 + 1,self.num_bands],
                 dtype=tf.float32),
     }
@@ -201,8 +201,10 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         ps = np.where(mask, np.log(ps**2), 10).astype('float32')
         ps_multi[:,:,i] = ps 
 
-    print(ps_multi.shape)
-    print(num_bands)
+    if num_bands == 1:
+        im_multi = np.squeeze(im_multi)
+        psf_multi = np.squeeze(psf_multi)
+        ps_multi = np.squeeze(ps_multi)
 
     serialized_output = {"image/encoded": [im_multi.tostring()],
             "image/format": ["raw"],

From 4a06be6a1794478999364b20c678946e6bef7cf5 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 15 Mar 2021 18:25:51 +0000
Subject: [PATCH 093/149] Fixed some parameters

---
 galaxy2galaxy/data_generators/cosmos.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index a0ad203..1ed31e7 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -574,14 +574,14 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
 
 @registry.register_problem
-class Attrs2imgCosmosMultiband64(Img2imgCosmos64):
+class Attrs2imgCosmosMultiband64(Img2imgCosmosMultiband64):
   """ Smaller version of the Img2imgCosmos problem, at half the pixel
   resolution
   """
   @property
   def num_bands(self):
     """Number of bands."""
-    return 2
+    return 1
 
   def eval_metrics(self):
     eval_metrics = [ ]
@@ -592,7 +592,7 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = 0.1
     p.img_len = 64
     p.example_per_shard = 1000
-    p.flux_ratio = [1,0.5]
+    p.flux_ratio = [1]
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}

From 3e747126b39724d049af3e7fbfc46b94c0f2e243 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 15 Mar 2021 19:40:41 +0000
Subject: [PATCH 094/149] Typo

---
 galaxy2galaxy/data_generators/cosmos.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 1ed31e7..04153b6 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -574,7 +574,7 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
 
 @registry.register_problem
-class Attrs2imgCosmosMultiband64(Img2imgCosmosMultiband64):
+class Attrs2imgCosmosMultiband64(Img2imgCosmosMultiband):
   """ Smaller version of the Img2imgCosmos problem, at half the pixel
   resolution
   """

From 428f2a75b6297f4776ad8c9a8bc19b61415dad90 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 16 Mar 2021 13:19:08 +0000
Subject: [PATCH 095/149] Fix the damn type bug

---
 galaxy2galaxy/data_generators/cosmos.py       | 4 ++--
 galaxy2galaxy/data_generators/galsim_utils.py | 6 +++---
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 04153b6..316a4e3 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -581,7 +581,7 @@ class Attrs2imgCosmosMultiband64(Img2imgCosmosMultiband):
   @property
   def num_bands(self):
     """Number of bands."""
-    return 1
+    return 2
 
   def eval_metrics(self):
     eval_metrics = [ ]
@@ -592,7 +592,7 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = 0.1
     p.img_len = 64
     p.example_per_shard = 1000
-    p.flux_ratio = [1]
+    p.flux_ratio = [1,0.5]
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 7cb94ee..94f0807 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -206,11 +206,11 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         psf_multi = np.squeeze(psf_multi)
         ps_multi = np.squeeze(ps_multi)
 
-    serialized_output = {"image/encoded": [im_multi.tostring()],
+    serialized_output = {"image/encoded": [im_multi.astype('float32').tostring()],
             "image/format": ["raw"],
-            "psf/encoded": [psf_multi.tostring()],
+            "psf/encoded": [psf_multi.astype('float32').tostring()],
             "psf/format": ["raw"],
-            "ps/encoded": [ps_multi.tostring()],
+            "ps/encoded": [ps_multi.astype('float32').tostring()],
             "ps/format": ["raw"]}
 
     # Adding the parameters provided

From 3e546651258b443c4a27be942559c32d48ce436c Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Mon, 22 Mar 2021 17:04:32 +0000
Subject: [PATCH 096/149] Data augmentation

---
 galaxy2galaxy/data_generators/candels.py      |  1 +
 galaxy2galaxy/data_generators/cosmos.py       | 11 ++++++-----
 galaxy2galaxy/data_generators/galsim_utils.py |  6 +++---
 3 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index acd3eed..373276f 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -515,6 +515,7 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
 
     # Clip to 1 the values of the image
     # image = tf.clip_by_value(image, -1, 1)
+    image = tf.keras.preprocessing.image.random_rotation(image,360,row_axis=0,col_axis=1,channel_axis=2)
 
     # Aggregate the conditions
     if hasattr(p, 'attributes'):
diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 316a4e3..6132b37 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -534,10 +534,10 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
       # Apply rotation so that the galaxy is at 0 PA
 #       if hasattr(p, "rotation") and p.rotation:
-      rotation_angle = galsim.Angle(-cat_param[ind]['sersic_beta'],
-                                      galsim.radians)
-      gal = gal.rotate(rotation_angle)
-      psf = psf.rotate(rotation_angle)
+      # rotation_angle = galsim.Angle(-cat_param[ind]['sersic_beta'],
+      #                                 galsim.radians)
+      # gal = gal.rotate(rotation_angle)
+      # psf = psf.rotate(rotation_angle)
 
       # We save the corresponding attributes for this galaxy
       if hasattr(p, 'attributes'):
@@ -561,6 +561,7 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
     p = self.get_hparams()
     image = example["inputs"]
 
+    image = tf.keras.preprocessing.image.random_rotation(image,360,row_axis=0,col_axis=1,channel_axis=2)
     # Clip to 1 the values of the image
     # image = tf.clip_by_value(image, -1, 1)
 
@@ -592,7 +593,7 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = 0.1
     p.img_len = 64
     p.example_per_shard = 1000
-    p.flux_ratio = [1,0.5]
+    p.flux_ratio = [1,1]
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 94f0807..7063df3 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -92,21 +92,21 @@ def example_reading_spec(self):
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
-                channels=self.num_bands,
+                channels=None,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "psf": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="psf/encoded",
                 format_key="psf/format",
-                channels=self.num_bands,
+                channels=None,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "ps": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="ps/encoded",
                 format_key="ps/format",
-                channels=self.num_bands,
+                channels=None,
                 shape=[p.img_len, p.img_len // 2 + 1,self.num_bands],
                 dtype=tf.float32),
     }

From e28959536f77545087f365c7aca11b58d1f7f0a7 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 23 Mar 2021 12:33:56 +0000
Subject: [PATCH 097/149] Removed useless things in candels and added random
 rotation to cosmos

---
 galaxy2galaxy/data_generators/candels.py      | 623 +-----------------
 galaxy2galaxy/data_generators/cosmos.py       |   2 +-
 galaxy2galaxy/data_generators/galsim_utils.py |  18 +
 3 files changed, 20 insertions(+), 623 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 373276f..c573444 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -48,627 +48,6 @@ def _resize_image(im, size):
   assert cropped.shape[0]==size and cropped.shape[1]==size, f"Wrong size! Still {cropped.shape}"
   return cropped
 
-@registry.register_problem
-class Img2imgCandelsOnefilter(astroimage_utils.AstroImageProblem):
-  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
-  """
-
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each.
-      .
-    """
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 0,
-    }]
-
-  # START: Subclass interface
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.img_len = 128
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "targets": None}
-    p.add_hparam("psf", None)
-    
-  @property
-  def num_bands(self):
-    """Number of bands."""
-    return 1
-
-  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
-    """ 
-    Generator yielding individual postage stamps.
-    """
-    
-    p = self.get_hparams()
-    
-    '''Load the catalogue containing every fields and every filter'''
-    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='egs   ')] = 'EGS'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='UDS   ')] = 'UDS'
-    
-    ''' Load the psfs for each filter and resize'''
-    psf = np.expand_dims(fits.open(tmp_dir + '/psfs/psf_acs_f606w.fits')[0].data, -1)
-    psf = _resize_image(psf, p.img_len)
-    # Step 2: Extract postage stamps, resize them to requested size
-    ''' Loop on the fields'''
-    for n_field, field in enumerate(['GDS', 'GDN', 'EGS', 'UDS', 'COSMOS']):
-        
-        print(f"\n generating{field}\n")
-        n_gal_creat = 0
-        index = 0
-        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-
-        ''' Loop on all the galaxies of the field '''
-        for gal in sub_cat['RB_ID']:
-                print(f"\n generating{gal}\n")
-
-                if gal == index :     # To take care of the redudency inside the cat
-                    continue
-                index = gal
-
-#             try:
-
-                tmp_file = glob.glob(os.path.join(tmp_dir, field, 'acs_f606w')+ '/galaxy_'+str(index)+'_*')[0]
-                im = np.zeros((128, 128, 1))
-                im[:, :, 0] = fits.open(tmp_file)[0].data
-
-                        
-                ''' Resize the image to the wanted size'''
-                im = _resize_image(im, p.img_len)
-                
-                ''' Load the wanted physical parameters of the galaxy '''
-                if hasattr(p, 'attributes'):
-                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
-
-                else:
-                    attributes=None
-                
-                ''' Add a flag corresponding to the field '''
-                field_info = np.asarray(n_field)
-
-                ''' Create the output to match T2T format '''
-                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
-                "image/format": ["raw"],
-                "psf/encoded": [psf.astype('float32').tostring()],
-                "psf/format": ["raw"],
-                "sigma_noise/encoded": [str(4.8e-3)],
-                "sigma_noise/format": ["raw"],
-                "field/encoded": [nb_field.astype('float32').tostring()],
-                "field/format": ["raw"]}
-
-                if attributes is not None:
-                    for k in attributes:
-                        serialized_output['attrs/'+k] = [attributes[k]]
-#                     serialized_output['attrs/field'] = [attributes['field']]
-                
-                ''' Increment the number of galaxy created on the shard '''
-                n_gal_creat += 1
-                
-                if n_gal_creat > p.example_per_shard:
-                    break
-                yield serialized_output
-#             except Exception:
-#                 print('No image corresponding to the Index')
-#                 n+=1
-
-  @property
-  def is_generate_per_split(self):
-    return False
-
-  def example_reading_spec(self):
-    """Define how data is serialized to file and read back.
-
-    Returns:
-      data_fields: A dictionary mapping data names to its feature type.
-      data_items_to_decoders: A dictionary mapping data names to TF Example
-         decoders, to be used when reading back TF examples from disk.
-    """
-    p = self.get_hparams()
-
-    data_fields = {
-        "image/encoded": tf.FixedLenFeature((), tf.string),
-        "image/format": tf.FixedLenFeature((), tf.string),
-
-        "psf/encoded": tf.FixedLenFeature((), tf.string),
-        "psf/format": tf.FixedLenFeature((), tf.string),
-
-        "sigma_noise/encoded": tf.FixedLenFeature((), tf.string),
-        "sigma_noise/format": tf.FixedLenFeature((), tf.string),
-        
-        "field/encoded": tf.FixedLenFeature((), tf.string),
-        "field/format": tf.FixedLenFeature((), tf.string)
-        }
-
-    # Adds additional attributes to be decoded as specified in the configuration
-    if hasattr(p, 'attributes'):
-        for k in p.attributes:
-            data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
-    data_items_to_decoders = {
-        "inputs": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="image/encoded",
-                format_key="image/format",
-                channels=self.num_bands,
-                shape=[p.img_len, p.img_len, self.num_bands],
-                dtype=tf.float32),
-
-        "psf": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="psf/encoded",
-                format_key="psf/format",
-                channels=self.num_bands,
-                shape=[p.img_len, p.img_len, self.num_bands],
-                dtype=tf.float32),
-
-        "sigma": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="sigma_noise/encoded",
-                format_key="sigma_noise/format",
-                channels=self.num_bands,
-                shape=[self.num_bands],
-                dtype=tf.float32),
-        
-        "field": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="field/encoded",
-                format_key="field/format",
-                shape=[1],
-                dtype=tf.float32),
-    }
-    if hasattr(p, 'attributes'):
-        for k in p.attributes:
-            data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
-#         data_items_to_decoders['field'] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/field')
-
-    return data_fields, data_items_to_decoders
-#     END: Subclass interface
-
-@registry.register_problem
-class Img2imgCandelsAll(astroimage_utils.AstroImageProblem):
-  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
-  """
-
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each.
-      .
-    """
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 0,
-    }]
-
-  # START: Subclass interface
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.img_len = 128
-    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1.5e-3]
-    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f814w', 'acs_f850l']
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "targets": None}
-    p.add_hparam("psf", None)
-
-  @property
-  def num_bands(self):
-    """Number of bands."""
-    p = self.get_hparams()
-    return len(p.filters)
-
-  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
-    """ 
-    Generator yielding individual postage stamps.
-    """
-    
-    p = self.get_hparams()
-    
-    '''Load the catalogue containing every fields and every filter'''
-    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='egs   ')] = 'EGS'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='UDS   ')] = 'UDS'
-    
-    ''' Load the psfs for each filter and resize'''
-    cube_psf = np.zeros((167, 167, len(p.filters)))
-    for i, filter in enumerate(p.filters):
-        cube_psf[:, :, i] = fits.open(tmp_dir + '/psfs/psf_' + filter +'.fits')[0].data
-                
-
-    # Step 2: Extract postage stamps, resize them to requested size
-    ''' Loop on the fields'''
-    for n_field, field in enumerate(['GDS', 'GDN', 'EGS', 'UDS', 'COSMOS']):
-        
-        print(f"\n generating{field}, {n_field}\n")
-        n_gal_creat = 0
-        index = 0
-        im = np.zeros((128, 128, len(p.filters)))
-        sigmas = np.zeros(len(p.filters)) + 10
-        im_psf = np.zeros((167, 167, len(p.filters)))
-        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-        
-        if field in ['GDS', 'GDN']:
-            filters_field = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850l']
-            bands = [0, 1, 2, 3, 4, 5, 7]
-        else :
-            filters_field = ['f125w', 'wfc3_f160w', 'acs_f606w', 'acs_f814w']
-            bands = [1, 2, 4, 6]
-
-        sigmas[bands] = p.sigmas[bands]
-        im_psf[bands] = cube_psf[bands]
-        ''' Loop on all the galaxies of the field '''
-        for gal in sub_cat['RB_ID']:
-                if gal == index :     # To take care of the redudency inside the cat
-                    continue
-                index = gal
-
-#             try:
-                ''' Loop on the filters '''
-                for n_filter, filter in enumerate(filters_field):
-                    try :
-                        ''' Open the image corresponding to the index of the current galaxy'''
-                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
-                        im[:, :, bands[n_filter]] = fits.open(tmp_file)[0].data
-                    except Exception:
-                        print('Galaxy not seen in every filter')
-                        continue
-                        
-                ''' Resize the image to the wanted size'''
-                im = _resize_image(im, p.img_len)
-                print(self.num_bands)
-                
-                ''' Load the wanted physical parameters of the galaxy '''
-                if hasattr(p, 'attributes'):
-                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
-
-                else:
-                    attributes=None
-                
-#                 ''' Create the power spectrum '''
-#                 ps = np.zeros((p.img_len, p.img_len//2+1, len(p.filters)))
-#                 for i in range(len(p.filters)):
-#                     ps[:, :, n_filter] = np.random.normal(0, p.sigmas[n_filter], (p.img_len, p.img_len // 2 + 1))
-                
-                ''' Add a flag corresponding to the field '''
-                field_info = np.asarray(n_field)
-
-                ''' Create the output to match T2T format '''
-                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
-                "image/format": ["raw"],
-                "psf/encoded": [im_psf.astype('float32').tostring()],
-                "psf/format": ["raw"],
-                "sigma_noise/encoded": [np.asarray(sigmas).astype('float32').tostring()],
-                "sigma_noise/format": ["raw"],
-                "field/encoded": [field_info.astype('float32').tostring()],
-                "field/format": ["raw"]}
-                
-                if attributes is not None:
-                    for k in attributes:
-                        serialized_output['attrs/'+k] = [attributes[k]]
-#                     serialized_output['attrs/field'] = [attributes['field']]
-                
-                ''' Increment the number of galaxy created on the shard '''
-                n_gal_creat += 1
-                
-                if n_gal_creat > p.example_per_shard:
-                    print('out')
-                    break
-                yield serialized_output
-#             except Exception:
-#                 print('No image corresponding to the Index')
-#                 n+=1
-
-
-@registry.register_problem
-class Img2imgCandelsGoods(astroimage_utils.AstroImageProblem):
-  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
-  """
-
-  @property
-  def dataset_splits(self):
-    """Splits of data to produce and number of output shards for each.
-      .
-    """
-    return [{
-        "split": problem.DatasetSplit.TRAIN,
-        "shards": 1,
-    }, {
-        "split": problem.DatasetSplit.EVAL,
-        "shards": 0,
-    }]
-
-  # START: Subclass interface
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.img_len = 128
-    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1.5e-3]
-    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850l']
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "targets": None}
-    p.add_hparam("psf", None)
-
-  @property
-  def num_bands(self):
-    """Number of bands."""
-    p = self.get_hparams()
-    return len(p.filters)
-
-  def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
-    """ 
-    Generator yielding individual postage stamps.
-    """
-    
-    p = self.get_hparams()
-    
-    '''Load the catalogue containing every fields and every filter'''
-    all_cat = Table.read(os.path.join(tmp_dir, 'CANDELS_morphology_v8_3dhst_galfit_ALLFIELDS.fit'))
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='gdn   ')] = 'GDN'
-    all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='GDS   ')] = 'GDS'
-    
-    ''' Load the psfs for each filter and resize'''
-    cube_psf = np.zeros((167, 167, len(p.filters)))
-    for i, filter in enumerate(p.filters):
-        cube_psf[:, :, i] = fits.open(tmp_dir + '/psfs/psf_' + filter +'.fits')[0].data
-    
-    scaling = p.pixel_scale/0.06
-    im_psf = resize(cube_psf, (np.ceil(128/scaling)+1, np.ceil(128/scaling)+1,len(p.filters)))
-    im_psf = _resize_image(cube_psf, p.img_len)
-    
-    im_tmp = np.zeros((128, 128, len(p.filters)))
-    sigmas = np.asarray(p.sigmas)
-
-    # Step 2: Extract postage stamps, resize them to requested size
-    ''' Loop on the two fields'''
-    for n_field, field in enumerate(['GDS']):#, 'GDN']):
-        
-        print(f"\n generating{field}, {n_field}\n")
-        n_gal_creat = 0
-        index = 0
-        
-        ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
-        sub_cat = all_cat[np.where(all_cat["FIELD_1"]==field)[0]]
-        
-        ''' Loop on all the galaxies of the field '''
-        for gal in sub_cat['RB_ID']:
-                if gal == index :     # To take care of the redudency inside the cat
-                    continue
-                index = gal
-
-#             try:
-                ''' Loop on the filters '''
-                for n_filter, filter in enumerate(p.filters):
-                    try :
-                        ''' Open the image corresponding to the index of the current galaxy'''
-                        tmp_file = glob.glob(os.path.join(tmp_dir, field, filter)+'/galaxy_'+str(index)+'_*')[0]
-                        if np.max(fits.open(tmp_file)[0].data) == 0.:
-                            sigmas[n_filter] = 10
-                        im_tmp[:, :, n_filter] = fits.open(tmp_file)[0].data
-
-                    except Exception:
-                        print('Galaxy not seen in every filter')
-                        continue
-                        
-                ''' Resize the image to the wanted size'''
-                im = resize(im_tmp, (np.ceil(128/scaling)+1, np.ceil(128/scaling)+1, len(p.filters)))
-                
-                im = _resize_image(im, p.img_len)
-                
-                ''' Load the wanted physical parameters of the galaxy '''
-                if hasattr(p, 'attributes'):
-                    attributes = {k: float(all_cat[k][index]) for k in p.attributes}
-
-                else:
-                    attributes=None
-                
-#                 ''' Create the power spectrum '''
-#                 ps = np.zeros((p.img_len, p.img_len//2+1, len(p.filters)))
-#                 for i in range(len(p.filters)):
-#                     ps[:, :, n_filter] = np.random.normal(0, p.sigmas[n_filter], (p.img_len, p.img_len // 2 + 1))
-                
-                ''' Add a flag corresponding to the field '''
-                field_info = np.asarray(n_field)
-
-                ''' Create the output to match T2T format '''
-                serialized_output = {"image/encoded": [im.astype('float32').tostring()],
-                "image/format": ["raw"],
-                "psf/encoded": [im_psf.astype('float32').tostring()],
-                "psf/format": ["raw"],
-                "sigma_noise/encoded": [sigmas.astype('float32').tostring()],
-                "sigma_noise/format": ["raw"],
-                "field/encoded": [field_info.astype('float32').tostring()],
-                "field/format": ["raw"]}
-                
-                if attributes is not None:
-                    for k in attributes:
-                        serialized_output['attrs/'+k] = [attributes[k]]
-                
-                ''' Increment the number of galaxy created on the shard '''
-                n_gal_creat += 1
-                
-                if n_gal_creat > p.example_per_shard:
-                    print('out')
-                    break
-                yield serialized_output
-#             except Exception:
-#                 print('No image corresponding to the Index')
-#                 n+=1
-
-  def preprocess_example(self, example, unused_mode, unused_hparams):
-    """ Preprocess the examples, can be used for further augmentation or
-    image standardization.
-    """
-    p = self.get_hparams()
-    image = example["inputs"]
-
-    # Clip to 1 the values of the image
-    # image = tf.clip_by_value(image, -1, 1)
-    image = tf.keras.preprocessing.image.random_rotation(image,360,row_axis=0,col_axis=1,channel_axis=2)
-
-    # Aggregate the conditions
-    if hasattr(p, 'attributes'):
-      example['attributes'] = tf.stack([example[k] for k in p.attributes])
-
-
-    example["inputs"] = image
-    example["targets"] = image
-    return example
-
-  def example_reading_spec(self):
-    """Define how data is serialized to file and read back.
-
-    Returns:
-      data_fields: A dictionary mapping data names to its feature type.
-      data_items_to_decoders: A dictionary mapping data names to TF Example
-         decoders, to be used when reading back TF examples from disk.
-    """
-    p = self.get_hparams()
-
-    data_fields = {
-        "image/encoded": tf.FixedLenFeature((), tf.string),
-        "image/format": tf.FixedLenFeature((), tf.string),
-
-        "psf/encoded": tf.FixedLenFeature((), tf.string),
-        "psf/format": tf.FixedLenFeature((), tf.string),
-
-        "sigma_noise/encoded": tf.FixedLenFeature((), tf.string),
-        "sigma_noise/format": tf.FixedLenFeature((), tf.string),
-
-        "field/encoded": tf.FixedLenFeature((), tf.string),
-        "field/format": tf.FixedLenFeature((), tf.string),
-    }
-
-    # Adds additional attributes to be decoded as specified in the configuration
-    if hasattr(p, 'attributes'):
-        for k in p.attributes:
-            data_fields['attrs/'+k] = tf.FixedLenFeature([], tf.float32, -1)
-    data_items_to_decoders = {
-        "inputs": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="image/encoded",
-                format_key="image/format",
-#                 channels=self.num_bands,
-                shape=[p.img_len, p.img_len, self.num_bands],
-                dtype=tf.float32),
-
-        "psf": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="psf/encoded",
-                format_key="psf/format",
-#                 channels=self.num_bands,
-                shape=[p.img_len, p.img_len, self.num_bands],
-                dtype=tf.float32),
-
-        "sigma": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="sigma_noise/encoded",
-                format_key="sigma_noise/format",
-#                 channels=self.num_bands,
-                shape=[self.num_bands],
-                dtype=tf.float32),
-        
-        "field": tf.contrib.slim.tfexample_decoder.Image(
-                image_key="field/encoded",
-                format_key="field/format",
-                shape=[1],
-                dtype=tf.float32),
-    }
-    if hasattr(p, 'attributes'):
-        for k in p.attributes:
-            data_items_to_decoders[k] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/'+k)
-#         data_items_to_decoders['field'] = tf.contrib.slim.tfexample_decoder.Tensor('attrs/field')
-
-    return data_fields, data_items_to_decoders
-#     END: Subclass interface
-    
-    
-  @property
-  def is_generate_per_split(self):
-    return False
-  
-@registry.register_problem
-class Attrs2imgCandelsOnefilter64Euclid(Img2imgCandelsOnefilter):
-  """
-  """
-
-  def eval_metrics(self):
-    eval_metrics = [ ]
-    return eval_metrics
-
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.pixel_scale = 0.1
-    p.img_len = 64
-    p.example_per_shard = 10
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "attributes":  modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "attributes": None,
-                    "targets": None}
-    p.attributes = ['mag', 're', 'q']
-    
-
-@registry.register_problem
-class Attrs2imgCandelsGoods64Euclid(Img2imgCandelsGoods):
-  """
-  """
-
-  def eval_metrics(self):
-    eval_metrics = [ ]
-    return eval_metrics
-
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.pixel_scale = 0.1
-    p.img_len = 64
-    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1.5e-3]
-    p.example_per_shard = 5
-    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f850l']
-
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "attributes":  modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "attributes": None,
-                    "targets": None}
-    p.attributes = ['mag', 're', 'q']
-    
-@registry.register_problem
-class Attrs2imgCandelsAll64Euclid(Img2imgCandelsAll):
-  """
-  """
-
-  def eval_metrics(self):
-    eval_metrics = [ ]
-    return eval_metrics
-
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.pixel_scale = 0.1
-    p.img_len = 64
-    p.sigmas = [6.7e-3, 5.4e-3, 4.0e-3, 2.5e-3, 4.8e-3, 3.4e-3, 1e-4, 1.5e-3]
-    p.example_per_shard = 5
-    p.filters = ['f105w', 'f125w', 'wfc3_f160w', 'acs_f435w', 'acs_f606w', 'acs_f775w', 'acs_f814w', 'acs_f850l']
-
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "attributes":  modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "attributes": None,
-                    "targets": None}
-    p.attributes = ['mag', 're', 'q']
-
-
 
 @registry.register_problem
 class Img2imgCandelsGoodsMultires(astroimage_utils.AstroImageProblem):
@@ -1000,7 +379,7 @@ def eval_metrics(self):
 
   def hparams(self, defaults, model_hparams):
     p = defaults
-    p.pixel_scale = {'high' : 0.06, 'low' : 0.06}
+    p.pixel_scale = {'high' : 0.1, 'low' : 0.1}
     p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
     p.img_len = 64
     p.sigmas = {"high" : [0.004094741966557142], "low" : [0.004017507500562]}
diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 6132b37..3f3429c 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -561,7 +561,7 @@ def preprocess_example(self, example, unused_mode, unused_hparams):
     p = self.get_hparams()
     image = example["inputs"]
 
-    image = tf.keras.preprocessing.image.random_rotation(image,360,row_axis=0,col_axis=1,channel_axis=2)
+    image = galsim_utils.tf_rotate(image)
     # Clip to 1 the values of the image
     # image = tf.clip_by_value(image, -1, 1)
 
diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 7063df3..1d9f636 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -279,3 +279,21 @@ def maybe_download_cosmos(target_dir, sample="25.2"):
     if do_remove:
         logger.info("Removing the tarball to save space")
         os.remove(target)
+
+def tf_rotate(input_image, min_angle = -np.pi/2, max_angle = np.pi/2):
+  '''
+  Tensorflow rotates the image randomly
+  : param input_image: image input
+  : param min_angle: minimum rotation angle
+  : param max? Angle: maximum rotation angle
+  : Return: rotated image
+  '''
+  distorted_image = tf.expand_dims(input_image, 0)
+  random_angles = tf.random.uniform(shape=(tf.shape(distorted_image)[0],), minval = min_angle , maxval = max_angle)
+  distorted_image = tf.contrib.image.transform(
+    distorted_image,
+    tf.contrib.image.angles_to_projective_transforms(
+      random_angles, tf.cast(tf.shape(distorted_image)[1], tf.float32), tf.cast(tf.shape(distorted_image)[2], tf.float32)
+    ))
+  rotate_image = tf.squeeze(distorted_image, [0])
+  return rotate_image
\ No newline at end of file

From f76c5a006ccfe567eef60499d70e2dd718eb2d81 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Mar 2021 12:00:03 +0000
Subject: [PATCH 098/149] Correct psf shape

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index e6d2ef8..340e264 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -153,7 +153,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
     gal = galsim.Convolve(gal, psf)
 
     im_multi = np.zeros((stamp_size,stamp_size,num_bands))
-    psf_multi = np.zeros((stamp_size,stamp_size,num_bands))
+    psf_multi = np.zeros((2*stamp_size,2*stamp_size//2+1,num_bands))
     ps_multi = np.zeros((stamp_size,stamp_size//2+1,num_bands))
     # Draw the Fourier domain image of the galaxy
     for i in range(num_bands):

From ca8aac74ed21338f092cf1ea8e3cb71426d57f1d Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Mar 2021 13:04:28 +0000
Subject: [PATCH 099/149] Try to solve issue with psf

---
 galaxy2galaxy/models/autoencoders_utils.py | 29 ++++++++++++----------
 1 file changed, 16 insertions(+), 13 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 438d6af..12adfe7 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -259,19 +259,22 @@ def make_model_spec():
 
   # Apply channel-wise convolution with the PSF if requested
   if hparams.apply_psf and 'psf' in features:
-    rec_padded = tf.pad(reconstr, [[0,0],
-                                  [0, int(hparams.psf_convolution_pad_factor*shape[1])],
-                                  [0, int(hparams.psf_convolution_pad_factor*shape[2])],
-                                  [0,0]])
-    rec_padded_t = tf.transpose(rec_padded, perm=[0, 3,  1, 2])
-    psf_padded = tf.pad(features['psf'], [[0,0],
-                                          [0, int(hparams.psf_convolution_pad_factor*shape[1])],
-                                          [0, int(hparams.psf_convolution_pad_factor*shape[2])],
-                                          [0,0]])
-    psf_padded_t = tf.transpose(psf_padded, perm=[0, 3,  1, 2])
-    reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded_t)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded_t)), tf.complex64))
-    reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
-    reconstr = reconstr[:, :shape[1], :shape[2], :]
+    # rec_padded = tf.pad(reconstr, [[0,0],
+    #                               [0, int(hparams.psf_convolution_pad_factor*shape[1])],
+    #                               [0, int(hparams.psf_convolution_pad_factor*shape[2])],
+    #                               [0,0]])
+    # rec_padded_t = tf.transpose(rec_padded, perm=[0, 3,  1, 2])
+    # psf_padded = tf.pad(features['psf'], [[0,0],
+    #                                       [0, int(hparams.psf_convolution_pad_factor*shape[1])],
+    #                                       [0, int(hparams.psf_convolution_pad_factor*shape[2])],
+    #                                       [0,0]])
+    # psf_padded_t = tf.transpose(psf_padded, perm=[0, 3,  1, 2])
+    # reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded_t)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded_t)), tf.complex64))
+    # reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
+    # reconstr = reconstr[:, :shape[1], :shape[2], :]
+    for i in shape[3]:
+      reconstr[...,i] = convolve(reconstr[...,i], tf.cast(features['psf'][...,i], tf.complex64),
+                          zero_padding_factor=1)
 
   # Losses.
   losses = {

From 914670a101b37666a4baeefa774ff932879d31e0 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Mar 2021 14:59:35 +0000
Subject: [PATCH 100/149] Try to solve issue with psf

---
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 12adfe7..4ea0912 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -272,7 +272,7 @@ def make_model_spec():
     # reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded_t)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded_t)), tf.complex64))
     # reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
     # reconstr = reconstr[:, :shape[1], :shape[2], :]
-    for i in shape[3]:
+    for i in range(shape[3]):
       reconstr[...,i] = convolve(reconstr[...,i], tf.cast(features['psf'][...,i], tf.complex64),
                           zero_padding_factor=1)
 

From d9a7b632a497bfab6e3fd9d5b4030baca0b9e057 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Mar 2021 17:38:28 +0000
Subject: [PATCH 101/149] Try to solve issue with psf

---
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 4ea0912..4759f60 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -273,7 +273,7 @@ def make_model_spec():
     # reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
     # reconstr = reconstr[:, :shape[1], :shape[2], :]
     for i in range(shape[3]):
-      reconstr[...,i] = convolve(reconstr[...,i], tf.cast(features['psf'][...,i], tf.complex64),
+      reconstr = convolve(reconstr, tf.cast(features['psf'], tf.complex64),
                           zero_padding_factor=1)
 
   # Losses.

From 56a0b4694f6d5957f63f6da6f02741b8aef987b0 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Mar 2021 18:51:24 +0000
Subject: [PATCH 102/149] Try to solve issue with psf

---
 galaxy2galaxy/models/autoencoders_utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 4759f60..55e8e21 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -273,8 +273,8 @@ def make_model_spec():
     # reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
     # reconstr = reconstr[:, :shape[1], :shape[2], :]
     for i in range(shape[3]):
-      reconstr = convolve(reconstr, tf.cast(features['psf'], tf.complex64),
-                          zero_padding_factor=1)
+      reconstr[...,i] = tf.squeeze(convolve(tf.expand_dims(reconstr[...,i],-1), tf.cast(features['psf'][...,i], tf.complex64),
+                          zero_padding_factor=1))
 
   # Losses.
   losses = {

From 3c5c10c0169f157fd73e9704cbfd0668594f4f7c Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Wed, 24 Mar 2021 19:46:11 +0000
Subject: [PATCH 103/149] Try to solve issue with psf

---
 galaxy2galaxy/models/autoencoders_utils.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 55e8e21..9a4f975 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -272,9 +272,11 @@ def make_model_spec():
     # reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded_t)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded_t)), tf.complex64))
     # reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
     # reconstr = reconstr[:, :shape[1], :shape[2], :]
+    output_list = []
     for i in range(shape[3]):
-      reconstr[...,i] = tf.squeeze(convolve(tf.expand_dims(reconstr[...,i],-1), tf.cast(features['psf'][...,i], tf.complex64),
-                          zero_padding_factor=1))
+      output_list.append(tf.squeeze(convolve(tf.expand_dims(reconstr[...,i],-1), tf.cast(features['psf'][...,i], tf.complex64),
+                          zero_padding_factor=1)))
+    reconstr = tf.stack(output_list)
 
   # Losses.
   losses = {

From 963a5d189951e6eef20af4cc2e50764bf97190e0 Mon Sep 17 00:00:00 2001
From: Thomas Sainrat <thomas.sainrat@student-cs.fr>
Date: Tue, 30 Mar 2021 12:23:26 +0100
Subject: [PATCH 104/149] Update

---
 galaxy2galaxy/data_generators/cosmos.py    | 2 +-
 galaxy2galaxy/models/autoencoders_utils.py | 9 +++++----
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index d1115fc..4c1379d 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -66,7 +66,7 @@ def num_bands(self):
 
   def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     """
-    Generates and yields poattrs2img_cosmos64_euclid_pixscale_1_wmorpho_2stage stamps obtained with GalSim.
+    Generates and yields postage stamps obtained with GalSim.
     """
     p = self.get_hparams()
 
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 9a4f975..aa86601 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -49,17 +49,17 @@ def loglikelihood_fn(xin, yin, features, hparams):
     return -pz
   elif hparams.likelihood_type == 'Pixel':
     # TODO: include per example noise std
-    pz = 0.5 * tf.reduce_sum(tf.abs(xin[:,:,:,0] - yin[...,0])**2, axis=[-1, -2]) / hparams.noise_rms**2 #/ size**2
+    pz = 0.5 * tf.reduce_sum(tf.abs(xin - yin)**2, axis=[-1, -2, -3]) / hparams.noise_rms**2 #/ size**2
     return -pz
   else:
     raise NotImplementedError
 
 def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
   """Helper for image summaries that are safe on TPU."""
-  if len(image_logits.get_shape()) != 4 or not image_logits.get_shape()[3] in [1,3,4]:
+  if len(image_logits.get_shape()) != 4:
     tf.logging.info("Not generating image summary, maybe not an image.")
     return
-  return tf.summary.image(name, pack_images(image_logits, rows, cols),
+  return tf.summary.image(name, pack_images(tf.expand_dims(image_logits[...,0],-1), rows, cols)
       max_outputs=max_outputs)
 
 def autoencoder_body(self, features):
@@ -276,7 +276,8 @@ def make_model_spec():
     for i in range(shape[3]):
       output_list.append(tf.squeeze(convolve(tf.expand_dims(reconstr[...,i],-1), tf.cast(features['psf'][...,i], tf.complex64),
                           zero_padding_factor=1)))
-    reconstr = tf.stack(output_list)
+    reconstr = tf.stack(output_list,axis=-1)
+    reconstr = tf.reshape(reconstr,shape)
 
   # Losses.
   losses = {

From a92377905979efe151f43d694ee55e08f44a44dc Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 12:29:03 +0100
Subject: [PATCH 105/149] Updated PSF for candels

---
 galaxy2galaxy/data_generators/candels.py   | 19 ++++++++++++++-----
 galaxy2galaxy/models/autoencoders_utils.py |  1 +
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c573444..5f6a883 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -115,13 +115,22 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     all_cat['FIELD_1'][np.where(all_cat['FIELD_1']=='COSMOS   ')] = 'COSMOS'
     
     ''' Load the psfs for each filter and resize'''
-    cube_psf = np.zeros((167, 167, band_num))
+    cube_psf = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, band_num))
+    bounds = _BoundsI(0, p.stamp_size//2, -p.stamp_size//2, p.stamp_size//2-1)
+
     k = 0
     for res in p.resolutions:
-        cube_psf_tmp = np.zeros((167, 167, len(p.filters[res])))
+        cube_psf_tmp = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
-            cube_psf_tmp[:, :, i] = fits.open(data_dir + '/psfs/psf_' + filt +'.fits')[0].data
-        cube_psf_tmp = resize(cube_psf_tmp, (167,167,len(p.filters[res])))
+            psf = galsim.InterpolatedImage(data_dir + '/psfs/psf_' + filt +'.fits')
+            imCp = psf.drawKImage(bounds=bounds,
+                             scale=2.*np.pi/(Nk * pixel_scale / interp_factor),
+                             recenter=False)
+
+            # Transform the psf array into proper format, remove the phase
+            im_psf = np.abs(np.fft.fftshift(imCp.array, axes=0)).astype('float32')
+            cube_psf_tmp[:, :, i] = im_psf
+        cube_psf_tmp = resize(cube_psf_tmp, (2*p.img_len, 2*p.img_len // 2 + 1,len(p.filters[res])))
         cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
 
         k += len(p.filters[res])
@@ -303,7 +312,7 @@ def example_reading_spec(self):
                 image_key="psf/encoded",
                 format_key="psf/format",
                 # channels=self.num_bands,
-                shape=[p.img_len, p.img_len, self.num_bands],
+                shape=[2*p.img_len, 2*p.img_len // 2 + 1, self.num_bands],
                 dtype=tf.float32),
 
         "ps": tf.contrib.slim.tfexample_decoder.Image(
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index aa86601..b9403cd 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -257,6 +257,7 @@ def make_model_spec():
   #tv = tf.reduce_sum(tf.sqrt(im_dx**2 + im_dy**2 + 1e-6), axis=[1,2,3])
   #tv = tf.reduce_mean(tv)
 
+  image_summary("without_psf",tf.reshape(reconstr, labels_shape))
   # Apply channel-wise convolution with the PSF if requested
   if hparams.apply_psf and 'psf' in features:
     # rec_padded = tf.pad(reconstr, [[0,0],

From d6e836e418fce4b00fd96449d6e9fa22026a77a2 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 12:30:23 +0100
Subject: [PATCH 106/149] missing coma

---
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index b9403cd..7c41409 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -59,7 +59,7 @@ def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
   if len(image_logits.get_shape()) != 4:
     tf.logging.info("Not generating image summary, maybe not an image.")
     return
-  return tf.summary.image(name, pack_images(tf.expand_dims(image_logits[...,0],-1), rows, cols)
+  return tf.summary.image(name, pack_images(tf.expand_dims(image_logits[...,0],-1), rows, cols),
       max_outputs=max_outputs)
 
 def autoencoder_body(self, features):

From 6954bc53d85690c71ea370b3847f3f30a282b4a5 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 12:38:01 +0100
Subject: [PATCH 107/149] Fix error

---
 galaxy2galaxy/data_generators/candels.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 5f6a883..c2ea110 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -134,9 +134,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
 
         k += len(p.filters[res])
-    
-    im_psf = _resize_image(cube_psf, p.img_len)
-    
+        
     sigmas = p.sigmas
 
     # Step 2: Extract postage stamps, resize them to requested size

From 5949104bf870c22dea087b0a277d1430a1528ac9 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 12:44:29 +0100
Subject: [PATCH 108/149] Fix psf

---
 galaxy2galaxy/data_generators/candels.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c2ea110..3a8c819 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -132,9 +132,10 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             cube_psf_tmp[:, :, i] = im_psf
         cube_psf_tmp = resize(cube_psf_tmp, (2*p.img_len, 2*p.img_len // 2 + 1,len(p.filters[res])))
         cube_psf[:,:,k:k+len(p.filters[res])] = cube_psf_tmp
-
         k += len(p.filters[res])
-        
+
+    psf = cube_psf
+    
     sigmas = p.sigmas
 
     # Step 2: Extract postage stamps, resize them to requested size
@@ -222,7 +223,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ''' Create the output to match T2T format '''
             serialized_output = {"image/encoded": [im.astype('float32').tostring()],
             "image/format": ["raw"],
-            "psf/encoded": [im_psf.astype('float32').tostring()],
+            "psf/encoded": [psf.astype('float32').tostring()],
             "psf/format": ["raw"],
             "ps/encoded": [ps.astype('float32').tostring()],
             "ps/format": ["raw"],

From f7bae0127de4ec5fa6a9fe8c8684ad5f5fc0f88c Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 13:02:00 +0100
Subject: [PATCH 109/149] Fix psf

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 3a8c819..8b62b10 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -116,7 +116,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     
     ''' Load the psfs for each filter and resize'''
     cube_psf = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, band_num))
-    bounds = _BoundsI(0, p.stamp_size//2, -p.stamp_size//2, p.stamp_size//2-1)
+    bounds = galsim.bounds._BoundsI(0, p.stamp_size//2, -p.stamp_size//2, p.stamp_size//2-1)
 
     k = 0
     for res in p.resolutions:

From 8cb2139b64498ff49851f05e22de46ac158550b4 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 13:03:09 +0100
Subject: [PATCH 110/149] Import galsim

---
 galaxy2galaxy/data_generators/candels.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 8b62b10..256f1bd 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -26,6 +26,7 @@
 import glob
 import os
 import sys
+import galsim
 from skimage.transform import resize,rescale
 from scipy.ndimage import binary_dilation  # type: ignore
 from astropy.table import Table

From 0b67d35018ed73b1ca8bbdc4514c38838091c73e Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 13:08:47 +0100
Subject: [PATCH 111/149] Fix typo

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 256f1bd..69017f3 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -117,7 +117,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     
     ''' Load the psfs for each filter and resize'''
     cube_psf = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, band_num))
-    bounds = galsim.bounds._BoundsI(0, p.stamp_size//2, -p.stamp_size//2, p.stamp_size//2-1)
+    bounds = galsim.bounds._BoundsI(0, p.img_len//2, -p.img_len//2, p.img_len//2-1)
 
     k = 0
     for res in p.resolutions:

From 89cf2d44f699b63705f23b1c881b9b9fb4456712 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 13:15:56 +0100
Subject: [PATCH 112/149]  Fix psf

---
 galaxy2galaxy/data_generators/candels.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 69017f3..d8c34d9 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -124,9 +124,13 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         cube_psf_tmp = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
             psf = galsim.InterpolatedImage(data_dir + '/psfs/psf_' + filt +'.fits')
+            interp_factor=2
+            padding_factor=1
+            Nk = p.img_len*interp_factor*padding_factor
+            bounds = _BoundsI(0, Nk//2, -Nk//2, Nk//2-1)
             imCp = psf.drawKImage(bounds=bounds,
-                             scale=2.*np.pi/(Nk * pixel_scale / interp_factor),
-                             recenter=False)
+                                 scale=2.*np.pi/(Nk * p.pixel_scale[res] / interp_factor),
+                                 recenter=False)
 
             # Transform the psf array into proper format, remove the phase
             im_psf = np.abs(np.fft.fftshift(imCp.array, axes=0)).astype('float32')

From 94ca1b047b5c0591a9b180cba585d8e424fcd503 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 3 May 2021 14:38:22 +0100
Subject: [PATCH 113/149] Fix psf

---
 galaxy2galaxy/data_generators/candels.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index d8c34d9..c13b05e 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -117,17 +117,16 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     
     ''' Load the psfs for each filter and resize'''
     cube_psf = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, band_num))
-    bounds = galsim.bounds._BoundsI(0, p.img_len//2, -p.img_len//2, p.img_len//2-1)
-
+    interp_factor=2
+    padding_factor=1
+    Nk = p.img_len*interp_factor*padding_factor
+    bounds = galsim.bounds._BoundsI(0, Nk//2, -Nk//2, Nk//2-1)
     k = 0
     for res in p.resolutions:
         cube_psf_tmp = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
             psf = galsim.InterpolatedImage(data_dir + '/psfs/psf_' + filt +'.fits')
-            interp_factor=2
-            padding_factor=1
-            Nk = p.img_len*interp_factor*padding_factor
-            bounds = _BoundsI(0, Nk//2, -Nk//2, Nk//2-1)
+            
             imCp = psf.drawKImage(bounds=bounds,
                                  scale=2.*np.pi/(Nk * p.pixel_scale[res] / interp_factor),
                                  recenter=False)

From a98ed519cc4d4ebe143f9c137c3859e6564d0ac3 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Wed, 5 May 2021 10:26:57 +0100
Subject: [PATCH 114/149] Added magnitude cut and fixed attributes

---
 galaxy2galaxy/data_generators/candels.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index c13b05e..326e806 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -150,6 +150,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
+    sub_cat = sub_cat[np.where(sub_cat['mag'] <= 25.3)]
     assert(task_id > -1)
     n_shards = self.dataset_splits[0]["shards"] + self.dataset_splits[1]["shards"]
     indexes = list(range(task_id*len(sub_cat)//n_shards,
@@ -200,7 +201,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
             ''' Load the wanted physical parameters of the galaxy '''
             if hasattr(p, 'attributes'):
-                attributes = {k: float(all_cat[k][index]) for k in p.attributes}
+                attributes = {k: float(sub_cat[k][m]) for k in p.attributes}
 
             else:
                 attributes=None
@@ -217,7 +218,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ps = np.transpose(ps,[1,2,0])
 
             ''' Add a flag corresponding to the field '''
-            field_info = np.asarray(1 if all_cat["FIELD_1"][m] == "GDS" else 0)
+            field_info = np.asarray(1 if sub_cat["FIELD_1"][m] == "GDS" else 0)
 
             sigmas_array = []
             for res in p.resolutions:

From 93419cbf1f0cfdb94fd5c1ba1120cf549d28cd55 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Wed, 5 May 2021 15:20:10 +0100
Subject: [PATCH 115/149] Fix psf scale

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 326e806..a9b615d 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -125,7 +125,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     for res in p.resolutions:
         cube_psf_tmp = np.zeros((2*p.img_len, 2*p.img_len // 2 + 1, len(p.filters[res])))
         for i, filt in enumerate(p.filters[res]):
-            psf = galsim.InterpolatedImage(data_dir + '/psfs/psf_' + filt +'.fits')
+            psf = galsim.InterpolatedImage(data_dir + '/psfs/psf_' + filt +'.fits',scale=0.06)
             
             imCp = psf.drawKImage(bounds=bounds,
                                  scale=2.*np.pi/(Nk * p.pixel_scale[res] / interp_factor),

From cc903f7c73cec4c634ca66bfe0b3091a20aea5de Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 6 May 2021 12:44:23 +0100
Subject: [PATCH 116/149] Remove images missing a galaxy (hopefully)

---
 galaxy2galaxy/data_generators/candels.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index a9b615d..67876fa 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -198,6 +198,12 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             
             im = _resize_image(im, p.img_len)
             
+            #Check that there is still a galaxy
+            img_s = im[:,:,0].byteswap().newbyteorder()
+            bkg = sep.Background(img_s)
+            cat_s = sep.extract(img_s-bkg,sigma_sex,err=bkg.globalrms)  
+            if len(cat_s) == 0:
+                raise ValueError('No galaxy detected in the field')
 
             ''' Load the wanted physical parameters of the galaxy '''
             if hasattr(p, 'attributes'):

From 54c5369efeab0b25469af0d4658bc82ef91fb419 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 6 May 2021 14:58:29 +0100
Subject: [PATCH 117/149] Correction

---
 galaxy2galaxy/data_generators/candels.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 67876fa..97c1b61 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -199,7 +199,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             im = _resize_image(im, p.img_len)
             
             #Check that there is still a galaxy
-            img_s = im[:,:,0].byteswap().newbyteorder()
+            img_s = im[:,:,0]
+            img_s = img_s.byteswap().newbyteorder()
             bkg = sep.Background(img_s)
             cat_s = sep.extract(img_s-bkg,sigma_sex,err=bkg.globalrms)  
             if len(cat_s) == 0:

From ec2e1fa030448168086950cae38e05bf9806f3be Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Fri, 7 May 2021 15:23:24 +0100
Subject: [PATCH 118/149] Removed rotation candels

---
 galaxy2galaxy/data_generators/candels.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 97c1b61..9345cc2 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -524,10 +524,11 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
             if blend_flux/blended_galaxy_flux > 0.5:
               raise ValueError('Blending suspected')
 
-    '''Rotate'''
-    PA = cat[find_central(cat)[0]][4]
-    img_rotate = rotate(cleaned, PA, reshape=False)
+    # '''Rotate'''
+    # PA = cat[find_central(cat)[0]][4]
+    # img_rotate = rotate(cleaned, PA, reshape=False)
     
+    img_rotate = cleaned
     '''Add noise'''
     background_mask = np.logical_and(np.logical_not(sex_seg==0),np.array(img,dtype=bool))
     if noise_level == None:

From 36e3e4b47239b5140fbd23a6f3404ffc74f4387c Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Fri, 7 May 2021 15:25:05 +0100
Subject: [PATCH 119/149] Fix sep

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 9345cc2..f511670 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -200,9 +200,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             
             #Check that there is still a galaxy
             img_s = im[:,:,0]
-            img_s = img_s.byteswap().newbyteorder()
+            img_s = img_s = img_s.copy(order='C')
             bkg = sep.Background(img_s)
-            cat_s = sep.extract(img_s-bkg,sigma_sex,err=bkg.globalrms)  
+            cat_s = sep.extract(img_s-bkg,2,err=bkg.globalrms)  
             if len(cat_s) == 0:
                 raise ValueError('No galaxy detected in the field')
 

From 6531333902cbe3c7e8fa462702395a9e5a8cff5b Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 11 May 2021 11:55:16 +0100
Subject: [PATCH 120/149] Changed noise power spectrum

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index f511670..ea77aab 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -218,7 +218,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             noise_im = np.zeros((p.img_len, p.img_len, band_num))
             for res in p.resolutions:
                 for n_filter in range(len(p.filters[res])):
-                    noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
+                    noise_im[:, :, n_filter+k] = np.random.normal(0, bkg.globalrms, (p.img_len, p.img_len))
                 k+=1
             noise_im = np.transpose(noise_im,[2,0,1])
             ps = np.abs(np.fft.rfft2(noise_im))

From 14ae526e4c3a5c8ab3203265695c47f45c60a35a Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 11 May 2021 17:13:32 +0100
Subject: [PATCH 121/149] Fix Nan error

---
 galaxy2galaxy/data_generators/candels.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index ea77aab..606eacc 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -218,7 +218,10 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             noise_im = np.zeros((p.img_len, p.img_len, band_num))
             for res in p.resolutions:
                 for n_filter in range(len(p.filters[res])):
-                    noise_im[:, :, n_filter+k] = np.random.normal(0, bkg.globalrms, (p.img_len, p.img_len))
+                    if bkg.globalrms > 0:
+                        noise_im[:, :, n_filter+k] = np.random.normal(0, bkg.globalrms, (p.img_len, p.img_len))
+                    else:
+                        noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
                 k+=1
             noise_im = np.transpose(noise_im,[2,0,1])
             ps = np.abs(np.fft.rfft2(noise_im))

From b4428551e4abf39b5e913f8d161b6c0252dc2255 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Wed, 12 May 2021 15:45:25 +0100
Subject: [PATCH 122/149] Added new images for tensorboard

---
 galaxy2galaxy/models/autoencoders_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 7c41409..6503a8b 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -300,5 +300,7 @@ def make_model_spec():
 
   image_summary("ae", reconstr)
   image_summary("input", labels)
+  image_summary("input2", features["inputs"])
+  
 
   return logits, losses

From b27ecfa347fb07dd3bd64355d5be0a5134e6c360 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 13 May 2021 10:22:00 +0100
Subject: [PATCH 123/149] All channels to image summary

---
 galaxy2galaxy/data_generators/candels.py   |  2 +-
 galaxy2galaxy/models/autoencoders_utils.py | 12 ++++++------
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 606eacc..aff37d2 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -213,7 +213,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             else:
                 attributes=None
             
-#                 ''' Create the power spectrum '''
+            ''' Create the power spectrum '''
             k = 0
             noise_im = np.zeros((p.img_len, p.img_len, band_num))
             for res in p.resolutions:
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 6503a8b..c6c9844 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -56,12 +56,14 @@ def loglikelihood_fn(xin, yin, features, hparams):
 
 def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
   """Helper for image summaries that are safe on TPU."""
-  if len(image_logits.get_shape()) != 4:
+  shape = image_logits.get_shape()
+  if len(shape) != 4:
     tf.logging.info("Not generating image summary, maybe not an image.")
     return
-  return tf.summary.image(name, pack_images(tf.expand_dims(image_logits[...,0],-1), rows, cols),
+  for i in range(shape[3]):
+      tf.summary.image(name+'i', pack_images(tf.expand_dims(image_logits[...,i],-1), rows, cols),
       max_outputs=max_outputs)
-
+  return 0
 def autoencoder_body(self, features):
   """ Customized body function for autoencoders acting on continuous images.
   This is based on tensor2tensor.models.research.AutoencoderBasic.body
@@ -299,8 +301,6 @@ def make_model_spec():
   logits = tf.reshape(reconstr, labels_shape)
 
   image_summary("ae", reconstr)
-  image_summary("input", labels)
-  image_summary("input2", features["inputs"])
-  
+  image_summary("input", labels)  
 
   return logits, losses

From ff48c72135b3ab1b37d3c79778c9ebc20a1df843 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Fri, 14 May 2021 10:22:55 +0100
Subject: [PATCH 124/149] Normalize

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index aff37d2..552a7e5 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -218,7 +218,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             noise_im = np.zeros((p.img_len, p.img_len, band_num))
             for res in p.resolutions:
                 for n_filter in range(len(p.filters[res])):
-                    if bkg.globalrms > 0:
+                    if False:
                         noise_im[:, :, n_filter+k] = np.random.normal(0, bkg.globalrms, (p.img_len, p.img_len))
                     else:
                         noise_im[:, :, n_filter+k] = np.random.normal(0, p.sigmas[res][n_filter], (p.img_len, p.img_len))
@@ -541,4 +541,4 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
     random_background = np.random.normal(scale=background_std, size=img_rotate.shape)
     rotated = np.where(img_rotate == 0, random_background, img_rotate)
 
-    return rotated
\ No newline at end of file
+    return rotated/np.max(rotated)
\ No newline at end of file

From 54b25910b44a9893992b9e90172622ce21b54750 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 17 May 2021 12:53:49 +0100
Subject: [PATCH 125/149] Rescale to correct flux

---
 galaxy2galaxy/data_generators/candels.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 552a7e5..a1ed9de 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -163,6 +163,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             continue
         index = gal
         print(index)
+        target_flux_main_band = 10**(-0.4*(sub_cat['mag'][m]-p.zeropoint))
 
         try:
             ''' Loop on the filters '''
@@ -180,8 +181,12 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     # if np.max(fits.open(tmp_file)[0].data) == 0.:
                     #     sigmas[res][n_filter] = 10
                     im_import = fits.open(tmp_file)[0].data
-                    im_tmp[:, :, n_filter] = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
+                    cleaned_image = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
 
+                    if res == p.resolutions[0] and n_filter == 0:
+                        flux_ratio = target_flux_main_band/np.sum(cleaned_image)
+
+                    im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
                     # except Exception:
                     #     print('Galaxy not seen in every filter')
                     #     continue
@@ -407,6 +412,7 @@ def hparams(self, defaults, model_hparams):
     p.img_len = 64
     p.sigmas = {"high" : [0.004094741966557142], "low" : [0.004017507500562]}
     p.filters = {"high" : ['acs_f606w'], "low" : ['wfc3_f160w']}
+    p.zeropoint = 26.49
     p.resolutions = ["high","low"]
     p.example_per_shard = 1000
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
@@ -541,4 +547,4 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
     random_background = np.random.normal(scale=background_std, size=img_rotate.shape)
     rotated = np.where(img_rotate == 0, random_background, img_rotate)
 
-    return rotated/np.max(rotated)
\ No newline at end of file
+    return rotated
\ No newline at end of file

From 8452ed4257a80f1354997797ecf8b659a3cd3ec7 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 1 Jun 2021 14:41:03 +0100
Subject: [PATCH 126/149] Added images for flow

---
 galaxy2galaxy/models/autoencoders_utils.py |  2 +-
 galaxy2galaxy/models/latent_flow.py        | 14 ++++++++++++++
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index c6c9844..89519af 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -61,7 +61,7 @@ def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
     tf.logging.info("Not generating image summary, maybe not an image.")
     return
   for i in range(shape[3]):
-      tf.summary.image(name+'i', pack_images(tf.expand_dims(image_logits[...,i],-1), rows, cols),
+      tf.summary.image(name+str(i), pack_images(tf.expand_dims(image_logits[...,i],-1), rows, cols),
       max_outputs=max_outputs)
   return 0
 def autoencoder_body(self, features):
diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index d1d9419..36f66ae 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -24,6 +24,18 @@
 tfb = tfp.bijectors
 tfd = tfp.distributions
 
+def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4, description=None):
+  """Helper for image summaries that are safe on TPU."""
+  shape = image_logits.get_shape()
+  if len(shape) != 4:
+    tf.logging.info("Not generating image summary, maybe not an image.")
+    return
+  for i in range(shape[3]):
+      tf.summary.image(name+str(i), pack_images(tf.expand_dims(image_logits[...,i],-1), rows, cols),
+      max_outputs=max_outputs,
+      description=description)
+  return 0
+
 class LatentFlow(t2t_model.T2TModel):
   """ Base class for latent flows
   This assumes that an already exported tensorflow hub autoencoder is provided
@@ -54,6 +66,8 @@ def body(self, features):
     x = features['inputs']
     cond = {k: features[k] for k in hparamsp.attributes}
 
+    image_summary("input",x,description=str(cond))
+
     # Load the encoder and decoder modules
     encoder = hub.Module(hparams.encoder_module, trainable=False)
 

From e458d000c57274b80a22fadf1dcaa574c763a72f Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 1 Jun 2021 14:54:52 +0100
Subject: [PATCH 127/149] Added images for flow

---
 galaxy2galaxy/models/latent_flow.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index 36f66ae..3bf7c6c 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -17,6 +17,7 @@
 
 from galaxy2galaxy.layers.flows import masked_autoregressive_conditional_template, ConditionalNeuralSpline, conditional_neural_spline_template, autoregressive_conditional_neural_spline_template
 from galaxy2galaxy.layers.tfp_utils import RealNVP, MaskedAutoregressiveFlow
+from galaxy2galaxy.layers.image_utils import pack_images
 
 import tensorflow as tf
 import tensorflow_hub as hub

From d649f7ea593d9d7a8d6fee20728cfe566a73a0f4 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 1 Jun 2021 15:10:29 +0100
Subject: [PATCH 128/149] Added images for flow

---
 galaxy2galaxy/models/latent_flow.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index 3bf7c6c..22180a3 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -25,7 +25,7 @@
 tfb = tfp.bijectors
 tfd = tfp.distributions
 
-def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4, description=None):
+def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4):
   """Helper for image summaries that are safe on TPU."""
   shape = image_logits.get_shape()
   if len(shape) != 4:
@@ -33,8 +33,7 @@ def image_summary(name, image_logits, max_outputs=1, rows=4, cols=4, description
     return
   for i in range(shape[3]):
       tf.summary.image(name+str(i), pack_images(tf.expand_dims(image_logits[...,i],-1), rows, cols),
-      max_outputs=max_outputs,
-      description=description)
+      max_outputs=max_outputs)
   return 0
 
 class LatentFlow(t2t_model.T2TModel):
@@ -67,7 +66,7 @@ def body(self, features):
     x = features['inputs']
     cond = {k: features[k] for k in hparamsp.attributes}
 
-    image_summary("input",x,description=str(cond))
+    image_summary("input",x)
 
     # Load the encoder and decoder modules
     encoder = hub.Module(hparams.encoder_module, trainable=False)

From e622bb58768c338909ebbccfdd9b3f5677354d40 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 3 Jun 2021 15:31:46 +0100
Subject: [PATCH 129/149] Changed normalization and excluded some placeholder
 values in catalog

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index a1ed9de..639d025 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -159,7 +159,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
     ''' Loop on all the galaxies of the field '''
     for m,gal in enumerate(sub_cat['RB_ID']):
-        if gal == index or gal == 15431:     # To take care of the redudency inside the cat
+        if gal == index or gal == 15431 or sub_cat["mag"][m] < 0:     # To take care of the redudency inside the cat
             continue
         index = gal
         print(index)
@@ -184,7 +184,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     cleaned_image = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
 
                     if res == p.resolutions[0] and n_filter == 0:
-                        flux_ratio = target_flux_main_band/np.sum(cleaned_image)
+                        flux_ratio = 1/100000
 
                     im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
                     # except Exception:

From a7e79e4e3716759fac43f8dd85c3e2044811a2b0 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Wed, 9 Jun 2021 12:59:57 +0100
Subject: [PATCH 130/149] Changed a few things

---
 galaxy2galaxy/data_generators/candels.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 639d025..e2bf107 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -184,8 +184,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     cleaned_image = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
 
                     if res == p.resolutions[0] and n_filter == 0:
-                        flux_ratio = 1/100000
+                        flux_ratio = target_flux_main_band/np.sum(cleaned_image) if np.sum(cleaned_image) != 0 else 1
 
+                    print(flux_ratio)
                     im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
                     # except Exception:
                     #     print('Galaxy not seen in every filter')
@@ -195,7 +196,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 new_size = np.ceil(128/scalings[res])+1
                 im_tmp = resize(im_tmp, (new_size, new_size, len(p.filters[res])))
                 ''' Resize the image to the highest resolution to get consistent array sizes'''
-                im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel=True)
+                im_tmp = rescale(im_tmp,p.pixel_scale[res]/target_pixel_scale,multichannel=True,preserve_range=True)
                 im_tmp = _resize_image(im_tmp,target_size)
 
                 im[:,:,k:k+len(p.filters[res])] = im_tmp

From ebc6ab6177a1afdf71afee34d3ae286642a8e210 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Wed, 9 Jun 2021 17:22:16 +0100
Subject: [PATCH 131/149] Added prints

---
 galaxy2galaxy/data_generators/candels.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index e2bf107..5a39ac0 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -169,7 +169,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             ''' Loop on the filters '''
             im = np.zeros((target_size, target_size, band_num))
 
-            k = 0
+            k = 0  
             for res in p.resolutions:
                 im_tmp = np.zeros((128, 128, len(p.filters[res])))
                 for n_filter, filt in enumerate(p.filters[res]):
@@ -188,6 +188,8 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
 
                     print(flux_ratio)
                     im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
+                    print(cleaned_image)
+                    print(np.max(im_tmp[:, :, n_filter]))
                     # except Exception:
                     #     print('Galaxy not seen in every filter')
                     #     continue

From 6d4912eecc6ad322011a5a3d99afabc2034645ff Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 10 Jun 2021 10:30:12 +0100
Subject: [PATCH 132/149] Removing the problematic galaxies (?)

---
 galaxy2galaxy/data_generators/candels.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 5a39ac0..9c81bb2 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -186,10 +186,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     if res == p.resolutions[0] and n_filter == 0:
                         flux_ratio = target_flux_main_band/np.sum(cleaned_image) if np.sum(cleaned_image) != 0 else 1
 
-                    print(flux_ratio)
                     im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
-                    print(cleaned_image)
-                    print(np.max(im_tmp[:, :, n_filter]))
+                    if np.max(cleaned_image) <= 10**(-5):
+                        raise ValueError("Very weak image")
                     # except Exception:
                     #     print('Galaxy not seen in every filter')
                     #     continue

From 91393c7fcb91596fe4f292ed058f84ac5859c640 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 17 Jun 2021 14:24:37 +0100
Subject: [PATCH 133/149] Added redshift as an attribute

---
 galaxy2galaxy/data_generators/candels.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 9c81bb2..7816fbe 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -423,7 +423,7 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "attributes": None,
                     "targets": None}
-    p.attributes = ['mag', 're', 'q']
+    p.attributes = ['mag', 're', 'q','ZPHOT']
 
 
 

From ba151159df726a8cb2f955b0ebd951e735c0ae91 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 17 Jun 2021 14:27:18 +0100
Subject: [PATCH 134/149] Normalization

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 7816fbe..bb1dd5a 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -184,10 +184,10 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     cleaned_image = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
 
                     if res == p.resolutions[0] and n_filter == 0:
-                        flux_ratio = target_flux_main_band/np.sum(cleaned_image) if np.sum(cleaned_image) != 0 else 1
+                        flux_ratio = 1/np.max(cleaned_image) if np.max(cleaned_image) != 0 else 1
 
                     im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
-                    if np.max(cleaned_image) <= 10**(-5):
+                    if np.max(cleaned_image) <= 5*10**(-3):
                         raise ValueError("Very weak image")
                     # except Exception:
                     #     print('Galaxy not seen in every filter')

From 85f2309c1c2ef7bf97311f9af8af7bddcbe6a8af Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 22 Jun 2021 10:30:09 +0100
Subject: [PATCH 135/149] Added problem with blue band

---
 galaxy2galaxy/data_generators/candels.py | 29 ++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index bb1dd5a..afa6f44 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -426,6 +426,35 @@ def hparams(self, defaults, model_hparams):
     p.attributes = ['mag', 're', 'q','ZPHOT']
 
 
+@registry.register_problem
+class Attrs2imgCandelsGoodsEuclid64Blue(Img2imgCandelsGoodsMultires):
+  """
+  """
+
+  def eval_metrics(self):
+    eval_metrics = [ ]
+    return eval_metrics
+
+  def hparams(self, defaults, model_hparams):
+    p = defaults
+    p.pixel_scale = {'high' : 0.1, 'low' : 0.1}
+    p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
+    p.img_len = 64
+    p.sigmas = {"high" : [0.004094741966557142], "low" : [0.004017507500562]}
+    p.filters = {"high" : ['acs_f435w'], "low" : ['wfc3_f160w']}
+    p.zeropoint = 26.49
+    p.resolutions = ["high","low"]
+    p.example_per_shard = 1000
+    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
+                  "attributes":  modalities.ModalityType.IDENTITY,
+                  "targets": modalities.ModalityType.IDENTITY}
+    p.vocab_size = {"inputs": None,
+                    "attributes": None,
+                    "targets": None}
+    p.attributes = ['mag', 're', 'q','ZPHOT']
+
+
+
 
 
 

From 1c69746b665c73df7caffbe4b96d1fa0a40b3aad Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 22 Jun 2021 16:41:49 +0100
Subject: [PATCH 136/149] Normalize COSMOS

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 340e264..ec72e75 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -169,6 +169,8 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         im = gal.drawImage(nx=stamp_size, ny=stamp_size, scale=pixel_scale,
                            method='no_pixel', use_true_center=False).array.astype('float32')
 
+        im = im.array/np.max(im.array)
+
         # Draw the Fourier domain image of the galaxy, using x1 zero padding,
         # and x2 subsampling
         interp_factor=2

From 0310f0c1ae377a76810d25e5098060e7296ffe24 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Thu, 24 Jun 2021 11:12:22 +0100
Subject: [PATCH 137/149] Correct mistake

---
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index ec72e75..3deb059 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -169,7 +169,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         im = gal.drawImage(nx=stamp_size, ny=stamp_size, scale=pixel_scale,
                            method='no_pixel', use_true_center=False).array.astype('float32')
 
-        im = im.array/np.max(im.array)
+        im = im/np.max(im)
 
         # Draw the Fourier domain image of the galaxy, using x1 zero padding,
         # and x2 subsampling

From 3d63331e579f66c80f8b78d011b201403a71568d Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 29 Jun 2021 11:10:30 +0100
Subject: [PATCH 138/149] Added possibility to change conditionning parameters
 for flow

---
 galaxy2galaxy/models/latent_flow.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index 22180a3..7aa0c89 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -61,10 +61,9 @@ def infer(self,
 
   def body(self, features):
     hparams = self.hparams
-    hparamsp = hparams.problem.get_hparams()
-
+    attributes = hparams.attributes
     x = features['inputs']
-    cond = {k: features[k] for k in hparamsp.attributes}
+    cond = {k: features[k] for k in attributes}
 
     image_summary("input",x)
 
@@ -77,7 +76,7 @@ def body(self, features):
     code_shape = [-1, code_shape[1].value, code_shape[2].value, code_shape[3].value]
 
     def get_flow(inputs, is_training=True):
-      y = tf.concat([tf.expand_dims(inputs[k], axis=1) for k in hparamsp.attributes] ,axis=1)
+      y = tf.concat([tf.expand_dims(inputs[k], axis=1) for k in attributes] ,axis=1)
       y = tf.layers.batch_normalization(y, name="y_norm", training=is_training)
       flow = self.normalizing_flow(y, latent_size)
       return flow
@@ -85,7 +84,7 @@ def get_flow(inputs, is_training=True):
     if hparams.mode == tf.estimator.ModeKeys.PREDICT:
       # Export the latent flow alone
       def flow_module_spec():
-        inputs_params = {k: tf.placeholder(tf.float32, shape=[None]) for k in hparamsp.attributes}
+        inputs_params = {k: tf.placeholder(tf.float32, shape=[None]) for k in attributes}
         random_normal = tf.placeholder(tf.float32, shape=[None, latent_size])
         flow = get_flow(inputs_params, is_training=False)
         samples = flow._bijector.forward(random_normal)
@@ -95,7 +94,7 @@ def flow_module_spec():
       flow_spec = hub.create_module_spec(flow_module_spec)
       flow = hub.Module(flow_spec, name='flow_module')
       hub.register_module_for_export(flow, "code_sampler")
-      cond['random_normal'] = tf.random_normal(shape=[tf.shape(cond[hparamsp.attributes[0]])[0] , latent_size])
+      cond['random_normal'] = tf.random_normal(shape=[tf.shape(cond[attributes[0]])[0] , latent_size])
       samples = flow(cond)
       return samples, {'loglikelihood': 0}
 
@@ -218,6 +217,8 @@ def init_once(x, name, trainable=False):
 def latent_flow():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
+  hparamsp = hparams.problem.get_hparams()
+  hparams.attributes = hparamsp.attributes
   hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
@@ -246,6 +247,8 @@ def latent_flow():
 def latent_flow_larger():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
+  hparamsp = hparams.problem.get_hparams()
+  hparams.attributes = hparamsp.attributes
   hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.1
   hparams.learning_rate_warmup_steps = 1000
@@ -273,6 +276,8 @@ def latent_flow_larger():
 def latent_flow_nsf():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
+  hparamsp = hparams.problem.get_hparams()
+  hparams.attributes = hparamsp.attributes
   hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.1
   hparams.learning_rate_warmup_steps = 1000

From 0f8bc4f6e6cee7d11c6a5455e28fc806ccf077bc Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 29 Jun 2021 11:33:04 +0100
Subject: [PATCH 139/149] Fixed attribute param

---
 galaxy2galaxy/models/latent_flow.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index 7aa0c89..2d7002a 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -217,8 +217,6 @@ def init_once(x, name, trainable=False):
 def latent_flow():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
-  hparamsp = hparams.problem.get_hparams()
-  hparams.attributes = hparamsp.attributes
   hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.0002
   hparams.learning_rate_warmup_steps = 500
@@ -240,6 +238,9 @@ def latent_flow():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
+  hparamsp = hparams.problem.get_hparams()
+  hparams.add_hparam("attributes",hparamsp.attributes)
+
   return hparams
 
 
@@ -247,8 +248,6 @@ def latent_flow():
 def latent_flow_larger():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
-  hparamsp = hparams.problem.get_hparams()
-  hparams.attributes = hparamsp.attributes
   hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.1
   hparams.learning_rate_warmup_steps = 1000
@@ -270,14 +269,15 @@ def latent_flow_larger():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
+  hparamsp = hparams.problem.get_hparams()
+  hparams.add_hparam("attributes",hparamsp.attributes)
+
   return hparams
 
 @registry.register_hparams
 def latent_flow_nsf():
   """Basic autoencoder model."""
   hparams = common_hparams.basic_params1()
-  hparamsp = hparams.problem.get_hparams()
-  hparams.attributes = hparamsp.attributes
   hparams.optimizer = "adam"
   hparams.learning_rate_constant = 0.1
   hparams.learning_rate_warmup_steps = 1000
@@ -301,4 +301,7 @@ def latent_flow_nsf():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
+  hparamsp = hparams.problem.get_hparams()
+  hparams.add_hparam("attributes",hparamsp.attributes)
+
   return hparams

From 1606ae2097ea143a73ef3e55f6f9d8428716ad47 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 29 Jun 2021 11:44:26 +0100
Subject: [PATCH 140/149] Fixed attribute param

---
 galaxy2galaxy/models/latent_flow.py | 13 +++++++------
 1 file changed, 7 insertions(+), 6 deletions(-)

diff --git a/galaxy2galaxy/models/latent_flow.py b/galaxy2galaxy/models/latent_flow.py
index 2d7002a..4d27032 100644
--- a/galaxy2galaxy/models/latent_flow.py
+++ b/galaxy2galaxy/models/latent_flow.py
@@ -62,6 +62,10 @@ def infer(self,
   def body(self, features):
     hparams = self.hparams
     attributes = hparams.attributes
+    if len(attributes[0]) == 0:
+      hparamsp = hparams.problem.get_hparams()
+      attributes = hparamsp.attributes
+    
     x = features['inputs']
     cond = {k: features[k] for k in attributes}
 
@@ -238,8 +242,7 @@ def latent_flow():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
-  hparamsp = hparams.problem.get_hparams()
-  hparams.add_hparam("attributes",hparamsp.attributes)
+  hparams.add_hparam("attributes",[""])
 
   return hparams
 
@@ -269,8 +272,7 @@ def latent_flow_larger():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
-  hparamsp = hparams.problem.get_hparams()
-  hparams.add_hparam("attributes",hparamsp.attributes)
+  hparams.add_hparam("attributes",[""])
 
   return hparams
 
@@ -301,7 +303,6 @@ def latent_flow_nsf():
   # hparams related to the PSF
   hparams.add_hparam("encode_psf", True) # Should we use the PSF at the encoder
 
-  hparamsp = hparams.problem.get_hparams()
-  hparams.add_hparam("attributes",hparamsp.attributes)
+  hparams.add_hparam("attributes",[""])
 
   return hparams

From 160e6ac0afc206622387704a080a5f8458149ad4 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Wed, 14 Jul 2021 11:16:29 +0100
Subject: [PATCH 141/149] Fixed changing the flow attributes and added random
 flux for cosmos multiband

---
 galaxy2galaxy/bin/g2g-exporter                | 7 ++++++-
 galaxy2galaxy/data_generators/candels.py      | 2 +-
 galaxy2galaxy/data_generators/cosmos.py       | 8 ++++++--
 galaxy2galaxy/data_generators/galsim_utils.py | 2 +-
 4 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/galaxy2galaxy/bin/g2g-exporter b/galaxy2galaxy/bin/g2g-exporter
index ee92961..53d851f 100755
--- a/galaxy2galaxy/bin/g2g-exporter
+++ b/galaxy2galaxy/bin/g2g-exporter
@@ -102,7 +102,12 @@ def main(_):
 
   estimator = create_estimator(run_config, hparams)
   hparams.img_len = problem.get_hparams().img_len
-  hparams.attributes = problem.get_hparams().attributes
+  try:
+    if len(hparams.attributes) == 0:
+      hparams.attributes = hparams.problem.get_hparams().attributes
+  except:
+    pass
+
 
   # Use tf hub to export any module that has been registered
   exporter = hub.LatestModuleExporter("tf_hub",
diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index afa6f44..868a06f 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -423,7 +423,7 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "attributes": None,
                     "targets": None}
-    p.attributes = ['mag', 're', 'q','ZPHOT']
+    p.attributes = ['mag','re', 'q','ZPHOT','F_IRR','F_SPHEROID','F_DISK']
 
 
 @registry.register_problem
diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 4c1379d..fa02520 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -531,12 +531,15 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
       else:
         attributes = None
 
+      flux_r = [1.0]
+      for i in range(1,self.num_bands):
+        flux_r.append(max(np.random.normal(p.flux_ratio_mean[i],p.flux_ratio_std[i]),0))
       # Utility function encodes the postage stamp for serialized features
       yield galsim_utils.draw_and_encode_stamp(gal, psf,
                                                stamp_size=p.img_len,
                                                pixel_scale=p.pixel_scale,
                                                attributes=attributes,
-                                               flux_r=p.flux_ratio,
+                                               flux_r=flux_r,
                                                num_bands=self.num_bands)
 
   def preprocess_example(self, example, unused_mode, unused_hparams):
@@ -578,7 +581,8 @@ def hparams(self, defaults, model_hparams):
     p.pixel_scale = 0.1
     p.img_len = 64
     p.example_per_shard = 1000
-    p.flux_ratio = [1,1]
+    p.flux_ratio_mean = [1,0.299]
+    p.flux_ratio_std = [0,1.038]
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "attributes":  modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
diff --git a/galaxy2galaxy/data_generators/galsim_utils.py b/galaxy2galaxy/data_generators/galsim_utils.py
index 3deb059..6643aae 100644
--- a/galaxy2galaxy/data_generators/galsim_utils.py
+++ b/galaxy2galaxy/data_generators/galsim_utils.py
@@ -169,7 +169,7 @@ def draw_and_encode_stamp(gal, psf, stamp_size, pixel_scale, num_bands = 1, flux
         im = gal.drawImage(nx=stamp_size, ny=stamp_size, scale=pixel_scale,
                            method='no_pixel', use_true_center=False).array.astype('float32')
 
-        im = im/np.max(im)
+        im = im/np.max(im) * flux_r[i]
 
         # Draw the Fourier domain image of the galaxy, using x1 zero padding,
         # and x2 subsampling

From ce4e7603a2a024fec38c7de8df0a43f4409fe5e4 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 20 Jul 2021 13:01:34 +0100
Subject: [PATCH 142/149] Rename class and fix detecting central galaxy

---
 galaxy2galaxy/data_generators/candels.py | 48 ++++--------------------
 1 file changed, 8 insertions(+), 40 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 868a06f..8fd09b8 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -51,7 +51,7 @@ def _resize_image(im, size):
 
 
 @registry.register_problem
-class Img2imgCandelsGoodsMultires(astroimage_utils.AstroImageProblem):
+class Img2imgCandelsMultires(astroimage_utils.AstroImageProblem):
   """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
   """
 
@@ -371,7 +371,7 @@ def is_generate_per_split(self):
 
 
 @registry.register_problem
-class Attrs2imgCandelsGoodsEuclid64(Img2imgCandelsGoodsMultires):
+class Attrs2imgCandelsEuclid64(Img2imgCandelsMultires):
   """
   """
 
@@ -399,7 +399,7 @@ def hparams(self, defaults, model_hparams):
 
 
 @registry.register_problem
-class Attrs2imgCandelsGoodsEuclid64Test(Img2imgCandelsGoodsMultires):
+class Attrs2imgCandelsEuclid64TwoBands(Img2imgCandelsMultires):
   """
   """
 
@@ -426,41 +426,9 @@ def hparams(self, defaults, model_hparams):
     p.attributes = ['mag','re', 'q','ZPHOT','F_IRR','F_SPHEROID','F_DISK']
 
 
-@registry.register_problem
-class Attrs2imgCandelsGoodsEuclid64Blue(Img2imgCandelsGoodsMultires):
-  """
-  """
-
-  def eval_metrics(self):
-    eval_metrics = [ ]
-    return eval_metrics
-
-  def hparams(self, defaults, model_hparams):
-    p = defaults
-    p.pixel_scale = {'high' : 0.1, 'low' : 0.1}
-    p.base_pixel_scale = {'high' : 0.06,'low' : 0.06}
-    p.img_len = 64
-    p.sigmas = {"high" : [0.004094741966557142], "low" : [0.004017507500562]}
-    p.filters = {"high" : ['acs_f435w'], "low" : ['wfc3_f160w']}
-    p.zeropoint = 26.49
-    p.resolutions = ["high","low"]
-    p.example_per_shard = 1000
-    p.modality = {"inputs": modalities.ModalityType.IDENTITY,
-                  "attributes":  modalities.ModalityType.IDENTITY,
-                  "targets": modalities.ModalityType.IDENTITY}
-    p.vocab_size = {"inputs": None,
-                    "attributes": None,
-                    "targets": None}
-    p.attributes = ['mag', 're', 'q','ZPHOT']
-
-
-
-
-
-
-
-def find_central(sex_cat):
-
+def find_central(sex_cat,center_coords=32):
+    """Find the central galaxy in a catalog provided by SExtractor
+    """
     n_detect = len(sex_cat)
     
     ''' Match the pred and true cat'''
@@ -469,8 +437,8 @@ def find_central(sex_cat):
     pred_pos[:, 1] = sex_cat['y']
 
     true_pos = np.zeros((1, 2))
-    true_pos[:, 0] = 64
-    true_pos[:, 1] = 64
+    true_pos[:, 0] = center_coords
+    true_pos[:, 1] = center_coords
 
     _, match_index = KDTree(pred_pos).query(true_pos)
     

From da7c3e5760e9c77e0ac4bdf244f111f6b4fa45c3 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 20 Jul 2021 13:07:37 +0100
Subject: [PATCH 143/149] Change hardcoded params to relative

---
 galaxy2galaxy/data_generators/candels.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 8fd09b8..0b4b431 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -503,6 +503,7 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
 
     ''' Sex for clean'''
     img = img.byteswap().newbyteorder()
+    im_size = img.shape[0]
     bkg = sep.Background(img)
    
     cat,sex_seg = sep.extract(img-bkg,sigma_sex,err=bkg.globalrms,segmentation_map=True)
@@ -510,9 +511,9 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
     if len(cat) == 0:
         raise ValueError('No galaxy detected in the field')
     
-    middle_pos = [cat[find_central(cat)[0]]['x'],cat[find_central(cat)[0]]['y']]
+    middle_pos = [cat[find_central(cat,im_shape[0]//2)[0]]['x'],cat[find_central(cat,im_shape[0]//2)[0]]['y']]
     
-    distance = np.sqrt((middle_pos[0]-64)**2 + (middle_pos[1]-64)**2)
+    distance = np.sqrt((middle_pos[0]-im_shape[0]//2)**2 + (middle_pos[1]-im_shape[0]//2)**2)
     if distance > 10 :
         raise ValueError('No galaxy detected in the center')
 
@@ -529,7 +530,7 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
         blended_galaxies = np.unique(sex_seg[loc])
         for blended_galaxy in blended_galaxies:
             blended_galaxy_flux = np.sum(img[np.where(sex_seg==blended_galaxy)])
-            if blend_flux/blended_galaxy_flux > 0.5:
+            if blend_flux/blended_galaxy_flux > 0.1:
               raise ValueError('Blending suspected')
 
     # '''Rotate'''

From aa324115dcd0e14561c1dceed4bcff8364ebe7ea Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 20 Jul 2021 15:25:51 +0100
Subject: [PATCH 144/149] Fixed error and minor adjustments

---
 galaxy2galaxy/data_generators/candels.py   | 48 +++++++++-------------
 galaxy2galaxy/models/autoencoders_utils.py | 13 ------
 2 files changed, 20 insertions(+), 41 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 0b4b431..b3875e8 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -52,7 +52,7 @@ def _resize_image(im, size):
 
 @registry.register_problem
 class Img2imgCandelsMultires(astroimage_utils.AstroImageProblem):
-  """ Base class for image problems created from CANDELS GOODS (North and South) fields, in 7 bands.
+  """ Base class for image problems with the CANDELS catalog, with multiresolution images.
   """
 
   @property
@@ -143,11 +143,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
     sigmas = p.sigmas
 
     # Step 2: Extract postage stamps, resize them to requested size
-    ''' Loop on the two fields'''
     n_gal_creat = 0
     index = 0
    
-    
     ''' Create a subcat containing only the galaxies (in every filters) of the current field'''
     sub_cat = all_cat[np.where(np.isin(list(all_cat["FIELD_1"]),["GDS","GDN","EGS","COSMOS","UDS"]))]
     sub_cat = sub_cat[np.where(sub_cat['mag'] <= 25.3)]
@@ -174,12 +172,9 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                 im_tmp = np.zeros((128, 128, len(p.filters[res])))
                 for n_filter, filt in enumerate(p.filters[res]):
                     print(filt)
-                    # try :
                     ''' Open the image corresponding to the index of the current galaxy'''
 
                     tmp_file = glob.glob(os.path.join(data_dir, sub_cat["FIELD_1"][m], filt)+'/galaxy_'+str(index)+'_*')[0]
-                    # if np.max(fits.open(tmp_file)[0].data) == 0.:
-                    #     sigmas[res][n_filter] = 10
                     im_import = fits.open(tmp_file)[0].data
                     cleaned_image = clean_rotate_stamp(im_import,sigma_sex=1.5)#,noise_level=p.sigmas[res][n_filter])
 
@@ -189,9 +184,6 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
                     im_tmp[:, :, n_filter] = cleaned_image * flux_ratio
                     if np.max(cleaned_image) <= 5*10**(-3):
                         raise ValueError("Very weak image")
-                    # except Exception:
-                    #     print('Galaxy not seen in every filter')
-                    #     continue
                         
                 ''' Resize the image to the low resolution'''
                 new_size = np.ceil(128/scalings[res])+1
@@ -205,7 +197,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             
             im = _resize_image(im, p.img_len)
             
-            #Check that there is still a galaxy
+            # Check that there is still a galaxy
             img_s = im[:,:,0]
             img_s = img_s = img_s.copy(order='C')
             bkg = sep.Background(img_s)
@@ -325,28 +317,24 @@ def example_reading_spec(self):
         "inputs": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="image/encoded",
                 format_key="image/format",
-                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len, self.num_bands],
                 dtype=tf.float32),
 
         "psf": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="psf/encoded",
                 format_key="psf/format",
-                # channels=self.num_bands,
                 shape=[2*p.img_len, 2*p.img_len // 2 + 1, self.num_bands],
                 dtype=tf.float32),
 
         "ps": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="ps/encoded",
                 format_key="ps/format",
-                # channels=self.num_bands,
                 shape=[p.img_len, p.img_len//2+1, self.num_bands],
                 dtype=tf.float32),
 
         "sigma_noise": tf.contrib.slim.tfexample_decoder.Image(
                 image_key="sigma_noise/encoded",
                 format_key="sigma_noise/format",
-                # channels=self.num_bands,
                 shape=[self.num_bands],
                 dtype=tf.float32),
         
@@ -372,7 +360,7 @@ def is_generate_per_split(self):
 
 @registry.register_problem
 class Attrs2imgCandelsEuclid64(Img2imgCandelsMultires):
-  """
+  """For generating images with the Euclid bands
   """
 
   def eval_metrics(self):
@@ -400,7 +388,7 @@ def hparams(self, defaults, model_hparams):
 
 @registry.register_problem
 class Attrs2imgCandelsEuclid64TwoBands(Img2imgCandelsMultires):
-  """
+  """ For generating two-band images (visible and infrared)
   """
 
   def eval_metrics(self):
@@ -426,7 +414,7 @@ def hparams(self, defaults, model_hparams):
     p.attributes = ['mag','re', 'q','ZPHOT','F_IRR','F_SPHEROID','F_DISK']
 
 
-def find_central(sex_cat,center_coords=32):
+def find_central(sex_cat,center_coords=64):
     """Find the central galaxy in a catalog provided by SExtractor
     """
     n_detect = len(sex_cat)
@@ -499,9 +487,11 @@ def mask_out_pixels(img, segmap, segval,
         
     return masked_img.astype(img.dtype), sources, background_mask, central_source, sources_except_central
 
-def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
+def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None, rotate_b=False, blend_threshold=0.1):
+    """Clean images by removing galaxies other than the central one.
+    """
 
-    ''' Sex for clean'''
+    # Detect galaxies with SExtractor
     img = img.byteswap().newbyteorder()
     im_size = img.shape[0]
     bkg = sep.Background(img)
@@ -511,9 +501,9 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
     if len(cat) == 0:
         raise ValueError('No galaxy detected in the field')
     
-    middle_pos = [cat[find_central(cat,im_shape[0]//2)[0]]['x'],cat[find_central(cat,im_shape[0]//2)[0]]['y']]
+    middle_pos = [cat[find_central(cat,im_size[0]//2)[0]]['x'],cat[find_central(cat,im_size[0]//2)[0]]['y']]
     
-    distance = np.sqrt((middle_pos[0]-im_shape[0]//2)**2 + (middle_pos[1]-im_shape[0]//2)**2)
+    distance = np.sqrt((middle_pos[0]-im_size[0]//2)**2 + (middle_pos[1]-im_size[0]//2)**2)
     if distance > 10 :
         raise ValueError('No galaxy detected in the center')
 
@@ -530,15 +520,17 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None):
         blended_galaxies = np.unique(sex_seg[loc])
         for blended_galaxy in blended_galaxies:
             blended_galaxy_flux = np.sum(img[np.where(sex_seg==blended_galaxy)])
-            if blend_flux/blended_galaxy_flux > 0.1:
+            if blend_flux/blended_galaxy_flux > blend_threshold:
               raise ValueError('Blending suspected')
 
-    # '''Rotate'''
-    # PA = cat[find_central(cat)[0]][4]
-    # img_rotate = rotate(cleaned, PA, reshape=False)
-    
-    img_rotate = cleaned
-    '''Add noise'''
+    # Rotate
+    if rotate_b:
+        PA = cat[find_central(cat)[0]][4]
+        img_rotate = rotate(cleaned, PA, reshape=False)
+    else:
+        img_rotate = cleaned
+
+    # Add noise
     background_mask = np.logical_and(np.logical_not(sex_seg==0),np.array(img,dtype=bool))
     if noise_level == None:
         background_std = np.std(img * background_mask)
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 89519af..06e24ee 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -262,19 +262,6 @@ def make_model_spec():
   image_summary("without_psf",tf.reshape(reconstr, labels_shape))
   # Apply channel-wise convolution with the PSF if requested
   if hparams.apply_psf and 'psf' in features:
-    # rec_padded = tf.pad(reconstr, [[0,0],
-    #                               [0, int(hparams.psf_convolution_pad_factor*shape[1])],
-    #                               [0, int(hparams.psf_convolution_pad_factor*shape[2])],
-    #                               [0,0]])
-    # rec_padded_t = tf.transpose(rec_padded, perm=[0, 3,  1, 2])
-    # psf_padded = tf.pad(features['psf'], [[0,0],
-    #                                       [0, int(hparams.psf_convolution_pad_factor*shape[1])],
-    #                                       [0, int(hparams.psf_convolution_pad_factor*shape[2])],
-    #                                       [0,0]])
-    # psf_padded_t = tf.transpose(psf_padded, perm=[0, 3,  1, 2])
-    # reconstr_t = tf.spectral.irfft2d(tf.spectral.rfft2d(rec_padded_t)*tf.cast(tf.abs(tf.spectral.rfft2d(psf_padded_t)), tf.complex64))
-    # reconstr = tf.transpose(reconstr_t, perm=[0, 2, 3, 1])
-    # reconstr = reconstr[:, :shape[1], :shape[2], :]
     output_list = []
     for i in range(shape[3]):
       output_list.append(tf.squeeze(convolve(tf.expand_dims(reconstr[...,i],-1), tf.cast(features['psf'][...,i], tf.complex64),

From 8d0b578d18beabef6a2c12d59edcae0c8df7ca11 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 20 Jul 2021 16:50:33 +0100
Subject: [PATCH 145/149] Correct typo

---
 galaxy2galaxy/data_generators/candels.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index b3875e8..9ae433f 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -501,9 +501,9 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None, rotate_b=False
     if len(cat) == 0:
         raise ValueError('No galaxy detected in the field')
     
-    middle_pos = [cat[find_central(cat,im_size[0]//2)[0]]['x'],cat[find_central(cat,im_size[0]//2)[0]]['y']]
+    middle_pos = [cat[find_central(cat,im_size//2)[0]]['x'],cat[find_central(cat,im_size//2)[0]]['y']]
     
-    distance = np.sqrt((middle_pos[0]-im_size[0]//2)**2 + (middle_pos[1]-im_size[0]//2)**2)
+    distance = np.sqrt((middle_pos[0]-im_size//2)**2 + (middle_pos[1]-im_size//2)**2)
     if distance > 10 :
         raise ValueError('No galaxy detected in the center')
 

From ee84d8c467c400eeefe5698d4c90cb0eae83718c Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 26 Jul 2021 10:55:32 +0100
Subject: [PATCH 146/149] Remove useless file

---
 .../data_generators/check_candels_emulation.py | 18 ------------------
 1 file changed, 18 deletions(-)
 delete mode 100644 galaxy2galaxy/data_generators/check_candels_emulation.py

diff --git a/galaxy2galaxy/data_generators/check_candels_emulation.py b/galaxy2galaxy/data_generators/check_candels_emulation.py
deleted file mode 100644
index 7a3091e..0000000
--- a/galaxy2galaxy/data_generators/check_candels_emulation.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from galaxy2galaxy import problems
-import tensorflow as tf
-import matplotlib.pyplot as plt
-import numpy as np
-Modes = tf.estimator.ModeKeys
-
-candels = problems.problem('attrs2img_candels_goods64_euclid')
-dset = candels.dataset(Modes.TRAIN, data_dir='/data57/hbretonniere/datasets/candels_for_g2g/test_candels')
-batch = dset.batch(1)
-
-print("data set :", dset, "\n \n")
-
-print("batch : ", batch, "\n \n")
-
-
-it = batch.make_one_shot_iterator().get_next()
-sess = tf.Session()
-batch_1 = sess.run(it)
\ No newline at end of file

From ff19206233c53974a1abb869bfbc6d7c01798f02 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Mon, 26 Jul 2021 11:29:34 +0100
Subject: [PATCH 147/149] Cleanup

---
 galaxy2galaxy/data_generators/cosmos.py    | 8 ++++----
 galaxy2galaxy/models/autoencoders.py       | 2 +-
 galaxy2galaxy/models/autoencoders_utils.py | 2 +-
 galaxy2galaxy/models/autoregressive.py     | 2 +-
 4 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index fa02520..4caca8c 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -119,11 +119,11 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         psf = gal.original_psf
 
       # Apply rotation so that the galaxy is at 0 PA
-#       if hasattr(p, "rotation") and p.rotation:
-      rotation_angle = galsim.Angle(-cat_param[ind]['sersic_beta'],
+      if hasattr(p, "rotation") and p.rotation:
+        rotation_angle = galsim.Angle(-np.random.rand()* 2 * np.pi,
                                       galsim.radians)
-      gal = gal.rotate(rotation_angle)
-      psf = psf.rotate(rotation_angle)
+        gal = gal.rotate(rotation_angle)
+        psf = psf.rotate(rotation_angle)
 
       # We save the corresponding attributes for this galaxy
       if hasattr(p, 'attributes'):
diff --git a/galaxy2galaxy/models/autoencoders.py b/galaxy2galaxy/models/autoencoders.py
index 2a3518e..333ba94 100644
--- a/galaxy2galaxy/models/autoencoders.py
+++ b/galaxy2galaxy/models/autoencoders.py
@@ -254,7 +254,7 @@ def continuous_autoencoder_residual_128():
   hparams = continuous_autoencoder_basic()
   hparams.optimizer = "Adafactor"
   hparams.clip_grad_norm = 1.0
-  hparams.learning_rate_constant = 0.35
+  hparams.learning_rate_constant = 0.25
   hparams.learning_rate_warmup_steps = 500
   hparams.learning_rate_schedule = "constant * linear_warmup * rsqrt_decay"
   hparams.num_hidden_layers = 7
diff --git a/galaxy2galaxy/models/autoencoders_utils.py b/galaxy2galaxy/models/autoencoders_utils.py
index 06e24ee..77e3457 100644
--- a/galaxy2galaxy/models/autoencoders_utils.py
+++ b/galaxy2galaxy/models/autoencoders_utils.py
@@ -288,6 +288,6 @@ def make_model_spec():
   logits = tf.reshape(reconstr, labels_shape)
 
   image_summary("ae", reconstr)
-  image_summary("input", labels)  
+  image_summary("input", labels)
 
   return logits, losses
diff --git a/galaxy2galaxy/models/autoregressive.py b/galaxy2galaxy/models/autoregressive.py
index 34ae1a3..9b83348 100644
--- a/galaxy2galaxy/models/autoregressive.py
+++ b/galaxy2galaxy/models/autoregressive.py
@@ -14,7 +14,7 @@
 from tensor2tensor.layers import modalities
 from tensor2tensor.utils import registry
 from tensor2tensor.utils import t2t_model
-# from pixel_cnn_pp.model import model_spec
+from pixel_cnn_pp.model import model_spec
 
 import tensorflow as tf
 import tensorflow_probability as tfp

From a878dd01241170e55a1f8bbcd09953dddd944ff7 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 27 Jul 2021 12:02:43 +0200
Subject: [PATCH 148/149] Cleaning

---
 galaxy2galaxy/data_generators/candels.py | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/galaxy2galaxy/data_generators/candels.py b/galaxy2galaxy/data_generators/candels.py
index 9ae433f..d8bcb9f 100644
--- a/galaxy2galaxy/data_generators/candels.py
+++ b/galaxy2galaxy/data_generators/candels.py
@@ -34,12 +34,6 @@
 from scipy.spatial import KDTree
 import sep
 
-
-# HSC default pixel scale
-_HSC_PIXEL_SCALE=0.168 #arcsec
-# Path to sql files for HSC samples
-_HSC_SAMPLE_SQL_DIR=os.path.join(os.path.dirname(os.path.realpath(__file__)), 'data')
-
 def _resize_image(im, size):
   centh = im.shape[0]/2
   centw = im.shape[1]/2
@@ -160,7 +154,6 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
         if gal == index or gal == 15431 or sub_cat["mag"][m] < 0:     # To take care of the redudency inside the cat
             continue
         index = gal
-        print(index)
         target_flux_main_band = 10**(-0.4*(sub_cat['mag'][m]-p.zeropoint))
 
         try:
@@ -171,7 +164,6 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
             for res in p.resolutions:
                 im_tmp = np.zeros((128, 128, len(p.filters[res])))
                 for n_filter, filt in enumerate(p.filters[res]):
-                    print(filt)
                     ''' Open the image corresponding to the index of the current galaxy'''
 
                     tmp_file = glob.glob(os.path.join(data_dir, sub_cat["FIELD_1"][m], filt)+'/galaxy_'+str(index)+'_*')[0]
@@ -539,4 +531,4 @@ def clean_rotate_stamp(img, eps=5, sigma_sex=2, noise_level=None, rotate_b=False
     random_background = np.random.normal(scale=background_std, size=img_rotate.shape)
     rotated = np.where(img_rotate == 0, random_background, img_rotate)
 
-    return rotated
\ No newline at end of file
+    return rotated

From de73237ecd3429f79f79fab155a05da4cb144a03 Mon Sep 17 00:00:00 2001
From: thuiop <thuiop@hotmail.fr>
Date: Tue, 27 Jul 2021 12:04:49 +0200
Subject: [PATCH 149/149] Remove useless changes

---
 galaxy2galaxy/data_generators/cosmos.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/galaxy2galaxy/data_generators/cosmos.py b/galaxy2galaxy/data_generators/cosmos.py
index 4caca8c..d6d4d68 100644
--- a/galaxy2galaxy/data_generators/cosmos.py
+++ b/galaxy2galaxy/data_generators/cosmos.py
@@ -51,7 +51,8 @@ def hparams(self, defaults, model_hparams):
     p = defaults
     p.pixel_scale = 0.03
     p.img_len = 64
-
+    p.example_per_shard = 1000
+    
     p.modality = {"inputs": modalities.ModalityType.IDENTITY,
                   "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": None,
@@ -118,7 +119,7 @@ def generator(self, data_dir, tmp_dir, dataset_split, task_id=-1):
       else:
         psf = gal.original_psf
 
-      # Apply rotation so that the galaxy is at 0 PA
+      # Apply random rotation if requested
       if hasattr(p, "rotation") and p.rotation:
         rotation_angle = galsim.Angle(-np.random.rand()* 2 * np.pi,
                                       galsim.radians)
@@ -228,7 +229,6 @@ def hparams(self, defaults, model_hparams):
                   "targets": modalities.ModalityType.IDENTITY}
     p.vocab_size = {"inputs": None,
                     "targets": None}
-    
 
 @registry.register_problem
 class Img2imgCosmos128(Img2imgCosmos):
@@ -589,4 +589,4 @@ def hparams(self, defaults, model_hparams):
     p.vocab_size = {"inputs": None,
                     "attributes": None,
                     "targets": None}
-    p.attributes = ['mag_auto', 'flux_radius', 'zphot']
\ No newline at end of file
+    p.attributes = ['mag_auto', 'flux_radius', 'zphot']