From 4666fb44188963eda85ed18336a2c29c756b2c69 Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Mon, 14 Aug 2017 13:02:53 -0700
Subject: [PATCH 1/8] tool to compute a pixel level flatfield from a list of
 preprocessed images obtained with the pixel flatfield slit

---
 bin/desi_compute_pixel_flatfield | 188 +++++++++++++++++++++++++++++++
 1 file changed, 188 insertions(+)
 create mode 100755 bin/desi_compute_pixel_flatfield

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
new file mode 100755
index 000000000..88b85d67a
--- /dev/null
+++ b/bin/desi_compute_pixel_flatfield
@@ -0,0 +1,188 @@
+#!/usr/bin/env python
+
+
+import sys,string
+import astropy.io.fits as pyfits
+import argparse
+import numpy as np
+import scipy.signal
+
+
+from desiutil.log import get_logger
+
+def median_image(image_filenames) :
+    """ Return a median of input images after rescaling each image
+    
+    Args:
+        image_filenames : list of preprocessed image path
+
+    Returns:
+        mimage : median image (2D np.array)
+    """
+    
+    log.debug("first median")
+    images=[]
+    for filename in image_filenames  :
+        h=pyfits.open(filename)
+        images.append(h[0].data)
+    mimage=np.median(images,axis=0)
+    log.debug("compute a scale per image")
+    smimage2=np.sum(mimage**2)
+    for i in range(len(images)) :
+        a=np.sum(images[i]*mimage)/smimage2
+        log.debug("scale %d = %f"%(i,a))
+        if a<=0 :
+            raise ValueError("scale = %f for image %s"%(a,image_filenames[i]))
+        images[i] /= a
+    mimage=np.median(images,axis=0)
+    return mimage
+
+def add_margins_to_image(image,margin) :
+    """ Adds a margin to the image with extrapolated values.
+    This is needed for Gaussian smoothing using FFTs
+
+    Args:
+        image : 2D array input image of shape (n0,n1)
+        margin : integer number >0
+    
+    Returns:
+        larger : larger image of shape (n0+2*margin,n1*2*margin)
+    """
+    
+    if margin<=0 :
+        raise ValueError("margin must be >0")
+    larger=np.zeros((image.shape[0]+2*margin,image.shape[1]+2*margin))
+    larger[margin:-margin,margin:-margin]=image
+    eps=10
+    larger[:margin+1,margin:-margin]=np.tile(np.median(image[:eps,:],axis=0),(margin+1,1))
+    larger[-margin-1:,margin:-margin]=np.tile(np.median(image[-eps:,:],axis=0),(margin+1,1))
+    larger[margin:-margin,:margin+1]=np.tile(np.median(image[:,:eps],axis=1),(margin+1,1)).T
+    larger[margin:-margin,-margin-1:]=np.tile(np.median(image[:,-eps:],axis=1),(margin+1,1)).T
+    larger[:margin,:margin]=np.median(larger[:margin,margin])
+    larger[-margin:,:margin]=np.median(larger[-margin:,margin])
+    larger[-margin:,-margin:]=np.median(larger[-margin:,-margin-1])
+    larger[:margin,-margin:]=np.median(larger[:margin,-margin-1])
+    return larger
+
+def gaussian_smoothing_1d_per_axis(image,sigma,npass=2) :
+    """Computes a smooth model of the input image using two
+    1D convolution with a Gaussian kernel of parameter sigma. 
+    Can do several passes.
+    
+    Args:
+        image : 2D array input image
+        sigma : float number (>0)
+        npass : integer number (>=1) 
+    
+    Returns:
+        model : 2D array image of same shape as image
+    """
+    log=get_logger()
+    hw=int(3*sigma)
+    tmp = add_margins_to_image(image,hw)
+    model = np.ones(tmp.shape)
+    x=(np.arange(2*hw+1)-hw)/sigma
+    kernel=np.zeros((2*hw+1,3))
+    kernel[:,1]=np.exp(-x**2/2)
+    kernel/=np.sum(kernel)
+    for p in range(npass) : # possibly do several passes
+        for a in range(2) : # convolve in 1d on each axis
+            log.debug("p=%d a=%d"%(p,a))
+            res=scipy.signal.fftconvolve(tmp,kernel,"same")
+            model *= res
+            tmp /= (res+(res==0))
+            kernel=kernel.T
+    return model[hw:-hw,hw:-hw]
+    
+
+def gaussian_smoothing_2d(image,sigma) :
+    """Computes a smooth model of the input image using one
+    2D convolution with a 2D Gaussian kernel of parameter sigma. 
+    
+    Args:
+        image : 2D array input image
+        sigma : float number (>0)
+    
+    Returns:
+        model : 2D array image of same shape as image
+    """
+    hw=int(3*sigma)
+    tmp = add_margins_to_image(image,hw)
+    model = np.ones(tmp.shape)
+    
+    x=np.tile((np.arange(2*hw+1)-hw)/sigma,(2*hw+1,1))
+    r2=x**2+x.T**2    
+    kernel=np.exp(-r2/2.)
+    kernel/=np.sum(kernel)
+    return scipy.signal.fftconvolve(tmp,kernel,"same")[hw:-hw,hw:-hw]
+
+    
+parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+description="Compute a pixel level flat field image from a set of preprocessed images obtained with the pixel flatfield slit."
+)
+
+parser.add_argument('-i','--images', type = str, nargs='*', default = None, required = True,
+                    help = 'path to input preprocessed image fits files, or a single median image')
+parser.add_argument('-o','--outfile', type = str, default = None, required = True,
+                    help = 'output flatfield image filename')
+parser.add_argument('--sigma', type = int, default = 60 , required = False, 
+                    help = "gaussian smoothing sigma (default is tuned on teststand data")
+parser.add_argument('--out-median', type = str, default = None , required = False, help = "save output median image (for development)")
+args        = parser.parse_args()
+log = get_logger()
+
+if len(args.images) == 1 :
+    log.info("read a single image")
+    image=pyfits.open(args.images[0])[0].data
+else :
+    log.info("compute a median of the input images")
+    image=median_image(args.images)
+
+log.info("first gaussian smoothing")
+model=gaussian_smoothing_1d_per_axis(image,args.sigma)
+
+nsloop=2
+for sloop in range(nsloop) :
+    log.info("compute mask using a temporary flat %d/%d"%(sloop,nsloop))
+    flat=(model>0)*image/(model+(model==0))
+    flat+=(model<=0)
+    flaterr=np.sqrt(1.+model*(model>0))/(model+(model==0))
+    rms=np.abs(flat-1)/flaterr
+    srms=gaussian_smoothing_2d(rms,10.)
+    mask=(srms>1.)|(rms>3.)
+    log.info("number of masked pixels=%d"%np.sum(mask.astype(int)))
+    if 1 :
+        log.info("replace masked pixels by model %d/%d"%(sloop,nsloop))
+        nloop=6
+        for loop in range(nloop) :
+            log.debug("pass %d/%d"%(loop,nloop))
+            tmp_image=image*(mask==0)+model*(mask>0)
+            model=gaussian_smoothing_1d_per_axis(tmp_image,args.sigma,npass=1)
+
+        
+log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
+n0=tmp_image.shape[0]//2
+n1=tmp_image.shape[1]//2
+model=np.zeros(tmp_image.shape)
+model[:n0,:n1]=gaussian_smoothing_1d_per_axis(tmp_image[:n0,:n1],args.sigma)
+model[n0:,:n1]=gaussian_smoothing_1d_per_axis(tmp_image[n0:,:n1],args.sigma)
+model[n0:,n1:]=gaussian_smoothing_1d_per_axis(tmp_image[n0:,n1:],args.sigma)
+model[:n0,n1:]=gaussian_smoothing_1d_per_axis(tmp_image[:n0,n1:],args.sigma)
+    
+
+flat=(model>0)*image/(model+(model==0))
+flat+=(model<=0)
+
+
+log.info("writing %s ..."%args.outfile)
+h=pyfits.HDUList([pyfits.PrimaryHDU(flat),pyfits.ImageHDU(model,name="MODEL"),pyfits.ImageHDU(mask.astype(int),name="MODMASK")])
+h[0].header["EXTNAME"]="FLAT"
+h[0].header["KERNSIG"]=args.sigma
+h.writeto(args.outfile,clobber=True)
+
+if args.out_median is not None :
+    log.info("writing median image %s ..."%args.out_median)
+    pyfits.writeto(args.out_median,image,clobber=True)
+
+log.info("done")
+

