Merge pull request #783 from desihub/ADM-pix-rands

Functionality to split random catalogs by HEALPixel

bin/rewrite_randoms_in_healpix

@@ -0,0 +1,32 @@
+#!/usr/bin/env python
+
+from desitarget.randoms import rewrite_randoms_in_hp
+from desiutil.log import get_logger
+log = get_logger()
+
+# ADM default nside for writing HEALPixel-split files.
+nside = 8
+
+from argparse import ArgumentParser
+ap = ArgumentParser(description='Rewrite randoms in files that are split by HEALPixel and that can use the desitarget.io functions')
+ap.add_argument("infile",
+                help="Filename of monolithic random catalog to rewrite split    \
+                across HEALPixels. For full functionality, include the entire   \
+                directory string to be searched for the data release number.")
+ap.add_argument("outdirectory",
+                help="Directory to which to write the files. The sub-directory  \
+                structure and filenames are built on-the-fly from the header of \
+                the infile.")
+ap.add_argument("--nside", type=int,
+                help="The NESTED HEALPixel resolution at which to write the     \
+                files. Defaults to [{}]".format(nside),
+                default=nside)
+ap.add_argument("--verbose", action='store_true',
+                help="Log extra information")
+ap.add_argument("--ntest", type=int,
+                help=" Read first `ntest` randoms instead of the full catalog.")
+
+ns = ap.parse_args()
+
+rewrite_randoms_in_hp(ns.infile, ns.outdirectory, ns.nside,
+                      verbose=ns.verbose, ntest=ns.ntest)

doc/changes.rst

@@ -5,11 +5,13 @@ desitarget Change Log
 2.3.1 (unreleased)
 ------------------
 
-* Fix a bug in the randoms due to inaccrate rounding of pixel coordinates [`PR #782`_].
-    * This change will affect brick-based values in randoms (e.g.,
-      NOBS, MASKBITS and PSFSIZE).
+* Functionality to split random catalogs by HEALPixel [`PR #783`_].
+    * Allows the io/reading utilities to be used on the resulting files.
+* More accurately round pixel coordinates for randoms [`PR #782`_].
+    * Alters brick-based values for randoms (e.g. NOBS/MASKBITS/PSFSIZE).
 
 .. _`PR #782`: https://github.com/desihub/desitarget/pull/782
+.. _`PR #783`: https://github.com/desihub/desitarget/pull/783
 
 2.3.0 (2021-12-14)
 ------------------

py/desitarget/io.py

