Merge branch 'main' of https://github.com/desihub/redrock

.readthedocs.yml

@@ -0,0 +1,23 @@
+# .readthedocs.yml
+# Read the Docs configuration file
+# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details
+
+# Required
+version: 2
+
+build:
+    os: "ubuntu-22.04"
+    tools:
+        python: "3.10"
+
+# Build documentation in the doc/ directory with Sphinx
+sphinx:
+    configuration: doc/conf.py
+
+# Optionally build your docs in additional formats such as PDF and ePub
+# formats: all
+
+# Optionally set the version of Python and requirements required to build your docs
+python:
+    install:
+        - requirements: doc/rtd-requirements.txt

README.rst

@@ -2,16 +2,18 @@
 redrock
 =======
 
-.. image:: https://travis-ci.org/desihub/redrock.svg?branch=master
-    :target: https://travis-ci.org/desihub/redrock
-    :alt: Travis Build Status
+|Actions Status| |Coveralls Status| |Documentation Status|
 
-.. image:: https://coveralls.io/repos/github/desihub/redrock/badge.svg?branch=master
-    :target: https://coveralls.io/github/desihub/redrock?branch=master
+.. |Actions Status| image:: https://github.com/desihub/redrock/workflows/CI/badge.svg
+    :target: https://github.com/desihub/redrock/actions
+    :alt: GitHub Actions CI Status
+
+.. |Coveralls Status| image:: https://coveralls.io/repos/desihub/redrock/badge.svg
+    :target: https://coveralls.io/github/desihub/redrock
     :alt: Test Coverage Status
 
-.. image:: https://readthedocs.org/projects/redrock/badge/?version=latest
-    :target: http://redrock.readthedocs.org/en/latest/
+.. |Documentation Status| image:: https://readthedocs.org/projects/redrock/badge/?version=latest
+    :target: https://redrock.readthedocs.io/en/latest/
     :alt: Documentation Status
 
 Introduction

doc/rtd-requirements.txt

@@ -1,2 +1,4 @@
-setuptools
-docutils<0.18
+setuptools>60
+Sphinx>6,<7
+sphinx-rtd-theme>1
+urllib3<2

py/redrock/fitz.py

@@ -107,7 +107,7 @@ def minfit(x, y):
     return (x0, xerr, y0, zwarn)
 
 
