def test_tractor_columns(self):
# ADM Gaia columns that get added on input.
from desitarget.gaiamatch import gaiadatamodel
from desitarget.gaiamatch import pop_gaia_coords, pop_gaia_columns
# ADM have to remove the GAIA_RA, GAIA_DEC columns used for matching.
gaiadatamodel = pop_gaia_coords(gaiadatamodel)
# ADM prior to DR8, we're also missing some other Gaia columns.
# gaiadatamodel = pop_gaia_columns(
# gaiadatamodel,
# ['REF_CAT', 'GAIA_PHOT_BP_RP_EXCESS_FACTOR',
# 'GAIA_ASTROMETRIC_SIGMA5D_MAX', 'GAIA_ASTROMETRIC_PARAMS_SOLVED']
# )
# ADM PHOTSYS gets added on input.
tscolumns = list(io.tsdatamodel.dtype.names) \
+ ['PHOTSYS'] \
+ list(gaiadatamodel.dtype.names)
tractorfile = io.list_tractorfiles(self.datadir)[0]
data = io.read_tractor(tractorfile)
self.assertEqual(set(data.dtype.names), set(tscolumns))
# columns = ['BX', 'BY']
columns = ['RA', 'DEC']
data = io.read_tractor(tractorfile, columns=columns)
self.assertEqual(set(data.dtype.names), set(columns))
data = io.read_tractor(tractorfile, columns=tuple(columns))
self.assertEqual(set(data.dtype.names), set(columns))
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_tractor_columns(self):
# ADM Gaia columns that get added on input.
from desitarget.gaiamatch import gaiadatamodel
from desitarget.gaiamatch import pop_gaia_coords, pop_gaia_columns
# ADM remove the GAIA_RA, GAIA_DEC columns used for matching.
gaiadatamodel = pop_gaia_coords(gaiadatamodel)
tractorfile = io.list_tractorfiles(self.datadir)[0]
data = io.read_tractor(tractorfile)
# ADM form the final data model in a manner that maintains
# ADM backwards-compatability with DR8.
if "FRACDEV" in data.dtype.names:
tsdatamodel = np.array(
[], dtype=io.basetsdatamodel.dtype.descr +
io.dr8addedcols.dtype.descr)
else:
tsdatamodel = np.array(
[], dtype=io.basetsdatamodel.dtype.descr +
io.dr9addedcols.dtype.descr)
# ADM PHOTSYS gets added on input.
tscolumns = list(tsdatamodel.dtype.names) \
+ ['PHOTSYS'] \
+ list(gaiadatamodel.dtype.names)
self.assertEqual(set(data.dtype.names), set(tscolumns))
# columns = ['BX', 'BY']
columns = ['RA', 'DEC']
data = io.read_tractor(tractorfile, columns=columns)
self.assertEqual(set(data.dtype.names), set(columns))
data = io.read_tractor(tractorfile, columns=tuple(columns))
self.assertEqual(set(data.dtype.names), set(columns))
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_tractor_columns(self):
tscolumns = io.tscolumns + ['BRICK_PRIMARY',]
tractorfile = io.list_tractorfiles(self.datadir)[0]
data = io.read_tractor(tractorfile)
self.assertEqual(set(data.dtype.names), set(tscolumns))
columns = ['BX', 'BY']
data = io.read_tractor(tractorfile, columns=columns)
self.assertEqual(set(data.dtype.names), set(columns))
data = io.read_tractor(tractorfile, columns=tuple(columns))
self.assertEqual(set(data.dtype.names), set(columns))
def test_tractor_columns(self):
tscolumns = list(io.tsdatamodel.dtype.names) + [
'BRICK_PRIMARY',
]
tractorfile = io.list_tractorfiles(self.datadir)[0]
data = io.read_tractor(tractorfile)
self.assertEqual(set(data.dtype.names), set(tscolumns))
# columns = ['BX', 'BY']
columns = ['RA', 'DEC']
data = io.read_tractor(tractorfile, columns=columns)
self.assertEqual(set(data.dtype.names), set(columns))
data = io.read_tractor(tractorfile, columns=tuple(columns))
self.assertEqual(set(data.dtype.names), set(columns))
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 main():
# Build a catalog of all the grz, 3-pass sources.
tracdir = '/project/projectdirs/cosmo/data/legacysurvey/dr2/tractor'
outfile = 'lensedecals-parent.fits'
ramin = 210
ramax = 240 # 250
rarange = np.arange(ramin, ramax+1, 1).astype(str)
infiles = [io.list_tractorfiles(os.path.join(tracdir, this)) for this in rarange]
infiles = np.concatenate(infiles).flatten()
#infiles = io.list_tractorfiles(tracdir)
print('Found {} tractor catalogs.'.format(len(infiles)))
cat = select_targets(infiles, verbose=True)#, numproc=1)
print('Selected {} objects.'.format(len(cat)))
io.write_targets(outfile, cat)
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 main():
# Build a catalog of all the grz, 3-pass sources.
tracdir = '/project/projectdirs/cosmo/data/legacysurvey/dr2/tractor'
outfile = 'lensedecals-parent.fits'
ramin = 210
ramax = 240 # 250
rarange = np.arange(ramin, ramax + 1, 1).astype(str)
infiles = [
io.list_tractorfiles(os.path.join(tracdir, this)) for this in rarange
]
infiles = np.concatenate(infiles).flatten()
#infiles = io.list_tractorfiles(tracdir)
print('Found {} tractor catalogs.'.format(len(infiles)))
cat = select_targets(infiles, verbose=True) #, numproc=1)
print('Selected {} objects.'.format(len(cat)))
io.write_targets(outfile, cat)
def test_list_tractorfiles(self):
files = io.list_tractorfiles(self.datadir)
self.assertEqual(len(files), 3)
for x in files:
self.assertTrue(os.path.basename(x).startswith('tractor'))
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')
# ADM All that is required is a link to the desitarget GAIA_DIR.
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)
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
def test_list_tractorfiles(self):
files = io.list_tractorfiles(self.datadir)
self.assertEqual(len(files), 3)
for x in files:
self.assertTrue(os.path.basename(x).startswith('tractor'))
self.assertTrue(os.path.basename(x).endswith('.fits'))
