Merge branch 'main' into abhijeet_per_camera_prior_gpu

desihub · Jan 3, 2024 · 00d89ba · 00d89ba
2 parents 8133f74 + ad32b87
commit 00d89ba
Show file tree

Hide file tree

Showing 15 changed files with 789 additions and 35 deletions.
diff --git a/.github/workflows/python-package.yml b/.github/workflows/python-package.yml
@@ -17,7 +17,7 @@ jobs:
             fail-fast: true
             matrix:
                 os: [ubuntu-latest]
-                python-version: ['3.8', '3.9']  # fitsio does not like Python < 3.7
+                python-version: ['3.9', '3.10']  # fitsio does not like Python < 3.7
                 # astropy-version: ['<3.0', '<4.1']  #, '<5.0']
         env:
             DESIUTIL_VERSION: 3.2.5

diff --git a/.readthedocs.yml b/.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
diff --git a/README.rst b/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

diff --git a/doc/changes.rst b/doc/changes.rst
@@ -2,12 +2,64 @@
 redrock Change Log
 ==================
 
-0.17.1 (unreleased)
+0.18.1 (unreleased)
 -------------------
 
+* Write test files to temporary directory (PR `#263`_).
+* Check template dimension so code works on a single template (PR `#264`_). 
+
+.. _`#263`: https://github.com/desihub/redrock/pull/263
+.. _`#264`: https://github.com/desihub/redrock/pull/264
+
+0.18.0 (2023-09-14)
+-------------------
+
+Significant speedups through GPU support, and bug fixes for archetype mode.
+
 * Update API documentation for completeness (PR `#229`_).
+* Fix to force mpprocs=1 when using GPU mode with multiprocessing (PR `#230`_).
+* Remove multiprocessing from rebin (PR `#232`_).
+* Bugfixes for crashes when MPI ranks > ntargets (PR `#233`_).
+* Refactored np.linalg.solve into solve_matrices (PR `#234`_).
+* Fix used of max_gpuprocs (PR `#236`_).
+* Bug fixes for archetype mode (PR `#238`_).
+* Added min/maxwave to templates.py that store wave[0] (PR `#239`_).
+* Fix archetype mode numba conflict (PR `#240`_).
+* Fix archetype mode with wrong input dir (PR `#244`_).
+* Refactor GPU to CPU copying for big speedup (PR `#245`_).
+* Fix unit test_gpu_zscan (PR `#247`_).
+* Update readthedocs config (PR `#252`_).
 
 .. _`#229`: https://github.com/desihub/redrock/pull/229
+.. _`#230`: https://github.com/desihub/redrock/pull/230
+.. _`#232`: https://github.com/desihub/redrock/pull/232
+.. _`#233`: https://github.com/desihub/redrock/pull/233
+.. _`#234`: https://github.com/desihub/redrock/pull/234
+.. _`#236`: https://github.com/desihub/redrock/pull/236
+.. _`#238`: https://github.com/desihub/redrock/pull/238
+.. _`#239`: https://github.com/desihub/redrock/pull/239
+.. _`#240`: https://github.com/desihub/redrock/pull/240
+.. _`#244`: https://github.com/desihub/redrock/pull/244
+.. _`#245`: https://github.com/desihub/redrock/pull/245
+.. _`#247`: https://github.com/desihub/redrock/pull/247
+.. _`#252`: https://github.com/desihub/redrock/pull/252
+
+0.17.9 (2023-12-01)
+-------------------
+
+This is a tag of the "opticaldepth" branch used by Allyson Brodzeller
+to re-process Iron production z>1.6 QSOs with a new HIZ template
+(redrock-templates tag 0.8.1) and updated optical depth coefficients
+to match Kamble et al. 2020 (Arxiv: 1904.01110).
+Although this tag was made after 0.18.0, it was branched prior to that
+tag and includes all of the 0.18.0 PRs except:
+
+* Refactor GPU to CPU copying for big speedup (PR `#245`_).
+* Update readthedocs config (PR `#252`_).
+
+This branch will not be merged into main as-is, but the updated coefficients
+will be included in a future tag to handle multiple optical depth models
+to support templates trained with different models.
 
 0.17.0 (2023-01-12)
 -------------------

