Adds testing of multiplicative biases against ngmix

autometacal/python/__init__.py

@@ -1,5 +1,5 @@
 import autometacal.python.datasets as datasets
 import autometacal.python.tf_ngmix as tf_ngmix
-from autometacal.python.metacal import generate_mcal_image, get_metacal_response
+from autometacal.python.metacal import generate_mcal_image, get_metacal_response, get_metacal_response_finitediff
 from autometacal.python.fitting import fit_multivariate_gaussian, get_ellipticity
 from autometacal.python.moments import get_moment_ellipticities, gaussian_moments

autometacal/python/metacal.py

@@ -4,16 +4,17 @@
 import galflow as gf
 import numpy as np
 
-def generate_mcal_image(gal_image,
-                        psf_image,
+def generate_mcal_image(gal_images,
+                        psf_images,
                         reconvolution_psf_image,
-                        g):
+                        g,
+                        padfactor=3):
   """ Generate a metacalibrated image given input and target PSFs.
 
   Args: 
-    gal_image: tf.Tensor or np.array
+    gal_images: tf.Tensor or np.array
       (batch_size, N, N ) image of galaxies
-    psf_image: tf.Tensor or np.array
+    psf_images: tf.Tensor or np.array
       (batch_size, N, N ) image of psf model
     reconvolution_psf_image: tf.Tensor
       (N, N ) tensor of reconvolution psf model
@@ -25,31 +26,47 @@ def generate_mcal_image(gal_image,
       shearing by g, and reconvolution with reconvolution_psf_image.
 
   """
-  #convert and cast
-  gal_image = tf.convert_to_tensor(gal_image, dtype=tf.float32)
-  psf_image = tf.convert_to_tensor(psf_image, dtype=tf.float32)
-  reconvolution_psf_image = tf.convert_to_tensor(reconvolution_psf_image, dtype=tf.float32)
+  #cast stuff as float32 tensors
+  gal_images = tf.convert_to_tensor(gal_images, dtype=tf.float32)  
+  psf_images = tf.convert_to_tensor(psf_images, dtype=tf.float32)  
+  reconvolution_psf_image = tf.convert_to_tensor(reconvolution_psf_image, dtype=tf.float32)  
+  g = tf.convert_to_tensor(g, dtype=tf.float32)  
 
+  #Get batch info
+  batch_size, nx, ny = gal_images.get_shape().as_list()  
+
+  #add pads in real space
+  fact = (padfactor - 1)//2 #how many image sizes to one direction
+  paddings = tf.constant([[0, 0,], [nx*fact, nx*fact], [ny*fact, ny*fact]])
+
+  padded_gal_images = tf.pad(gal_images,paddings)
+  padded_psf_images = tf.pad(psf_images,paddings)
+  padded_reconvolution_psf_image = tf.pad(reconvolution_psf_image,paddings)
+
+  #Convert galaxy images to k space
+  im_shift = tf.signal.ifftshift(padded_gal_images,axes=[1,2]) # The ifftshift is to remove the phase for centered objects
+  im_complex = tf.cast(im_shift, tf.complex64)
+  im_fft = tf.signal.fft2d(im_complex)
+  imk = tf.signal.fftshift(im_fft, axes=[1,2])#the fftshift is to put the 0 frequency at the center of the k image
 
-  g = tf.convert_to_tensor(g, dtype=tf.float32)
-  batch_size, nx, ny = gal_image.get_shape().as_list()
+  #Convert psf images to k space  
+  psf_complex = tf.cast(padded_psf_images, tf.complex64)
+  psf_fft = tf.signal.fft2d(psf_complex)
+  psf_fft_abs = tf.abs(psf_fft)
+  psf_fft_abs_complex = tf.cast(psf_fft_abs,tf.complex64)
+  kpsf = tf.signal.fftshift(psf_fft_abs_complex,axes=[1,2])
 
-  # Convert input stamps to k space
-  # The ifftshift is to remove the phase for centered objects
-  # the fftshift is to put the 0 frequency at the center of the k image
-  imk = tf.signal.fftshift(tf.signal.fft2d(tf.cast(tf.signal.ifftshift(gal_image,axes=[1,2]),
-                                                   tf.complex64)),axes=[1,2])
-  # Note the abs here, to remove the phase of the PSF
-  kpsf = tf.cast(tf.abs(tf.signal.fft2d(tf.cast(psf_image, tf.complex64))), 
-                 tf.complex64) 
-  kpsf = tf.signal.fftshift(kpsf,axes=[1,2])
-  krpsf = tf.cast(tf.abs(tf.signal.fft2d(tf.cast(reconvolution_psf_image,tf.complex64))), tf.complex64)
-  krpsf = tf.signal.fftshift(krpsf,axes=[1,2])
+  #Convert reconvolution psf image to k space 
+  rpsf_complex = tf.cast(padded_reconvolution_psf_image, tf.complex64)
+  rpsf_fft =  tf.signal.fft2d(rpsf_complex)
+  rpsf_fft_abs = tf.abs(rpsf_fft)
+  psf_fft_abs_complex = tf.cast(rpsf_fft_abs,tf.complex64)
+  krpsf = tf.signal.fftshift(psf_fft_abs_complex,axes=[1,2])
 
   # Compute Fourier mask for high frequencies
   # careful, this is not exactly the correct formula for fftfreq