From e3d8dd4a02bb94410eab9402f827dbbe6a92532d Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Mon, 14 Aug 2017 16:07:19 -0700
Subject: [PATCH 2/8] accounting for trace orientation when defining 1D
 gaussian kernels

---
 bin/desi_compute_pixel_flatfield | 208 ++++++++++++++++++++++++++++---
 1 file changed, 188 insertions(+), 20 deletions(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index 88b85d67a..42220c1fe 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -6,10 +6,67 @@ import astropy.io.fits as pyfits
 import argparse
 import numpy as np
 import scipy.signal
-
+import specter.psf
 
 from desiutil.log import get_logger
 
+def grid_from_psf(filename) :
+    try :
+        psftype=pyfits.open(filename)[0].header["PSFTYPE"]
+    except KeyError :
+        psftype=""
+
+    psf=None
+    
+    if psftype=="GAUSS-HERMITE" :
+        psf=specter.psf.GaussHermitePSF(filename)
+    elif psftype=="SPOTGRID" :
+        psf=specter.psf.SpotGridPSF(filename)
+    
+    if psf is None :
+        raise ValueError("cannot read PSFTYPE=%s"%psftype)
+    
+    # make a grid of points from this PSF
+    x=[]
+    y=[]
+    f=[]
+    w=[]
+    wstep=200.
+    waves=np.linspace(psf.wmin,psf.wmax,int((psf.wmax-psf.wmin)/wstep))
+    fibers=np.arange(psf.nspec)
+    for fiber in fibers :
+        for wave in waves :
+            tx,ty = psf.xy(fiber,wave)
+            x.append(tx)
+            y.append(ty)
+            f.append(fiber)
+            w.append(wave)
+    
+    x=np.array(x)
+    y=np.array(y)
+    dxdy=np.zeros(x.size).astype(float)
+    dydx=np.zeros(x.size).astype(float)
+
+    # use this grid of points to determine dx/dy along wave and dy/dx along fiber
+    for fiber in fibers :
+        mask=np.where(f==fiber)[0]
+        tx=x[mask]
+        ty=y[mask]
+        i=np.argsort(ty)
+        dx=np.gradient(tx[i])
+        dy=np.gradient(ty[i])
+        dxdy[mask[i]]=dx/dy
+    
+    for wave in waves :
+        mask=np.where(w==wave)[0]
+        tx=x[mask]
+        ty=y[mask]
+        i=np.argsort(tx)
+        dx=np.gradient(tx[i])
+        dy=np.gradient(ty[i])
+        dydx[mask[i]]=dy/dx
+    return x,y,dxdy,dydx
+
 def median_image(image_filenames) :
     """ Return a median of input images after rescaling each image
     
@@ -64,7 +121,7 @@ def add_margins_to_image(image,margin) :
     larger[:margin,-margin:]=np.median(larger[:margin,-margin-1])
     return larger
 
-def gaussian_smoothing_1d_per_axis(image,sigma,npass=2) :
+def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0,dydx=0) :
     """Computes a smooth model of the input image using two
     1D convolution with a Gaussian kernel of parameter sigma. 
     Can do several passes.
@@ -77,23 +134,120 @@ def gaussian_smoothing_1d_per_axis(image,sigma,npass=2) :
     Returns:
         model : 2D array image of same shape as image
     """
+    
     log=get_logger()
     hw=int(3*sigma)
     tmp = add_margins_to_image(image,hw)
     model = np.ones(tmp.shape)
-    x=(np.arange(2*hw+1)-hw)/sigma
-    kernel=np.zeros((2*hw+1,3))
-    kernel[:,1]=np.exp(-x**2/2)
-    kernel/=np.sum(kernel)
+
+    # single Gaussian profile
+    u=(np.arange(2*hw+1)-hw)
+    prof=np.exp(-u**2/sigma**2/2.)
+    prof/=np.sum(prof)
+    
+    # two kernels along two axes
+    # 
+    kernels=[]
+
+    # axis 0
+    if dxdy==0 :
+        kernel=np.zeros((2*hw+1,3))
+        kernel[:,1]=prof
+        kernels.append(kernel)
+    else :
+        # non trivial because of axis orientations
+        dxmax=int(hw*np.abs(dxdy))+1
+        n0=2*hw+1
+        n1=2*dxmax+1
+    
+        x=u*dxdy
+        x2d=np.tile(x,(n1,1)).T
+        prof2d=np.tile(prof,(n1,1)).T
+        v2d=np.tile(np.arange(n1)-dxmax,(n0,1))
+        weight=1.-np.abs(x2d-v2d)
+        weight*=(weight>0)
+        kernel=prof2d*weight
+        kernels.append(kernel)
+    
+    
+    # axis 1
+    if dydx==0 :
+        kernel=np.zeros((3,2*hw+1))
+        kernel[1,:]=prof
+        kernels.append(kernel)
+    else :
+        dymax=int(hw*np.abs(dydx))+1
+        n0=2*dymax+1
+        n1=2*hw+1
+    
+        y=u*dydx
+        y2d=np.tile(y,(n0,1))
+        prof2d=np.tile(prof,(n0,1))
+        v2d=np.tile(np.arange(n0)-dymax,(n1,1)).T
+        weight=1.-np.abs(y2d-v2d)
+        weight*=(weight>0)
+        kernel=prof2d*weight
+        kernels.append(kernel)
+    
+    
     for p in range(npass) : # possibly do several passes
         for a in range(2) : # convolve in 1d on each axis
             log.debug("p=%d a=%d"%(p,a))
-            res=scipy.signal.fftconvolve(tmp,kernel,"same")
+            res=scipy.signal.fftconvolve(tmp,kernels[a],"same")
             model *= res
             tmp /= (res+(res==0))
-            kernel=kernel.T
     return model[hw:-hw,hw:-hw]
     
+def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nblocks=5) :
+    """Computes a smooth model of the input image using two
+    1D convolution with a Gaussian kernel of parameter sigma. 
+    Can do several passes.
+    
+    Args:
+        image : 2D array input image
+        sigma : float number (>0)
+        npass : integer number (>=1) 
+    
+    Returns:
+        model : 2D array image of same shape as image
+    """
+    if x is None or nblocks==1 :
+        return gaussian_smoothing_1d_per_axis(image,sigma,npass=npass,dxdy=0,dydx=0)
+    
+    # defining blocks where trace directions are averaged
+    b0size=image.shape[0]//nblocks
+    b1size=image.shape[1]//nblocks
+    
+    b0=b0size*np.arange(nblocks)
+    e0=b0+b0size
+    e0[-1]=image.shape[0]
+    b1=b1size*np.arange(nblocks)
+    e1=b1+b1size
+    e1[-1]=image.shape[1]
+    
+    # average trace direction per block
+    bdxdy=np.zeros((nblocks,nblocks))
+    bdydx=np.zeros((nblocks,nblocks))
+    for i in range(nblocks) :
+        for j in range(nblocks) :
+            mask=(y>b0[i])&(y<=e0[i])&(x>b1[j])&(x<=e1[j])
+            
+            if np.sum(mask)>0 :                
+                bdxdy[i,j]=np.mean(dxdy[mask])
+                bdydx[i,j]=np.mean(dydx[mask])
+            
+            log.debug("Block %d,%d n psf points=%d dxdy=%f dydx=%f"%(i,j,np.sum(mask),bdxdy[i,j],bdydx[i,j]))
+    
+    model=np.zeros(image.shape)
+    for i in range(nblocks) :
+        for j in range(nblocks) : 
+            # fast, but with edge issues
+            #model[b0[i]:e0[i],b1[j]:e1[j]] = gaussian_smoothing_1d_per_axis(image[b0[i]:e0[i],b1[j]:e1[j]],sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])
+            # slower, but without edge issues
+            model = gaussian_smoothing_1d_per_axis(image,sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])[b0[i]:e0[i],b1[j]:e1[j]]
+            
+    return model
+    
 
 def gaussian_smoothing_2d(image,sigma) :
     """Computes a smooth model of the input image using one
@@ -128,6 +282,8 @@ parser.add_argument('-o','--outfile', type = str, default = None, required = Tru
 parser.add_argument('--sigma', type = int, default = 60 , required = False, 
                     help = "gaussian smoothing sigma (default is tuned on teststand data")
 parser.add_argument('--out-median', type = str, default = None , required = False, help = "save output median image (for development)")
+parser.add_argument('--psf', type = str, default = None , required = False, help = "use traces in this PSF to orient 1D convolutions")
+
 args        = parser.parse_args()
 log = get_logger()
 
@@ -138,8 +294,17 @@ else :
     log.info("compute a median of the input images")
     image=median_image(args.images)
 
+if args.psf :
+    log.info("get trace coordinates from psf %s"%args.psf)
+    x,y,dxdy,dydx=grid_from_psf(args.psf)
+else :
+    x=None
+    y=None
+    dxdy=None
+    dydx=None
+
 log.info("first gaussian smoothing")
-model=gaussian_smoothing_1d_per_axis(image,args.sigma)
+model=gaussian_smoothing_1d_with_tilted_axes(image,args.sigma,2,x,y,dxdy,dydx)
 
 nsloop=2
 for sloop in range(nsloop) :
@@ -157,17 +322,20 @@ for sloop in range(nsloop) :
         for loop in range(nloop) :
             log.debug("pass %d/%d"%(loop,nloop))
             tmp_image=image*(mask==0)+model*(mask>0)
-            model=gaussian_smoothing_1d_per_axis(tmp_image,args.sigma,npass=1)
-
-        
-log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
-n0=tmp_image.shape[0]//2
-n1=tmp_image.shape[1]//2
-model=np.zeros(tmp_image.shape)
-model[:n0,:n1]=gaussian_smoothing_1d_per_axis(tmp_image[:n0,:n1],args.sigma)
-model[n0:,:n1]=gaussian_smoothing_1d_per_axis(tmp_image[n0:,:n1],args.sigma)
-model[n0:,n1:]=gaussian_smoothing_1d_per_axis(tmp_image[n0:,n1:],args.sigma)
-model[:n0,n1:]=gaussian_smoothing_1d_per_axis(tmp_image[:n0,n1:],args.sigma)
+            model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,args.sigma,1,x,y,dxdy,dydx)
+
+if 0 :        
+    log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
+    n0=tmp_image.shape[0]//2
+    n1=tmp_image.shape[1]//2
+    model=np.zeros(tmp_image.shape)
+    model[:n0,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,:n1],args.sigma,2,x,y,dxdy,dydx)
+    model[n0:,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,:n1],args.sigma,2,x,y,dxdy,dydx)
+    model[n0:,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,n1:],args.sigma,2,x,y,dxdy,dydx)
+    model[:n0,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,n1:],args.sigma,2,x,y,dxdy,dydx)
+else :
+    log.info("last gaussian smoothing")
+    model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,args.sigma,2,x,y,dxdy,dydx)
     
 
 flat=(model>0)*image/(model+(model==0))

