From 31710adc736760c8a38824c447cddc595911d82d Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 03:02:07 -0700
Subject: [PATCH 01/18] move reference spectra transforms into dedicated
 function. Also use different psfs to convolve the spectrum

---
 py/desispec/trace_shifts.py | 176 +++++++++++++++++++++---------------
 1 file changed, 105 insertions(+), 71 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 4ec99fe05..6d773434b 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -618,6 +618,84 @@ def compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image
 
     return ox,oy,odx,oex,of,ol
 
+def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
+    log = get_logger()
+    
+    # trim ref_spectrum
+    i = (ref_wave >= wave[0]) & (ref_wave <= wave[-1])
+    ref_wave = ref_wave[i]
+    ref_spectrum = ref_spectrum[i]
+
+    # check wave is linear or make it linear
+    if (np.abs((ref_wave[1] - ref_wave[0]) - (ref_wave[-1] - ref_wave[-2])) >
+    0.0001 * (ref_wave[1]-ref_wave[0])):
+        log.info("reference spectrum wavelength is not on a linear grid, resample it")
+        dwave = np.min(np.gradient(ref_wave))
+        tmp_wave = np.linspace(ref_wave[0], ref_wave[-1],
+                               int((ref_wave[-1]-ref_wave[0])/dwave))
+        ref_spectrum = resample_flux(tmp_wave, ref_wave, ref_spectrum)
+        ref_wave = tmp_wave
+    
+    n_wave_bins = 1
+    fiber_step = max(nfibers//2, 1)
+    wave_bins = np.linspace(wave[0], wave[-1], n_wave_bins+1)
+    wave_bins = wave_bins[:-1] + .5 * np.diff(wave_bins)
+    spectra = []
+    for central_wave in wave_bins:
+        i = np.searchsorted(ref_wave, central_wave)
+        log.warning(' xx %s %s'%(central_wave,i))
+        angstrom_hwidth = 3 # psf half width
+        dwave = ref_wave[i+1] - ref_wave[i]
+        hw = int(angstrom_hwidth / dwave) + 1
+        wave_range = ref_wave[i - hw:i + hw + 1]
+        kernels = []
+        #for fiber in range(0, nfibers, fiber_step):
+        for fiber in [250]:
+            # compute psf at most significant line of ref_spectrum
+            x, y = psf.xy(fiber, wave_range)
+            # x = np.tile(x[hw] + np.linspace(-hw, hw, 2 * hw + 1),
+            #             (y.size, 1))
+            # x = x + np.linspace(-hw,hw,2*hw+1)
+            # x = x[:,None] + np.linspace(-hw, hw, 2*hw+1)[None,:]
+            x=np.tile(x[hw]+np.arange(-hw,hw+1)*(y[-1]-y[0])/(2*hw+1),(y.size,1))
+            y = np.tile(y, (2 * hw + 1, 1)).T
+
+            kernel2d = psf._value(x, y, fiber, central_wave)
+            kernel1d = np.sum(kernel2d, axis=1)
+            kernels.append(kernel1d)
+        kernels = np.mean(kernels, axis=0)
+        
+        ref_spectrum = fftconvolve(ref_spectrum, kernels, mode='same')
+        spectra.append(ref_spectrum)
+    log.info("convolve reference spectrum using PSF")
+    spectra = np.array(spectra)
+    spectra_pos = np.searchsorted(wave_bins, ref_wave)
+    result = ref_spectrum * 0.
+    ''' Edges of the spectrum '''
+    result[spectra_pos == 0] = spectra[0][spectra_pos==0]
+    result[spectra_pos == n_wave_bins] = spectra[-1][spectra_pos == n_wave_bins]
+    inside = (spectra_pos > 0) & (spectra_pos < n_wave_bins)
+    weight = (ref_wave[inside]-wave_bins[spectra_pos[inside]-1])/(
+        wave_bins[spectra_pos[inside]]-wave_bins[spectra_pos[inside]-1])
+    ''' weight is zero on left edge one on right'''
+    result[inside] = (1-weight) * spectra[spectra_pos[inside] - 1, inside]+(
+        weight * spectra[spectra_pos[inside],inside])
+    assert len(ref_wave)==len(result)
+    ref_spectrum = result
+    
+    # resample input spectrum
+    log.info("resample convolved reference spectrum")
+    ref_spectrum = resample_flux(wave, ref_wave , ref_spectrum)
+
+    log.info("absorb difference of calibration")
+    x = (wave - wave[wave.size//2]) / 50.
+    kernel = np.exp(- x**2/2)
+    f1 = fftconvolve(mflux, kernel, mode='same')
+    f2 = fftconvolve(ref_spectrum, kernel, mode='same')
+    # We scale by a constant factor
+    scale = (f1 * f2).sum() / (f2 * f2).sum()
+    ref_spectrum *= scale
+    return ref_wave, ref_spectrum
 
 def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
                                         spectrum_filename, degyy=2, width=7,
@@ -698,86 +776,42 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
     mivar = np.minimum(mivar, 1. / mad_factor**2 / np.median(np.abs(mflux))**2)
     # do not allow negatives
     mflux[mflux <  0] = 0
-
-
-    # trim ref_spectrum
-    i=(ref_wave>=wave[0])&(ref_wave<=wave[-1])
-    ref_wave=ref_wave[i]
-    ref_spectrum=ref_spectrum[i]
-
-    # check wave is linear or make it linear
-    if np.abs((ref_wave[1]-ref_wave[0])-(ref_wave[-1]-ref_wave[-2]))>0.0001*(ref_wave[1]-ref_wave[0]) :
-        log.info("reference spectrum wavelength is not on a linear grid, resample it")
-        dwave = np.min(np.gradient(ref_wave))
-        tmp_wave = np.linspace(ref_wave[0],ref_wave[-1],int((ref_wave[-1]-ref_wave[0])/dwave))
-        ref_spectrum = resample_flux(tmp_wave, ref_wave , ref_spectrum)
-        ref_wave = tmp_wave
-
-    i=np.argmax(ref_spectrum)
-    central_wave_for_psf_evaluation  = ref_wave[i]
-    fiber_for_psf_evaluation = (flux.shape[0]//2)
-    try :
-        # compute psf at most significant line of ref_spectrum
-        dwave=ref_wave[i+1]-ref_wave[i]
-        hw=int(3./dwave)+1 # 3A half width
-        wave_range = ref_wave[i-hw:i+hw+1]
-        x,y=psf.xy(fiber_for_psf_evaluation,wave_range)
-        x=np.tile(x[hw]+np.arange(-hw,hw+1)*(y[-1]-y[0])/(2*hw+1),(y.size,1))
-        y=np.tile(y,(2*hw+1,1)).T
-        kernel2d=psf._value(x,y,fiber_for_psf_evaluation,central_wave_for_psf_evaluation)
-        kernel1d=np.sum(kernel2d,axis=1)
-        log.info("convolve reference spectrum using PSF at fiber %d and wavelength %dA"%(fiber_for_psf_evaluation,central_wave_for_psf_evaluation))
-        ref_spectrum=fftconvolve(ref_spectrum,kernel1d, mode='same')
-    except :
-        log.warning("couldn't convolve reference spectrum: %s %s"%(sys.exc_info()[0],sys.exc_info()[1]))
-
-
-
-    # resample input spectrum
-    log.info("resample convolved reference spectrum")
-    ref_spectrum = resample_flux(wave, ref_wave , ref_spectrum)
-
-    log.info("absorb difference of calibration")
-    x = (wave - wave[wave.size//2]) / 50.
-    kernel = np.exp(- x**2/2)
-    f1 = fftconvolve(mflux, kernel, mode='same')
-    f2 = fftconvolve(ref_spectrum, kernel, mode='same')
-    # We scale by a constant factor
-    scale = (f1 * f2).sum() / (f2 * f2).sum()
-    ref_spectrum *= scale
     
+    ref_wave, ref_spectrum = _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, len(ivar))
+
     log.info("fit shifts on wavelength bins")
     # define bins
-    n_wavelength_bins = degyy+4
-    y_for_dy=np.array([])
-    dy=np.array([])
-    ey=np.array([])
-    wave_for_dy=np.array([])
-
+    n_wavelength_bins = degyy + 4
+    y_for_dy = np.array([])
+    dy = np.array([])
+    ey = np.array([])
+    wave_for_dy = np.array([])
+    fiber_for_psf_evaluation = flux.shape[0] //2
+    wavelength_bins = np.linspace(wave[0], wave[-1], n_wavelength_bins+1)
     for b in range(n_wavelength_bins) :
-        wmin=wave[0]+((wave[-1]-wave[0])/n_wavelength_bins)*b
-        if b<n_wavelength_bins-1 :
-            wmax=wave[0]+((wave[-1]-wave[0])/n_wavelength_bins)*(b+1)
-        else :
-            wmax=wave[-1]
-        ok=(wave>=wmin)&(wave<=wmax)
-        sw= np.sum(mflux[ok]*(mflux[ok]>0))
-        if sw==0 :
+        wmin, wmax = [wavelength_bins[_] for _ in [b, b + 1]]
+        ok = (wave >= wmin) & (wave <= wmax)
+        flux_weight = np.sum(mflux[ok] * (mflux[ok] > 0))
+        log.warning("%s %s "%(b, flux_weight))
+        if flux_weight == 0 :
             continue
         dwave,err = compute_dy_from_spectral_cross_correlation(mflux[ok],
                 wave[ok], ref_spectrum[ok], ivar=mivar[ok], hw=10.,
                 prior_width_dy=prior_width_dy)
-        bin_wave  = np.sum(mflux[ok]*(mflux[ok]>0)*wave[ok])/sw
-        x,y=psf.xy(fiber_for_psf_evaluation,bin_wave)
-        eps=0.1
-        x,yp=psf.xy(fiber_for_psf_evaluation,bin_wave+eps)
-        dydw=(yp-y)/eps
-        if err*dydw<1 :
-            dy=np.append(dy,-dwave*dydw)
-            ey=np.append(ey,err*dydw)
-            wave_for_dy=np.append(wave_for_dy,bin_wave)
+        bin_wave  = np.sum(mflux[ok] * (mflux[ok] > 0) * wave[ok]) / flux_weight
+        # flux weighted wavelength of the center
+        # Computing the derivative dy/dwavelength
+        x, y = psf.xy(fiber_for_psf_evaluation, bin_wave)
+        eps =  0.1
+        x, yp = psf.xy(fiber_for_psf_evaluation, bin_wave+eps)
+        dydw = (yp - y) / eps 
+        log.warning("%s %s %s"%(dwave, err, dydw))
+        if err * dydw < 1 :
+            dy = np.append(dy, -dwave * dydw)
+            ey = np.append(ey, err*dydw)
+            wave_for_dy = np.append(wave_for_dy,bin_wave)
             y_for_dy=np.append(y_for_dy,y)
-            log.info("wave = %fA , y=%d, measured dwave = %f +- %f A"%(bin_wave,y,dwave,err))
+            log.info("wave = %fA , y=%d, measured dwave = %f +- %f A"%(bin_wave, y, dwave, err))
 
     if False : # we don't need this for now
         try :

From b35478e5df7bd3b4e43c3ef6a0fa682f197f8737 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 07:24:58 -0700
Subject: [PATCH 02/18] additional updates

---
 py/desispec/trace_shifts.py | 63 ++++++++++++++++++++-----------------
 1 file changed, 35 insertions(+), 28 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 6d773434b..fb749b32d 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -619,12 +619,16 @@ def compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image
     return ox,oy,odx,oex,of,ol
 
 def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
+    """
+    Prepare the reference spectrum to be used for wavelegth offset 
+    determination
+    """
     log = get_logger()
     
     # trim ref_spectrum
-    i = (ref_wave >= wave[0]) & (ref_wave <= wave[-1])
-    ref_wave = ref_wave[i]
-    ref_spectrum = ref_spectrum[i]
+    subset = (ref_wave >= wave[0]) & (ref_wave <= wave[-1])
+    ref_wave = ref_wave[subset]
+    ref_spectrum = ref_spectrum[subset]
 
     # check wave is linear or make it linear
     if (np.abs((ref_wave[1] - ref_wave[0]) - (ref_wave[-1] - ref_wave[-2])) >
@@ -636,53 +640,56 @@ def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
         ref_spectrum = resample_flux(tmp_wave, ref_wave, ref_spectrum)
         ref_wave = tmp_wave
     
-    n_wave_bins = 1
-    fiber_step = max(nfibers//2, 1)
+    n_wave_bins = 20 # how many points along wavelength to use to get psf
+    n_representative_fibers = 20 # how many fibers to use to get psf
+    fiber_list = np.unique(np.linspace(0, nfibers - 1,
+                                       n_representative_fibers).astype(int))
     wave_bins = np.linspace(wave[0], wave[-1], n_wave_bins+1)
     wave_bins = wave_bins[:-1] + .5 * np.diff(wave_bins)
     spectra = []
-    for central_wave in wave_bins:
-        i = np.searchsorted(ref_wave, central_wave)
-        log.warning(' xx %s %s'%(central_wave,i))
-        angstrom_hwidth = 3 # psf half width
-        dwave = ref_wave[i+1] - ref_wave[i]
+    ref_spectrum0 = ref_spectrum * 1
+    # original before convolution
+    angstrom_hwidth = 3 # psf half width
+    
+    for central_wave0 in wave_bins:
+        ipos = np.searchsorted(ref_wave, central_wave0)
+        central_wave = ref_wave[ipos] # actual value from the grid
+        dwave = ref_wave[ipos + 1] - ref_wave[ipos]
         hw = int(angstrom_hwidth / dwave) + 1
-        wave_range = ref_wave[i - hw:i + hw + 1]
+        wave_range = ref_wave[ipos - hw:ipos + hw + 1]
         kernels = []
-        #for fiber in range(0, nfibers, fiber_step):
-        for fiber in [250]:
-            # compute psf at most significant line of ref_spectrum
+        for fiber in fiber_list:
             x, y = psf.xy(fiber, wave_range)
-            # x = np.tile(x[hw] + np.linspace(-hw, hw, 2 * hw + 1),
-            #             (y.size, 1))
-            # x = x + np.linspace(-hw,hw,2*hw+1)
-            # x = x[:,None] + np.linspace(-hw, hw, 2*hw+1)[None,:]
-            x=np.tile(x[hw]+np.arange(-hw,hw+1)*(y[-1]-y[0])/(2*hw+1),(y.size,1))
+
+            x = x[:,None] + np.linspace(-hw, hw, 2*hw+1)[None,:]
+            # original code below but I don't understand y[-1]-y[0] part
+            # x=np.tile(x[hw]+np.arange(-hw,hw+1)*(y[-1]-y[0])/(2*hw+1),(y.size,1))
             y = np.tile(y, (2 * hw + 1, 1)).T
 
             kernel2d = psf._value(x, y, fiber, central_wave)
             kernel1d = np.sum(kernel2d, axis=1)
             kernels.append(kernel1d)
         kernels = np.mean(kernels, axis=0)
-        
-        ref_spectrum = fftconvolve(ref_spectrum, kernels, mode='same')
+        # average across fibers
+        ref_spectrum = fftconvolve(ref_spectrum0, kernels, mode='same')
         spectra.append(ref_spectrum)
     log.info("convolve reference spectrum using PSF")
     spectra = np.array(spectra)
     spectra_pos = np.searchsorted(wave_bins, ref_wave)
     result = ref_spectrum * 0.
-    ''' Edges of the spectrum '''
+    # Edges of the spectrum
     result[spectra_pos == 0] = spectra[0][spectra_pos==0]
     result[spectra_pos == n_wave_bins] = spectra[-1][spectra_pos == n_wave_bins]
     inside = (spectra_pos > 0) & (spectra_pos < n_wave_bins)
-    weight = (ref_wave[inside]-wave_bins[spectra_pos[inside]-1])/(
-        wave_bins[spectra_pos[inside]]-wave_bins[spectra_pos[inside]-1])
-    ''' weight is zero on left edge one on right'''
-    result[inside] = (1-weight) * spectra[spectra_pos[inside] - 1, inside]+(
+    weight = (ref_wave[inside] - wave_bins[spectra_pos[inside] - 1])/(
+        wave_bins[spectra_pos[inside]] - wave_bins[spectra_pos[inside] - 1])
+    # weight is zero on left edge one on right
+    
+    # linearly stitching the spectra convolved with different kernel
+    result[inside] = (1 - weight) * spectra[spectra_pos[inside] - 1, inside]+(
         weight * spectra[spectra_pos[inside],inside])
-    assert len(ref_wave)==len(result)
     ref_spectrum = result
-    
+
     # resample input spectrum
     log.info("resample convolved reference spectrum")
     ref_spectrum = resample_flux(wave, ref_wave , ref_spectrum)

From 61c241d93848798d3c26dbd5e3f2cf8390f6a86f Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 08:29:48 -0700
Subject: [PATCH 03/18] move code into separate function to be reused

---
 py/desispec/trace_shifts.py | 70 ++++++++++++++++++++-----------------
 1 file changed, 37 insertions(+), 33 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index fb749b32d..63916cfa5 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -704,6 +704,42 @@ def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
     ref_spectrum *= scale
     return ref_wave, ref_spectrum
 
+
+def _continuum_subtract_median(flux0, ivar, continuum_win = 17):
+    # here we get rid of continuum by applying a median filter 
+    continuum_foot = np.abs(np.arange(-continuum_win,continuum_win))>continuum_win /2.
+    flux = flux0 * 1 # we will modify flux
+    # we only keep emission lines and get rid of continuum
+    for ii in range(flux.shape[0]):
+        flux[ii] = flux[ii] - median_filter(flux[ii], footprint=continuum_foot)
+
+    # boolean mask of fibers with good data
+    good_fibers = (np.sum(ivar>0, axis=1) > 0)
+    num_good_fibers = np.sum(good_fibers)
+
+    # median flux used as internal spectral reference
+    mflux = np.median(flux[good_fibers], axis=0)
+
+    # we use data variance and MAD from different spectra
+    # to assign variance to a spectrum (1.48 is MAD factor,
+    # pi/2 is a factor from Stddev[median(N(0,1))]
+    mad_factor = 1.48
+    mad = np.maximum(np.median(np.abs(flux[good_fibers] - mflux[None, :]),
+                               axis=0), 1e-100)
+    # I prevent it from being zero to avoid the warning below
+    # The exact value does not matter as we're comparing to actual
+    # median(ivar)
+    mivar = np.minimum(
+        np.median(ivar[good_fibers], axis=0) ,
+        1./mad_factor**2 / mad**2) * num_good_fibers * (2. / np.pi)
+    # finally use use the MAD of the background subtracted spectra to
+    # assign further variance limit
+    # this is sort of "effective" noise in the continuum subtracted spectrum
+    mivar = np.minimum(mivar, 1. / mad_factor**2 / np.median(np.abs(mflux))**2)
+    # do not allow negatives
+    mflux[mflux <  0] = 0
+    return mflux, mivar, flux
+
 def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
                                         spectrum_filename, degyy=2, width=7,
                                         prior_width_dy=0.1):
@@ -750,40 +786,8 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
 
     flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=2)
 
-    # here we get rid of continuum by applying a median filter 
-    continuum_win = 17
-    continuum_foot = np.abs(np.arange(-continuum_win,continuum_win))>continuum_win /2.
-
-    # we only keep emission lines and get rid of continuum
-    for ii in range(flux.shape[0]):
-        flux[ii] = flux[ii] - median_filter(flux[ii], footprint=continuum_foot)
-
-    # boolean mask of fibers with good data
-    good_fibers = (np.sum(ivar>0, axis=1) > 0)
-    num_good_fibers = np.sum(good_fibers)
-
-    # median flux used as internal spectral reference
-    mflux = np.median(flux[good_fibers], axis=0)
+    mflux, mivar, flux = _continuum_subtract_median(flux, ivar, continuum_win = 17)
 
-    # we use data variance and MAD from different spectra
-    # to assign variance to a spectrum (1.48 is MAD factor,
-    # pi/2 is a factor from Stddev[median(N(0,1))]
-    mad_factor = 1.48
-    mad = np.maximum(np.median(np.abs(flux[good_fibers] - mflux[None, :]),
-                               axis=0), 1e-100)
-    # I prevent it from being zero to avoid the warning below
-    # The exact value does not matter as we're comparing to actual
-    # median(ivar)
-    mivar = np.minimum(
-        np.median(ivar[good_fibers], axis=0) ,
-        1./mad_factor**2 / mad**2) * num_good_fibers * (2. / np.pi)
-    # finally use use the MAD of the background subtracted spectra to
-    # assign further variance limit
-    # this is sort of "effective" noise in the continuum subtracted spectrum
-    mivar = np.minimum(mivar, 1. / mad_factor**2 / np.median(np.abs(mflux))**2)
-    # do not allow negatives
-    mflux[mflux <  0] = 0
-    
     ref_wave, ref_spectrum = _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, len(ivar))
 
     log.info("fit shifts on wavelength bins")

From 5a1b754a3a902a647413a8cac3469a7fd69eaf60 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 08:36:13 -0700
Subject: [PATCH 04/18] log the internal offsets

---
 py/desispec/trace_shifts.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 63916cfa5..ce03315ab 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -363,13 +363,15 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
             sw=np.sum(ivar[fiber,ok]*flux[fiber,ok]*(flux[fiber,ok]>0))
             if sw<=0 :
                 continue
+            block_wave = np.sum(ivar[fiber,ok]*flux[fiber,ok]*(flux[fiber,ok]>0)*wave[ok])/sw
 
             dwave,err = compute_dy_from_spectral_cross_correlation(flux[fiber,ok],wave[ok],reference_flux[ok],ivar=ivar[fiber,ok]*reference_flux[ok],hw=3., calibrate=True)
+            if fiber %10==0 :
+                log.info("Wavelength offset %f for fiber #%03d at wave %f "%(dwave, fiber, block_wave))
 
             if err > 1 :
                 continue
 
-            block_wave = np.sum(ivar[fiber,ok]*flux[fiber,ok]*(flux[fiber,ok]>0)*wave[ok])/sw
             rw = legx(block_wave,wavemin,wavemax)
             tx = legval(rw,xcoef[fiber])
             ty = legval(rw,ycoef[fiber])

From 7f500c506e6054524f317363da4bf3f0e6e6147a Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 09:13:46 -0700
Subject: [PATCH 05/18] subtract continuum when doing internal calirbation

---
 py/desispec/trace_shifts.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index ce03315ab..cc46ff767 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -367,7 +367,7 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
 
             dwave,err = compute_dy_from_spectral_cross_correlation(flux[fiber,ok],wave[ok],reference_flux[ok],ivar=ivar[fiber,ok]*reference_flux[ok],hw=3., calibrate=True)
             if fiber %10==0 :
-                log.info("Wavelength offset %f for fiber #%03d at wave %f "%(dwave, fiber, block_wave))
+                log.info("Wavelength offset %f +/- %f for fiber #%03d at wave %f "%(dwave, err, fiber, block_wave))
 
             if err > 1 :
                 continue
@@ -424,10 +424,12 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
     flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=4)
 
     # boolean mask of fibers with good data
-    good_fibers = (np.sum(ivar>0, axis=1) > 0)
+    # good_fibers = (np.sum(ivar>0, axis=1) > 0)
 
     # median flux of good fibers used as internal spectral reference
-    mflux=np.median(flux[good_fibers],axis=0)
+    # mflux=np.median(flux[good_fibers],axis=0)
+
+    mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
 
     # measure y shifts
     wavemin = xytraceset.wavemin
@@ -818,7 +820,6 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
         eps =  0.1
         x, yp = psf.xy(fiber_for_psf_evaluation, bin_wave+eps)
         dydw = (yp - y) / eps 
-        log.warning("%s %s %s"%(dwave, err, dydw))
         if err * dydw < 1 :
             dy = np.append(dy, -dwave * dydw)
             ey = np.append(ey, err*dydw)

From d5c29cf19c9ab6c854076bf43dd5cffe96cae1eb Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 09:36:43 -0700
Subject: [PATCH 06/18] make the continuum subtraction in self-calibration
 optional

---
 py/desispec/trace_shifts.py | 17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index cc46ff767..10b16f9dd 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -415,21 +415,22 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
     """
 
     log=get_logger()
+    continuum_subtract = True # use continuum subtracted spectra to self-calibrate
 
     # boxcar extraction
-
     qframe = qproc_boxcar_extraction(xytraceset, image, fibers=fibers, width=7)
 
     # resampling on common finer wavelength grid
     flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=4)
+    if continuum_subtract:
+        mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
+    else:
+        # boolean mask of fibers with good data
+        good_fibers = (np.sum(ivar>0, axis=1) > 0)
+        
+        # median flux of good fibers used as internal spectral reference
+        mflux=np.median(flux[good_fibers],axis=0)
 
-    # boolean mask of fibers with good data
-    # good_fibers = (np.sum(ivar>0, axis=1) > 0)
-
-    # median flux of good fibers used as internal spectral reference
-    # mflux=np.median(flux[good_fibers],axis=0)
-
-    mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
 
     # measure y shifts
     wavemin = xytraceset.wavemin

From 389f1e1b7890a4a37d657f66e08a7dfe6bbb9d58 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 11:09:25 -0700
Subject: [PATCH 07/18] use actual variance not variance times flux

---
 py/desispec/trace_shifts.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 10b16f9dd..b779d846a 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -365,7 +365,10 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
                 continue
             block_wave = np.sum(ivar[fiber,ok]*flux[fiber,ok]*(flux[fiber,ok]>0)*wave[ok])/sw
 
-            dwave,err = compute_dy_from_spectral_cross_correlation(flux[fiber,ok],wave[ok],reference_flux[ok],ivar=ivar[fiber,ok]*reference_flux[ok],hw=3., calibrate=True)
+            dwave,err = compute_dy_from_spectral_cross_correlation(flux[fiber,ok], wave[ok],
+                                                                   reference_flux[ok],
+                                                                   ivar=ivar[fiber,ok],
+                                                                   hw=3., calibrate=True)
             if fiber %10==0 :
                 log.info("Wavelength offset %f +/- %f for fiber #%03d at wave %f "%(dwave, err, fiber, block_wave))
 

From 958739599b2a5ae3dcb02c0c56bb3234e149b023 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 12:17:51 -0700
Subject: [PATCH 08/18] update window of medianing

---
 py/desispec/trace_shifts.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index b779d846a..aaccda645 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -425,6 +425,7 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
 
     # resampling on common finer wavelength grid
     flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=4)
+    flux0 = flux * 1 # for debugging 
     if continuum_subtract:
         mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
     else:
@@ -433,8 +434,7 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
         
         # median flux of good fibers used as internal spectral reference
         mflux=np.median(flux[good_fibers],axis=0)
-
-
+    
     # measure y shifts
     wavemin = xytraceset.wavemin
     wavemax = xytraceset.wavemax
@@ -713,9 +713,9 @@ def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
     return ref_wave, ref_spectrum
 
 
-def _continuum_subtract_median(flux0, ivar, continuum_win = 17):
+def _continuum_subtract_median(flux0, ivar, continuum_win = 51):
     # here we get rid of continuum by applying a median filter 
-    continuum_foot = np.abs(np.arange(-continuum_win,continuum_win))>continuum_win /2.
+    continuum_foot = np.abs(np.arange(-continuum_win,continuum_win+1))>continuum_win /4.
     flux = flux0 * 1 # we will modify flux
     # we only keep emission lines and get rid of continuum
     for ii in range(flux.shape[0]):
@@ -794,7 +794,7 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
 
     flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=2)
 
-    mflux, mivar, flux = _continuum_subtract_median(flux, ivar, continuum_win = 17)
+    mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
 
     ref_wave, ref_spectrum = _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, len(ivar))
 

From 764d4132524dd57d54a96034486444834e133066 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 14:08:16 -0700
Subject: [PATCH 09/18] deal with the fact that different functions oversample
 differently therefore the filtering size needs to depend on that

---
 py/desispec/trace_shifts.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index aaccda645..b580350c1 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -424,10 +424,12 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
     qframe = qproc_boxcar_extraction(xytraceset, image, fibers=fibers, width=7)
 
     # resampling on common finer wavelength grid
-    flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=4)
+    oversampling = 4 # The reason why we oversample is unclear to me 
+    flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=oversampling)
     flux0 = flux * 1 # for debugging 
     if continuum_subtract:
-        mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
+        mflux, mivar, flux = _continuum_subtract_median(flux, ivar, continuum_win = oversampling * 9)
+        flux[flux<0] = 0
     else:
         # boolean mask of fibers with good data
         good_fibers = (np.sum(ivar>0, axis=1) > 0)
@@ -713,9 +715,9 @@ def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
     return ref_wave, ref_spectrum
 
 
-def _continuum_subtract_median(flux0, ivar, continuum_win = 51):
+def _continuum_subtract_median(flux0, ivar, continuum_win = 17):
     # here we get rid of continuum by applying a median filter 
-    continuum_foot = np.abs(np.arange(-continuum_win,continuum_win+1))>continuum_win /4.
+    continuum_foot = np.abs(np.arange(-continuum_win,continuum_win+1))>continuum_win /2.
     flux = flux0 * 1 # we will modify flux
     # we only keep emission lines and get rid of continuum
     for ii in range(flux.shape[0]):
@@ -791,10 +793,10 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
     qframe = qproc_boxcar_extraction(xytraceset, image, fibers=fibers, width=7)
 
     # resampling on common finer wavelength grid
+    oversampling = 2 #
+    flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=oversampling)
 
-    flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=2)
-
-    mflux, mivar, flux = _continuum_subtract_median(flux, ivar)
+    mflux, mivar, flux = _continuum_subtract_median(flux, ivar, continuum_win=oversampling*9)
 
     ref_wave, ref_spectrum = _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, len(ivar))
 

From ee2037036a281804aee0d1762a76eca2c593e102 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 14:50:52 -0700
Subject: [PATCH 10/18] make sure arcs are still processed with continuum as
 before

---
 py/desispec/scripts/trace_shifts.py | 4 +++-
 py/desispec/trace_shifts.py         | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/py/desispec/scripts/trace_shifts.py b/py/desispec/scripts/trace_shifts.py
index e12a09ef1..b1ec93455 100644
--- a/py/desispec/scripts/trace_shifts.py
+++ b/py/desispec/scripts/trace_shifts.py
@@ -140,8 +140,10 @@ def fit_trace_shifts(image, args):
             internal_wavelength_calib = False
         elif flavor == "arc" :
             internal_wavelength_calib = True
+            subtract_continuum = False
             args.arc_lamps = True
         else :
+            subtract_continuum = True
             internal_wavelength_calib = True
             args.sky = True
         log.info("wavelength calib, internal={}, sky={} , arc_lamps={}".format(internal_wavelength_calib,args.sky,args.arc_lamps))
@@ -196,7 +198,7 @@ def fit_trace_shifts(image, args):
         x_for_dx,y_for_dx,dx,ex,fiber_for_dx,wave_for_dx = compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image=image, fibers=fibers, width=args.width, deg=args.degxy,image_rebin=args.ccd_rows_rebin)
         if internal_wavelength_calib :
             # measure y shifts
-            x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy = compute_dy_using_boxcar_extraction(tset, image=image, fibers=fibers, width=args.width)
+            x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy = compute_dy_using_boxcar_extraction(tset, image=image, fibers=fibers, width=args.width, subtract_continuum=subtract_continuum)
             mdy = np.median(dy)
             log.info("Subtract median(dy)={}".format(mdy))
             dy -= mdy # remove median, because this is an internal calibration
diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index b580350c1..48ee843e3 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -393,7 +393,8 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
 
     return x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy
 
-def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy=2) :
+def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy=2,
+                                           continuum_subtract = False):
     """
     Measures y offsets (internal wavelength calibration) from a preprocessed image and a trace set using a cross-correlation of boxcar extracted spectra.
     Uses boxcar_extraction , resample_boxcar_frame , compute_dy_from_spectral_cross_correlations_of_frame
@@ -406,6 +407,7 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
         fibers : 1D np.array of int (default is all fibers, the first fiber is always = 0)
         width  : int, extraction boxcar width, default is 7
         degyy  : int, degree of polynomial fit of shifts as a function of y, used to reject outliers.
+        continuum_subtract : bool if true subtract continuum before cross-correlation
 
     Returns:
         x  : 1D array of x coordinates on CCD (axis=1 in numpy image array, AXIS=0 in FITS, cross-dispersion axis = fiber number direction)
@@ -418,7 +420,6 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
     """
 
     log=get_logger()
-    continuum_subtract = True # use continuum subtracted spectra to self-calibrate
 
     # boxcar extraction
     qframe = qproc_boxcar_extraction(xytraceset, image, fibers=fibers, width=7)

From be4d31654a1538813a126774925247e881b1d589 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Wed, 9 Oct 2024 17:38:48 -0700
Subject: [PATCH 11/18] fix variable name

---
 py/desispec/scripts/trace_shifts.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/py/desispec/scripts/trace_shifts.py b/py/desispec/scripts/trace_shifts.py
index b1ec93455..fdc8f7470 100644
--- a/py/desispec/scripts/trace_shifts.py
+++ b/py/desispec/scripts/trace_shifts.py
@@ -140,22 +140,22 @@ def fit_trace_shifts(image, args):
             internal_wavelength_calib = False
         elif flavor == "arc" :
             internal_wavelength_calib = True
-            subtract_continuum = False
             args.arc_lamps = True
         else :
-            subtract_continuum = True
             internal_wavelength_calib = True
             args.sky = True
         log.info("wavelength calib, internal={}, sky={} , arc_lamps={}".format(internal_wavelength_calib,args.sky,args.arc_lamps))
 
     spectrum_filename = args.spectrum
     if args.sky :
+        continuum_subtract = True
         srch_file = "data/spec-sky.dat"
         if not resources.files('desispec').joinpath(srch_file).is_file():
             log.error("Cannot find sky spectrum file {:s}".format(srch_file))
             raise RuntimeError("Cannot find sky spectrum file {:s}".format(srch_file))
         spectrum_filename = resources.files('desispec').joinpath(srch_file)
     elif args.arc_lamps :
+        continuum_subtract = False
         srch_file = "data/spec-arc-lamps.dat"
         if not resources.files('desispec').joinpath(srch_file).is_file():
             log.error("Cannot find arc lamps spectrum file {:s}".format(srch_file))
@@ -198,7 +198,7 @@ def fit_trace_shifts(image, args):
         x_for_dx,y_for_dx,dx,ex,fiber_for_dx,wave_for_dx = compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image=image, fibers=fibers, width=args.width, deg=args.degxy,image_rebin=args.ccd_rows_rebin)
         if internal_wavelength_calib :
             # measure y shifts
-            x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy = compute_dy_using_boxcar_extraction(tset, image=image, fibers=fibers, width=args.width, subtract_continuum=subtract_continuum)
+            x_for_dy,y_for_dy,dy,ey,fiber_for_dy,wave_for_dy = compute_dy_using_boxcar_extraction(tset, image=image, fibers=fibers, width=args.width, continuum_subtract=continuum_subtract)
             mdy = np.median(dy)
             log.info("Subtract median(dy)={}".format(mdy))
             dy -= mdy # remove median, because this is an internal calibration

From a4050eb9190fc1d2d2618a33e9ac09e009692ec4 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Tue, 15 Oct 2024 05:48:38 -0700
Subject: [PATCH 12/18] switch to f-formatting in a couple of instances

---
 py/desispec/trace_shifts.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 48ee843e3..09b83b090 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -370,7 +370,7 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
                                                                    ivar=ivar[fiber,ok],
                                                                    hw=3., calibrate=True)
             if fiber %10==0 :
-                log.info("Wavelength offset %f +/- %f for fiber #%03d at wave %f "%(dwave, err, fiber, block_wave))
+                log.info(f"Wavelength offset {dwave} +/- {err} for fiber {fiber:03d} at wave {block_wave}")
 
             if err > 1 :
                 continue
@@ -832,7 +832,7 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
             ey = np.append(ey, err*dydw)
             wave_for_dy = np.append(wave_for_dy,bin_wave)
             y_for_dy=np.append(y_for_dy,y)
-            log.info("wave = %fA , y=%d, measured dwave = %f +- %f A"%(bin_wave, y, dwave, err))
+            log.info(f"wave = {bin_wave}A , y={y}, measured dwave = {dwave} +- {err} A")
 
     if False : # we don't need this for now
         try :

From 2f6e80ad2108bd1c603972ad32d3e612b4cec58e Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Sat, 26 Oct 2024 06:57:29 -0700
Subject: [PATCH 13/18] use correct weights for central wavelength

---
 py/desispec/trace_shifts.py | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 09b83b090..1e977bbb6 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -360,16 +360,17 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
             else :
                 wmax=wave[-1]
             ok=(wave>=wmin)&(wave<=wmax)
-            sw=np.sum(ivar[fiber,ok]*flux[fiber,ok]*(flux[fiber,ok]>0))
-            if sw<=0 :
+            flux_weight = ivar[fiber,ok] * flux[fiber,ok]**2 * (flux[fiber,ok]>0)
+            flux_weight_sum = np.sum(flux_weight)
+            if cur_weights_sum <= 0 :
                 continue
-            block_wave = np.sum(ivar[fiber,ok]*flux[fiber,ok]*(flux[fiber,ok]>0)*wave[ok])/sw
+            block_wave = np.sum(flux_weight * wave[ok]) / flux_weight_sum
 
             dwave,err = compute_dy_from_spectral_cross_correlation(flux[fiber,ok], wave[ok],
                                                                    reference_flux[ok],
                                                                    ivar=ivar[fiber,ok],
                                                                    hw=3., calibrate=True)
-            if fiber %10==0 :
+            if fiber % 10==0 :
                 log.info(f"Wavelength offset {dwave} +/- {err} for fiber {fiber:03d} at wave {block_wave}")
 
             if err > 1 :
@@ -813,14 +814,15 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
     for b in range(n_wavelength_bins) :
         wmin, wmax = [wavelength_bins[_] for _ in [b, b + 1]]
         ok = (wave >= wmin) & (wave <= wmax)
-        flux_weight = np.sum(mflux[ok] * (mflux[ok] > 0))
+        flux_weight = mflux[ok]**2 * (mflux[ok] > 0) * mivar[ok]
+        flux_weight_sum = np.sum(flux_weight)
         log.warning("%s %s "%(b, flux_weight))
-        if flux_weight == 0 :
+        if flux_weight_sum == 0 :
             continue
         dwave,err = compute_dy_from_spectral_cross_correlation(mflux[ok],
                 wave[ok], ref_spectrum[ok], ivar=mivar[ok], hw=10.,
                 prior_width_dy=prior_width_dy)
-        bin_wave  = np.sum(mflux[ok] * (mflux[ok] > 0) * wave[ok]) / flux_weight
+        bin_wave  = np.sum(flux_weight * wave[ok]) / flux_weight_sum
         # flux weighted wavelength of the center
         # Computing the derivative dy/dwavelength
         x, y = psf.xy(fiber_for_psf_evaluation, bin_wave)

From 0f03f5c8557e0429a9cec716a24253530b1218f1 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Sat, 2 Nov 2024 07:50:20 -0700
Subject: [PATCH 14/18] update variable names get rid of useless warning

---
 py/desispec/trace_shifts.py | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 1e977bbb6..c4f762b54 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -360,11 +360,11 @@ def compute_dy_from_spectral_cross_correlations_of_frame(flux, ivar, wave , xcoe
             else :
                 wmax=wave[-1]
             ok=(wave>=wmin)&(wave<=wmax)
-            flux_weight = ivar[fiber,ok] * flux[fiber,ok]**2 * (flux[fiber,ok]>0)
-            flux_weight_sum = np.sum(flux_weight)
-            if cur_weights_sum <= 0 :
+            flux_weights = ivar[fiber,ok] * flux[fiber,ok]**2 * (flux[fiber,ok]>0)
+            flux_weights_sum = np.sum(flux_weights)
+            if flux_weights_sum <= 0 :
                 continue
-            block_wave = np.sum(flux_weight * wave[ok]) / flux_weight_sum
+            block_wave = np.sum(flux_weights * wave[ok]) / flux_weights_sum
 
             dwave,err = compute_dy_from_spectral_cross_correlation(flux[fiber,ok], wave[ok],
                                                                    reference_flux[ok],
@@ -814,15 +814,14 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
     for b in range(n_wavelength_bins) :
         wmin, wmax = [wavelength_bins[_] for _ in [b, b + 1]]
         ok = (wave >= wmin) & (wave <= wmax)
-        flux_weight = mflux[ok]**2 * (mflux[ok] > 0) * mivar[ok]
-        flux_weight_sum = np.sum(flux_weight)
-        log.warning("%s %s "%(b, flux_weight))
-        if flux_weight_sum == 0 :
+        flux_weights = mflux[ok]**2 * (mflux[ok] > 0) * mivar[ok]
+        flux_weights_sum = np.sum(flux_weights)
+        if flux_weights_sum == 0 :
             continue
         dwave,err = compute_dy_from_spectral_cross_correlation(mflux[ok],
                 wave[ok], ref_spectrum[ok], ivar=mivar[ok], hw=10.,
                 prior_width_dy=prior_width_dy)
-        bin_wave  = np.sum(flux_weight * wave[ok]) / flux_weight_sum
+        bin_wave  = np.sum(flux_weights * wave[ok]) / flux_weights_sum
         # flux weighted wavelength of the center
         # Computing the derivative dy/dwavelength
         x, y = psf.xy(fiber_for_psf_evaluation, bin_wave)

From aa1cf83dfd143942a21bafca952182a361e37f0b Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Mon, 4 Nov 2024 14:28:20 +0000
Subject: [PATCH 15/18] add functions docstrings

---
 py/desispec/trace_shifts.py | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index c4f762b54..b7b969aa5 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -632,8 +632,19 @@ def compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image
 
 def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
     """
-    Prepare the reference spectrum to be used for wavelegth offset 
-    determination
+    Prepare the reference spectrum to be used for wavelength offset 
+    determination. Here we convolve it to the right LSF and rescale it 
+    to match the measured flux.
+
+    Arguments:
+        ref_wave: np.array of wavelengths 
+        ref_spectrum: np.array of reference spectrum flux
+        psf: PSF object
+        wave: np.array wavelength of extracted spectra
+        mflux: np.array flux of extracted spectra
+        nfibers: int 
+    Returns:
+        ref_wave, ref_spectrum: tuple of wavelength and flux arrays
     """
     log = get_logger()
     
@@ -718,6 +729,20 @@ def _prepare_ref_spectrum(ref_wave, ref_spectrum, psf, wave, mflux, nfibers):
 
 
 def _continuum_subtract_median(flux0, ivar, continuum_win = 17):
+    """
+    Compute the median spectrum after continuum subtraction
+    
+    Arguments:
+        flux0: (nfibers, npix) np.array of fluxes
+        ivar: (nfibers, npix)-shaped np.array of inverser variances
+        continuum_win: integer. How-many pixels around are used to get continuum. 
+            Here we use the 1d annulus from continuum_win/2 to continuum_win
+    
+    Returns:
+        mfux: npix np.array of median spectrum
+        mivar: npix np.array with ivar of the median spectrum 
+        flux: (nfibers, npix) continuum subtracted original flux array
+    """
     # here we get rid of continuum by applying a median filter 
     continuum_foot = np.abs(np.arange(-continuum_win,continuum_win+1))>continuum_win /2.
     flux = flux0 * 1 # we will modify flux

From b3c5ba30af89c05b931463c1cf064345b50b0bb5 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Tue, 5 Nov 2024 00:24:32 +0000
Subject: [PATCH 16/18] get rid of comment on oversampling fix a few typos in
 comments

---
 py/desispec/trace_shifts.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index b7b969aa5..2a142e029 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -218,7 +218,7 @@ def compute_dy_from_spectral_cross_correlation(flux, wave, refflux, ivar=None,
     A relative flux calibration of the two spectra is done internally.
 
     Args:
-        flux    : 1D array of spectral flux as a function of wavelenght
+        flux    : 1D array of spectral flux as a function of wavelength
         wave    : 1D array of wavelength (in Angstrom)
         refflux : 1D array of reference spectral flux
 
@@ -426,7 +426,7 @@ def compute_dy_using_boxcar_extraction(xytraceset, image, fibers, width=7, degyy
     qframe = qproc_boxcar_extraction(xytraceset, image, fibers=fibers, width=7)
 
     # resampling on common finer wavelength grid
-    oversampling = 4 # The reason why we oversample is unclear to me 
+    oversampling = 4
     flux, ivar, wave = resample_boxcar_frame(qframe.flux, qframe.ivar, qframe.wave, oversampling=oversampling)
     flux0 = flux * 1 # for debugging 
     if continuum_subtract:
@@ -477,8 +477,8 @@ def compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image
     Measure x offsets from a preprocessed image and a trace set
 
     Args:
-        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelenght to XCCD
-        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelenght to YCCD
+        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to XCCD
+        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to YCCD
         wavemin : float
         wavemax : float. wavemin and wavemax are used to define a reduced variable legx(wave,wavemin,wavemax)=2*(wave-wavemin)/(wavemax-wavemin)-1
                   used to compute the traces, xccd=legval(legx(wave,wavemin,wavemax),xtrace[fiber])
@@ -799,7 +799,7 @@ def shift_ycoef_using_external_spectrum(psf, xytraceset, image, fibers,
         prior_width_dy: float with of the Gaussian prior on dy
 
     Returns:
-        ycoef  : 2D np.array of same shape as input, with modified Legendre coefficents for each fiber to convert wavelenght to YCCD
+        ycoef  : 2D np.array of same shape as input, with modified Legendre coefficients for each fiber to convert wavelength to YCCD
 
     """
     log = get_logger()
@@ -1294,8 +1294,8 @@ def recompute_legendre_coefficients(xcoef,ycoef,wavemin,wavemax,degxx,degxy,degy
     Modifies legendre coefficients of an input trace set using polynomial coefficents (as defined by the routine monomials)
 
     Args:
-        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelenght to XCCD
-        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelenght to YCCD
+        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to XCCD
+        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to YCCD
         wavemin : float
         wavemax : float. wavemin and wavemax are used to define a reduced variable legx(wave,wavemin,wavemax)=2*(wave-wavemin)/(wavemax-wavemin)-1
                   used to compute the traces, xccd=legval(legx(wave,wavemin,wavemax),xtrace[fiber])

From b08cb4eacc3efc80d11a02ca2642e1e2ff40c025 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Tue, 5 Nov 2024 00:44:21 +0000
Subject: [PATCH 17/18] few more typos

---
 py/desispec/trace_shifts.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/py/desispec/trace_shifts.py b/py/desispec/trace_shifts.py
index 2a142e029..738a280f7 100644
--- a/py/desispec/trace_shifts.py
+++ b/py/desispec/trace_shifts.py
@@ -477,8 +477,8 @@ def compute_dx_from_cross_dispersion_profiles(xcoef,ycoef,wavemin,wavemax, image
     Measure x offsets from a preprocessed image and a trace set
 
     Args:
-        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to XCCD
-        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to YCCD
+        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficients for each fiber to convert wavelength to XCCD
+        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficients for each fiber to convert wavelength to YCCD
         wavemin : float
         wavemax : float. wavemin and wavemax are used to define a reduced variable legx(wave,wavemin,wavemax)=2*(wave-wavemin)/(wavemax-wavemin)-1
                   used to compute the traces, xccd=legval(legx(wave,wavemin,wavemax),xtrace[fiber])
@@ -1291,11 +1291,11 @@ def polynomial_fit(z,ez,xx,yy,degx,degy) :
 
 def recompute_legendre_coefficients(xcoef,ycoef,wavemin,wavemax,degxx,degxy,degyx,degyy,dx_coeff,dy_coeff) :
     """
-    Modifies legendre coefficients of an input trace set using polynomial coefficents (as defined by the routine monomials)
+    Modifies legendre coefficients of an input trace set using polynomial coefficients (as defined by the routine monomials)
 
     Args:
-        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to XCCD
-        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficents for each fiber to convert wavelength to YCCD
+        xcoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficients for each fiber to convert wavelength to XCCD
+        ycoef : 2D np.array of shape (nfibers,ncoef) containing Legendre coefficients for each fiber to convert wavelength to YCCD
         wavemin : float
         wavemax : float. wavemin and wavemax are used to define a reduced variable legx(wave,wavemin,wavemax)=2*(wave-wavemin)/(wavemax-wavemin)-1
                   used to compute the traces, xccd=legval(legx(wave,wavemin,wavemax),xtrace[fiber])
@@ -1307,8 +1307,8 @@ def recompute_legendre_coefficients(xcoef,ycoef,wavemin,wavemax,degxx,degxy,degy
         dy_coeff : 1D np.array of polynomial coefficients of size (degyx*degyy) as defined by the routine monomials.
 
     Returns:
-        xcoef : 2D np.array of shape (nfibers,ncoef) with modified Legendre coefficents
-        ycoef : 2D np.array of shape (nfibers,ncoef) with modified Legendre coefficents
+        xcoef : 2D np.array of shape (nfibers,ncoef) with modified Legendre coefficients
+        ycoef : 2D np.array of shape (nfibers,ncoef) with modified Legendre coefficients
     """
     wave=np.linspace(wavemin,wavemax,100)
     nfibers=xcoef.shape[0]

From ece6b599c191bad7d7effe2bb63a7524a0bce617 Mon Sep 17 00:00:00 2001
From: "Sergey E. Koposov" <skoposov@ed.ac.uk>
Date: Sat, 9 Nov 2024 11:12:18 -0800
Subject: [PATCH 18/18] make the selection more readable

---
 py/desispec/scripts/trace_shifts.py | 19 ++++++++-----------
 1 file changed, 8 insertions(+), 11 deletions(-)

diff --git a/py/desispec/scripts/trace_shifts.py b/py/desispec/scripts/trace_shifts.py
index fdc8f7470..58ace0d5d 100644
--- a/py/desispec/scripts/trace_shifts.py
+++ b/py/desispec/scripts/trace_shifts.py
@@ -218,17 +218,14 @@ def fit_trace_shifts(image, args):
     degyy=args.degyy
 
     # if any quadrant is masked, reduce to a single offset
-    hy=image.pix.shape[0]//2
-    hx=image.pix.shape[1]//2
-    allgood=True
-    allgood &= (np.sum((x_for_dx<hx)&(y_for_dx<hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dx<hx)&(y_for_dx>hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dx>hx)&(y_for_dx<hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dx>hx)&(y_for_dx>hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dy<hx)&(y_for_dy<hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dy<hx)&(y_for_dy>hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dy>hx)&(y_for_dy<hy))>0) # some data in this quadrant
-    allgood &= (np.sum((x_for_dy>hx)&(y_for_dy>hy))>0) # some data in this quadrant
+    hy = image.pix.shape[0] // 2
+    hx = image.pix.shape[1] // 2
+    allgood = True
+    for _curx, _cury in [(x_for_dx,y_for_dx),(x_for_dy, y_for_dy)]:
+        for curxop in [np.less, np.greater]:
+            for curyop in [np.less, np.greater]:
+                allgood &= np.any(curxop(_curx, hx) & curyop(_cury, hy))
+        # some data in this quadrant
     if not allgood :
         log.warning("No shift data for at least one quadrant of the CCD, falls back to deg=0 shift")
         degxx=0