def get_current_results():
thisdir = os.path.split(os.path.abspath(__file__))[0]
psffile = thisdir + '/../etc/psfnight-r0.fits'
psf = specter.psf.load_psf(psffile)
np.random.seed(0)
ny, nx = psf.npix_y, psf.npix_x
image = np.random.normal(loc=0, scale=1, size=(ny, nx))
imageivar = np.ones_like(image)
#- Spectra and wavelengths to extract
nspec = 10
wave = np.arange(psf.wmin_all, psf.wmin_all + 200, 1)
#- Wake up code
flux, ivar, R = ex2d(image, imageivar, psf, 0, nspec, wave)
#- Now do it for real
t0 = time.time()
flux, ivar, R = ex2d(image,
imageivar,
psf,
0,
nspec,
wave,
bundlesize=nspec)
runtime = time.time() - t0
results = dict(flux=flux,
ivar=ivar,
R=R,
runtime=runtime,
nspec=nspec,
wave=wave)
return results
def run_pa(self,
input_image,
psf,
specmin,
nspec,
wave,
regularize=None,
ndecorr=True,
bundlesize=25,
wavesize=50,
outfile=None,
fibers=None,
fibermap=None):
import specter
from specter.extract import ex2d
from desispec.frame import Frame as fr
flux, ivar, Rdata = ex2d(input_image.pix,
input_image.ivar * (input_image.mask == 0),
psf,
specmin,
nspec,
wave,
regularize=regularize,
ndecorr=ndecorr,
bundlesize=bundlesize,
wavesize=wavesize)
#- Augment input image header for output
input_image.meta['NSPEC'] = (nspec, 'Number of spectra')
input_image.meta['WAVEMIN'] = (wave[0], 'First wavelength [Angstroms]')
input_image.meta['WAVEMAX'] = (wave[-1], 'Last wavelength [Angstroms]')
input_image.meta['WAVESTEP'] = (wave[1] - wave[0],
'Wavelength step size [Angstroms]')
input_image.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
#input_image.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
#input_image.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = fr(wave,
flux,
ivar,
resolution_data=Rdata,
fibers=fibers,
meta=input_image.meta,
fibermap=fibermap)
if outfile is not None: #- writing to a frame file if needed.
from desispec import io
io.write_frame(outfile, frame)
log.info("wrote frame output file %s" % outfile)
return frame
def wrap_ex2d(icpu, qin, qout):
p = psutil.Process(os.getpid())
p.cpu_affinity([
icpu,
])
while True:
orig_cpu = get_cpu()
i, specmin, nspec, wave = qin.get()
results = ex2d(image, imageivar, psf, specmin, nspec, wave)
# A = np.random.uniform(size=(500,500))
# results = np.linalg.svd(A.T.dot(A))
# results = 1
# time.sleep(np.random.uniform(0,1))
final_cpu = get_cpu()
qout.put((i, results, icpu, orig_cpu, final_cpu))
def run_pa(self,input_image,psf,specmin,nspec,wave,regularize=None,ndecorr=True,bundlesize=25,wavesize=50, outfile=None,fibers=None,fibermap=None):
import specter
from specter.extract import ex2d
from desispec.frame import Frame as fr
flux,ivar,Rdata=ex2d(input_image.pix,input_image.ivar*(input_image.mask==0),psf,specmin,nspec,wave,regularize=regularize,ndecorr=ndecorr,bundlesize=bundlesize,wavesize=wavesize)
#- Augment input image header for output
input_image.meta['NSPEC'] = (nspec, 'Number of spectra')
input_image.meta['WAVEMIN'] = (wave[0], 'First wavelength [Angstroms]')
input_image.meta['WAVEMAX'] = (wave[-1], 'Last wavelength [Angstroms]')
input_image.meta['WAVESTEP']= (wave[1]-wave[0], 'Wavelength step size [Angstroms]')
input_image.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
#input_image.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
#input_image.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = fr(wave, flux, ivar, resolution_data=Rdata,fibers=fibers, meta=input_image.meta, fibermap=fibermap)
if outfile is not None: #- writing to a frame file if needed.
from desispec import io
io.write_frame(outfile,frame)
log.info("wrote frame output file %s"%outfile)
return frame
opts.psf = thisdir + '/../etc/psfnight-r0.fits'
assert os.path.exists(opts.psf)
psf = specter.psf.load_psf(opts.psf)
#- Create fake noisy image
np.random.seed(opts.seed)
ny, nx = psf.npix_y, psf.npix_x
image = np.random.normal(loc=0, scale=1, size=(ny, nx))
imageivar = np.ones_like(image)
#- Spectra and wavelengths to extract
wave = np.arange(psf.wmin_all, psf.wmin_all + opts.numwave, 1)
#- Wake up the code in case there is library loading overhead
flux, ivar, R = ex2d(image, imageivar, psf, 0, 2, wave[0:10])
print('Running on {}/{}'.format(platform.node(), platform.processor()))
if 'OMP_NUM_THREADS' not in os.environ:
print('$OMP_NUM_THREADS not set')
else:
print('$OMP_NUM_THREADS={}'.format(os.getenv('OMP_NUM_THREADS')))
print('nspec rate')
for n in nspec:
t0 = time.time()
flux, ivar, R = ex2d(image,
imageivar,
psf,
0,
n,
wave,
thisdir = os.path.split(knltest.__file__)[0]
opts.psf = thisdir + '/../etc/psfnight-r0.fits'
assert os.path.exists(opts.psf)
psf = specter.psf.load_psf(opts.psf)
#- Create fake noisy image
ny, nx = psf.npix_y, psf.npix_x
image = np.random.normal(loc=0, scale=1, size=(ny, nx))
imageivar = np.ones_like(image)
#- Spectra and wavelengths to extract
w = np.linspace(psf.wmin_all, psf.wmax_all, 2000)
#- Wake up the code in case there is library loading overhead
flux, ivar, R = ex2d(image, imageivar, psf, 0, 2, w[0:10])
#- Setup sub extractions
extract_args = list()
iarg = 0
for specmin in range(0, psf.nspec, opts.bundlesize):
for i in range(0, len(w), opts.numwave):
x = (iarg, specmin, opts.bundlesize, w[i:i + opts.numwave])
extract_args.append(x)
iarg += 1
#- add comm.barrier() here to make sure everyone is woken up?
comm.barrier()
if comm.rank == 0:
t2 = time.time()
print('setup time {:.1f}'.format(t2 - t1))
def main_mpi(args, comm=None):
log = get_logger()
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if comm is None:
img = io.read_image(input_file)
else:
if comm.rank == 0:
img = io.read_image(input_file)
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
# get spectral range
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')]
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0))
psf_wavemax = np.min(
psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError(
'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.
format(wstart, psf_wavemin))
if psf_wavemax < wstop:
raise ValueError(
'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.
format(wstop, psf_wavemax))
# Now we divide our spectra into bundles
bundlesize = args.bundlesize
checkbundles = set()
checkbundles.update(
np.floor_divide(np.arange(specmin, specmax),
bundlesize * np.ones(nspec)).astype(int))
bundles = sorted(checkbundles)
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b + 1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mynbundle = int(nbundle // nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle *
(rank - leftover))
if rank == 0:
#- Print parameters
log.info("extract: input = {}".format(input_file))
log.info("extract: psf = {}".format(psf_file))
log.info("extract: specmin = {}".format(specmin))
log.info("extract: nspec = {}".format(nspec))
log.info("extract: wavelength = {},{},{}".format(wstart, wstop, dw))
log.info("extract: nwavestep = {}".format(args.nwavestep))
log.info("extract: regularize = {}".format(args.regularize))
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError(
"extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.normpath(os.path.dirname(outroot))
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
failcount = 0
for b in range(myfirstbundle, myfirstbundle + mynbundle):
outbundle = "{}_{:02d}.fits".format(outroot, b)
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
log.info('extract: Rank {} starting {} spectra {}:{} at {}'.format(
rank,
os.path.basename(input_file),
bspecmin[b],
bspecmin[b] + bnspec[b],
time.asctime(),
))
sys.stdout.flush()
#- The actual extraction
try:
results = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
bspecmin[b],
bnspec[b],
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX
mask[chi2pix > 100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
if fibermap is not None:
bfibermap = fibermap[bspecmin[b] - specmin:bspecmin[b] +
bnspec[b] - specmin]
else:
bfibermap = None
bfibers = fibers[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] -
specmin]
frame = Frame(wave,
flux,
ivar,
mask=mask,
resolution_data=Rdata,
fibers=bfibers,
meta=img.meta,
fibermap=bfibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(outbundle, frame, units='photon/bin')
if args.model is not None:
from astropy.io import fits
fits.writeto(outmodel,
results['modelimage'],
header=frame.meta)
log.info('extract: Done {} spectra {}:{} at {}'.format(
os.path.basename(input_file), bspecmin[b],
bspecmin[b] + bnspec[b], time.asctime()))
sys.stdout.flush()
except:
# Log the error and increment the number of failures
log.error(
"extract: FAILED bundle {}, spectrum range {}:{}".format(
b, bspecmin[b], bspecmin[b] + bnspec[b]))
exc_type, exc_value, exc_traceback = sys.exc_info()
lines = traceback.format_exception(exc_type, exc_value,
exc_traceback)
log.error(''.join(lines))
failcount += 1
sys.stdout.flush()
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
if rank == 0:
mergeopts = ['--output', args.output, '--force', '--delete']
mergeopts.extend(
["{}_{:02d}.fits".format(outroot, b) for b in bundles])
mergeargs = mergebundles.parse(mergeopts)
mergebundles.main(mergeargs)
if args.model is not None:
model = None
for b in bundles:
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
if model is None:
model = fits.getdata(outmodel)
else:
#- TODO: test and warn if models overlap for pixels with
#- non-zero values
model += fits.getdata(outmodel)
os.remove(outmodel)
fits.writeto(args.model, model)
def main(args):
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
print('Starting {} spectra {}:{} at {}'.format(
os.path.basename(input_file), specmin, specmin + nspec,
time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')]
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0))
psf_wavemax = np.min(
psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError(
'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.
format(wstart, psf_wavemin))
if psf_wavemax < wstop:
raise ValueError(
'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.
format(wstop, psf_wavemax))
#- Print parameters
print("""\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file,
psf=psf_file,
output=args.output,
wstart=wstart,
wstop=wstop,
dw=dw,
specmin=specmin,
nspec=nspec,
regularize=args.regularize))
#- The actual extraction
results = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
specmin,
nspec,
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX
mask[chi2pix > 100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave,
flux,
ivar,
mask=mask,
resolution_data=Rdata,
fibers=fibers,
meta=img.meta,
fibermap=fibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(args.output, frame, units='photon/bin')
if args.model is not None:
from astropy.io import fits
fits.writeto(args.model,
results['modelimage'],
header=frame.meta,
clobber=True)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin + nspec,
time.asctime()))
def main(args):
if args.mpi:
from mpi4py import MPI
comm = MPI.COMM_WORLD
return main_mpi(args, comm)
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
if args.fibermap_index is not None :
fibermin = args.fibermap_index
else :
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * 500 + specmin
print('Starting {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
else:
try:
fibermap = io.read_fibermap(args.input)
except (AttributeError, IOError, KeyError):
fibermap = None
#- Trim fibermap to matching fiber range and create fibers array
if fibermap:
ii = np.in1d(fibermap['FIBER'], np.arange(fibermin, fibermin+nspec))
fibermap = fibermap[ii]
fibers = fibermap['FIBER']
else:
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = [float(tmp) for tmp in args.wavelength.split(',')]
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.7
if args.heliocentric_correction :
heliocentric_correction_factor = heliocentric_correction_multiplicative_factor(img.meta)
wstart /= heliocentric_correction_factor
wstop /= heliocentric_correction_factor
dw /= heliocentric_correction_factor
else :
heliocentric_correction_factor = 1.
wave = np.arange(wstart, wstop+dw/2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(list(range(specmin, specmax)), y=0))
psf_wavemax = np.min(psf.wavelength(list(range(specmin, specmax)), y=psf.npix_y-1))
if psf_wavemin-5 > wstart:
raise ValueError('Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin))
if psf_wavemax+5 < wstop:
raise ValueError('Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax))
#- Print parameters
print("""\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file, psf=psf_file, output=args.output,
wstart=wstart, wstop=wstop, dw=dw,
specmin=specmin, nspec=nspec,
regularize=args.regularize))
#- The actual extraction
results = ex2d(img.pix, img.ivar*(img.mask==0), psf, specmin, nspec, wave,
regularize=args.regularize, ndecorr=args.decorrelate_fibers,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True, nsubbundles=args.nsubbundles,psferr=args.psferr)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX
mask[chi2pix>100.0] |= specmask.BAD2DFIT
if heliocentric_correction_factor != 1 :
#- Apply heliocentric correction factor to the wavelength
#- without touching the spectra, that is the whole point
wave *= heliocentric_correction_factor
wstart *= heliocentric_correction_factor
wstop *= heliocentric_correction_factor
dw *= heliocentric_correction_factor
img.meta['HELIOCOR'] = heliocentric_correction_factor
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata,
fibers=fibers, meta=img.meta, fibermap=fibermap,
chi2pix=chi2pix)
#- Add unit
# In specter.extract.ex2d one has flux /= dwave
# to convert the measured total number of electrons per
# wavelength node to an electron 'density'
frame.meta['BUNIT'] = 'count/Angstrom'
#- Add scores to frame
if not args.no_scores :
compute_and_append_frame_scores(frame,suffix="RAW")
#- Write output
io.write_frame(args.output, frame)
if args.model is not None:
from astropy.io import fits
fits.writeto(args.model, results['modelimage'], header=frame.meta, overwrite=True)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
def wrap_ex2d(args):
specmin, nspec, ww, xyrange = args
xmin, xmax, ymin, ymax = xyrange
return ex2d(image[ymin:ymax, xmin:xmax], image_ivar[ymin:ymax, xmin:xmax],
psf, specmin, nspec, ww, xyrange)
#- Load point spread function model
if opts.psf is None:
thisdir = os.path.split(knltest.__file__)[0]
opts.psf = thisdir + '/../etc/psfnight-r0.fits'
assert os.path.exists(opts.psf)
psf = specter.psf.load_psf(opts.psf)
#- Create fake noisy image
ny, nx = psf.npix_y, psf.npix_x
image = np.random.normal(loc=0, scale=1, size=(ny, nx))
imageivar = np.ones_like(image)
#- Spectra and wavelengths to extract
w = np.arange(psf.wmin_all, psf.wmin_all + opts.numwave, 1)
args = [image, imageivar, psf]
kwargs = dict(specmin=0, nspec=opts.numspec, wavelengths=w)
#- Wake up the code in case there is library loading overhead
flux, ivar, R = ex2d(image, imageivar, psf, 0, 2, w[0:10])
print('Running on {}/{} with {} logical cores'.format(
platform.node(), platform.processor(), multiprocessing.cpu_count()))
print('{} spectra x {} wavelengths extracted'.format(opts.numspec,
opts.numwave))
print('OMP_NUM_THREADS time')
for n in ntest:
os.environ['OMP_NUM_THREADS'] = str(n)
t = knltest.timeit(ex2d, args, kwargs)
print("{:3} {:5.1f}".format(n, t), flush=True)
def wrap_ex2d(x):
i, specmin, nspec, wave = x
return ex2d(image, imageivar, psf, specmin, nspec, wave)
def main_mpi(args, comm=None):
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if comm is None:
img = io.read_image(input_file)
else:
if comm.rank == 0:
img = io.read_image(input_file)
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
# get spectral range
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(
psf.wavelength(range(specmin, specmax), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(
wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(
wstop, psf_wavemax)
# Now we divide our spectra into bundles
bundlesize = args.bundlesize
checkbundles = set()
checkbundles.update(
np.floor_divide(np.arange(specmin, specmax),
bundlesize * np.ones(nspec)).astype(int))
bundles = sorted(list(checkbundles))
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b + 1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mynbundle = int(nbundle / nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle *
(rank - leftover))
if rank == 0:
#- Print parameters
print "extract: input = {}".format(input_file)
print "extract: psf = {}".format(psf_file)
print "extract: specmin = {}".format(specmin)
print "extract: nspec = {}".format(nspec)
print "extract: wavelength = {},{},{}".format(wstart, wstop, dw)
print "extract: nwavestep = {}".format(args.nwavestep)
print "extract: regularize = {}".format(args.regularize)
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError(
"extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.dirname(outroot)
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
failcount = 0
for b in range(myfirstbundle, myfirstbundle + mynbundle):
outbundle = "{}_{:02d}.fits".format(outroot, b)
print('extract: Starting {} spectra {}:{} at {}'.format(
os.path.basename(input_file), bspecmin[b], bspecmin[b] + bnspec[b],
time.asctime()))
#- The actual extraction
try:
flux, ivar, Rdata = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
bspecmin[b],
bnspec[b],
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose)
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
bfibermap = fibermap[bspecmin[b] - specmin:bspecmin[b] +
bnspec[b] - specmin]
bfibers = fibers[bspecmin[b] - specmin:bspecmin[b] + bnspec[b] -
specmin]
frame = Frame(wave,
flux,
ivar,
resolution_data=Rdata,
fibers=bfibers,
meta=img.meta,
fibermap=bfibermap)
#- Write output
io.write_frame(outbundle, frame)
print('extract: Done {} spectra {}:{} at {}'.format(
os.path.basename(input_file), bspecmin[b],
bspecmin[b] + bnspec[b], time.asctime()))
except:
failcount += 1
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
if rank == 0:
opts = ['--output', args.output, '--force', '--delete']
opts.extend(["{}_{:02d}.fits".format(outroot, b) for b in bundles])
args = merge.parse(opts)
merge.main(args)
def main(args):
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
print('Starting {} spectra {}:{} at {}'.format(
os.path.basename(input_file), specmin, specmin + nspec,
time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin + nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin + nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop + dw / 2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(
psf.wavelength(range(specmin, specmax), y=psf.npix_y - 1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(
wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(
wstop, psf_wavemax)
#- Print parameters
print """\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file,
psf=psf_file,
output=args.output,
wstart=wstart,
wstop=wstop,
dw=dw,
specmin=specmin,
nspec=nspec,
regularize=args.regularize)
#- The actual extraction
flux, ivar, Rdata = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
specmin,
nspec,
wave,
regularize=args.regularize,
ndecorr=True,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose)
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP'] = (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__,
'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave,
flux,
ivar,
resolution_data=Rdata,
fibers=fibers,
meta=img.meta,
fibermap=fibermap)
#- Write output
io.write_frame(args.output, frame)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin + nspec,
time.asctime()))
def main_mpi(args, comm=None):
psf_file = args.psf
input_file = args.input
# these parameters are interpreted as the *global* spec range,
# to be divided among processes.
specmin = args.specmin
nspec = args.nspec
#- Load input files and broadcast
# FIXME: after we have fixed the serialization
# of the PSF, read and broadcast here, to reduce
# disk contention.
img = None
if comm is None:
img = io.read_image(input_file)
else:
if comm.rank == 0:
img = io.read_image(input_file)
img = comm.bcast(img, root=0)
psf = load_psf(psf_file)
# get spectral range
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin+nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop+dw/2.0, dw)
nwave = len(wave)
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(psf.wavelength(range(specmin, specmax), y=psf.npix_y-1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax)
# Now we divide our spectra into bundles
bundlesize = args.bundlesize
checkbundles = set()
checkbundles.update(np.floor_divide(np.arange(specmin, specmax), bundlesize*np.ones(nspec)).astype(int))
bundles = sorted(list(checkbundles))
nbundle = len(bundles)
bspecmin = {}
bnspec = {}
for b in bundles:
if specmin > b * bundlesize:
bspecmin[b] = specmin
else:
bspecmin[b] = b * bundlesize
if (b+1) * bundlesize > specmax:
bnspec[b] = specmax - bspecmin[b]
else:
bnspec[b] = bundlesize
# Now we assign bundles to processes
nproc = 1
rank = 0
if comm is not None:
nproc = comm.size
rank = comm.rank
mynbundle = int(nbundle // nproc)
myfirstbundle = 0
leftover = nbundle % nproc
if rank < leftover:
mynbundle += 1
myfirstbundle = rank * mynbundle
else:
myfirstbundle = ((mynbundle + 1) * leftover) + (mynbundle * (rank - leftover))
if rank == 0:
#- Print parameters
print "extract: input = {}".format(input_file)
print "extract: psf = {}".format(psf_file)
print "extract: specmin = {}".format(specmin)
print "extract: nspec = {}".format(nspec)
print "extract: wavelength = {},{},{}".format(wstart, wstop, dw)
print "extract: nwavestep = {}".format(args.nwavestep)
print "extract: regularize = {}".format(args.regularize)
# get the root output file
outpat = re.compile(r'(.*)\.fits')
outmat = outpat.match(args.output)
if outmat is None:
raise RuntimeError("extraction output file should have .fits extension")
outroot = outmat.group(1)
outdir = os.path.normpath(os.path.dirname(outroot))
if rank == 0:
if not os.path.isdir(outdir):
os.makedirs(outdir)
if comm is not None:
comm.barrier()
failcount = 0
for b in range(myfirstbundle, myfirstbundle+mynbundle):
outbundle = "{}_{:02d}.fits".format(outroot, b)
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
print('extract: Rank {} starting {} spectra {}:{} at {}'.format(
rank, os.path.basename(input_file),
bspecmin[b], bspecmin[b]+bnspec[b], time.asctime(),
) )
#- The actual extraction
try:
results = ex2d(img.pix, img.ivar*(img.mask==0), psf, bspecmin[b],
bnspec[b], wave, regularize=args.regularize, ndecorr=True,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX
mask[chi2pix>100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
if fibermap is not None:
bfibermap = fibermap[bspecmin[b]-specmin:bspecmin[b]+bnspec[b]-specmin]
else:
bfibermap = None
bfibers = fibers[bspecmin[b]-specmin:bspecmin[b]+bnspec[b]-specmin]
frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata,
fibers=bfibers, meta=img.meta, fibermap=bfibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(outbundle, frame)
if args.model is not None:
from astropy.io import fits
fits.writeto(outmodel, results['modelimage'], header=frame.meta)
print('extract: Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
bspecmin[b], bspecmin[b]+bnspec[b], time.asctime()))
except:
failcount += 1
if comm is not None:
failcount = comm.allreduce(failcount)
if failcount > 0:
# all processes throw
raise RuntimeError("some extraction bundles failed")
if rank == 0:
mergeopts = [
'--output', args.output,
'--force',
'--delete'
]
mergeopts.extend([ "{}_{:02d}.fits".format(outroot, b) for b in bundles ])
mergeargs = mergebundles.parse(mergeopts)
mergebundles.main(mergeargs)
if args.model is not None:
model = None
for b in bundles:
outmodel = "{}_model_{:02d}.fits".format(outroot, b)
if model is None:
model = fits.getdata(outmodel)
else:
#- TODO: test and warn if models overlap for pixels with
#- non-zero values
model += fits.getdata(outmodel)
os.remove(outmodel)
fits.writeto(args.model, model)
def main(args):
psf_file = args.psf
input_file = args.input
specmin = args.specmin
nspec = args.nspec
#- Load input files
psf = load_psf(psf_file)
img = io.read_image(input_file)
if nspec is None:
nspec = psf.nspec
specmax = specmin + nspec
camera = img.meta['CAMERA'].lower() #- b0, r1, .. z9
spectrograph = int(camera[1])
fibermin = spectrograph * psf.nspec + specmin
print('Starting {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
if args.fibermap is not None:
fibermap = io.read_fibermap(args.fibermap)
fibermap = fibermap[fibermin:fibermin+nspec]
fibers = fibermap['FIBER']
else:
fibermap = None
fibers = np.arange(fibermin, fibermin+nspec, dtype='i4')
#- Get wavelength grid from options
if args.wavelength is not None:
wstart, wstop, dw = map(float, args.wavelength.split(','))
else:
wstart = np.ceil(psf.wmin_all)
wstop = np.floor(psf.wmax_all)
dw = 0.5
wave = np.arange(wstart, wstop+dw/2.0, dw)
nwave = len(wave)
bundlesize = args.bundlesize
#- Confirm that this PSF covers these wavelengths for these spectra
psf_wavemin = np.max(psf.wavelength(range(specmin, specmax), y=0))
psf_wavemax = np.min(psf.wavelength(range(specmin, specmax), y=psf.npix_y-1))
if psf_wavemin > wstart:
raise ValueError, 'Start wavelength {:.2f} < min wavelength {:.2f} for these fibers'.format(wstart, psf_wavemin)
if psf_wavemax < wstop:
raise ValueError, 'Stop wavelength {:.2f} > max wavelength {:.2f} for these fibers'.format(wstop, psf_wavemax)
#- Print parameters
print """\
#--- Extraction Parameters ---
input: {input}
psf: {psf}
output: {output}
wavelength: {wstart} - {wstop} AA steps {dw}
specmin: {specmin}
nspec: {nspec}
regularize: {regularize}
#-----------------------------\
""".format(input=input_file, psf=psf_file, output=args.output,
wstart=wstart, wstop=wstop, dw=dw,
specmin=specmin, nspec=nspec,
regularize=args.regularize)
#- The actual extraction
results = ex2d(img.pix, img.ivar*(img.mask==0), psf, specmin, nspec, wave,
regularize=args.regularize, ndecorr=True,
bundlesize=bundlesize, wavesize=args.nwavestep, verbose=args.verbose,
full_output=True)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction']>0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction']==1.0] |= specmask.ALLBADPIX
mask[chi2pix>100.0] |= specmask.BAD2DFIT
#- Augment input image header for output
img.meta['NSPEC'] = (nspec, 'Number of spectra')
img.meta['WAVEMIN'] = (wstart, 'First wavelength [Angstroms]')
img.meta['WAVEMAX'] = (wstop, 'Last wavelength [Angstroms]')
img.meta['WAVESTEP']= (dw, 'Wavelength step size [Angstroms]')
img.meta['SPECTER'] = (specter.__version__, 'https://github.com/desihub/specter')
img.meta['IN_PSF'] = (_trim(psf_file), 'Input spectral PSF')
img.meta['IN_IMG'] = (_trim(input_file), 'Input image')
frame = Frame(wave, flux, ivar, mask=mask, resolution_data=Rdata,
fibers=fibers, meta=img.meta, fibermap=fibermap,
chi2pix=chi2pix)
#- Write output
io.write_frame(args.output, frame)
if args.model is not None:
from astropy.io import fits
fits.writeto(args.model, results['modelimage'], header=frame.meta, clobber=True)
print('Done {} spectra {}:{} at {}'.format(os.path.basename(input_file),
specmin, specmin+nspec, time.asctime()))
def _extract_and_save(img, psf, bspecmin, bnspec, specmin, wave, raw_wave,
fibers, fibermap, outbundle, outmodel, bundlesize, args,
log):
'''
Performs the main extraction and saving of extracted frames found in the body of the
main loop. Refactored to be callable by both MPI and non-MPI versions of the code.
This should be viewed as a shorthand for the following commands.
'''
results = ex2d(img.pix,
img.ivar * (img.mask == 0),
psf,
bspecmin,
bnspec,
wave,
regularize=args.regularize,
ndecorr=args.decorrelate_fibers,
bundlesize=bundlesize,
wavesize=args.nwavestep,
verbose=args.verbose,
full_output=True,
nsubbundles=args.nsubbundles)
flux = results['flux']
ivar = results['ivar']
Rdata = results['resolution_data']
chi2pix = results['chi2pix']
mask = np.zeros(flux.shape, dtype=np.uint32)
mask[results['pixmask_fraction'] > 0.5] |= specmask.SOMEBADPIX
mask[results['pixmask_fraction'] == 1.0] |= specmask.ALLBADPIX
mask[chi2pix > 100.0] |= specmask.BAD2DFIT
if fibermap is not None:
bfibermap = fibermap[bspecmin - specmin:bspecmin + bnspec - specmin]
else:
bfibermap = None
bfibers = fibers[bspecmin - specmin:bspecmin + bnspec - specmin]
#- Save the raw wavelength, not the corrected one (if corrected)
frame = Frame(raw_wave,
flux,
ivar,
mask=mask,
resolution_data=Rdata,
fibers=bfibers,
meta=img.meta,
fibermap=bfibermap,
chi2pix=chi2pix)
#- Add unit
# In specter.extract.ex2d one has flux /= dwave
# to convert the measured total number of electrons per
# wavelength node to an electron 'density'
frame.meta['BUNIT'] = 'electron/Angstrom'
#- Add scores to frame
if not args.no_scores:
compute_and_append_frame_scores(frame, suffix="RAW")
mark_extraction = time.time()
#- Write output
io.write_frame(outbundle, frame)
if args.model is not None:
fits.writeto(outmodel, results['modelimage'], header=frame.meta)
log.info('extract: Done {} spectra {}:{} at {}'.format(
os.path.basename(args.input), bspecmin, bspecmin + bnspec,
time.asctime()))
sys.stdout.flush()
return mark_extraction