From 7fc6a101c3825d77c722b9833e666fdde053108b Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Tue, 15 Aug 2017 15:25:37 -0700
Subject: [PATCH 3/8] debugged version with tilted axis, still working on CCD
 edges ...

---
 bin/desi_compute_pixel_flatfield | 217 +++++++++++++++++++------------
 1 file changed, 136 insertions(+), 81 deletions(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index 42220c1fe..f8c347a78 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -94,34 +94,30 @@ def median_image(image_filenames) :
     mimage=np.median(images,axis=0)
     return mimage
 
-def add_margins_to_image(image,margin) :
-    """ Adds a margin to the image with extrapolated values.
-    This is needed for Gaussian smoothing using FFTs
+def convolve2d(image,k) :
+    
+    if len(k.shape) != 2 or len(image.shape) != 2:
+        raise ValueError("kernel and image should have 2 dimensions")
+    for d in range(2) :
+        if k.shape[d]<=1 or k.shape[d]-(k.shape[d]//2)*2 != 1 :
+            raise ValueError("kernel dimensions should both be odd and >1, and input as shape %s"%str(k.shape))
+    m0=k.shape[0]//2
+    m1=k.shape[1]//2
+    eps0=m0
+    eps1=m1
+    tmp=np.zeros((image.shape[0]+2*m0,image.shape[1]+2*m1))
+    tmp[m0:-m0,m1:-m1]=image
+    tmp[:m0+1,m1:-m1]=np.tile(np.median(image[:eps0,:],axis=0),(m0+1,1))
+    tmp[-m0-1:,m1:-m1]=np.tile(np.median(image[-eps0:,:],axis=0),(m0+1,1))
+    tmp[m0:-m0,:m1+1]=np.tile(np.median(image[:,:eps1],axis=1),(m1+1,1)).T
+    tmp[m0:-m0,-m1-1:]=np.tile(np.median(image[:,-eps1:],axis=1),(m1+1,1)).T
+    tmp[:m0,:m1]=np.median(tmp[:m0,m1])
+    tmp[-m0:,:m1]=np.median(tmp[-m0:,m1])
+    tmp[-m0:,-m1:]=np.median(tmp[-m0:,-m1-1])
+    tmp[:m0,-m1:]=np.median(tmp[:m0,-m1-1])
+    return scipy.signal.fftconvolve(tmp,k,"valid")
 
-    Args:
-        image : 2D array input image of shape (n0,n1)
-        margin : integer number >0
-    
-    Returns:
-        larger : larger image of shape (n0+2*margin,n1*2*margin)
-    """
-    
-    if margin<=0 :
-        raise ValueError("margin must be >0")
-    larger=np.zeros((image.shape[0]+2*margin,image.shape[1]+2*margin))
-    larger[margin:-margin,margin:-margin]=image
-    eps=10
-    larger[:margin+1,margin:-margin]=np.tile(np.median(image[:eps,:],axis=0),(margin+1,1))
-    larger[-margin-1:,margin:-margin]=np.tile(np.median(image[-eps:,:],axis=0),(margin+1,1))
-    larger[margin:-margin,:margin+1]=np.tile(np.median(image[:,:eps],axis=1),(margin+1,1)).T
-    larger[margin:-margin,-margin-1:]=np.tile(np.median(image[:,-eps:],axis=1),(margin+1,1)).T
-    larger[:margin,:margin]=np.median(larger[:margin,margin])
-    larger[-margin:,:margin]=np.median(larger[-margin:,margin])
-    larger[-margin:,-margin:]=np.median(larger[-margin:,-margin-1])
-    larger[:margin,-margin:]=np.median(larger[:margin,-margin-1])
-    return larger
-
-def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0,dydx=0) :
+def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0.,dydx=0.) :
     """Computes a smooth model of the input image using two
     1D convolution with a Gaussian kernel of parameter sigma. 
     Can do several passes.
@@ -137,8 +133,8 @@ def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0,dydx=0) :
     
     log=get_logger()
     hw=int(3*sigma)
-    tmp = add_margins_to_image(image,hw)
-    model = np.ones(tmp.shape)
+    tmp = image.copy()
+    model = np.ones(tmp.shape).astype(float)
 
     # single Gaussian profile
     u=(np.arange(2*hw+1)-hw)
@@ -155,48 +151,45 @@ def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0,dydx=0) :
         kernel[:,1]=prof
         kernels.append(kernel)
     else :
-        # non trivial because of axis orientations
-        dxmax=int(hw*np.abs(dxdy))+1
-        n0=2*hw+1
-        n1=2*dxmax+1
-    
         x=u*dxdy
-        x2d=np.tile(x,(n1,1)).T
-        prof2d=np.tile(prof,(n1,1)).T
-        v2d=np.tile(np.arange(n1)-dxmax,(n0,1))
-        weight=1.-np.abs(x2d-v2d)
-        weight*=(weight>0)
-        kernel=prof2d*weight
+        i=(x+0.5*(x>0)-0.5*(x<0)).astype(int)
+        j=np.arange(2*hw+1)
+        hwb=max(1,np.max(np.abs(i)))
+        kernel=np.zeros((2*hw+1,2*hwb+1))
+        kernel[j,i+hwb]=prof
         kernels.append(kernel)
     
-    
     # axis 1
     if dydx==0 :
         kernel=np.zeros((3,2*hw+1))
         kernel[1,:]=prof
         kernels.append(kernel)
     else :
-        dymax=int(hw*np.abs(dydx))+1
-        n0=2*dymax+1
-        n1=2*hw+1
-    
         y=u*dydx
-        y2d=np.tile(y,(n0,1))
-        prof2d=np.tile(prof,(n0,1))
-        v2d=np.tile(np.arange(n0)-dymax,(n1,1)).T
-        weight=1.-np.abs(y2d-v2d)
-        weight*=(weight>0)
-        kernel=prof2d*weight
+        j=(y+0.5*(y>0)-0.5*(y<0)).astype(int)
+        i=np.arange(2*hw+1)
+        hwb=max(1,np.max(np.abs(j)))
+        kernel=np.zeros((2*hwb+1,2*hw+1))
+        kernel[j+hwb,i]=prof
         kernels.append(kernel)
     
-    
     for p in range(npass) : # possibly do several passes
         for a in range(2) : # convolve in 1d on each axis
-            log.debug("p=%d a=%d"%(p,a))
-            res=scipy.signal.fftconvolve(tmp,kernels[a],"same")
+            #log.debug("p=%d a=%d"%(p,a))
+            res=convolve2d(tmp,kernels[a])
             model *= res
             tmp /= (res+(res==0))
-    return model[hw:-hw,hw:-hw]
+
+    if 1 : # add 2D smoothing (does not help)
+        x=np.tile((np.arange(2*hw+1)-hw)/sigma,(2*hw+1,1))
+        r2=x**2+x.T**2    
+        kernel2d=np.exp(-r2/2.)
+        kernel2d/=np.sum(kernel2d)
+        res = convolve2d(tmp,kernel2d)
+        model *= res
+    
+    return model
+
     
 def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nblocks=5) :
     """Computes a smooth model of the input image using two