-def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, use_gpu=False):
+def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=False):
     """Refines redshift measurement around up to nminima minima.
 
     TODO:
@@ -116,8 +116,8 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, use_gpu
     Args:
         zchi2 (array): chi^2 values for each redshift.
         redshifts (array): the redshift values.
-        spectra (list): list of Spectrum objects at different wavelengths
-            grids.
+        target (Target): the target for this fit which includes a list
+            of Spectrum objects at different wavelength grids.
         template (Template): the template for this fit.
         nminima (int): the number of minima to consider.
         use_gpu (bool): use GPU or not
@@ -132,6 +132,7 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, use_gpu
         import cupy as cp
 
     nbasis = template.nbasis
+    spectra = target.spectra
 
     # Build dictionary of wavelength grids
     dwave = { s.wavehash:s.wave for s in spectra }
@@ -146,10 +147,13 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, use_gpu
 
     (weights, flux, wflux) = spectral_data(spectra)
     if (use_gpu):
-        #Copy arrays to GPU
-        weights = cp.asarray(weights)
-        flux = cp.asarray(flux)
-        wflux = cp.asarray(wflux)
+        #Get CuPy arrays of weights, flux, wflux
+        #These are created on the first call of gpu_spectral_data() for a
+        #target and stored.  They are retrieved on subsequent calls.
+        (gpuweights, gpuflux, gpuwflux) = target.gpu_spectral_data()
+        # Build dictionaries of wavelength bin edges, min/max, and centers
+        gpuedges = { s.wavehash:(s.gpuedges, s.minedge, s.maxedge) for s in spectra }
+        gpudwave = { s.wavehash:s.gpuwave for s in spectra }
 
     results = list()
     #Define nz here instead of hard-coding length 15 and then defining nz as
@@ -171,12 +175,25 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, use_gpu
         ilo = max(0, imin-1)
         ihi = min(imin+1, len(zchi2)-1)
         zz = np.linspace(redshifts[ilo], redshifts[ihi], nz)
+        if (use_gpu):
+            #Create a redshift grid on the GPU as well
+            gpuzz = cp.asarray(zz)
 
         zzchi2 = np.zeros(nz, dtype=np.float64)
         zzcoeff = np.zeros((nz, nbasis), dtype=np.float64)
 
+        #Calculate xmin and xmax from template and zz array on CPU and
+        #pass as scalars
+        xmin = template.minwave*(1+zz.max())
+        xmax = template.maxwave*(1+zz.min())
+
         #Use batch mode for rebin_template, transmission_Lyman, and calc_zchi2
-        binned = rebin_template(template, zz, dwave, use_gpu=use_gpu)
+        if (use_gpu):
+            #Use gpuedges already calculated and on GPU
+            binned = rebin_template(template, gpuzz, dedges=gpuedges, use_gpu=use_gpu, xmin=xmin, xmax=xmax)
+        else:
+            #Use numpy CPU arrays
+            binned = rebin_template(template, zz, dwave, use_gpu=use_gpu, xmin=xmin, xmax=xmax)
         # Correct spectra for Lyman-series
         for k in list(dwave.keys()):
             #New algorithm accepts all z as an array and returns T, a 2-d
@@ -188,23 +205,41 @@ def fitz(zchi2, redshifts, spectra, template, nminima=3, archetype=None, use_gpu
                 continue
             #Vectorize multiplication
             binned[k] *= T[:,:,None]
-        (zzchi2, zzcoeff) = calc_zchi2_batch(spectra, binned, weights, flux, wflux, nz, nbasis, use_gpu=use_gpu)
+        if (use_gpu):
+            #Use gpu arrays for weights, flux, wflux
+            (zzchi2, zzcoeff) = calc_zchi2_batch(spectra, binned, gpuweights, gpuflux, gpuwflux, nz, nbasis, use_gpu=use_gpu)
+        else:
+            #Use numpy CPU arrays for weights, flux, wflux 
+            (zzchi2, zzcoeff) = calc_zchi2_batch(spectra, binned, weights, flux, wflux, nz, nbasis, use_gpu=use_gpu)
 
         #- fit parabola to 3 points around minimum
         i = min(max(np.argmin(zzchi2),1), len(zz)-2)
         zmin, sigma, chi2min, zwarn = minfit(zz[i-1:i+2], zzchi2[i-1:i+2])
 
         try:
-            binned = rebin_template(template, np.array([zmin]), dwave, use_gpu=use_gpu)
+            #Calculate xmin and xmax from template and pass as scalars 
+            xmin = template.minwave*(1+zmin)
+            ximax = template.maxwave*(1+zmin)
+            if (use_gpu):
+                #Use gpuedges already calculated and on GPU
+                binned = rebin_template(template, cp.array([zmin]), dedges=gpuedges, use_gpu=use_gpu, xmin=xmin, xmax=xmax)
+            else:
+                binned = rebin_template(template, np.array([zmin]), dwave, use_gpu=use_gpu, xmin=xmin, xmax=xmax)
             for k in list(dwave.keys()):
-                T = transmission_Lyman(np.array([zmin]),dwave[k], use_gpu=use_gpu)
+                if (use_gpu):
+                    #Copy binned[k] back to CPU to perform next steps on CPU
+                    #because faster with only 1 redshift
+                    binned[k] = binned[k].get()
+                #Use CPU always
+                T = transmission_Lyman(np.array([zmin]),dwave[k], use_gpu=False)
                 if (T is None):
                     #Return value of None means that wavelenght regime
                     #does not overlap Lyman transmission - continue here
                     continue
                 #Vectorize multiplication
                 binned[k] *= T[:,:,None]
-            (chi2, coeff) = calc_zchi2_batch(spectra, binned, weights, flux, wflux, 1, nbasis, use_gpu=use_gpu)
+            #Use CPU always with one redshift
+            (chi2, coeff) = calc_zchi2_batch(spectra, binned, weights, flux, wflux, 1, nbasis, use_gpu=False)
             coeff = coeff[0,:]
         except ValueError as err:
             if zmin<redshifts[0] or redshifts[-1]<zmin:

py/redrock/rebin.py

@@ -209,7 +209,7 @@ def _trapz_rebin_batch(x, y, edges, myz, results, redshifted_x):
     return
 
 
-def trapz_rebin(x, y, xnew=None, edges=None, myz=None, use_gpu=False, xmin=None, xmax=None):
+def trapz_rebin(x, y, xnew=None, edges=None, myz=None, use_gpu=False, xmin=None, xmax=None, edge_min=None, edge_max=None):
     """Rebin y(x) flux density using trapezoidal integration between bin edges
     Optionally use GPU helper method to rebin in batch, see trapz_rebin_batch_gpu
     Note - current return array shape is (nz, nbins, nbasis).  Changing to