diff --git a/doc/rtd-requirements.txt b/doc/rtd-requirements.txt
@@ -1,2 +1,4 @@
-setuptools
-docutils<0.18
+setuptools>60
+Sphinx>6,<7
+sphinx-rtd-theme>1
+urllib3<2
diff --git a/py/redrock/_version.py b/py/redrock/_version.py
@@ -1 +1 @@
-__version__ = '0.17.0.dev758'
+__version__ = '0.18.0.dev797'
diff --git a/py/redrock/data/DLAcoeff.txt b/py/redrock/data/DLAcoeff.txt
@@ -0,0 +1,39 @@
+    2  1215.670 1.61698E-04 5.38995E-05
+    3  1025.720 1.54539E-04 5.15129E-05
+    4   972.537 1.49767E-04 4.99222E-05
+    5   949.743 1.46031E-04 4.86769E-05
+    6   937.803 1.42893E-04 4.76312E-05
+    7   930.748 1.40159E-04 4.67196E-05
+    8   926.226 1.37714E-04 4.59048E-05
+    9   923.150 1.35495E-04 4.51650E-05
+   10   920.963 1.33452E-04 4.44841E-05
+   11   919.352 1.31561E-04 4.38536E-05
+   12   918.129 1.29785E-04 4.32617E-05
+   13   917.181 1.28117E-04 4.27056E-05
+   14   916.429 1.26540E-04 4.21799E-05
+   15   915.824 1.25041E-04 4.16804E-05
+   16   915.329 1.23614E-04 4.12046E-05
+   17   914.919 1.22248E-04 4.07494E-05
+   18   914.576 1.20938E-04 4.03127E-05
+   19   914.286 1.19681E-04 3.98938E-05
+   20   914.039 1.18469E-04 3.94896E-05
+   21   913.826 1.17298E-04 3.90995E-05
+   22   913.641 1.16167E-04 3.87225E-05
+   23   913.480 1.15071E-04 3.83572E-05
+   24   913.339 1.14011E-04 3.80037E-05
+   25   913.215 1.12983E-04 3.76609E-05
+   26   913.104 1.11972E-04 3.73241E-05
+   27   913.006 1.11002E-04 3.70005E-05
+   28   912.918 1.10051E-04 3.66836E-05
+   29   912.839 1.09125E-04 3.63749E-05
+   30   912.768 1.08220E-04 3.60734E-05
+   31   912.703 1.07337E-04 3.57789E-05
+   32   912.645 1.06473E-04 3.54909E-05
+   33   912.592 1.05629E-04 3.52096E-05
+   34   912.543 1.04802E-04 3.49340E-05
+   35   912.499 1.03991E-04 3.46636E-05
+   36   912.458 1.03198E-04 3.43994E-05
+   37   912.420 1.02420E-04 3.41402E-05
+   38   912.385 1.01657E-04 3.38856E-05
+   39   912.353 1.00908E-04 3.36359E-05
+   40   912.324 1.00168E-04 3.33895E-05
diff --git a/py/redrock/data/LAFcoeff.txt b/py/redrock/data/LAFcoeff.txt
@@ -0,0 +1,39 @@
+    2  1215.670 1.68976E-02 2.35379E-03 1.02611E-04
+    3  1025.720 4.69229E-03 6.53625E-04 2.84940E-05
+    4   972.537 2.23898E-03 3.11884E-04 1.35962E-05
+    5   949.743 1.31901E-03 1.83735E-04 8.00974E-06
+    6   937.803 8.70656E-04 1.21280E-04 5.28707E-06
+    7   930.748 6.17843E-04 8.60640E-05 3.75186E-06
+    8   926.226 4.60924E-04 6.42055E-05 2.79897E-06
+    9   923.150 3.56887E-04 4.97135E-05 2.16720E-06
+   10   920.963 2.84278E-04 3.95992E-05 1.72628E-06
+   11   919.352 2.31771E-04 3.22851E-05 1.40743E-06
+   12   918.129 1.92348E-04 2.67936E-05 1.16804E-06
+   13   917.181 1.62155E-04 2.25878E-05 9.84689E-07
+   14   916.429 1.38498E-04 1.92925E-05 8.41033E-07
+   15   915.824 1.19611E-04 1.66615E-05 7.26340E-07
+   16   915.329 1.04314E-04 1.45306E-05 6.33446E-07
+   17   914.919 9.17397E-05 1.27791E-05 5.57091E-07
+   18   914.576 8.12784E-05 1.13219E-05 4.93564E-07
+   19   914.286 7.25069E-05 1.01000E-05 4.40299E-07
+   20   914.039 6.50549E-05 9.06198E-06 3.95047E-07
+   21   913.826 5.86816E-05 8.17421E-06 3.56345E-07
+   22   913.641 5.31918E-05 7.40949E-06 3.23008E-07
+   23   913.480 4.84261E-05 6.74563E-06 2.94068E-07
+   24   913.339 4.42740E-05 6.16726E-06 2.68854E-07
+   25   913.215 4.06311E-05 5.65981E-06 2.46733E-07
+   26   913.104 3.73821E-05 5.20723E-06 2.27003E-07
+   27   913.006 3.45377E-05 4.81102E-06 2.09731E-07
+   28   912.918 3.19891E-05 4.45601E-06 1.94255E-07
+   29   912.839 2.97110E-05 4.13867E-06 1.80421E-07
+   30   912.768 2.76635E-05 3.85346E-06 1.67987E-07
+   31   912.703 2.58178E-05 3.59636E-06 1.56779E-07
+   32   912.645 2.41479E-05 3.36374E-06 1.46638E-07
+   33   912.592 2.26347E-05 3.15296E-06 1.37450E-07
+   34   912.543 2.12567E-05 2.96100E-06 1.29081E-07
+   35   912.499 1.99967E-05 2.78549E-06 1.21430E-07
+   36   912.458 1.88476E-05 2.62543E-06 1.14452E-07
+   37   912.420 1.77928E-05 2.47850E-06 1.08047E-07
+   38   912.385 1.68222E-05 2.34330E-06 1.02153E-07
+   39   912.353 1.59286E-05 2.21882E-06 9.67268E-08
+   40   912.324 1.50996E-05 2.10334E-06 9.16925E-08
diff --git a/py/redrock/external/desi.py b/py/redrock/external/desi.py
@@ -593,6 +593,15 @@ def rrdesi(options=None, comm=None):
     parser.add_argument("--chi2-scan", type=str, default=None,
         required=False, help="Load the chi2-scan from the input file")
 