@@ -1445,6 +1445,96 @@ def write_gfas(targdir, data, indir=None, indir2=None, nside=None,
     return len(data), filename
 
 
+def write_randoms_in_hp(randoms, outdirname, nside, pixlist, hpx=None,
+                        verbose=True, infile=None, hdr=None):
+    """
+    Rewrite a random catalog split by HEALPixel.
+
+    Parameters
+    ----------
+    randoms : :class:`~numpy.array`
+        Entries a from random catalog as made by, e.g.,
+        :func:`desitarget.randoms.select_randoms()`.
+    outdirname : :class:`str`
+        Output directory to which to write the random catalog (the file
+        names are constructed on the fly).
+    nside : :class:`int`
+        (NESTED) HEALPixel nside that corresponds to `pixlist`.
+    hpx : :class:`~numpy.array`, optional, defaults to ``None``
+        The HEALPixel occupied by each of `randoms` in the NESTED scheme
+        at `nside`. If `hpx` isn't passed it will be calculated from the
+        required inputs. But, for large `randoms` array it may be quicker
+        to calculate it in advance. Must contain one HEALPixel number for
+        each row in `randoms`.
+    pixlist : :class:`list` or `int`
+        HEALPixels at `nside` at which to write the split catalogs.
+    verbose : :class:`bool`, optional, defaults to ``True``
+        If ``True`` then log target and file information.
+    infile : :class:`str`, optional, defaults to ``None``
+        Name of input random catalog to add to the output file header.
+    hdr : :class:`dict`, optional, defaults to ``None``
+        Dictionary to write as the header of the output file. Note that
+        `hdr` is a kwarg so will be modified-in-place.
+
+    Returns
+    -------
+    :class:`int`
+        The number of randoms that were written to file.
+    :class:`str`
+        The name of the file to which randoms were written.
+
+    Notes
+    -------
+        - A new column HPXPIXEL is added to `randoms`. HPXNSIDE=`nside`
+          and FILENSID=`nside` and "HPXNEST"=``True`` are added to the
+          output header (or overwritten, if they already exist).
+    """
+    t0 = time()
+
+    # ADM set up output file header.
+    if hdr is None:
+        hdr = {}
+    if infile is not None:
+        hdr["INFILE"] = infile
+    hdr["HPXNEST"] = True
+    hdr["HPXNSIDE"] = nside
+    hdr["FILENSID"] = nside
+
+    # ADM in case an integer was passed.
+    pixlist = np.atleast_1d(pixlist)
+
+    # ADM calculate hpx if it wasn't passed.
+    if hpx is None:
+        theta, phi = np.radians(90-randoms["DEC"]), np.radians(randoms["RA"])
+        hpx = hp.ang2pix(nside, theta, phi, nest=True)
+
+    # ADM write the catalogs...
+    nt, fn = 0, None
+    for pix in pixlist:
+        inpix = hpx == pix
+        if np.any(inpix):
+            # ADM set up a new array restricted to just the pixel of
+            # ADM interest and add the HPXPIXEL column.
+            nrows = np.sum(inpix)
+            dt = randoms.dtype.descr + [('HPXPIXEL', '>i8')]
+            outran = np.zeros(nrows, dtype=dt)
+            outran["HPXPIXEL"] = pix
+            for col in randoms.dtype.names:
+                outran[col] = randoms[col][inpix]
+            # ADM add the pixel to the output file hdr. dict is in case
+            # ADM hdr was passed as a FITSHDR, which can't be copied.
+            outhdr = dict(hdr).copy()
+            outhdr["FILEHPX"] = pix
+            # ADM write out the randoms in this pixel.
+            nt, fn = write_randoms(outdirname, outran, indir=infile,
+                                   nsidefile=nside, hpxlist=pix, extra=outhdr)
+            if verbose:
+                log.info('{} targets written to {}...t={:.1f}s'.format(
+                    nt, fn, time()-t0))
+
+    return nt, fn
+
+
 def write_randoms(targdir, data, indir=None, hdr=None, nside=None, supp=False,
                   nsidefile=None, hpxlist=None, resolve=True, north=None,
                   extra=None):
@@ -1516,7 +1606,13 @@ def write_randoms(targdir, data, indir=None, hdr=None, nside=None, supp=False,
             drstring = 'dr'+str(drint)
             depend.setdep(hdr, 'photcat', drstring)
         except (ValueError, IndexError, AttributeError):
-            drint = None
+            # ADM can also try reading the DR from a passed header.
+            try:
+                drint = int(extra["DR"])
+            except KeyError:
+                drint = None
+    else:
+        drint = None
 
     # ADM whether this is a north-specific or south-specific file.
     region = None
@@ -1544,6 +1640,8 @@ def write_randoms(targdir, data, indir=None, hdr=None, nside=None, supp=False,
 
     # ADM add the extra keywords to the header.
     hdr["DR"] = drint
+    if drint is None:
+        drint = "X"
     if extra is not None:
         for key in extra:
             hdr[key] = extra[key]
@@ -2422,7 +2520,9 @@ def find_target_files(targdir, dr='X', flavor="targets", survey="main",
     seed : :class:`int` or `str`, optional
         If `seed` is not ``None``, then it is added to the file name just
         before the ".fits" extension (i.e. "-8.fits" for `seed` of 8).
-        Only relevant if `flavor` is "randoms".
+        Only relevant if `flavor` is "randoms". If `seed` is not ``None``
+        and `hp` is not ``None`` then `seed` is added to the file name
+        AFTER `hp` (i.e. "-8-hp-44.fits" for `seed` of 8 and `hp` of 44).
     region : :class:`int`, optional
         If `region` is not ``None``, then it is added to the directory
         name after `resolve`. Only relevant if `flavor` is "randoms".
@@ -2573,7 +2673,11 @@ def find_target_files(targdir, dr='X', flavor="targets", survey="main",
         # ADM note that these clauses won't do anything
         # ADM unless a file name was already constructed.
         if seed is not None:
-            fn = fn.replace(".{}".format(ender), "-{}.{}".format(seed, ender))
+            if hp is None:
+                fn = fn.replace(".{}".format(ender),
+                                "-{}.{}".format(seed, ender))
+            else:
+                fn = fn.replace("hp", "{}-hp".format(seed))
         if supp:
             justfn = os.path.basename(fn)
             justfn = justfn.replace("randoms", "randoms-outside")
@@ -2788,7 +2892,7 @@ def read_target_files(filename, columns=None, rows=None, header=False,
     start = time()
     # ADM start with some checking that this is a target file.
     targtypes = ["TARGETS", "GFA_TARGETS", "SKY_TARGETS",
-                 "MASKS", "MTL", "SCND_TARGETS"]
+                 "MASKS", "MTL", "SCND_TARGETS", "RANDOMS"]
     # ADM read in the FITS extension info.
     f = fitsio.FITS(filename)
     if len(f) != 2:

py/desitarget/randoms.py

@@ -113,6 +113,62 @@ def finalize_randoms(randoms):
     return finalize(randoms, dt, dt*0, dt*0, randoms=True)
 
 
+def rewrite_randoms_in_hp(fn, outdn, nside, verbose=True, ntest=None):
+    """Rewrite randoms in HEALPixels, in files that can use io functions.
+
+    Parameters
+    ----------
+    fn : :class:`str`
+        The filename of a monolithic random catalog to be rewritten split
+        across HEALPixels. For full functionality, include the entire
+        directory string to be searched for the data release number.
+    outdn : :class:`str`
+        Output directory to which to write the files. The sub-directory
+        structure and filenames are built on-the-fly from the header of
+        `fn` (which must correspond to a typical random catalog).
+    nside : :class:`int`
+        The NESTED HEALPixel resolution at which to write files.
+    verbose : :class:`bool`, optional, defaults to ``True``
+        If ``True`` then log extra information.
+    ntest : :class:`int`, optional, default to ``None``
+        If passed (and not ``None``) then read the first `ntest` randoms
+        instead of the entire catalog. Useful for testing.
+
+    Returns
+    -------
+    Nothing, but rewrites the randoms in `fn` split across HEALPixels.
+
+    Notes
+    -----
+    - Really just a simple wrapper on :func:`io.write_randoms_in_hp()`
+      that also reads the input file.
+    """
+    t0 = time()
+    from desitarget.io import write_randoms_in_hp
+
+    # ADM read in the random catalog (and header).
+    if ntest is None:
+        randoms, hdr = fitsio.read(fn, header=True)
+    else:
+        randoms, hdr = fitsio.read(fn, header=True, rows=np.arange(ntest))
+
+    if verbose:
+        log.info('Read {} randoms from {}...t={:.1f}m'
+                 .format(len(randoms), fn, (time()-t0)/60.))
+
+    # ADM determine all possible pixels at the passed nside.
+    pixlist = np.arange(hp.nside2npix(nside))
+
+    # ADM write the catalog back out split-by-healpixel
+    _, _ = write_randoms_in_hp(randoms, outdn, nside, pixlist,
+                               verbose=verbose, infile=fn, hdr=hdr)
+
+    if verbose:
+        log.info('Done...t={:.1f}m'.format((time()-t0)/60.))
+
+    return
+
+
 def add_default_mtl(randoms, seed):
     """Add default columns that are added by MTL.