@@ -214,10 +207,12 @@ def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nbloc
     if x is None or nblocks==1 :
         return gaussian_smoothing_1d_per_axis(image,sigma,npass=npass,dxdy=0,dydx=0)
     
+    
     # defining blocks where trace directions are averaged
     b0size=image.shape[0]//nblocks
     b1size=image.shape[1]//nblocks
     
+    # blocks begin and end coordinates
     b0=b0size*np.arange(nblocks)
     e0=b0+b0size
     e0[-1]=image.shape[0]
@@ -225,6 +220,18 @@ def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nbloc
     e1=b1+b1size
     e1[-1]=image.shape[1]
     
+    # blocks begin and end coordinates with margins
+    hw=int(3*sigma)
+    b0m = b0-hw    
+    e0m = e0+hw
+    b1m = b1-hw
+    e1m = e1+hw
+    b0m[b0m<0]=0
+    e0m[e0m>image.shape[0]]=image.shape[0]
+    b1m[b1m<0]=0
+    e1m[e1m>image.shape[1]]=image.shape[1]
+    
+    
     # average trace direction per block
     bdxdy=np.zeros((nblocks,nblocks))
     bdydx=np.zeros((nblocks,nblocks))
@@ -236,16 +243,20 @@ def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nbloc
                 bdxdy[i,j]=np.mean(dxdy[mask])
                 bdydx[i,j]=np.mean(dydx[mask])
             
-            log.debug("Block %d,%d n psf points=%d dxdy=%f dydx=%f"%(i,j,np.sum(mask),bdxdy[i,j],bdydx[i,j]))
+            #log.debug("Block %d,%d n psf points=%d dxdy=%f dydx=%f"%(i,j,np.sum(mask),bdxdy[i,j],bdydx[i,j]))
     
     model=np.zeros(image.shape)
     for i in range(nblocks) :
-        for j in range(nblocks) : 
-            # fast, but with edge issues
-            #model[b0[i]:e0[i],b1[j]:e1[j]] = gaussian_smoothing_1d_per_axis(image[b0[i]:e0[i],b1[j]:e1[j]],sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])
-            # slower, but without edge issues
-            model = gaussian_smoothing_1d_per_axis(image,sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])[b0[i]:e0[i],b1[j]:e1[j]]
+        for j in range(nblocks) :
+            log.info("calling gaussian_smoothing_1d_per_axis for block (%d,%d)"%(i,j))
+            block_sigma = sigma
+            if False :
+                if (i==0 and j==0) or (i==0 and j==(nblocks-1)) or  (i==(nblocks-1) and j==(nblocks-1)) or (i==(nblocks-1) and j==0) :
+                    block_sigma = sigma/2.
+                    log.info("Using lower sigma for edge blocks (%d,%d) = %f"%(i,j,block_sigma))
             
+            model[b0[i]:e0[i],b1[j]:e1[j]] = gaussian_smoothing_1d_per_axis(image[b0m[i]:e0m[i],b1m[j]:e1m[j]],sigma=block_sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])[b0[i]-b0m[i]:,b1[j]-b1m[j]:][:e0[i]-b0[i],:e1[j]-b1[j]]
+    
     return model
     
 