+    parser.add_argument("--zscan-galaxy", type=str, default='-0.005,1.7,3e-4',
+        required=False, help="Redshift scan parameters for galaxies: zmin,zmax,dz")
+
+    parser.add_argument("--zscan-qso", type=str, default='0.05,6.0,5e-4',
+        required=False, help="Redshift scan parameters for QSO: zmin,zmax,dz")
+
+    parser.add_argument("--zscan-star", type=str, default='-0.002,0.00201,4e-5',
+        required=False, help="Redshift scan parameters for stars: zmin,zmax,dz")
+
     parser.add_argument("-n", "--ntargets", type=int,
         required=False, help="the number of targets to process in each file")
 
@@ -868,7 +877,12 @@ def rrdesi(options=None, comm=None):
         # Read the template data
         # Pass both use_gpu (this proc) and args.gpu (if any proc is using GPU)
         dtemplates = load_dist_templates(dwave, templates=args.templates,
-            comm=comm, mp_procs=mpprocs, redistribute=redistribute_templates, use_gpu=use_gpu, gpu_mode=args.gpu)
+                                         zscan_galaxy=args.zscan_galaxy,
+                                         zscan_qso=args.zscan_qso,
+                                         zscan_star=args.zscan_star,
+                                         comm=comm, mp_procs=mpprocs,
+                                         redistribute=redistribute_templates,
+                                         use_gpu=use_gpu, gpu_mode=args.gpu)
 
         # Compute the redshifts, including both the coarse scan and the
         # refinement.  This function only returns data on the rank 0 process.

diff --git a/py/redrock/fitz.py b/py/redrock/fitz.py
@@ -253,7 +253,7 @@ def fitz(zchi2, redshifts, target, template, nminima=3, archetype=None, use_gpu=
         try:
             #Calculate xmin and xmax from template and pass as scalars 
             xmin = template.minwave*(1+zmin)
-            ximax = template.maxwave*(1+zmin)
+            xmax = 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)

diff --git a/py/redrock/rebin.py b/py/redrock/rebin.py
@@ -123,7 +123,7 @@ def centers2edges(centers):
 # of this code have already been tested and shown to perform no better
 # than numba on Intel haswell and KNL architectures.
 
-@numba.jit
+@numba.jit(nopython=True)
 def _trapz_rebin_1d(x, y, edges, results):
     '''
     Numba-friendly version of trapezoidal rebinning
@@ -175,7 +175,7 @@ def _trapz_rebin_1d(x, y, edges, results):
 
     return
 
-@numba.jit
+@numba.jit(nopython=True)
 def _trapz_rebin_batch(x, y, edges, myz, results, redshifted_x):
     '''
     Numba-friendly version of trapezoidal rebinning

