def test_config():
with pytest.raises(RuntimeError):
config = utils.read_config('xx.yaml')
cwd = os.getcwd()
os.chdir(path)
try:
config = utils.read_config()
finally:
os.chdir(cwd)
def proc_many(files,
oprefix,
fig_prefix,
config=None,
nthreads=1,
overwrite=True):
"""
Process many spectral files
Parameters:
-----------
mask: string
The filename mask with spectra, i.e path/*fits
oprefix: string
The prefix where the table with measurements will be stored
fig_prefix: string
The prfix where the figures will be stored
"""
config = utils.read_config(config)
if nthreads > 1:
parallel = True
else:
parallel = False
if parallel:
pool = mp.Pool(nthreads)
for f in files:
res = []
fname = f.split('/')[-1]
ofname = oprefix + 'outtab_' + fname
if (not overwrite) and os.path.exists(ofname):
print('skipping, products already exist', f)
continue
if parallel:
for i in range(nthreads):
res.append(
pool.apply_async(proc_weave_wrapper,
(f, fig_prefix, config, i, nthreads)))
tabs = []
for r in res:
tabs.append(r.get())
tabs = ([_ for _ in tabs if _ is not None])
if len(tabs) == 0:
continue
tabs = atpy.vstack(tabs)
tabs.write(ofname, overwrite=True)
else:
tabs = proc_weave_wrapper(f, fig_prefix, config, 0, 1)
if tabs is not None:
tabs.write(ofname, overwrite=True)
if parallel:
pool.close()
pool.join()
def doone(seed,
sn=100,
nlam=400,
config_name='tests/test.yaml',
doplot=False,
resol=1000):
if config_name != si.fname:
si.config = utils.read_config(config_name)
si.fname = config_name
config = si.config
st = np.random.get_state()
np.random.seed(seed)
# read data
lamcen = 5000.
#lamcen_air = lamcen / (1.0 + 2.735182E-4 + 131.4182 / lamcen**2 +
# 2.76249E8 / lamcen**4)
v0 = np.random.normal(0, 300)
slope = (np.random.uniform(-2, 2))
wresol = (lamcen / resol / 2.35)
lam = np.linspace(4600, 5400, nlam)
#spec0 = (1 - 0.02 * np.exp(-0.5 * ((lam - lamcen1) / w)**2))*lam**slope
teff = np.random.uniform(3000, 12000)
feh = np.random.uniform(-2, 0)
alpha = np.random.uniform(0, 1)
logg = np.random.uniform(0, 5)
c = 299792.458
lam1 = lam / np.sqrt((1 + v0 / c) / (1 - v0 / c))
spec0 = mktemps.getspec(lam1, teff, logg, feh, alpha,
wresol=wresol) * lam**slope
spec0 = spec0 / np.median(spec0) * 10**np.random.uniform(-3, 3)
espec = spec0 / sn
spec = np.random.normal(spec0, espec)
# construct specdata object
specdata = [spec_fit.SpecData('config1', lam, spec, espec)]
options = {'npoly': 10}
paramDict0 = {'logg': 2.5, 'teff': 5000, 'feh': -1, 'alpha': 0.5}
#, 'vsini': 0.0}
# fixParam = ['vsini']
res = vel_fit.process(
specdata,
paramDict0,
# fixParam=fixParam,
config=config,
options=options)
ret = (v0, res['vel'], res['vel_err'])
if doplot:
plt.figure(figsize=(6, 2), dpi=300)
plt.plot(specdata[0].lam, specdata[0].spec, 'k-')
plt.plot(specdata[0].lam, res['yfit'][0], 'r-')
plt.tight_layout()
plt.savefig('plot_accuracy_test.png')
#1/0
np.random.set_state(st)
return ret
def test_fake_grid():
config = utils.read_config(path + '/test.yaml')
# read data
lam = np.linspace(4600, 5400, 800)
v0 = np.random.normal(0, 100)
lam1 = lam / (1 + v0 / 3e5)
resol = 1000.
lamcen = 5000
w = lamcen / resol / 2.35
spec0 = getspec(lam1, 5000, 2, -1, 0.2, wresol=w)
espec = spec0 * 0.01
spec = np.random.normal(spec0, espec)
# construct specdata object
specdata = [spec_fit.SpecData('config1_grid', lam, spec, espec)]
options = {'npoly': 15}
paramDict0 = {
'logg': 2,
'teff': 5000,
'feh': -0.2,
'alpha': 0.2,
'vsini': 0.1
}
fixParam = [] #'vsini']
res = vel_fit.process(specdata,
paramDict0,
fixParam=fixParam,
config=config,
options=options)
print(res['vel'] - v0, res['vel_err'])
print(res['param'])
if True:
plt.figure(figsize=(6, 2), dpi=300)
plt.plot(specdata[0].lam, specdata[0].spec, 'k-')
plt.plot(specdata[0].lam, res['yfit'][0], 'r-')
plt.tight_layout()
plt.savefig(path + '/plot_test_fit_fake_grid.png')
def proc_mws(args):
"""
Process DESI MWS spectral files
"""
# N_processor=64 # Number of processors on Haswell
parser = argparse.ArgumentParser()
parser.add_argument('--nthreads',
help='Number of processor per node',
type=int,
default=1)
parser.add_argument('--input_dir',
help='Output directory of files to process',
type=str,
default=None,
required=True)
parser.add_argument('--output_dir',
help='Output directory for the data tables',
type=str,
default=None,
required=True)
parser.add_argument(
'--output_script_dir',
help='Output directory for the slurm scripts and ferre input files',
type=str,
default=None,
required=True)
parser.add_argument(
'--commissioning_folder',
help='The folder (named by date) that stores commissioning data',
type=str,
default=None,
required=False)
parser.add_argument('--report_dir',
help='directory of the report files',
type=str,
default=None,
required=True)
parser.add_argument('--allobjects',
help='Fit all objects, not just MWS_TARGET',
action='store_true',
default=False)
parser.add_argument('--whole_spectra64',
help='Process the whole spectra64 directory',
action='store_true',
default=False)
args = parser.parse_args(args)
path = args.input_dir
report_dir = args.report_dir
mwonly = not args.allobjects
whole_spectra64 = args.whole_spectra64
out_path = args.output_dir
out_script_path = args.output_script_dir
nthreads = args.nthreads
commissioning = args.commissioning_folder
config_path = os.environ['rvspecfit_config'] + "/config.yaml"
python_path = os.environ['DESI_MWS_root'] + "/piferre"
config = utils.read_config(config_path)
#-------- Read in latest previous report and generate a new report -----------
latest_report, last_expid, prev_date = check_latest_report(report_dir)
print('Reading in latest report:', latest_report)
print('Previous largest expid:', last_expid)
n_proc_now, pix_list, sdir_list, expid_range, prev_date, report_data = check_spectra64(
report_dir, path, out_path, mwonly, last_expid, prev_date,
whole_spectra64, commissioning)
min_expid = last_expid
#== Store all the scripts in a folder named with Today's date
#yr= datetime.date.today().year
#month=datetime.date.today().month
#day=datetime.date.today().day
#now=str(yr*10000+month*100+day)
time_now = datetime.datetime.now()
now = str(time_now.year * 10000 + time_now.month * 100 +
time_now.day) + '-' + str(100 + time_now.hour)[1:] + str(
100 + time_now.minute)[1:]
#== Create directories if not existing.
for i, pixel in enumerate(pix_list):
sdir = sdir_list[i]
if not os.path.exists(os.path.join(out_path, sdir)):
os.mkdir(os.path.join(out_path, sdir))
if not os.path.exists(os.path.join(out_script_path, now)):
os.mkdir(os.path.join(out_script_path, now))
if not os.path.exists(os.path.join(out_script_path, now, sdir)):
os.mkdir(os.path.join(out_script_path, now, sdir))
if not os.path.exists(os.path.join(out_script_path, now, sdir, pixel)):
os.mkdir(os.path.join(out_script_path, now, sdir, pixel))
if not os.path.exists(os.path.join(out_path, sdir, pixel)):
os.mkdir(os.path.join(out_path, sdir, pixel))
# Determining how many jobs to be submitted
pix_gp_list, sdir_gp_list, fn_sum_list, job_ind, sort_index = cal_node_n(
n_proc_now, pix_list, sdir_list, nthreads)
generate_reports(path, report_dir, expid_range, prev_date, report_data,
job_ind, sort_index)
#== Writing slurm scripts and lists of input files (x.xxx) for each job (group jobs according to fiber numbers and processors per node)
n_node = len(pix_gp_list) # requested node number = pixel group number
print('Number of jobs to be submitted:', n_node + len(pix_list))
for i in range(n_node):
file_ind = str(1000 + i)[1:]
#== Generating files that store lists of input files
f = open(os.path.join(out_script_path, now, 'x.' + file_ind), 'w')
sub_pixels = pix_gp_list[i]
sub_sdirs = sdir_gp_list[i]
for pixel, sdir in zip(sub_pixels, sub_sdirs):
entry = path + '/' + str(sdir) + '/' + str(pixel)
f.write(entry + "/spectra-64-" + str(pixel) + ".fits \n")
f.close()
suffix = 'rvspecfit_' + file_ind
#== Generating slurm scripts
print('Generating rvspecfit slurm scripts....', i, '/', n_node)
n_fiber = fn_sum_list[i]
write_slurm_tot(out_script_path, out_path, file_ind, path,
sdir_gp_list[i], pix_gp_list[i], now, min_expid,
n_fiber, nthreads, suffix, whole_spectra64, mwonly)
print('Writing shell scripts....', i, '/', n_node)
write_script_tot_gp(os.path.join(out_script_path),
out_path,
path,
sdir_gp_list[i],
pix_gp_list[i],
now,
file_ind,
min_expid,
nthreads=nthreads,
suffix=suffix,
mwonly=mwonly)
print('Generating ferre slurm scripts....')
for sdir, pixel, n_fiber in zip(sdir_list, pix_list, n_proc_now):
print('pixel number:', pixel)
cmd = "python3 -c \"import sys; sys.path.insert(0, '" + python_path + "'); from piferre import write_slurm; write_slurm(\'" + str(
sdir) + "\',\'" + str(pixel) + "\', \'" + str(
out_path) + "\'," + str(
n_fiber) + " , script_path='" + os.path.join(
out_script_path, now) + "', ngrids=9, nthreads=4)\"\n"
#print('cmd=',cmd)
err = subprocess.call(cmd, shell=True)
#== Executing the shell script and submitting jobs
run_scripts(os.path.join(out_script_path, now))
def proc_many(files,
output_dir,
output_tab_prefix,
output_mod_prefix,
figure_dir=None,
figure_prefix=None,
config_fname=None,
nthreads=1,
fit_targetid=None,
objtypes=None,
minsn=-1e9,
doplot=True,
expid_range=None,
overwrite=False,
skipexisting=False,
fitarm=None,
cmdline=None,
zbest_select=False,
ccfinit=True,
ccf_continuum_normalize=True,
process_status_file=None,
npoly=None,
throw_exceptions=None):
"""
Process many spectral files
Parameters
-----------
files: strings
The files with spectra
oprefix: string
The prefix where the table with measurements will be stored
figure_dir: string
The director where the figures will be stored
figure_prefix: string
The prefix of figure filenames
config_fname: string
The name of the config file
fit_targetid: integer or None
The targetid to fit (the rest will be ignored)
objtypes: lists
list of DESI_TARGET regular expressions
doplot: bool
Plotting
minsn: real
THe min S/N to fit
cmdline: string
The command line used in execution of rvspecfit
expid_range: tuple
None or a tule of two numbers for the range of expids to fit
skipexisting: bool
if True do not process anything if output files exist
fitarm: list
the list of arms/spec configurations to fit (can be None)
npoly: integer
the degree of the polynomial used for continuum
process_status_file: str
The filename where we'll put status of the fitting
"""
override = dict(ccf_continuum_normalize=ccf_continuum_normalize)
config = utils.read_config(config_fname, override)
assert (config is not None)
assert ('template_lib' in config)
if nthreads > 1:
parallel = True
else:
parallel = False
if process_status_file is not None:
update_process_status_file(process_status_file,
None,
None,
None,
None,
start=True)
if parallel:
poolEx = concurrent.futures.ProcessPoolExecutor(nthreads)
else:
poolEx = FakeExecutor()
res = []
for f in files:
fname = f.split('/')[-1]
assert (len(f.split('/')) > 2)
# we need that because we use the last two directories in the path
# to create output directory structure
# i.e. input file a/b/c/d/e/f/g.fits will produce output file in
# output_prefix/e/f/xxx.fits
fdirs = f.split('/')
folder_path = output_dir + '/' + fdirs[-3] + '/' + fdirs[-2] + '/'
os.makedirs(folder_path, exist_ok=True)
logging.debug(f'Making folder {folder_path}')
if figure_dir is not None:
figure_path = figure_dir + '/' + fdirs[-3] + '/' + fdirs[-2] + '/'
os.makedirs(figure_path, exist_ok=True)
cur_figure_prefix = figure_path + '/' + figure_prefix
logging.debug(f'Making folder {figure_path}')
else:
cur_figure_prefix = None
tab_ofname = folder_path + output_tab_prefix + '_' + fname
mod_ofname = folder_path + output_mod_prefix + '_' + fname
if (skipexisting and os.path.exists(tab_ofname)
and os.path.exists(mod_ofname)):
logging.info('skipping, products already exist %s' % f)
if process_status_file is not None:
update_process_status_file(process_status_file, f,
ProcessStatus.EXISTING, -1, 0)
continue
args = (f, tab_ofname, mod_ofname, cur_figure_prefix, config)
kwargs = dict(fit_targetid=fit_targetid,
objtypes=objtypes,
doplot=doplot,
minsn=minsn,
expid_range=expid_range,
overwrite=overwrite,
poolex=poolEx,
fitarm=fitarm,
cmdline=cmdline,
zbest_select=zbest_select,
process_status_file=process_status_file,
npoly=npoly,
ccfinit=ccfinit,
throw_exceptions=throw_exceptions)
proc_desi_wrapper(*args, **kwargs)
if parallel:
try:
poolEx.shutdown(wait=True)
except KeyboardInterrupt:
for r in res:
r.cancel()
poolEx.shutdown(wait=False)
raise
logging.info("Successfully finished processing")
def test_fits():
config = utils.read_config(path + '/config.yaml')
# read data
dat = pyfits.getdata(path + '/data/spec-0266-51602-0031.fits')
err = dat['ivar']
err = 1. / err**.5
err[~np.isfinite(err)] = 1e40
# construct specdata object
specdata = [
spec_fit.SpecData('sdss1', 10**dat['loglam'], dat['flux'], err)
]
rot_params = None
resols_params = None
params_list = [[4000, 3, -1, 0], [5000, 3, -1, 0], [6000, 2, -2, 0],
[5500, 5, 0, 0]]
vel_grid = np.linspace(-600, 600, 1000)
options = {'npoly': 10}
t1 = time.time()
res = spec_fit.find_best(specdata,
vel_grid,
params_list,
rot_params,
resols_params,
options=options,
config=config)
t2 = time.time()
bestv, bestpar, bestchi, vel_err = [
res[_] for _ in ['best_vel', 'best_param', 'best_chi', 'vel_err']
]
assert (np.abs(bestv - 15) < 15)
param0 = vel_fit.firstguess(specdata, options=options, config=config)
resfull = vel_fit.process(specdata,
param0,
resols_params,
options=options,
config=config)
chisquare = np.mean(
((specdata[0].spec - resfull['yfit'][0]) / specdata[0].espec)**2)
assert (chisquare < 1.2)
assert (np.abs(resfull['vel'] - 6) < 10)
rot_params = (300, )
plt.clf()
ret = spec_fit.get_chisq(specdata,
bestv,
bestpar,
rot_params,
resols_params,
options=options,
config=config,
full_output=True)
plt.plot(specdata[0].lam, specdata[0].spec, 'k')
plt.plot(specdata[0].lam, ret['models'][0], 'r')
plt.savefig(path + '/plot_sdss_test1.png')
# Test the fit with the resolution matrix
rot_params = None
resol_mat = spec_fit.construct_resol_mat(specdata[0].lam, 50)
resols_params = {'sdss1': resol_mat}
ret = spec_fit.get_chisq(specdata,
bestv,
bestpar,
rot_params,
resol_params=resols_params,
options=options,
config=config,
full_output=True)
plt.clf()
plt.plot(specdata[0].lam, specdata[0].spec, 'k')
plt.plot(specdata[0].lam, ret['models'][0], 'r')
plt.savefig(path + '/plot_sdss_test2.png')
resol_mat = spec_fit.construct_resol_mat(specdata[0].lam, 50)
specdata = [
spec_fit.SpecData('sdss1',
10**dat['loglam'],
dat['flux'],
err,
resolution=resol_mat)
]
ret = spec_fit.get_chisq(specdata,
bestv,
bestpar,
rot_params,
options=options,
config=config,
full_output=True)
plt.clf()
plt.plot(specdata[0].lam, specdata[0].spec, 'k')
plt.plot(specdata[0].lam, ret['models'][0], 'r')
plt.savefig(path + '/plot_sdss_test3.png')
ret = spec_fit.get_chisq_continuum(specdata, options=options)
def test_fit():
config = utils.read_config(path + '/config.yaml')
# read data
dat = pyfits.getdata(path + '/data/spec-0266-51602-0031.fits')
err = dat['ivar']
err = 1. / err**.5
err[~np.isfinite(err)] = 1e40
# construct specdata object
specdata = [
spec_fit.SpecData('sdss1', 10**dat['loglam'], dat['flux'], err)
]
options = {'npoly': 15}
paramDict0 = {
'logg': 2,
'teff': 5000,
'feh': -1,
'alpha': 0.2,
'vsini': 19
}
fixParam = ['vsini']
paramDict0 = {
'logg': 2,
'teff': 5000,
'feh': -1,
'alpha': 0.2,
'vsini': 19
}
fixParam = ['vsini']
# fit with fixed vssini
res = vel_fit.process(specdata,
paramDict0,
fixParam=fixParam,
config=config,
options=options)
paramDict0 = {
'logg': 2,
'teff': 5000,
'feh': -1,
'alpha': 0.2,
'vsini': 19
}
# fit witout fixin
fixParam = []
res = vel_fit.process(specdata,
paramDict0,
fixParam=fixParam,
config=config,
options=options)
options = {'npoly': 15}
# first guess fit
xres0 = vel_fit.firstguess(specdata, config=config)
# ccf
res = fitter_ccf.fit(specdata, config)
paramDict0 = res['best_par']
fixParam = []
if res['best_vsini'] is not None:
paramDict0['vsini'] = res['best_vsini']
res1 = vel_fit.process(specdata,
paramDict0,
fixParam=fixParam,
config=config,
options=options)
print(res1)
plt.figure(figsize=(6, 2), dpi=300)
plt.plot(specdata[0].lam, specdata[0].spec, 'k-')
plt.plot(specdata[0].lam, res1['yfit'][0], 'r-')
plt.tight_layout()
plt.savefig(path + '/plot_test_fit_sdss.png')
# test priors
res2 = vel_fit.process(specdata,
paramDict0,
fixParam=fixParam,
config=config,
options=options,
priors={'teff': (9000, 50)})
# test rbf
options['rbf_continuum'] = False
res2 = vel_fit.process(specdata,
paramDict0,
fixParam=fixParam,
config=config,
options=options)
class si:
fname = 'tests/test.yaml'
config = utils.read_config(fname)
def test_interp():
conf = utils.read_config(path + '/config.yaml')
interp = spec_inter.getInterpolator('sdss1', conf)
interp.eval({'teff': 5000, 'logg': 1, 'feh': -1, 'alpha': 0.3})