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_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 collect_bright_stars(
bands,
maglim,
numproc=4,
rootdirname='/global/project/projectdirs/cosmo/data/legacysurvey/dr3.1/sweep/3.1',
outfilename=None):
"""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
Returns
-------
:class:`recarray`
The structure of bright stars from the sweeps limited in the passed band(s) to the
passed maglim(s).
"""
#ADM set up default logger
from lvmutil.log import get_logger
log = get_logger()
#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()
log.info('Collecting bright stars from sweeps...')
def _update_status(result):
'''wrapper function for the critical reduction operation,
that occurs on the main parallel process'''
if nfiles % 25 == 0:
elapsed = time() - t0
rate = nfiles / elapsed
log.info(
'{}/{} 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
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('lvmtarget.test', 't')
cls.tractorfiles = sorted(io.list_tractorfiles(cls.datadir))
cls.sweepfiles = sorted(io.list_sweepfiles(cls.datadir))