@@ -258,14 +258,22 @@ def trapz_rebin(x, y, xnew=None, edges=None, myz=None, use_gpu=False, xmin=None,
         myz (array): (optional) redshift array to rebin in batch,
             applying redshifts on-the-fly to x
         use_gpu (boolean): whether or not to use GPU algorithm 
-        xmin (float): (optional) minimum x-value - x[0] will be used if
-            omitted - this is useful to avoid GPU to CPU copying in the case
-            where x is a CuPy array and providing the scalar value results
-            in a large speed gain
-        xmax (float) (optional) maximum x-value - x[-1] will be used if
-            omitted - this is useful to avoid GPU to CPU copying in the case
-            where x is a CuPy array and providing the scalar value results
-            in a large speed gain
+        xmin (float): (optional) minimum x-value - x[0]*(1+myz.max()) will be
+            used if omitted - this is useful to avoid GPU to CPU copying in
+            the case where x is a CuPy array and providing the scalar value
+            results in a large speed gain
+        xmax (float) (optional) maximum x-value - x[-1]*(1+myz.min()) will be
+            used if omitted - this is useful to avoid GPU to CPU copying in
+            the case where x is a CuPy array and providing the scalar value
+            results in a large speed gain
+        edge_min (float): (optional) minimum new bin edge - edges[0] will be
+            used if omitted - this is useful to avoid GPU to CPU copying in
+            the case where edges is precomputed on the GPU as a CuPy array
+            and providing the scalar value results in a large speed gain
+        edge_max (float): (optional) maximum new bin edge - edges[-1] will be
+            used if omitted - this is useful to avoid GPU to CPU copying in
+            the case where edges is precomputed on the GPU as a CuPy array
+            and providing the scalar value results in a large speed gain
 
     Returns:
         array: integrated results with len(results) = len(edges)-1
@@ -283,15 +291,21 @@ def trapz_rebin(x, y, xnew=None, edges=None, myz=None, use_gpu=False, xmin=None,
         ValueError: if edges are outside the range of x or if len(x) != len(y)
 
     """
+    if use_gpu:
+        import cupy as cp
     if edges is None:
         edges = centers2edges(xnew)
+    elif use_gpu:
+        if type(edges) != cp.ndarray:
+            edges = np.asarray(edges)
     else:
         edges = np.asarray(edges)
     nbins = len(edges)-1
-    if xmin is None:
-        xmin = x[0]
-    if xmax is None:
-        xmax = x[-1]
+    #Set edge_min and edge_max if not passed as input
+    if edge_min is None:
+        edge_min = edges[0]
+    if edge_max is None:
+        edge_max = edges[-1]
 
     #Use these booleans to determine output array shape based on input
     scalar_z = False
@@ -303,17 +317,28 @@ def trapz_rebin(x, y, xnew=None, edges=None, myz=None, use_gpu=False, xmin=None,
     elif (np.isscalar(myz)):
         myz = np.array([myz], dtype=np.float64)
         scalar_z = True
-    myz = np.asarray(myz, dtype=np.float64)
+    elif (use_gpu):
+        if (type(myz) != cp.ndarray):
+            myz = np.asarray(myz, dtype=np.float64)
+    else:
+        myz = np.asarray(myz, dtype=np.float64)
     nz = len(myz)
     if (nz == 0):
         #Empty myz array
         return np.zeros((0,nbins, 1), dtype=np.float64)
 
+    #Set xmin and xmax based on x and myz if not passed as input
     #Must multiply x by 1+z for comparison, only need to look at max/min cases
-    #if (edges[0] < x[0]*(1+myz.max()) or edges[-1] > x[-1]*(1+myz.min())):
-    if (edges[0] < xmin*(1+myz.max()) or edges[-1] > xmax*(1+myz.min())):
-        raise ValueError('edges must be within input x range')
+    if xmin is None:
+        xmin = x[0]*(1+myz.max())
+    if xmax is None:
+        xmax = x[-1]*(1+myz.min())
 
+    #Use edge_min, xmin, edge_max, and xmax for comparison - these are
+    #passed as scalars when x and edges are on GPU which results in
+    #large speed gain by minimizing GPU to CPU copying of single values
+    if (edge_min < xmin or edge_max > xmax):
+        raise ValueError('edges must be within input x range')
 
     if (not use_gpu and scalar_z and len(y.shape) == 1):
         #Special case, call _trapz_rebin_1d directly
@@ -400,19 +425,19 @@ def _trapz_rebin_batch_gpu(x, y, edges, myz, result_shape):
     #Divide edges output wavelength array by (1+myz) to get 2d array
     #of input wavelengths at each boundary in edges and
     #use serachsorted to find index for each boundary
-    e2d = cp.asarray(edges/(1+myz[:,None]))
+    e2d = edges/(1+myz[:,None])
     idx = cp.searchsorted(x, e2d, side='right')#.astype(cp.int32)
 
     # Load CUDA kernel
     cp_module = cp.RawModule(code=cuda_source)
     batch_trapz_rebin_kernel = cp_module.get_function('batch_trapz_rebin')
 
     #Array sizes
-    nbin = cp.int32(len(edges)-1)
-    nz = cp.int32(len(myz))
-    nbasis = cp.int32(y.shape[0])
+    nbin = len(edges)-1
+    nz = len(myz)
+    nbasis = y.shape[0]
     n = nbin*nbasis*nz
-    nt = cp.int32(len(x))
+    nt = len(x)
     if (result_shape is None):
       result_shape = (nz, nbasis, nbin)
 
@@ -425,18 +450,33 @@ def _trapz_rebin_batch_gpu(x, y, edges, myz, result_shape):
 
     return result
 
-def rebin_template(template, myz, dwave, use_gpu=False):
+def rebin_template(template, myz, dwave=None, dedges=None, use_gpu=False, xmin=None, xmax=None):
     """Rebin a template to a set of wavelengths.
 
     Given a template and a single redshift - or an array of redshifts,
     rebin the template to a set of wavelength arrays.
+    The wavelengths can be passed as centers (dwave) or edges (dedges)
+    of the wavelength bins.
 
     Args:
         template (Template): the template object
         myz (float or array of float): the redshift(s)
         dwave (dict): the keys are the "wavehash" and the values
             are a 1D array containing the wavelength grid.
+        dedges (dict): the keys are the "wavehash" and the values
+            can be either a 1D array containing the bin edges
+            of the wavelength grid (computed from centers2edges)
+            or a 3-element tuple with this 1D array as a CuPy array
+            and the min and max values as scalars in CPU memory.
         use_gpu (bool): whether or not to use the GPU algorithm
+        xmin (float): (optional) minimum x-value - x[0]*(1+myz.max()) will be
+            used if omitted - this is useful to avoid GPU to CPU copying in
+            the case where x is a CuPy array and providing the scalar value
+            results in a large speed gain
+        xmax (float) (optional) maximum x-value - x[-1]*(1+myz.min()) will be
+            used if omitted - this is useful to avoid GPU to CPU copying in
+            the case where x is a CuPy array and providing the scalar value
+            results in a large speed gain
 
     Returns:
         dict:  The rebinned template for every basis function and wavelength
@@ -446,9 +486,32 @@ def rebin_template(template, myz, dwave, use_gpu=False):
     """
     nbasis = template.flux.shape[0]  #- number of template basis vectors
     result = dict()
+    #calculate xmin and xmax if not given
+    if (xmin is None):
+        xmin = template.minwave*(1+myz.max())
+    if (xmax is None):
+        xmax = template.maxwave*(1+myz.min())
     #rebin all z and all bases in batch in parallel
     #and return dict of 3-d numpy / cupy arrays 
-    for hs, wave in dwave.items():
-        result[hs] = trapz_rebin(template.wave, template.flux, xnew=wave, myz=myz, use_gpu=use_gpu, xmin=template.minwave, xmax=template.maxwave)
-
+    if dwave is not None:
+        #Handle the case where dwave is passed
+        for hs, wave in dwave.items():
+            if (use_gpu and template.gpuwave is not None):
+                #Use gpuwave and gpuflux arrays already on GPU
+                result[hs] = trapz_rebin(template.gpuwave, template.gpuflux, xnew=wave, myz=myz, use_gpu=use_gpu, xmin=xmin, xmax=xmax)
+            else:
+                result[hs] = trapz_rebin(template.wave, template.flux, xnew=wave, myz=myz, use_gpu=use_gpu, xmin=xmin, xmax=xmax)
+    elif dedges is not None:
+        #Handle the case where dedges is passed because edges is already computed
+        minedge = None
+        maxedge = None
+        for hs, edges in dedges.items():
+            #Check if edges is a 1-d array or a tuple also containing scalar min/max values
+            if type(edges) is tuple:
+                (edges, minedge, maxedge) = edges
+            if (use_gpu and template.gpuwave is not None):
+                #Use gpuwave and gpuflux arrays already on GPU
+                result[hs] = trapz_rebin(template.gpuwave, template.gpuflux, edges=edges, myz=myz, use_gpu=use_gpu, xmin=xmin, xmax=xmax, edge_min=minedge, edge_max=maxedge)
+            else:
+                result[hs] = trapz_rebin(template.wave, template.flux, edges=edges, myz=myz, use_gpu=use_gpu, xmin=xmin, xmax=xmax, edge_min=minedge, edge_max=maxedge)
     return result