diff --git a/py/redrock/templates.py b/py/redrock/templates.py
@@ -31,7 +31,8 @@ class Template(object):
 
     """
     def __init__(self, filename=None, spectype=None, redshifts=None,
-        wave=None, flux=None, subtype=None):
+                 wave=None, flux=None, subtype=None,
+                 zscan_galaxy=None, zscan_qso=None, zscan_star=None):
 
         if filename is not None:
             fx = None
@@ -73,15 +74,23 @@ def __init__(self, filename=None, spectype=None, redshifts=None,
             self._rrtype = hdr['RRTYPE'].strip().upper()
             if old_style_templates:
                 if self._rrtype == 'GALAXY':
-                    # redshifts = 10**np.arange(np.log10(1+0.005),
-                    # np.log10(1+2.0), 1.5e-4) - 1
-                    self._redshifts = 10**np.arange(np.log10(1-0.005),
-                        np.log10(1+1.7), 3e-4) - 1
+                    if zscan_galaxy is not None:
+                        zmin, zmax, dz = zscan_galaxy.split(',')
+                        self._redshifts = 10**np.arange(np.log10(1+float(zmin)), np.log10(1+float(zmax)), float(dz)) - 1
+                    else:
+                        self._redshifts = 10**np.arange(np.log10(1-0.005), np.log10(1+1.7), 3e-4) - 1
                 elif self._rrtype == 'STAR':
-                    self._redshifts = np.arange(-0.002, 0.00201, 4e-5)
+                    if zscan_star is not None:
+                        zmin, zmax, dz = zscan_star.split(',')
+                        self._redshifts = np.arange(float(zmin), float(zmax), float(dz))
+                    else:
+                        self._redshifts = np.arange(-0.002, 0.00201, 4e-5)
                 elif self._rrtype == 'QSO':
-                    self._redshifts = 10**np.arange(np.log10(1+0.05),
-                        np.log10(1+6.0), 5e-4) - 1
+                    if zscan_qso is not None:
+                        zmin, zmax, dz = zscan_qso.split(',')
+                        self._redshifts = 10**np.arange(np.log10(1+float(zmin)), np.log10(1+float(zmax)), float(dz)) - 1
+                    else:
+                        self._redshifts = 10**np.arange(np.log10(1+0.05), np.log10(1+6.0), 5e-4) - 1
                 else:
                     raise ValueError("Unknown redshift range to use for "
                         "template type {}".format(self._rrtype))
@@ -446,7 +455,9 @@ def cycle(self):
         return True
 
 
-def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1, redistribute=False, use_gpu=False, gpu_mode=False):
+def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1,
+                        zscan_galaxy=None, zscan_qso=None, zscan_star=None,
+                        redistribute=False, use_gpu=False, gpu_mode=False):
     """Read and distribute templates from disk.
 
     This reads one or more template files from disk and distributes them among
@@ -509,7 +520,8 @@ def load_dist_templates(dwave, templates=None, comm=None, mp_procs=1, redistribu
     template_data = list()
     if (comm is None) or (comm.rank == 0):
         for t in template_files:
-            template_data.append(Template(filename=t))
+            template_data.append(Template(filename=t, zscan_galaxy=zscan_galaxy,
+                                          zscan_star=zscan_star, zscan_qso=zscan_qso))
 
     if comm is not None:
         template_data = comm.bcast(template_data, root=0)

diff --git a/py/redrock/test/test_io.py b/py/redrock/test/test_io.py
@@ -1,6 +1,7 @@
 from __future__ import division, print_function
 
-import os
+import os, tempfile
+from shutil import rmtree
 import unittest
 from uuid import uuid1
 import numpy as np
@@ -18,14 +19,18 @@ class TestIO(unittest.TestCase):
     #- Create unique test filename in a subdirectory
     @classmethod
     def setUpClass(cls):
-        cls.testfile = 'test-{uuid}.h5'.format(uuid=uuid1())
+        cls.testDir = tempfile.mkdtemp()        
+        cls.testfile = os.path.join(cls.testDir, 'test-{uuid}.h5'.format(uuid=uuid1()))
 
     #- Cleanup test files if they exist
     @classmethod
     def tearDownClass(cls):
         if os.path.exists(cls.testfile):
             os.remove(cls.testfile)
 
+        if os.path.exists(cls.testDir):
+            rmtree(cls.testDir)
+
     def setUp(self):
         #- remove testfile if leftover from a previous test
         if os.path.exists(self.testfile):
@@ -90,9 +95,3 @@ def test_zscan_io(self):
                     self.assertTrue(np.all(d1==d2), 'data mismatch {}/{}/{}'.format(targetid, spectype, key))
 
 
-def test_suite():
-    """Allows testing of only this module with the command::
-
-        python setup.py test -m <modulename>
-    """
-    return unittest.defaultTestLoader.loadTestsFromName(__name__)