@@ -261,14 +272,11 @@ def gaussian_smoothing_2d(image,sigma) :
         model : 2D array image of same shape as image
     """
     hw=int(3*sigma)
-    tmp = add_margins_to_image(image,hw)
-    model = np.ones(tmp.shape)
-    
     x=np.tile((np.arange(2*hw+1)-hw)/sigma,(2*hw+1,1))
     r2=x**2+x.T**2    
     kernel=np.exp(-r2/2.)
     kernel/=np.sum(kernel)
-    return scipy.signal.fftconvolve(tmp,kernel,"same")[hw:-hw,hw:-hw]
+    return convolve2d(image,kernel)
 
     
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -280,7 +288,13 @@ parser.add_argument('-i','--images', type = str, nargs='*', default = None, requ
 parser.add_argument('-o','--outfile', type = str, default = None, required = True,
                     help = 'output flatfield image filename')
 parser.add_argument('--sigma', type = int, default = 60 , required = False, 
-                    help = "gaussian smoothing sigma (default is tuned on teststand data")
+                    help = "gaussian smoothing sigma (default is tuned on teststand data)")
+parser.add_argument('--npass1', type = int, default = 2 , required = False, 
+                    help = "number of iterations in gaussian smoothing 1")
+parser.add_argument('--npass2', type = int, default = 6 , required = False, 
+                    help = "number of iterations in gaussian smoothing 2")
+parser.add_argument('--nblocks', type = int, default = 4 , required = False, 
+                    help = "number of blocks along one axis (total number of blocks is the square) where the trace orientations are averaged (to use in combination with --psf option otherwise ignored)")
 parser.add_argument('--out-median', type = str, default = None , required = False, help = "save output median image (for development)")
 parser.add_argument('--psf', type = str, default = None , required = False, help = "use traces in this PSF to orient 1D convolutions")
 
@@ -303,41 +317,82 @@ else :
     dxdy=None
     dydx=None
 
+
+
 log.info("first gaussian smoothing")
-model=gaussian_smoothing_1d_with_tilted_axes(image,args.sigma,2,x,y,dxdy,dydx)
+model=gaussian_smoothing_1d_with_tilted_axes(image,sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
 
 nsloop=2
 for sloop in range(nsloop) :
-    log.info("compute mask using a temporary flat %d/%d"%(sloop,nsloop))
+    log.info("compute mask using a temporary flat %d/%d"%(sloop+1,nsloop))
+    
+    log.info("estimate variance")
+    r_model=model.ravel()
+    ok=(r_model>-100)&(r_model<5000)
+    h0,bins=np.histogram(r_model[ok],bins=15)
+    hx,junk=np.histogram(r_model[ok],weights=r_model[ok],bins=bins)
+    valid=np.where(h0>50)[0]# cut on min number of pixels per flux bin
+    if valid.size < 4 :
+        log.warning("cannot measure variance, use guess")
+        var = 1.+model*(model>0)
+    else :
+        flux=hx[valid]/h0[valid]
+        var_of_flux=np.zeros(valid.size)
+        r_res=image.ravel()-r_model
+        for i,j in enumerate(valid) :
+            var_of_flux[i] = (1.48*np.median(np.abs(r_res[(r_model>=bins[j])&(r_model<bins[j+1])])))**2
+        c=np.polyfit(flux,var_of_flux,1)
+        var = c[0]*model+c[1]
+        log.info("var = %f + %f*flux"%(c[1],c[0]))
+        if c[1]<1. :
+            c[1]=1
+            log.warning("forcing var(flux=0)=%f"%c[1])
+        #import matplotlib.pyplot as plt
+        #pol=np.poly1d(c)
+        #err = var_of_flux*np.sqrt(2./h0[valid])
+        #plt.errorbar(flux,var_of_flux,err,fmt="o")
+        #plt.plot(flux,pol(flux),"-")
+        #plt.show()
+    
     flat=(model>0)*image/(model+(model==0))
     flat+=(model<=0)
-    flaterr=np.sqrt(1.+model*(model>0))/(model+(model==0))
+    flaterr=np.sqrt(var)/(model+(model==0))
     rms=np.abs(flat-1)/flaterr
     srms=gaussian_smoothing_2d(rms,10.)
     mask=(srms>1.)|(rms>3.)
+    
+    # DEBUG
+    pyfits.writeto("mask-%d.fits"%sloop,mask.astype(int),clobber=True)
+    pyfits.writeto("tmp-flat-%d.fits"%sloop,flat,clobber=True)
+    
     log.info("number of masked pixels=%d"%np.sum(mask.astype(int)))
     if 1 :
-        log.info("replace masked pixels by model %d/%d"%(sloop,nsloop))
-        nloop=6
-        for loop in range(nloop) :
-            log.debug("pass %d/%d"%(loop,nloop))
+        log.info("replace masked pixels by model %d/%d"%(sloop+1,nsloop))
+        for loop in range(args.npass2) :
+            log.debug("pass %d/%d"%(loop+1,args.npass2))
             tmp_image=image*(mask==0)+model*(mask>0)
-            model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,args.sigma,1,x,y,dxdy,dydx)
+            model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
 
-if 0 :        
+if True :        
     log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
     n0=tmp_image.shape[0]//2
     n1=tmp_image.shape[1]//2
     model=np.zeros(tmp_image.shape)
-    model[:n0,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,:n1],args.sigma,2,x,y,dxdy,dydx)
-    model[n0:,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,:n1],args.sigma,2,x,y,dxdy,dydx)
-    model[n0:,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,n1:],args.sigma,2,x,y,dxdy,dydx)
-    model[:n0,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,n1:],args.sigma,2,x,y,dxdy,dydx)
+    if args.nblocks == 1 :
+        nblocks=1
+    else :
+        nblocks=args.nblocks//2
+    model[:n0,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,:n1],sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    model[n0:,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,:n1],sigma=args.sigma,npass=args.npass1,x=x,y=y-n0,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    model[n0:,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,n1:],sigma=args.sigma,npass=args.npass1,x=x-n1,y=y-n0,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    model[:n0,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,n1:],sigma=args.sigma,npass=args.npass1,x=x-n1,y=y,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
 else :
-    log.info("last gaussian smoothing")
-    model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,args.sigma,2,x,y,dxdy,dydx)
+    log.info("last gaussian smoothing (not per quadrant)")
+    model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
     
 
+
+
 flat=(model>0)*image/(model+(model==0))
 flat+=(model<=0)
 

From c3e6e320186e172c4d2b30512aff6cbf9fe939a1 Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Tue, 15 Aug 2017 16:22:33 -0700
Subject: [PATCH 4/8] try to limit the masked area and add more help

---
 bin/desi_compute_pixel_flatfield | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index f8c347a78..341015d43 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -280,7 +280,11 @@ def gaussian_smoothing_2d(image,sigma) :
 
     
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-description="Compute a pixel level flat field image from a set of preprocessed images obtained with the pixel flatfield slit."
+description="""Computes a pixel level flat field image from a set of preprocessed images obtained with the flatfield slit.
+A median image is computed if several preprocessed images are given. 
+The method consists in iteratively dividing the median or input image by a smoothed version of the same image, flat(n+1) = flat(n)/smoothing(flat(n)).
+The smoothing consists in 1D FFT Gaussian convolution along the wavelength dispersion axis or the fiber axis alternatively. A masking of outliers is performed to avoid tails around CCD defects. The trace orientations is obtained from an input PSF file, and the orientations are averaged in blocks (the number of blocks is a parameter). Optionally the modeling can be performed per CCD amplifier to correct for gain mismatch (possibly due to non-linearities). The method fails in areas where the illumination pattern varies in both directions at a scale smaller than the sigma value, but the sigma value is also limited by the maximum size of the CCD defects to be captured in the flat.
+"""
 )
 
 parser.add_argument('-i','--images', type = str, nargs='*', default = None, required = True,
@@ -297,6 +301,7 @@ parser.add_argument('--nblocks', type = int, default = 4 , required = False,
                     help = "number of blocks along one axis (total number of blocks is the square) where the trace orientations are averaged (to use in combination with --psf option otherwise ignored)")
 parser.add_argument('--out-median', type = str, default = None , required = False, help = "save output median image (for development)")
 parser.add_argument('--psf', type = str, default = None , required = False, help = "use traces in this PSF to orient 1D convolutions")
+parser.add_argument('--per-amplifier', action = 'store_true', default = None , required = False, help = "solve model per amplifier if gains are uncertain or non-linarities")
 
 args        = parser.parse_args()
 log = get_logger()
@@ -361,6 +366,23 @@ for sloop in range(nsloop) :
     srms=gaussian_smoothing_2d(rms,10.)
     mask=(srms>1.)|(rms>3.)
     
+    # remove from mask regions with important model gradient
+    grad=np.zeros(model.shape)
+    grad[1:-1]=np.abs(model[2:]-model[:-2])
+    grad[:,1:-1]+=np.abs(model[:,2:]-model[:,:-2])
+    grad/=(model+(model==0))
+    mask[grad>0.2]=0
+    
+    # do not mask edges
+    mask[0,:]=0 ; mask[:,0]=1 ; mask[-1,:]=1 ; mask[:,-1]=1
+    if args.per_amplifier :
+        log.info("do not mask amp. boundaries")
+        mask[image.shape[0]//2-2:image.shape[0]//2+3]=0
+        mask[:,image.shape[1]//2-2:image.shape[1]//2+3]=0
+    
+    
+    
+
     # DEBUG
     pyfits.writeto("mask-%d.fits"%sloop,mask.astype(int),clobber=True)
     pyfits.writeto("tmp-flat-%d.fits"%sloop,flat,clobber=True)
@@ -373,7 +395,7 @@ for sloop in range(nsloop) :
             tmp_image=image*(mask==0)+model*(mask>0)
             model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
 
-if True :        
+if args.per_amplifier :        
     log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
     n0=tmp_image.shape[0]//2
     n1=tmp_image.shape[1]//2

From 6f468045b9939a53a64cfebade7f2bfcbce1f848 Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Wed, 16 Aug 2017 07:14:16 -0700
Subject: [PATCH 5/8] improve masking

---
 bin/desi_compute_pixel_flatfield | 26 +++++++++++++++-----------
 1 file changed, 15 insertions(+), 11 deletions(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index 341015d43..e3d60bd34 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -295,7 +295,7 @@ parser.add_argument('--sigma', type = int, default = 60 , required = False,
                     help = "gaussian smoothing sigma (default is tuned on teststand data)")
 parser.add_argument('--npass1', type = int, default = 2 , required = False, 
                     help = "number of iterations in gaussian smoothing 1")
-parser.add_argument('--npass2', type = int, default = 6 , required = False, 
+parser.add_argument('--npass2', type = int, default = 3 , required = False, 
                     help = "number of iterations in gaussian smoothing 2")
 parser.add_argument('--nblocks', type = int, default = 4 , required = False, 
                     help = "number of blocks along one axis (total number of blocks is the square) where the trace orientations are averaged (to use in combination with --psf option otherwise ignored)")
@@ -327,7 +327,7 @@ else :
 log.info("first gaussian smoothing")
 model=gaussian_smoothing_1d_with_tilted_axes(image,sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
 
-nsloop=2
+nsloop=3 # we have to do twice the masking to remove from mask the tails around CCD defects
 for sloop in range(nsloop) :
     log.info("compute mask using a temporary flat %d/%d"%(sloop+1,nsloop))
     
@@ -362,9 +362,9 @@ for sloop in range(nsloop) :
     flat=(model>0)*image/(model+(model==0))
     flat+=(model<=0)
     flaterr=np.sqrt(var)/(model+(model==0))
-    rms=np.abs(flat-1)/flaterr
+    rms=(flat-1)/flaterr
     srms=gaussian_smoothing_2d(rms,10.)
-    mask=(srms>1.)|(rms>3.)
+    mask=(np.abs(srms)>1.)|(np.abs(rms)>4.)
     
     # remove from mask regions with important model gradient
     grad=np.zeros(model.shape)
@@ -388,13 +388,17 @@ for sloop in range(nsloop) :
     pyfits.writeto("tmp-flat-%d.fits"%sloop,flat,clobber=True)
     
     log.info("number of masked pixels=%d"%np.sum(mask.astype(int)))
-    if 1 :
-        log.info("replace masked pixels by model %d/%d"%(sloop+1,nsloop))
-        for loop in range(args.npass2) :
-            log.debug("pass %d/%d"%(loop+1,args.npass2))
-            tmp_image=image*(mask==0)+model*(mask>0)
-            model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
-
+    log.info("replace masked pixels by model %d/%d"%(sloop+1,nsloop))
+    tmp_image=image.copy()
+    
+    for loop in range(args.npass2) :
+        log.info("iteration to fill mask %d/%d"%(loop+1,args.npass2))
+        a=0.5 # a=0 is a slow iterative solution, a>0 is faster but needs tuning
+        tmp_image = image*(mask==0)+(model+a*(model-tmp_image))*(mask>0)
+        model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
+        if loop == args.npass2-1 :
+            tmp_image=image*(mask==0)+(model+a*(model-tmp_image))*(mask>0)
+    
 if args.per_amplifier :        
     log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
     n0=tmp_image.shape[0]//2

From 2f12844e1d38f0b86c75003170d4168b63b58ff7 Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Wed, 16 Aug 2017 09:36:03 -0700
Subject: [PATCH 6/8] compute ivar of median and use it in convolutions

---
 bin/desi_compute_pixel_flatfield | 163 +++++++++++++------------------
 1 file changed, 68 insertions(+), 95 deletions(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index e3d60bd34..fb44444d7 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -79,9 +79,12 @@ def median_image(image_filenames) :
     
     log.debug("first median")
     images=[]
+    ivars=[]
     for filename in image_filenames  :
         h=pyfits.open(filename)
         images.append(h[0].data)
+        ivars.append(h["IVAR"].data)
+    
     mimage=np.median(images,axis=0)
     log.debug("compute a scale per image")
     smimage2=np.sum(mimage**2)
@@ -91,10 +94,19 @@ def median_image(image_filenames) :
         if a<=0 :
             raise ValueError("scale = %f for image %s"%(a,image_filenames[i]))
         images[i] /= a
+        ivars[i] *= a**2
     mimage=np.median(images,axis=0)
-    return mimage
+    ivar=np.sum(ivars,axis=0)*(2./np.pi) # penalty factor for median
+    return mimage,ivar
+
+def convolve2d(image,k,weight=None) :
 
-def convolve2d(image,k) :
+    if weight is not None :
+        if weight.shape != image.shape :
+            raise ValueError("weight and image should have same shape")
+        sw=convolve2d(weight,k,None)
+        swim=convolve2d(weight*image,k,None)
+        return swim/(sw+(sw==0))
     
     if len(k.shape) != 2 or len(image.shape) != 2:
         raise ValueError("kernel and image should have 2 dimensions")
@@ -117,7 +129,7 @@ def convolve2d(image,k) :
     tmp[:m0,-m1:]=np.median(tmp[:m0,-m1-1])
     return scipy.signal.fftconvolve(tmp,k,"valid")
 
-def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0.,dydx=0.) :
+def gaussian_smoothing_1d_per_axis(image,ivar,sigma,npass=2,dxdy=0.,dydx=0.) :
     """Computes a smooth model of the input image using two
     1D convolution with a Gaussian kernel of parameter sigma. 
     Can do several passes.