-  kx, ky = tf.meshgrid(tf.linspace(-0.5,0.5,nx),
-                       tf.linspace(-0.5,0.5,ny))
+  kx, ky = tf.meshgrid(tf.linspace(-0.5,0.5,padfactor*nx),
+                       tf.linspace(-0.5,0.5,padfactor*ny))
   mask = tf.cast(tf.math.sqrt(kx**2 + ky**2) <= 0.5, dtype='complex64')
   mask = tf.expand_dims(mask, axis=0)
 
@@ -65,29 +82,95 @@ def generate_mcal_image(gal_image,
   # Compute inverse Fourier transform
   img = tf.math.real(tf.signal.fftshift(im_reconv))
 
-  return img
+  return img[:,fact*nx:-fact*nx,fact*ny:-fact*ny]
 
-def get_metacal_response(gal_image,
-                         psf_image,
+def get_metacal_response(gal_images,
+                         psf_images,
                          reconvolution_psf_image,
                          method):
   """
   Convenience function to compute the shear response
   """  
-  gal_image = tf.convert_to_tensor(gal_image, dtype=tf.float32)
-  psf_image = tf.convert_to_tensor(psf_image, dtype=tf.float32)
+  gal_images = tf.convert_to_tensor(gal_images, dtype=tf.float32)
+  psf_images = tf.convert_to_tensor(psf_images, dtype=tf.float32)
   reconvolution_psf_image = tf.convert_to_tensor(reconvolution_psf_image, dtype=tf.float32)
-  batch_size, nx, ny = gal_image.get_shape().as_list()
+  batch_size, _ , _ = gal_images.get_shape().as_list()
   g = tf.zeros([batch_size, 2])
   with tf.GradientTape() as tape:
     tape.watch(g)
     # Measure ellipticity under metacal
-    e = method(generate_mcal_image(gal_image,
-                                   psf_image,
+    e = method(generate_mcal_image(gal_images,
+                                   psf_images,
                                    reconvolution_psf_image,
                                    g))
 
   # Compute response matrix
+
   R = tape.batch_jacobian(e, g)
   return e, R
 
+
+def get_metacal_response_finitediff(gal_image,psf_image,reconv_psf_image,step,method):
+  """
+  Gets shear response as a finite difference operation, 
+  instead of automatic differentiation.
+  """
+
+  img0s = generate_mcal_image(
+    gal_image,
+    psf_image,
+    reconv_psf_image,
+    [[0,0]]
+  ) 
+  img1p = generate_mcal_image(
+    gal_image,
+    psf_image,
+    reconv_psf_image,
+    [[step,0]]
+  ) 
+  img1m = generate_mcal_image(
+    gal_image,
+    psf_image,
+    reconv_psf_image,
+    [[-step,0]]
+  ) 
+  img2p = generate_mcal_image(
+    gal_image,
+    psf_image,
+    reconv_psf_image,
+    [[0,step]]
+  ) 
+  img2m = generate_mcal_image(
+    gal_image,
+    psf_image,
+    reconv_psf_image,
+    [[0,-step]]
+  ) 
+
+  g0s = method(img0s)
+  g1p = method(img1p)
+  g1m = method(img1m)
+  g2p = method(img2p)
+  g2m = method(img2m)
+
+  R11 = (g1p[:,0]-g1m[:,0])/(2*step)
+  R21 = (g1p[:,1]-g1m[:,1])/(2*step) 
+  R12 = (g2p[:,0]-g2m[:,0])/(2*step)
+  R22 = (g2p[:,1]-g2m[:,1])/(2*step)
+
+  #N. B.:The matrix is correct. 
+  #The transposition will swap R12 with R21 across a batch correctly.
+  R = tf.transpose(tf.convert_to_tensor(
+    [[R11,R21],
+     [R12,R22]],dtype=tf.float32)
+  )
+
+  ellip_dict = {
+    'noshear':g0s,
+    '1p':g1p,
+    '1m':g1m,
+    '2p':g2p,
+    '2m':g2m,    
+  } 
+
+  return ellip_dict, R

autometacal/python/moments.py

@@ -29,7 +29,7 @@ def get_moment_ellipticities(images, scale, fwhm, **kwargs):
 
   q1 = Q11 - Q22
   q2 = 2*Q12
-  T= Q11  + Q22 + 2*tf.math.sqrt(tf.math.abs(Q11*Q22-Q12*Q12)) # g convention 
+  T= Q11  + Q22 # e convention 
 
   g1 = q1/T
   g2 = q2/T