Merge pull request #200 from desihub/mock-scaling

fix isLRG vs. isLRG_colors

bin/mpi_select_mock_targets

@@ -120,16 +120,22 @@ else:
 
 params = comm.bcast(params, root=0)
 
+#- Generate a different random seed for each pixel
+np.random.seed(args.seed)
+randseeds = np.random.randint(2**31, size=len(pixels))
+
 #- Split the pixels into groups for each rank
 iedges = np.linspace(0, len(pixels), size+1, dtype=int)
 rankpix = pixels[iedges[rank]:iedges[rank+1]]
+rankseeds = randseeds[iedges[rank]:iedges[rank+1]]
 log.info('rank {} processes {} pixels {}'.format(rank, iedges[rank+1]-iedges[rank], rankpix))
 
 if len(rankpix) > 0:
-    #- Process in groups of 4 to balance mock reading overhead
-    #- vs. writing partial progress in case job times out
-    n = 4
+    #- Process one pixel at a time to avoid blowing out memory, but structure
+    #- it in a way that we could expand to multiple pixels per call if/when
+    #- we use less memory.
+    n = 1
     for i in range(0, len(rankpix), n):
-        targets_truth(params, args.output_dir, seed=args.seed, nside=args.nside,
-                      nproc=args.nproc, verbose=args.verbose, clobber=True,
+        targets_truth(params, args.output_dir, seed=rankseeds[i], nside=args.nside,
+                      nproc=args.nproc, verbose=args.verbose,
                       healpixels=rankpix[i:i+n])

doc/changes.rst

@@ -5,7 +5,12 @@ desitarget Change Log
 0.14.1 (unreleased)
 -------------------
 
-* no changes yet
+* Changes for mocks [`PR #200`_]:
+  * Fix isLRG vs. isLRG_colors
+  * Correct random seeds when processing pix in parallel
+  * Misc other small bug fixes
+
+.. _`PR #200`: https://github.com/desihub/desitarget/pull/200
 
 0.14.0 (2017-07-10)
 -------------------

py/desitarget/mock/build.py

@@ -350,13 +350,13 @@ def get_spectra_onebrick(target_name, mockformat, thisbrick, brick_info, Spectra
     truth = empty_truth_table(nobj)
 
     for key in ('RA', 'DEC', 'BRICKNAME'):
-        targets[key] = source_data[key][onbrick]
+        targets[key][:] = source_data[key][onbrick]
 
     for band, depthkey in zip((1, 2, 4), ('DEPTH_G', 'DEPTH_R', 'DEPTH_Z')):
         targets['DECAM_DEPTH'][:, band] = brick_info[depthkey][brickindx]
     for band, depthkey in zip((1, 2, 4), ('GALDEPTH_G', 'GALDEPTH_R', 'GALDEPTH_Z')):
         targets['DECAM_GALDEPTH'][:, band] = brick_info[depthkey][brickindx]
-    targets['EBV'] = brick_info['EBV'][brickindx]
+    targets['EBV'][:] = brick_info['EBV'][brickindx]
 
     # Use the point-source depth for point sources, although this should really
     # be tied to the morphology.
@@ -380,21 +380,21 @@ def get_spectra_onebrick(target_name, mockformat, thisbrick, brick_info, Spectra
     if 'SHAPEEXP_R' in source_data.keys(): # not all target types have shape information
         for key in ('SHAPEEXP_R', 'SHAPEEXP_E1', 'SHAPEEXP_E2',
                     'SHAPEDEV_R', 'SHAPEDEV_E1', 'SHAPEDEV_E2'):
-            targets[key] = source_data[key][onbrick]
+            targets[key][:] = source_data[key][onbrick]
 
     for key, source_key in zip( ['MOCKID', 'SEED', 'TEMPLATETYPE', 'TEMPLATESUBTYPE', 'TRUESPECTYPE'],
                                 ['MOCKID', 'SEED', 'TEMPLATETYPE', 'TEMPLATESUBTYPE', 'TRUESPECTYPE'] ):
         if isinstance(source_data[source_key], np.ndarray):
-            truth[key] = source_data[source_key][onbrick]
+            truth[key][:] = source_data[source_key][onbrick]
         else:
-            truth[key] = np.repeat(source_data[source_key], nobj)
+            truth[key][:] = np.repeat(source_data[source_key], nobj)
 
     # Sky targets are a special case without redshifts.
     if target_name == 'sky':
         select_targets_function(targets, truth)
         return [targets, truth]
 
-    truth['TRUEZ'] = source_data['Z'][onbrick]
+    truth['TRUEZ'][:] = source_data['Z'][onbrick]
 
     # For FAINTSTAR targets, preselect stars that are going to pass target
     # selection cuts without actually generating spectra, in order to save
@@ -454,7 +454,7 @@ def get_spectra_onebrick(target_name, mockformat, thisbrick, brick_info, Spectra
 
     for key in ('TEMPLATEID', 'MAG', 'DECAM_FLUX', 'WISE_FLUX',
                 'OIIFLUX', 'HBETAFLUX', 'TEFF', 'LOGG', 'FEH'):
-        truth[key] = meta[key]
+        truth[key][:] = meta[key]
 
     # Perturb the photometry based on the variance on this brick and apply
     # target selection.
@@ -555,7 +555,7 @@ def targets_truth(params, output_dir='.', realtargets=None, seed=None, verbose=F
         log = get_logger(DEBUG)
     else:
         log = get_logger()
-    
+
     if params is None:
         log.fatal('Required params input not given!')
         raise ValueError
@@ -802,13 +802,13 @@ def _update_spectra_status(result):
     targetid = rand.randint(2**62, size=ntarget + nsky)
     subpriority = rand.uniform(0.0, 1.0, size=ntarget + nsky)
 
-    truth['TARGETID'] = targetid[:ntarget]
-    targets['TARGETID'] = targetid[:ntarget]
-    targets['SUBPRIORITY'] = subpriority[:ntarget]
+    truth['TARGETID'][:] = targetid[:ntarget]
+    targets['TARGETID'][:] = targetid[:ntarget]
+    targets['SUBPRIORITY'][:] = subpriority[:ntarget]
 
     if nsky > 0:
-        skytargets['TARGETID'] = targetid[ntarget:ntarget+nsky]
-        skytargets['SUBPRIORITY'] = subpriority[ntarget:ntarget+nsky]
+        skytargets['TARGETID'][:] = targetid[ntarget:ntarget+nsky]
+        skytargets['SUBPRIORITY'][:] = subpriority[ntarget:ntarget+nsky]
 
     # Write the final catalogs out by healpixel.
     if seed is None:

py/desitarget/mock/io.py

@@ -23,7 +23,7 @@
 from desimodel.footprint import radec2pix
 
 from desiutil.log import get_logger, DEBUG
-log = get_logger(DEBUG)
+log = get_logger()
 
 """How to distribute 52 user bits of targetid.
 

py/desitarget/mock/selection.py

@@ -39,7 +39,7 @@ def _star_select(self, targets, truth):
         targets.
 
         """ 
-        from desitarget.cuts import isBGS_faint, isELG, isLRG, isQSO_colors
+        from desitarget.cuts import isBGS_faint, isELG, isLRG_colors, isQSO_colors
 
         gflux = targets['DECAM_FLUX'][..., 1]
         rflux = targets['DECAM_FLUX'][..., 2]
@@ -69,7 +69,8 @@ def _star_select(self, targets, truth):
         truth['CONTAM_TARGET'] |= (elg != 0) * self.contam_mask.ELG_CONTAM
 
         # Select stellar contaminants for LRG targets.
-        lrg = isLRG(rflux=rflux, zflux=zflux, w1flux=w1flux)
+        lrg = isLRG_colors(gflux=gflux, rflux=rflux, zflux=zflux,
+                           w1flux=w1flux, w2flux=w2flux)
         targets['DESI_TARGET'] |= (lrg != 0) * self.desi_mask.LRG
         targets['DESI_TARGET'] |= (lrg != 0) * self.desi_mask.LRG_SOUTH
         for oo in self.desi_mask.LRG.obsconditions.split('|'):
@@ -197,15 +198,16 @@ def faintstar_select(self, targets, truth, boss_std=None):
 
     def lrg_select(self, targets, truth, boss_std=None):
         """Select LRG targets."""
-        from desitarget.cuts import isLRG, isQSO_colors
+        from desitarget.cuts import isLRG_colors, isQSO_colors
 
         gflux = targets['DECAM_FLUX'][..., 1]
         rflux = targets['DECAM_FLUX'][..., 2]
         zflux = targets['DECAM_FLUX'][..., 4]
         w1flux = targets['WISE_FLUX'][..., 0]
         w2flux = targets['WISE_FLUX'][..., 1]
 
-        lrg = isLRG(rflux=rflux, zflux=zflux, w1flux=w1flux)
+        lrg = isLRG_colors(gflux=gflux, rflux=rflux, zflux=zflux,
+                           w1flux=w1flux, w2flux=w2flux)
 
         targets['DESI_TARGET'] |= (lrg != 0) * self.desi_mask.LRG
         targets['DESI_TARGET'] |= (lrg != 0) * self.desi_mask.LRG_SOUTH