@@ -134,6 +146,7 @@ def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0.,dydx=0.) :
     log=get_logger()
     hw=int(3*sigma)
     tmp = image.copy()
+    tmpivar = ivar.copy()
     model = np.ones(tmp.shape).astype(float)
 
     # single Gaussian profile
@@ -176,22 +189,25 @@ def gaussian_smoothing_1d_per_axis(image,sigma,npass=2,dxdy=0.,dydx=0.) :
     for p in range(npass) : # possibly do several passes
         for a in range(2) : # convolve in 1d on each axis
             #log.debug("p=%d a=%d"%(p,a))
-            res=convolve2d(tmp,kernels[a])
+            res=convolve2d(tmp,kernels[a],weight=tmpivar)
             model *= res
+            tmpivar *= res**2 # ?
+            #tmpivar *= tmp**2 # ?            
             tmp /= (res+(res==0))
+            
 
-    if 1 : # add 2D smoothing (does not help)
+    if 0 : # add 2D smoothing (does not help)
         x=np.tile((np.arange(2*hw+1)-hw)/sigma,(2*hw+1,1))
         r2=x**2+x.T**2    
         kernel2d=np.exp(-r2/2.)
         kernel2d/=np.sum(kernel2d)
-        res = convolve2d(tmp,kernel2d)
+        res = convolve2d(tmp,kernel2d,weight=tmpivar)
         model *= res
     
     return model
 
     
-def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nblocks=5) :
+def gaussian_smoothing_1d_with_tilted_axes(image,ivar,sigma,npass,x,y,dxdy,dydx,nblocks=5) :
     """Computes a smooth model of the input image using two
     1D convolution with a Gaussian kernel of parameter sigma. 
     Can do several passes.
@@ -205,7 +221,7 @@ def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nbloc
         model : 2D array image of same shape as image
     """
     if x is None or nblocks==1 :
-        return gaussian_smoothing_1d_per_axis(image,sigma,npass=npass,dxdy=0,dydx=0)
+        return gaussian_smoothing_1d_per_axis(image,ivar,sigma,npass=npass,dxdy=0,dydx=0)
     
     
     # defining blocks where trace directions are averaged
@@ -255,12 +271,12 @@ def gaussian_smoothing_1d_with_tilted_axes(image,sigma,npass,x,y,dxdy,dydx,nbloc
                     block_sigma = sigma/2.
                     log.info("Using lower sigma for edge blocks (%d,%d) = %f"%(i,j,block_sigma))
             
-            model[b0[i]:e0[i],b1[j]:e1[j]] = gaussian_smoothing_1d_per_axis(image[b0m[i]:e0m[i],b1m[j]:e1m[j]],sigma=block_sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])[b0[i]-b0m[i]:,b1[j]-b1m[j]:][:e0[i]-b0[i],:e1[j]-b1[j]]
+            model[b0[i]:e0[i],b1[j]:e1[j]] = gaussian_smoothing_1d_per_axis(image[b0m[i]:e0m[i],b1m[j]:e1m[j]],ivar[b0m[i]:e0m[i],b1m[j]:e1m[j]],sigma=block_sigma,npass=npass,dxdy=bdxdy[i,j],dydx=bdydx[i,j])[b0[i]-b0m[i]:,b1[j]-b1m[j]:][:e0[i]-b0[i],:e1[j]-b1[j]]
     
     return model
     
 
