def test_readwrite_tractor(self): tractorfile = io.list_tractorfiles(self.datadir)[0] sweepfile = io.list_sweepfiles(self.datadir)[0] data = io.read_tractor(sweepfile) self.assertEqual(len(data), 6) #- test data has 6 objects per file data = io.read_tractor(tractorfile) self.assertEqual(len(data), 6) #- test data has 6 objects per file data, hdr = io.read_tractor(tractorfile, header=True) self.assertEqual(len(data), 6) #- test data has 6 objects per file #ADM check PHOTSYS got added in writing targets io.write_targets(self.testfile, data, indir=self.datadir) #ADM use fits read wrapper in io to correctly handle whitespace d2, h2 = io.whitespace_fits_read(self.testfile, header=True) self.assertEqual( list(data.dtype.names) + ["PHOTSYS"], list(d2.dtype.names)) #ADM check PHOTSYS and HPXPIXEL got added writing targets with NSIDE request io.write_targets(self.testfile, data, nside=64, indir=self.datadir) #ADM use fits read wrapper in io to correctly handle whitespace d2, h2 = io.whitespace_fits_read(self.testfile, header=True) self.assertEqual( list(data.dtype.names) + ["HPXPIXEL", "PHOTSYS"], list(d2.dtype.names)) for column in data.dtype.names: self.assertTrue(np.all(data[column] == d2[column]))
def test_readwrite_tractor(self): tractorfile = io.list_tractorfiles(self.datadir)[0] sweepfile = io.list_sweepfiles(self.datadir)[0] data = io.read_tractor(sweepfile) self.assertEqual(len(data), 6) # - test data has 6 objects per file data = io.read_tractor(tractorfile) self.assertEqual(len(data), 6) # - test data has 6 objects per file data, hdr = io.read_tractor(tractorfile, header=True) self.assertEqual(len(data), 6) # - test data has 6 objects per file # ADM check that input and output columns are the same. _, filename = io.write_targets(self.testdir, data, indir=self.datadir) # ADM use fits read wrapper in io to correctly handle whitespace. d2, h2 = io.whitespace_fits_read(filename, header=True) self.assertEqual(list(data.dtype.names), list(d2.dtype.names)) # ADM check HPXPIXEL got added writing targets with NSIDE request. _, filename = io.write_targets(self.testdir, data, nside=64, indir=self.datadir) # ADM use fits read wrapper in io to correctly handle whitespace. d2, h2 = io.whitespace_fits_read(filename, header=True) self.assertEqual( list(data.dtype.names) + ["HPXPIXEL"], list(d2.dtype.names)) for column in data.dtype.names: kind = data[column].dtype.kind # ADM whitespace_fits_read() doesn't convert string data # ADM types identically for every version of fitsio. if kind == 'U' or kind == 'S': self.assertTrue( np.all( data[column] == d2[column].astype(data[column].dtype))) else: self.assertTrue(np.all(data[column] == d2[column]))
def test_fix_dr1(self): '''test the DR1 TYPE dype fix (make everything S4)''' # - First, break it files = io.list_sweepfiles(self.datadir) objects = io.read_tractor(files[0]) dt = objects.dtype.descr for i in range(len(dt)): if dt[i][0] == 'TYPE': dt[i] = ('TYPE', 'S10') break badobjects = objects.astype(np.dtype(dt)) newobjects = io.fix_tractor_dr1_dtype(badobjects) self.assertEqual(newobjects['TYPE'].dtype, np.dtype('S4'))
def test_fix_dr1(self): '''test the DR1 TYPE dype fix (make everything S4)''' #- First, break it files = io.list_sweepfiles(self.datadir) objects = io.read_tractor(files[0]) dt = objects.dtype.descr for i in range(len(dt)): if dt[i][0] == 'TYPE': dt[i] = ('TYPE', 'S10') break badobjects = objects.astype(np.dtype(dt)) newobjects = io.fix_tractor_dr1_dtype(badobjects) self.assertEqual(newobjects['TYPE'].dtype, np.dtype('S4'))
def test_readwrite_tractor(self): tractorfile = io.list_tractorfiles(self.datadir)[0] sweepfile = io.list_sweepfiles(self.datadir)[0] data = io.read_tractor(sweepfile) data = io.read_tractor(tractorfile) self.assertEqual(len(data), 6) #- test data has 6 objects per file data, hdr = io.read_tractor(tractorfile, header=True) self.assertEqual(len(data), 6) #- test data has 6 objects per file io.write_targets(self.testfile, data, indir=self.datadir) d2, h2 = fits.getdata(self.testfile, header=True) self.assertEqual(h2['DEPVER02'], self.datadir) self.assertEqual(data.dtype.names, d2.dtype.names) for column in data.dtype.names: self.assertTrue(np.all(data[column] == d2[column]))
def setUpClass(cls): cls.datadir = resource_filename('desitarget.test', 't') cls.tractorfiles = sorted(io.list_tractorfiles(cls.datadir)) cls.sweepfiles = sorted(io.list_sweepfiles(cls.datadir)) # ADM find which HEALPixels are covered by test sweeps files. cls.nside = 32 pixlist = [] for fn in cls.sweepfiles: objs = fitsio.read(fn) theta, phi = np.radians(90-objs["DEC"]), np.radians(objs["RA"]) pixels = hp.ang2pix(cls.nside, theta, phi, nest=True) pixlist.append(pixels) cls.pix = np.unique(pixlist) # ADM set up the GAIA_DIR environment variable. cls.gaiadir_orig = os.getenv("GAIA_DIR") os.environ["GAIA_DIR"] = resource_filename('desitarget.test', 't4')
def test_list_sweepfiles(self): files = io.list_sweepfiles(self.datadir) self.assertEqual(len(files), 3) for x in files: self.assertTrue(os.path.basename(x).startswith('sweep')) self.assertTrue(os.path.basename(x).endswith('.fits'))
def setUpClass(cls): cls.datadir = resource_filename('desitarget.test', 't') cls.tractorfiles = sorted(io.list_tractorfiles(cls.datadir)) cls.sweepfiles = sorted(io.list_sweepfiles(cls.datadir)) cls.cmxdir = resource_filename('desitarget.test', 't3')
import os import fitsio import numpy as np from time import time from pkg_resources import resource_filename from desitarget.gaiamatch import find_gaia_files from desitarget import io start = time() # ADM choose the Gaia files to cover the same object # ADM locations as the sweeps/tractor files. datadir = resource_filename('desitarget.test', 't') tractorfiles = sorted(io.list_tractorfiles(datadir)) sweepfiles = sorted(io.list_sweepfiles(datadir)) # ADM read in each of the relevant Gaia files. gaiafiles = [] for fn in sweepfiles + tractorfiles: objs = fitsio.read(fn, columns=["RA", "DEC"]) gaiafiles.append(find_gaia_files(objs, neighbors=False)) gaiafiles = np.unique(gaiafiles) # ADM loop through the Gaia files and write out some rows # ADM to the "t4" unit test directory. if not os.path.exists("t4"): os.makedirs(os.path.join("t4", "healpix")) for fn in gaiafiles: objs = fitsio.read(fn) outfile = os.path.join("t4", "healpix", os.path.basename(fn))
def setUpClass(cls): # py/desitarget/test -> etc/datadir thisdir, thisfile = os.path.split(__file__) cls.datadir = os.path.abspath(thisdir+'/../../../') + '/etc/testdata' cls.tractorfiles = io.list_tractorfiles(cls.datadir) cls.sweepfiles = io.list_sweepfiles(cls.datadir)
def collect_bright_stars( bands, maglim, numproc=4, rootdirname='/global/project/projectdirs/cosmo/data/legacysurvey/dr3.1/sweep/3.1', outfilename=None, verbose=False): """Extract a structure from the sweeps containing only bright stars in a given band to a given magnitude limit Parameters ---------- bands : :class:`str` A magnitude band from the sweeps, e.g., "G", "R", "Z". Can pass multiple bands as string, e.g. "GRZ", in which case maglim has to be a list of the same length as the string maglim : :class:`float` The upper limit in that magnitude band for which to assemble a list of bright stars. Can pass a list of magnitude limits, in which case bands has to be a string of the same length (e.g., "GRZ" for [12.3,12.7,12.6] numproc : :class:`int`, optional Number of processes over which to parallelize rootdirname : :class:`str`, optional, defaults to dr3 Root directory containing either sweeps or tractor files...e.g. for dr3 this might be /global/project/projectdirs/cosmo/data/legacysurvey/dr3/sweeps/dr3.1 outfilename : :class:`str`, optional, defaults to not writing anything to file (FITS) File name to which to write the output structure of bright stars verbose : :class:`bool`, optional Send to write progress to screen Returns ------- :class:`recarray` The structure of bright stars from the sweeps limited in the passed band(s) to the passed maglim(s). """ #ADM use io.py to retrieve list of sweeps or tractor files infiles = io.list_sweepfiles(rootdirname) if len(infiles) == 0: infiles = io.list_tractorfiles(rootdirname) if len(infiles) == 0: raise IOError( 'No sweep or tractor files found in {}'.format(rootdirname)) #ADM force the input maglim to be a list (in case a single value was passed) if type(maglim) == type(16) or type(maglim) == type(16.): maglim = [maglim] #ADM set bands to uppercase if passed as lower case bands = bands.upper() #ADM the band names as a flux array instead of a string bandnames = np.array(["FLUX_" + band for band in bands]) if len(bandnames) != len(maglim): raise IOError( 'bands has to be the same length as maglim and {} does not equal {}' .format(len(bands), len(maglim))) #ADM change input magnitude(s) to a flux to test against fluxlim = 10.**((22.5 - np.array(maglim)) / 2.5) #ADM parallel formalism from this step forward is stolen from cuts.select_targets #ADM function to grab the bright stars from a given file def _get_bright_stars(filename): '''Retrieves bright stars from a sweeps/Tractor file''' objs = io.read_tractor(filename) #ADM write the fluxes as an array instead of as named columns fluxes = objs[bandnames].view( objs[bandnames].dtype[0]).reshape(objs[bandnames].shape + (-1, )) #ADM Retain rows for which ANY band is brighter than maglim w = np.where(np.any(fluxes > fluxlim, axis=1)) if len(w[0]) > 0: return objs[w] #ADM counter for how many files have been processed #ADM critical to use np.ones because a numpy scalar allows in place modifications # c.f https://www.python.org/dev/peps/pep-3104/ totfiles = np.ones((), dtype='i8') * len(infiles) nfiles = np.ones((), dtype='i8') t0 = time() if verbose: print('Collecting bright stars from sweeps...') def _update_status(result): '''wrapper function for the critical reduction operation, that occurs on the main parallel process''' if verbose and nfiles % 25 == 0: elapsed = time() - t0 rate = nfiles / elapsed print('{}/{} files; {:.1f} files/sec; {:.1f} total mins elapsed'. format(nfiles, totfiles, rate, elapsed / 60.)) nfiles[...] += 1 #this is an in-place modification return result #ADM did we ask to parallelize, or not? if numproc > 1: pool = sharedmem.MapReduce(np=numproc) with pool: starstruc = pool.map(_get_bright_stars, infiles, reduce=_update_status) else: starstruc = [] for file in infiles: starstruc.append(_update_status(_get_bright_stars(file))) #ADM note that if there were no bright stars in a file then #ADM the _get_bright_stars function will have returned NoneTypes #ADM so we need to filter those out starstruc = [x for x in starstruc if x is not None] if len(starstruc) == 0: raise IOError( 'There are no stars brighter than {} in {} in files in {} with which to make a mask' .format(str(maglim), bands, rootdirname)) #ADM concatenate all of the output recarrays starstruc = np.hstack(starstruc) #ADM if the name of a file for output is passed, then write to it if outfilename is not None: fitsio.write(outfilename, starstruc, clobber=True) return starstruc
import fitsio import numpy as np import matplotlib.pyplot as plt import sys, os sys.path.insert(0, '/global/homes/q/qmxp55/DESI/BGS_paper/') from desitarget.cuts import select_targets from desitarget import io #from desitarget.sv1.sv1_targetmask import desi_mask, bgs_mask, mws_mask #nside = io.desitarget_nside() sweep_dir_dr8south = '/global/project/projectdirs/cosmo/data/legacysurvey/dr8/south/sweep/8.0' sweep_dir_dr8north = '/global/project/projectdirs/cosmo/data/legacysurvey/dr8/north/sweep/8.0' dest = '/global/cscratch1/sd/qmxp55/desitarget_output/sweeps-dr8_south_r20.fits' infilessouth = io.list_sweepfiles(sweep_dir_dr8south) infilesnorth = io.list_sweepfiles(sweep_dir_dr8north) #infiles = np.concatenate((infilessouth, infilesnorth)).tolist() infiles = infilessouth # import time start = time.time() def flux_to_mag(flux): mag = 22.5 - 2.5 * np.log10(flux) return mag targets = select_targets(infiles,