truetype to match spectype, not sourcetype

doc/mock_example/input.yaml

@@ -1,6 +1,12 @@
 #mock target configuration file
 output_dir: /tmp
 run_label:  GX_100PC_MXXL_JT2P4Y_JULY16_B
+subset:
+    ra_dec_cut: False
+    min_ra: 200
+    max_ra: 220
+    min_dec: 0
+    max_dec: 20
 sources:
     SKY: {
         number_density: 1400.0,

py/desitarget/mock/build.py

@@ -54,14 +54,33 @@ def targets_truth(params):
         source_path = params['sources'][source_name]['root_mock_dir']
         source_dict = params['sources'][source_name]
 
-
         print('type: {} format: {}'.format(source_name, source_format))
         function = 'read_'+source_format
         if 'mock_name' in source_dict.keys():
             mock_name = source_dict['mock_name']
         else:
             mock_name = None
         result = getattr(mockio, function)(source_path, source_name, mock_name=mock_name)
+
+        if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
+            print('Trimming {} to RA,dec subselection'.format(source_name))
+            ii  = (result['RA']  >= params['subset']['min_ra']) & \
+                  (result['RA']  <= params['subset']['max_ra']) & \
+                  (result['DEC'] >= params['subset']['min_dec']) & \
+                  (result['DEC'] <= params['subset']['max_dec'])
+
+            #- Trim RA,DEC,Z, ... columns to subselection
+            #- Different types of mocks have different metadata, so assume
+            #- that any ndarray of the same length as number of targets should
+            #- be trimmed.
+            ntargets = len(result['RA'])
+            for key in result:
+                if isinstance(result[key], np.ndarray) and len(result[key]) == ntargets:
+                    result[key] = result[key][ii]
+
+            #- Add min/max ra/dec to source_dict for use in density estimates
+            source_dict.update(params['subset'])
+
         source_data_all[source_name] = result
 
     print('loaded {} mock sources'.format(len(source_data_all)))
@@ -86,7 +105,7 @@ def targets_truth(params):
     bgs_target_total = np.empty(0, dtype='i8')
     mws_target_total = np.empty(0, dtype='i8')
     true_type_total = np.empty(0, dtype='S10')
-    true_subtype_total = np.empty(0, dtype='S10')
+    source_type_total = np.empty(0, dtype='S10')
     obsconditions_total = np.empty(0, dtype='uint16')
 
 
@@ -112,8 +131,20 @@ def targets_truth(params):
         # define names that go into Truth
         n = len(source_data['RA'][ii])
         if source_name not in ['STD_FSTAR', 'SKY']:
-            true_type = np.zeros(n, dtype='S10'); true_type[:] = source_name
-            true_subtype = np.zeros(n, dtype='S10');true_subtype[:] = source_name
+            true_type_map = {
+                'ELG': 'GALAXY',
+                'LRG': 'GALAXY',
+                'BGS': 'GALAXY',
+                'QSO': 'QSO',
+                'STD_FSTAR': 'STAR',
+                'MWS_MAIN': 'STAR',
+                'MWS_WD': 'STAR',
+                'SKY': 'SKY',
+            }
+            source_type = np.zeros(n, dtype='S10')
+            source_type[:] = source_name
+            true_type = np.zeros(n, dtype='S10')
+            true_type[:] = true_type_map[source_name]
 
 
         #define obsconditions
@@ -152,7 +183,7 @@ def targets_truth(params):
             bgs_target_total = np.append(bgs_target_total, bgs_target)
             mws_target_total = np.append(mws_target_total, mws_target)
             true_type_total = np.append(true_type_total, true_type)
-            true_subtype_total = np.append(true_subtype_total, true_subtype)
+            source_type_total = np.append(source_type_total, source_type)
             obsconditions_total = np.append(obsconditions_total, source_obsconditions)
 
 
@@ -173,83 +204,64 @@ def targets_truth(params):
     print('Great total of {} targets {} stdstars {} sky pos'.format(n_target, n_star, n_sky))
     targetid = np.random.randint(2**62, size=n)
 
-    # make a subselection in 
-    if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
-        ii_stars = (ra_stars > params['subset']['min_ra']) & (ra_stars< params['subset']['max_ra'])
-        ii_stars &= (dec_stars > params['subset']['min_dec']) & (dec_stars< params['subset']['max_dec'])
-
-        ii_sky = (ra_sky > params['subset']['min_ra']) & (ra_sky< params['subset']['max_ra'])
-        ii_sky &= (dec_sky > params['subset']['min_dec']) & (dec_sky< params['subset']['max_dec'])
-
-        ii_targets = (ra_total > params['subset']['min_ra']) & (ra_total < params['subset']['max_ra'])
-        ii_targets &= (dec_total > params['subset']['min_dec']) & (dec_total < params['subset']['max_dec'])
-        print('IDs to select subset ready')
-
-    #write to disk
+    # write to disk
     if 'STD_FSTAR' in source_defs.keys():
         subprior = np.random.uniform(0., 1., size=n_star)
-        brickname = desispec.brick.brickname(ra_stars, dec_stars)
         #write the Std Stars to disk
         print('Started writing StdStars file')
         stars_filename = os.path.join(params['output_dir'], 'stdstars.fits')
         stars = Table()
         stars['TARGETID'] = targetid[n_target:n_target+n_star]
-        stars['BRICKNAME'] = brickname
         stars['RA'] = ra_stars
         stars['DEC'] = dec_stars
         stars['DESI_TARGET'] = desi_target_stars
         stars['BGS_TARGET'] = bgs_target_stars
         stars['MWS_TARGET'] = mws_target_stars
         stars['SUBPRIORITY'] = subprior
         stars['OBSCONDITIONS'] = obsconditions_stars
-        if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
-            stars = stars[ii_stars]
-            print('subsetting in std_stars data done')
+        # if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
+        #     stars = stars[ii_stars]
+        #     print('subsetting in std_stars data done')
+        brickname = desispec.brick.brickname(stars['RA'], stars['DEC'])
+        stars['BRICKNAME'] = brickname
         stars.write(stars_filename, overwrite=True)
         print('Finished writing stdstars file')
 
     if 'SKY' in source_defs.keys():
         subprior = np.random.uniform(0., 1., size=n_sky)
-        brickname = desispec.brick.brickname(ra_sky, dec_sky)
         #write the Std Stars to disk
         print('Started writing sky to file')
         sky_filename = os.path.join(params['output_dir'], 'sky.fits')
         sky = Table()
         sky['TARGETID'] = targetid[n_target+n_star:n_target+n_star+n_sky]
-        sky['BRICKNAME'] = brickname
         sky['RA'] = ra_sky
         sky['DEC'] = dec_sky
         sky['DESI_TARGET'] = desi_target_sky
         sky['BGS_TARGET'] = bgs_target_sky
         sky['MWS_TARGET'] = mws_target_sky
         sky['SUBPRIORITY'] = subprior
         sky['OBSCONDITIONS'] = obsconditions_sky
-        if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
-            sky = sky[ii_sky]
-            print('subsetting in sky data done')
+        brickname = desispec.brick.brickname(sky['RA'], sky['DEC'])
+        sky['BRICKNAME'] = brickname
         sky.write(sky_filename, overwrite=True)
         print('Finished writing sky file')
 
     if n_target > 0:
         subprior = np.random.uniform(0., 1., size=n_target)
-        brickname = desispec.brick.brickname(ra_total, dec_total)
-
-    # write the Targets to disk
+        # write the Targets to disk
         print('Started writing Targets file')
         targets_filename = os.path.join(params['output_dir'], 'targets.fits')
         targets = Table()
         targets['TARGETID'] = targetid[0:n_target]
-        targets['BRICKNAME'] = brickname
         targets['RA'] = ra_total
         targets['DEC'] = dec_total
         targets['DESI_TARGET'] = desi_target_total
         targets['BGS_TARGET'] = bgs_target_total
         targets['MWS_TARGET'] = mws_target_total
         targets['SUBPRIORITY'] = subprior
         targets['OBSCONDITIONS'] = obsconditions_total
-        if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
-            targets = targets[ii_targets]
-            print('subsetting in targets data done')
+        brickname = desispec.brick.brickname(targets['RA'], targets['DEC'])
+        targets['BRICKNAME'] = brickname
         targets.write(targets_filename, overwrite=True)
         print('Finished writing Targets file')
 
@@ -262,21 +274,17 @@ def targets_truth(params):
         mtl_table.write(mtl_filename, overwrite=True)
         print('Finished writing mtl file')
 
-
-    # write the Truth to disk
+        # write the Truth to disk
         print('Started writing Truth file')
         truth_filename = os.path.join(params['output_dir'], 'truth.fits')
         truth = Table()
         truth['TARGETID'] = targetid[0:n_target]
-        truth['BRICKNAME'] = brickname
         truth['RA'] = ra_total
         truth['DEC'] = dec_total
         truth['TRUEZ'] = z_total
         truth['TRUETYPE'] = true_type_total
-        truth['TRUESUBTYPE'] = true_subtype_total
-        if ('subset' in params.keys()) & (params['subset']['ra_dec_cut']==True):
-            truth = truth[ii_targets]
-            print('subsetting in truth data done')
+        truth['SOURCETYPE'] = source_type_total
+        truth['BRICKNAME'] = brickname
         truth.write(truth_filename, overwrite=True)
         print('Finished writing Truth file')
 

py/desitarget/mock/io.py

@@ -62,7 +62,7 @@ def _load_mock_mws_file(filename):
         'SDSS[grz]_obs': :class: `numpy.ndarray`
              Apparent magnitudes in SDSS grz bands, including extinction.
     """
-
+    print('Reading '+filename)
     C_LIGHT = 299792.458
     desitarget.io.check_fitsio_version()
     data = fitsio.read(filename,
@@ -213,15 +213,23 @@ def read_galaxia(mock_dir, target_type, mock_name=None):
 
     # Read each file
     print('Reading individual mock files')
-    target_list = list()
-    file_list   = list()
-    nfiles      = 0
-    for mock_file in iter_mock_files:
-        nfiles += 1
-        data_this_file = _load_mock_mws_file(mock_file)
-        target_list.append(data_this_file)
-        file_list.append(mock_file)
-        print('read file {} {}'.format(nfiles, mock_file))
+    # target_list = list()
+    # file_list   = list()
+    # nfiles      = 0
+    # for mock_file in iter_mock_files:
+    #     nfiles += 1
+    #     data_this_file = _load_mock_mws_file(mock_file)
+    #     target_list.append(data_this_file)
+    #     file_list.append(mock_file)
+    #     print('read file {} {}'.format(nfiles, mock_file))
+
+    file_list = list(iter_mock_files)
+    nfiles = len(file_list)
+    import multiprocessing
+    ncpu = max(1, multiprocessing.cpu_count() // 2)
+    print('using {} parallel readers'.format(ncpu))
+    p = multiprocessing.Pool(ncpu)
+    target_list = p.map(_load_mock_mws_file, file_list)
 
     print('Read {} files'.format(nfiles))
 

py/desitarget/mock/selection.py

@@ -318,7 +318,7 @@ def build_mock_target(root_mock_dir='', output_dir='',
 
 
 
-def estimate_density(ra, dec):
+def estimate_density(ra, dec, bounds=(170, 190, 0, 35)):
     """Estimate the number density from a small patch
     
     Args:
@@ -327,15 +327,20 @@ def estimate_density(ra, dec):
         dec: array_like
             An array with Dec positions.
 
+    Options:
+        bounds: (min_ra, max_ra, min_dec, max_dec) to use for density
+            estimation, assuming complete coverage within those bounds [deg]
+
     Returns:
         density: float
            Object number density computed over a small patch.
     """
     density = 0.0 
 
-    footprint_area = 20. * 35.* np.sin(35. * np.pi/180.)/(35. * np.pi/180.)
-    smalldata = ra[(ra>170.) & (ra<190.) & (dec>0.) & (dec<35.)]
-    n_in = len(smalldata)
+    min_ra, max_ra, min_dec, max_dec = bounds
+    footprint_area = (max_ra-min_ra) * (np.sin(max_dec*np.pi/180.) - np.sin(min_dec*np.pi/180.)) * 180 / np.pi
+
+    n_in = np.count_nonzero((ra>=min_ra) & (ra<max_ra) & (dec>=min_dec) & (dec<max_dec))
     density = n_in/footprint_area
     if(n_in==0):
         density = 1E-6
@@ -352,12 +357,17 @@ def ndens_select(data, source_name, **kwargs):
     dec = data['DEC']
     z = data['Z']
 
-    if ('min_z' in data) & ('max_z' in data):
+    if ('min_z' in kwargs) & ('max_z' in kwargs):
         in_z = ((z>=kwargs['min_z']) & (z<=kwargs['max_z']))
     else:
         in_z = z>0.0
 
-    mock_dens = estimate_density(ra[in_z], dec[in_z])
+    try:
+        bounds = kwargs['min_ra'], kwargs['max_ra'], kwargs['min_dec'], kwargs['max_dec']
+        mock_dens = estimate_density(ra[in_z], dec[in_z], bounds=bounds)
+    except KeyError:
+        mock_dens = estimate_density(ra[in_z], dec[in_z])
+
     frac_keep = min(kwargs['number_density']/mock_dens , 1.0)
     if mock_dens < kwargs['number_density']:
         print("WARNING: mock cannot achieve the goal density for source {} Goal {}. Mock {}".format(source_name, kwargs['number_density'], mock_dens))