-def gaussian_smoothing_2d(image,sigma) :
+def gaussian_smoothing_2d(image,ivar,sigma) :
     """Computes a smooth model of the input image using one
     2D convolution with a 2D Gaussian kernel of parameter sigma. 
     
@@ -276,7 +292,7 @@ def gaussian_smoothing_2d(image,sigma) :
     r2=x**2+x.T**2    
     kernel=np.exp(-r2/2.)
     kernel/=np.sum(kernel)
-    return convolve2d(image,kernel)
+    return convolve2d(image,kernel,weight=ivar)
 
     
 parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter,
@@ -292,11 +308,17 @@ parser.add_argument('-i','--images', type = str, nargs='*', default = None, requ
 parser.add_argument('-o','--outfile', type = str, default = None, required = True,
                     help = 'output flatfield image filename')
 parser.add_argument('--sigma', type = int, default = 60 , required = False, 
-                    help = "gaussian smoothing sigma (default is tuned on teststand data)")
-parser.add_argument('--npass1', type = int, default = 2 , required = False, 
-                    help = "number of iterations in gaussian smoothing 1")
-parser.add_argument('--npass2', type = int, default = 3 , required = False, 
-                    help = "number of iterations in gaussian smoothing 2")
+                    help = "gaussian filtering sigma")
+parser.add_argument('--niter_filter', type = int, default = 2 , required = False, 
+                    help = "number of iterations in gaussian filtering")
+parser.add_argument('--niter-mask', type = int, default = 3 , required = False, 
+                    help = "number of iterations for mask evaluation")
+parser.add_argument('--sigma-mask', type = int, default = 10 , required = False, 
+                    help = "gaussian smoothing sigma for the mask")
+parser.add_argument('--nsig-mask', type = float, default = 4.0 , required = False, 
+                    help = "# sigma cut for mask")
+parser.add_argument('--nsig-smooth-mask', type = float, default = 1.5 , required = False, 
+                    help = "# sigma cut for mask after smoothing")
 parser.add_argument('--nblocks', type = int, default = 4 , required = False, 
                     help = "number of blocks along one axis (total number of blocks is the square) where the trace orientations are averaged (to use in combination with --psf option otherwise ignored)")
 parser.add_argument('--out-median', type = str, default = None , required = False, help = "save output median image (for development)")
@@ -308,10 +330,12 @@ log = get_logger()
 
 if len(args.images) == 1 :
     log.info("read a single image")
-    image=pyfits.open(args.images[0])[0].data
+    h=pyfits.open(args.images[0])
+    image=h[0].data
+    ivar=h["IVAR"].data
 else :
     log.info("compute a median of the input images")
-    image=median_image(args.images)
+    image,ivar=median_image(args.images)
 
 if args.psf :
     log.info("get trace coordinates from psf %s"%args.psf)
@@ -322,99 +346,47 @@ else :
     dxdy=None
     dydx=None
 
+log.warning("need to add trim option")
+
 
 
 log.info("first gaussian smoothing")
-model=gaussian_smoothing_1d_with_tilted_axes(image,sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
-
-nsloop=3 # we have to do twice the masking to remove from mask the tails around CCD defects
-for sloop in range(nsloop) :
-    log.info("compute mask using a temporary flat %d/%d"%(sloop+1,nsloop))
-    
-    log.info("estimate variance")
-    r_model=model.ravel()
-    ok=(r_model>-100)&(r_model<5000)
-    h0,bins=np.histogram(r_model[ok],bins=15)
-    hx,junk=np.histogram(r_model[ok],weights=r_model[ok],bins=bins)
-    valid=np.where(h0>50)[0]# cut on min number of pixels per flux bin
-    if valid.size < 4 :
-        log.warning("cannot measure variance, use guess")
-        var = 1.+model*(model>0)
-    else :
-        flux=hx[valid]/h0[valid]
-        var_of_flux=np.zeros(valid.size)
-        r_res=image.ravel()-r_model
-        for i,j in enumerate(valid) :
-            var_of_flux[i] = (1.48*np.median(np.abs(r_res[(r_model>=bins[j])&(r_model<bins[j+1])])))**2
-        c=np.polyfit(flux,var_of_flux,1)
-        var = c[0]*model+c[1]
-        log.info("var = %f + %f*flux"%(c[1],c[0]))
-        if c[1]<1. :
-            c[1]=1
-            log.warning("forcing var(flux=0)=%f"%c[1])
-        #import matplotlib.pyplot as plt
-        #pol=np.poly1d(c)
-        #err = var_of_flux*np.sqrt(2./h0[valid])
-        #plt.errorbar(flux,var_of_flux,err,fmt="o")
-        #plt.plot(flux,pol(flux),"-")
-        #plt.show()
-    
+model=gaussian_smoothing_1d_with_tilted_axes(image,ivar,sigma=args.sigma,npass=args.niter_filter,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
+
+# we have to do several times the masking to remove from mask the tails around CCD defects
+for sloop in range(args.niter_mask) :
+    log.info("compute mask using a temporary flat and refit %d/%d"%(sloop+1,args.niter_mask))
     flat=(model>0)*image/(model+(model==0))
     flat+=(model<=0)
-    flaterr=np.sqrt(var)/(model+(model==0))
-    rms=(flat-1)/flaterr
-    srms=gaussian_smoothing_2d(rms,10.)
-    mask=(np.abs(srms)>1.)|(np.abs(rms)>4.)
-    
-    # remove from mask regions with important model gradient
-    grad=np.zeros(model.shape)
-    grad[1:-1]=np.abs(model[2:]-model[:-2])
-    grad[:,1:-1]+=np.abs(model[:,2:]-model[:,:-2])
-    grad/=(model+(model==0))
-    mask[grad>0.2]=0
-    
-    # do not mask edges
-    mask[0,:]=0 ; mask[:,0]=1 ; mask[-1,:]=1 ; mask[:,-1]=1
+    rms=(flat-1)*(np.sqrt(ivar)*model*(model>0))
+    srms=gaussian_smoothing_2d(rms,ivar=None,sigma=args.sigma_mask)
+    mask=(np.abs(srms)>args.nsig_smooth_mask)|(np.abs(rms)>args.nsig_mask)
     if args.per_amplifier :
-        log.info("do not mask amp. boundaries")
+        log.info("do not mask amp. boundaries to get a correct stitching")
         mask[image.shape[0]//2-2:image.shape[0]//2+3]=0
         mask[:,image.shape[1]//2-2:image.shape[1]//2+3]=0
-    
-    
-    
+    if sloop<(args.niter_mask-1) :
+        model=gaussian_smoothing_1d_with_tilted_axes(image,ivar=ivar*(mask==0),sigma=args.sigma,npass=1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
+ 
+mivar=(mask==0)*ivar   
 
-    # DEBUG
-    pyfits.writeto("mask-%d.fits"%sloop,mask.astype(int),clobber=True)
-    pyfits.writeto("tmp-flat-%d.fits"%sloop,flat,clobber=True)
-    
-    log.info("number of masked pixels=%d"%np.sum(mask.astype(int)))
-    log.info("replace masked pixels by model %d/%d"%(sloop+1,nsloop))
-    tmp_image=image.copy()
-    
-    for loop in range(args.npass2) :
-        log.info("iteration to fill mask %d/%d"%(loop+1,args.npass2))
-        a=0.5 # a=0 is a slow iterative solution, a>0 is faster but needs tuning
-        tmp_image = image*(mask==0)+(model+a*(model-tmp_image))*(mask>0)
-        model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
-        if loop == args.npass2-1 :
-            tmp_image=image*(mask==0)+(model+a*(model-tmp_image))*(mask>0)
-    
 if args.per_amplifier :        
     log.info("last gaussian smoothing, per quadrant, to remove residual gain difference")
-    n0=tmp_image.shape[0]//2
-    n1=tmp_image.shape[1]//2
-    model=np.zeros(tmp_image.shape)
+    n0=image.shape[0]//2
+    n1=image.shape[1]//2
+    model=np.zeros(image.shape)
     if args.nblocks == 1 :
         nblocks=1
     else :
         nblocks=args.nblocks//2
-    model[:n0,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,:n1],sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
-    model[n0:,:n1]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,:n1],sigma=args.sigma,npass=args.npass1,x=x,y=y-n0,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
-    model[n0:,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[n0:,n1:],sigma=args.sigma,npass=args.npass1,x=x-n1,y=y-n0,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
-    model[:n0,n1:]=gaussian_smoothing_1d_with_tilted_axes(tmp_image[:n0,n1:],sigma=args.sigma,npass=args.npass1,x=x-n1,y=y,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    
+    model[:n0,:n1]=gaussian_smoothing_1d_with_tilted_axes(image[:n0,:n1],mivar[:n0,:n1],sigma=args.sigma,npass=args.niter_filter,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    model[n0:,:n1]=gaussian_smoothing_1d_with_tilted_axes(image[n0:,:n1],mivar[n0:,:n1],sigma=args.sigma,npass=args.niter_filter,x=x,y=y-n0,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    model[n0:,n1:]=gaussian_smoothing_1d_with_tilted_axes(image[n0:,n1:],mivar[n0:,n1:],sigma=args.sigma,npass=args.niter_filter,x=x-n1,y=y-n0,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
+    model[:n0,n1:]=gaussian_smoothing_1d_with_tilted_axes(image[:n0,n1:],mivar[:n0,n1:],sigma=args.sigma,npass=args.niter_filter,x=x-n1,y=y,dxdy=dxdy,dydx=dydx,nblocks=nblocks)
 else :
     log.info("last gaussian smoothing (not per quadrant)")
-    model=gaussian_smoothing_1d_with_tilted_axes(tmp_image,sigma=args.sigma,npass=args.npass1,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
+    model=gaussian_smoothing_1d_with_tilted_axes(image,mivar,sigma=args.sigma,npass=args.niter_filter,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
     
 
 
@@ -431,7 +403,8 @@ h.writeto(args.outfile,clobber=True)
 
 if args.out_median is not None :
     log.info("writing median image %s ..."%args.out_median)
-    pyfits.writeto(args.out_median,image,clobber=True)
+    h=pyfits.HDUList([pyfits.PrimaryHDU(image),pyfits.ImageHDU(ivar,name="IVAR")])
+    h.writeto(args.out_median,clobber=True)
 
 log.info("done")
 

From 1c1484305b8391b7d5904b14a3084e1b276d1124 Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Wed, 16 Aug 2017 09:47:38 -0700
Subject: [PATCH 7/8] add/update docstring

---
 bin/desi_compute_pixel_flatfield | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index fb44444d7..05f8aab36 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -11,6 +11,17 @@ import specter.psf
 from desiutil.log import get_logger
 
 def grid_from_psf(filename) :
+    """ Return a list of spots coordinates along traces and trace orientations
+    
+    Args:
+        filename : a fits file withe a specter PSF 
+
+    Returns:
+        x     : 1D array, x_ccd coordinate of spots (fiber axis, axis=1 in np.array)
+        y     : 1D array, y_ccd coordinate of spots (wavelength axis, axis=0 in np.array)
+        dx/dy : 1D array, dx/dy for curves of constant fiber coordinate
+        dy/dx : 1D array, dy/dx for curves of constant wavelength        
+    """
     try :
         psftype=pyfits.open(filename)[0].header["PSFTYPE"]
     except KeyError :
@@ -75,6 +86,7 @@ def median_image(image_filenames) :
 
     Returns:
         mimage : median image (2D np.array)
+        ivar   : ivar of median
     """
     
     log.debug("first median")
@@ -100,7 +112,16 @@ def median_image(image_filenames) :
     return mimage,ivar
 
 def convolve2d(image,k,weight=None) :
-
+""" Return a 2D convolution of image with kernel k, optionally with a weight image
+    
+    Args:
+        image : 2D np.array image
+        k : 2D np.array kernel, each dimension must be odd and greater than 1
+    Options:
+        weight : 2D np.array of same shape as image
+    Returns:
+        cimage : 2D np.array convolved image of same shape as input image
+    """
     if weight is not None :
         if weight.shape != image.shape :
             raise ValueError("weight and image should have same shape")

From 115c691bb56649c949260ef374791aec6dac0e3b Mon Sep 17 00:00:00 2001
From: Julien Guy <jguy@lbl.gov>
Date: Wed, 16 Aug 2017 10:19:18 -0700
Subject: [PATCH 8/8] add rectangular triming (could be refined)

---
 bin/desi_compute_pixel_flatfield | 48 +++++++++++++++++++++++++++-----
 1 file changed, 41 insertions(+), 7 deletions(-)

diff --git a/bin/desi_compute_pixel_flatfield b/bin/desi_compute_pixel_flatfield
index 05f8aab36..3bf30c0d5 100755
--- a/bin/desi_compute_pixel_flatfield
+++ b/bin/desi_compute_pixel_flatfield
@@ -42,7 +42,7 @@ def grid_from_psf(filename) :
     y=[]
     f=[]
     w=[]
-    wstep=200.
+    wstep=50.
     waves=np.linspace(psf.wmin,psf.wmax,int((psf.wmax-psf.wmin)/wstep))
     fibers=np.arange(psf.nspec)
     for fiber in fibers :
@@ -112,7 +112,7 @@ def median_image(image_filenames) :
     return mimage,ivar
 
 def convolve2d(image,k,weight=None) :
-""" Return a 2D convolution of image with kernel k, optionally with a weight image
+    """ Return a 2D convolution of image with kernel k, optionally with a weight image
     
     Args:
         image : 2D np.array image
@@ -345,6 +345,7 @@ parser.add_argument('--nblocks', type = int, default = 4 , required = False,
 parser.add_argument('--out-median', type = str, default = None , required = False, help = "save output median image (for development)")
 parser.add_argument('--psf', type = str, default = None , required = False, help = "use traces in this PSF to orient 1D convolutions")
 parser.add_argument('--per-amplifier', action = 'store_true', default = None , required = False, help = "solve model per amplifier if gains are uncertain or non-linarities")
+parser.add_argument('--no-trim', action = 'store_true', default = None , required = False, help = "do not compute flat only in CCD area covered by traces (to use in combination with --psf option otherwise ignored)")
 
 args        = parser.parse_args()
 log = get_logger()
@@ -367,9 +368,28 @@ else :
     dxdy=None
     dydx=None
 
-log.warning("need to add trim option")
-
-
+xmin=0
+xmax=image.shape[1]
+ymin=0
+ymax=image.shape[0]
+
+original_image_shape=image.shape
+if args.psf and ( not args.no_trim ) :
+    xmin=int(np.min(x))
+    xmax=int(np.max(x))+1
+    xmin-=10
+    xmax+=10
+    xmin=max(0,xmin)
+    xmax=min(xmax,image.shape[1])
+    ymin=int(np.min(y))
+    ymax=int(np.max(y))+1
+    ymin-=10
+    ymax+=10
+    ymin=max(0,ymin)
+    ymax=min(ymax,image.shape[0])
+    image=image[ymin:ymax,xmin:xmax]
+    ivar=ivar[ymin:ymax,xmin:xmax]
+    log.info("trimed image %s -> %s"%(str(original_image_shape),str(image.shape)))
 
 log.info("first gaussian smoothing")
 model=gaussian_smoothing_1d_with_tilted_axes(image,ivar,sigma=args.sigma,npass=args.niter_filter,x=x,y=y,dxdy=dxdy,dydx=dydx,nblocks=args.nblocks)
@@ -414,8 +434,22 @@ else :
 
 flat=(model>0)*image/(model+(model==0))
 flat+=(model<=0)
-
-
+if args.psf and (not args.no_trim ):
+    log.info("restore image size after triming")
+    
+    tmp_flat=np.ones(original_image_shape)
+    tmp_flat[ymin:ymax,xmin:xmax]=flat    
+    flat=tmp_flat
+    
+    tmp_model=np.zeros(original_image_shape)
+    tmp_model[ymin:ymax,xmin:xmax]=model    
+    model=tmp_model
+    
+    mask += (ivar==0)*(mask==0)
+    tmp_mask=np.ones(original_image_shape)
+    tmp_mask[ymin:ymax,xmin:xmax]=mask
+    mask=tmp_mask
+    
 log.info("writing %s ..."%args.outfile)
 h=pyfits.HDUList([pyfits.PrimaryHDU(flat),pyfits.ImageHDU(model,name="MODEL"),pyfits.ImageHDU(mask.astype(int),name="MODMASK")])
 h[0].header["EXTNAME"]="FLAT"