def _on_select(self, xmin, xmax): wvln = self._ui['textbox'].text().strip() if wvln == '': self._ui['textbox'].setText('Error: please enter a wavelength value before selecting') return wave_val = float(wvln) # if line_list specified, find closest line from list: if self.line_list is not None: absdiff = np.abs(self.line_list - wave_val) idx = absdiff.argmin() if absdiff[idx] > 1.: logger.error("Couldn't find precise line corresponding to " "input {:.3f}".format(wave_val)) return logger.info("Snapping input wavelength {:.3f} to line list " "value {:.3f}".format(wave_val, self.line_list[idx])) wave_val = self.line_list[idx] self._done_wavel_idx.append(idx) # line_props, line_cov = self.get_line_props(xmin, xmax) line_props,_ = self.get_line_props(xmin, xmax) if line_props is None: return self.draw_line_marker(line_props, wave_val, xmin, xmax) self.fig.suptitle('') plt.draw() self.fig.canvas.draw() self._map_dict['wavel'].append(wave_val) self._map_dict['pixel'].append(line_props['x0'])
def main(db_path, run_name, data_root_path=None, filename=None, overwrite=False, pool=None): if pool is None: pool = schwimmbad.SerialPool() # connect to the database engine = db_connect(db_path) # engine.echo = True logger.debug("Connected to database at '{}'".format(db_path)) # create a new session for interacting with the database session = Session() root_path, _ = path.split(db_path) if data_root_path is None: data_root_path = root_path plot_path = path.join(root_path, 'plots', run_name) if not path.exists(plot_path): os.makedirs(plot_path, exist_ok=True) # TODO: there might be some bugs here... n_lines = session.query(SpectralLineInfo).count() Halpha = session.query(SpectralLineInfo)\ .filter(SpectralLineInfo.name == 'Halpha').one() OI_lines = session.query(SpectralLineInfo)\ .filter(SpectralLineInfo.name.contains('[OI]')).all() if filename is None: # grab all unfinished sources observations = session.query(Observation).join(Run)\ .filter(Run.name == run_name).all() else: # only process the observation corresponding to this filename observations = session.query(Observation).join(Run)\ .filter(Run.name == run_name)\ .filter(Observation.filename_raw == filename).all() for obs in observations: measurements = session.query(SpectralLineMeasurement)\ .join(Observation)\ .filter(Observation.id == obs.id).all() if len(measurements) == n_lines and not overwrite: logger.debug('All line measurements already complete for object ' '{0} in file {1}'.format(obs.object, obs.filename_raw)) continue # Read the spectrum data and get wavelength solution filebase, _ = path.splitext(obs.filename_1d) filename_1d = obs.path_1d(data_root_path) spec = Table.read(filename_1d) logger.debug('Loaded 1D spectrum for object {0} from file {1}'.format( obs.object, filename_1d)) # Extract region around Halpha x, (flux, ivar) = extract_region( spec['wavelength'], center=Halpha.wavelength.value, width=100, arrs=[spec['source_flux'], spec['source_ivar']]) # We start by doing maximum likelihood estimation to fit the line, then # use the best-fit parameters to initialize an MCMC run. # TODO: need to figure out if it's emission or absorption...for now just # assume absorption absorp_emiss = -1. lf = VoigtLineFitter(x, flux, ivar, absorp_emiss=absorp_emiss) lf.fit() fit_pars = lf.get_gp_mean_pars() if (not lf.success or abs(fit_pars['x0'] - Halpha.wavelength.value) > 16. or # 16 Å = ~700 km/s abs(fit_pars['amp']) < 10): # minimum amplitude - MAGIC NUMBER # TODO: should try again with emission line logger.error('absorption line has tiny amplitude! did ' 'auto-determination of absorption/emission fail?') # TODO: what now? continue fig = lf.plot_fit() fig.savefig(path.join(plot_path, '{}_maxlike.png'.format(filebase)), dpi=256) plt.close(fig) # ---------------------------------------------------------------------- # Run `emcee` instead to sample over GP model parameters: if fit_pars['std_G'] < 1E-2: lf.gp.freeze_parameter('mean:ln_std_G') initial = np.array(lf.gp.get_parameter_vector()) if initial[4] < -10: # TODO: ??? initial[4] = -8. if initial[5] < -10: # TODO: ??? initial[5] = -8. ndim, nwalkers = len(initial), 64 sampler = emcee.EnsembleSampler(nwalkers, ndim, log_probability, pool=pool, args=(lf.gp, flux)) logger.debug("Running burn-in...") p0 = initial + 1e-6 * np.random.randn(nwalkers, ndim) p0, lp, _ = sampler.run_mcmc(p0, 128) logger.debug("Running 2nd burn-in...") sampler.reset() p0 = p0[lp.argmax()] + 1e-3 * np.random.randn(nwalkers, ndim) p0, lp, _ = sampler.run_mcmc(p0, 512) logger.debug("Running production...") sampler.reset() pos, lp, _ = sampler.run_mcmc(p0, 1024) fit_kw = dict() for i, par_name in enumerate(lf.gp.get_parameter_names()): if 'kernel' in par_name: continue # remove 'mean:' par_name = par_name[5:] # skip bg if par_name.startswith('bg'): continue samples = sampler.flatchain[:, i] if par_name.startswith('ln_'): par_name = par_name[3:] samples = np.exp(samples) MAD = np.median(np.abs(samples - np.median(samples))) fit_kw[par_name] = np.median(samples) fit_kw[par_name + '_error'] = 1.5 * MAD # convert to ~stddev # remove all previous line measurements q = session.query(SpectralLineMeasurement).join(Observation)\ .filter(Observation.id == obs.id) if q.count() > 0: for meas in q.all(): session.delete(meas) session.commit() slm = SpectralLineMeasurement(**fit_kw) slm.info = Halpha slm.observation = obs session.add(slm) session.commit() # -------------------------------------------------------------------- # plot MCMC traces fig, axes = plt.subplots(2, 4, figsize=(18, 6)) for i in range(sampler.dim): for walker in sampler.chain[..., i]: axes.flat[i].plot(walker, marker='', drawstyle='steps-mid', alpha=0.2) axes.flat[i].set_title(lf.gp.get_parameter_names()[i], fontsize=12) fig.tight_layout() fig.savefig(path.join(plot_path, '{}_mcmc_trace.png'.format(filebase)), dpi=256) plt.close(fig) # -------------------------------------------------------------------- # -------------------------------------------------------------------- # plot samples fig, axes = plt.subplots(3, 1, figsize=(10, 10), sharex=True) samples = sampler.flatchain for s in samples[np.random.randint(len(samples), size=32)]: lf.gp.set_parameter_vector(s) lf.plot_fit(axes=axes, fit_alpha=0.2) fig.tight_layout() fig.savefig(path.join(plot_path, '{}_mcmc_fits.png'.format(filebase)), dpi=256) plt.close(fig) # -------------------------------------------------------------------- # -------------------------------------------------------------------- # corner plot fig = corner.corner( sampler.flatchain[::10, :], labels=[x.split(':')[1] for x in lf.gp.get_parameter_names()]) fig.savefig(path.join(plot_path, '{}_corner.png'.format(filebase)), dpi=256) plt.close(fig) # -------------------------------------------------------------------- # compute centroids for sky lines sky_centroids = [] for j, sky_line in enumerate(OI_lines): wvln = sky_line.wavelength.value x, (flux, ivar) = extract_region( spec['wavelength'], center=wvln, width=32., # angstroms arrs=[spec['background_flux'], spec['background_ivar']]) lf = GaussianLineFitter(x, flux, ivar, absorp_emiss=1.) # all emission lines try: lf.fit() fit_pars = lf.get_gp_mean_pars() except Exception as e: logger.warn("Failed to fit sky line {0}:\n{1}".format( sky_line, e)) lf.success = False fit_pars = lf.get_init() # OMG this is the biggest effing hack fit_pars['amp'] = 0. fit_pars['bg_coef'] = None fit_pars['x0'] = 0. # HACK: hackish signal-to-noise max_ = fit_pars['amp'] / np.sqrt(2 * np.pi * fit_pars['std']**2) SNR = max_ / np.median(1 / np.sqrt(ivar)) if (not lf.success or abs(fit_pars['x0'] - wvln) > 4 or fit_pars['amp'] < 10 or fit_pars['std'] > 4 or SNR < 2.5): # failed x0 = np.nan * u.angstrom title = 'f****d' fit_pars['amp'] = 0. else: x0 = fit_pars['x0'] * u.angstrom title = '{:.2f}'.format(fit_pars['amp']) if lf.success: fig = lf.plot_fit() fig.suptitle(title, y=0.95) fig.subplots_adjust(top=0.8) fig.savefig(path.join( plot_path, '{}_maxlike_sky_{:.0f}.png'.format(filebase, wvln)), dpi=256) plt.close(fig) # store the sky line measurements fit_pars['std_G'] = fit_pars.pop('std') # HACK fit_pars.pop('bg_coef') # HACK slm = SpectralLineMeasurement(**fit_pars) slm.info = sky_line slm.observation = obs session.add(slm) session.commit() sky_centroids.append(x0) sky_centroids = u.Quantity(sky_centroids) logger.info('{} [{}]: x0={x0:.3f} σ={err:.3f}\n--------'.format( obs.object, filebase, x0=fit_kw['x0'], err=fit_kw['x0_error'])) session.commit() pool.close()
def main(night_path, skip_list_file, mask_file, overwrite=False, plot=False): """ See argparse block at bottom of script for description of parameters. """ night_path = path.realpath(path.expanduser(night_path)) if not path.exists(night_path): raise IOError("Path '{}' doesn't exist".format(night_path)) logger.info("Reading data from path: {}".format(night_path)) base_path, night_name = path.split(night_path) data_path, run_name = path.split(base_path) output_path = path.realpath( path.join(data_path, 'processed', run_name, night_name)) os.makedirs(output_path, exist_ok=True) logger.info("Saving processed files to path: {}".format(output_path)) if plot: # if we're making plots plot_path = path.realpath(path.join(output_path, 'plots')) logger.debug("Will make and save plots to: {}".format(plot_path)) os.makedirs(plot_path, exist_ok=True) else: plot_path = None # check for files to skip (e.g., saturated or errored exposures) if skip_list_file is not None: # a file containing a list of filenames to skip with open(skip_list_file, 'r') as f: skip_list = [x.strip() for x in f if x.strip()] else: skip_list = None # look for pixel mask file if mask_file is not None: with open( mask_file, 'r' ) as f: # load YAML file specifying pixel masks for nearby sources pixel_mask_spec = yaml.load(f.read()) else: pixel_mask_spec = None # generate the raw image file collection to process ic = GlobImageFileCollection(night_path, skip_filenames=skip_list) logger.info("Frames to process:") logger.info("- Bias frames: {}".format( len(ic.files_filtered(imagetyp='BIAS')))) logger.info("- Flat frames: {}".format( len(ic.files_filtered(imagetyp='FLAT')))) logger.info("- Comparison lamp frames: {}".format( len(ic.files_filtered(imagetyp='COMP')))) logger.info("- Object frames: {}".format( len(ic.files_filtered(imagetyp='OBJECT')))) # HACK: ic = GlobImageFileCollection(night_path, skip_filenames=skip_list) # ============================ # Create the master bias frame # ============================ # overscan region of the CCD, using FITS index notation oscan_fits_section = "[{}:{},:]".format(oscan_idx, oscan_idx + oscan_size) master_bias_file = path.join(output_path, 'master_bias.fits') if not os.path.exists(master_bias_file) or overwrite: # get list of overscan-subtracted bias frames as 2D image arrays bias_list = [] for hdu, fname in ic.hdus(return_fname=True, imagetyp='BIAS'): logger.debug('Processing Bias frame: {0}'.format(fname)) ccd = CCDData.read(path.join(ic.location, fname), unit='adu') ccd = ccdproc.gain_correct(ccd, gain=ccd_gain) ccd = ccdproc.subtract_overscan(ccd, overscan=ccd[:, oscan_idx:]) ccd = ccdproc.trim_image(ccd, fits_section="[1:{},:]".format(oscan_idx)) bias_list.append(ccd) # combine all bias frames into a master bias frame logger.info("Creating master bias frame") master_bias = ccdproc.combine(bias_list, method='average', clip_extrema=True, nlow=1, nhigh=1, error=True) master_bias.write(master_bias_file, overwrite=True) else: logger.info("Master bias frame file already exists: {}".format( master_bias_file)) master_bias = CCDData.read(master_bias_file) if plot: # TODO: this assumes vertical CCD assert master_bias.shape[0] > master_bias.shape[1] aspect_ratio = master_bias.shape[1] / master_bias.shape[0] fig, ax = plt.subplots(1, 1, figsize=(10, 12 * aspect_ratio)) vmin, vmax = zscaler.get_limits(master_bias.data) cs = ax.imshow(master_bias.data.T, origin='bottom', cmap=cmap, vmin=max(0, vmin), vmax=vmax) ax.set_title('master bias frame [zscale]') fig.colorbar(cs) fig.tight_layout() fig.savefig(path.join(plot_path, 'master_bias.png')) plt.close(fig) # ============================ # Create the master flat field # ============================ # HACK: ic = GlobImageFileCollection(night_path, skip_filenames=skip_list) master_flat_file = path.join(output_path, 'master_flat.fits') if not os.path.exists(master_flat_file) or overwrite: # create a list of flat frames flat_list = [] for hdu, fname in ic.hdus(return_fname=True, imagetyp='FLAT'): logger.debug('Processing Flat frame: {0}'.format(fname)) ccd = CCDData.read(path.join(ic.location, fname), unit='adu') ccd = ccdproc.gain_correct(ccd, gain=ccd_gain) ccd = ccdproc.ccd_process(ccd, oscan=oscan_fits_section, trim="[1:{},:]".format(oscan_idx), master_bias=master_bias) flat_list.append(ccd) # combine into a single master flat - use 3*sigma sigma-clipping logger.info("Creating master flat frame") master_flat = ccdproc.combine(flat_list, method='average', sigma_clip=True, low_thresh=3, high_thresh=3) master_flat.write(master_flat_file, overwrite=True) # TODO: make plot if requested? else: logger.info("Master flat frame file already exists: {}".format( master_flat_file)) master_flat = CCDData.read(master_flat_file) if plot: # TODO: this assumes vertical CCD assert master_flat.shape[0] > master_flat.shape[1] aspect_ratio = master_flat.shape[1] / master_flat.shape[0] fig, ax = plt.subplots(1, 1, figsize=(10, 12 * aspect_ratio)) vmin, vmax = zscaler.get_limits(master_flat.data) cs = ax.imshow(master_flat.data.T, origin='bottom', cmap=cmap, vmin=max(0, vmin), vmax=vmax) ax.set_title('master flat frame [zscale]') fig.colorbar(cs) fig.tight_layout() fig.savefig(path.join(plot_path, 'master_flat.png')) plt.close(fig) # ===================== # Process object frames # ===================== # HACK: ic = GlobImageFileCollection(night_path, skip_filenames=skip_list) logger.info("Beginning object frame processing...") for hdu, fname in ic.hdus(return_fname=True, imagetyp='OBJECT'): new_fname = path.join(output_path, 'p_{}'.format(fname)) # ------------------------------------------- # First do the simple processing of the frame # ------------------------------------------- logger.debug("Processing '{}' [{}]".format(hdu.header['OBJECT'], fname)) if path.exists(new_fname) and not overwrite: logger.log(1, "\tAlready processed! {}".format(new_fname)) ext = SourceCCDExtractor(filename=path.join( ic.location, new_fname), plot_path=plot_path, zscaler=zscaler, cmap=cmap, **ccd_props) nccd = ext.ccd # HACK: F**K this is a bad hack ext._filename_base = ext._filename_base[2:] else: # process the frame! ext = SourceCCDExtractor(filename=path.join(ic.location, fname), plot_path=plot_path, zscaler=zscaler, cmap=cmap, unit='adu', **ccd_props) _pix_mask = pixel_mask_spec.get( fname, None) if pixel_mask_spec is not None else None nccd = ext.process_raw_frame(pixel_mask_spec=_pix_mask, master_bias=master_bias, master_flat=master_flat) nccd.write(new_fname, overwrite=overwrite) # ------------------------------------------- # Now do the 1D extraction # ------------------------------------------- fname_1d = path.join(output_path, '1d_{0}'.format(fname)) if path.exists(fname_1d) and not overwrite: logger.log(1, "\tAlready extracted! {}".format(fname_1d)) continue else: logger.debug("\tExtracting to 1D") # first step is to fit a voigt profile to a middle-ish row to determine LSF lsf_p = ext.get_lsf_pars() # MAGIC NUMBER try: tbl = ext.extract_1d(lsf_p) except Exception as e: logger.error('Failed! {}: {}'.format(e.__class__.__name__, str(e))) continue hdu0 = fits.PrimaryHDU(header=nccd.header) hdu1 = fits.table_to_hdu(tbl) hdulist = fits.HDUList([hdu0, hdu1]) hdulist.writeto(fname_1d, overwrite=overwrite) del ext # ============================== # Process comparison lamp frames # ============================== # HACK: ic = GlobImageFileCollection(night_path, skip_filenames=skip_list) logger.info("Beginning comp. lamp frame processing...") for hdu, fname in ic.hdus(return_fname=True, imagetyp='COMP'): new_fname = path.join(output_path, 'p_{}'.format(fname)) logger.debug("\tProcessing '{}'".format(hdu.header['OBJECT'])) if path.exists(new_fname) and not overwrite: logger.log(1, "\tAlready processed! {}".format(new_fname)) ext = CompCCDExtractor(filename=path.join(ic.location, new_fname), plot_path=plot_path, zscaler=zscaler, cmap=cmap, **ccd_props) nccd = ext.ccd # HACK: F**K this is a bad hack ext._filename_base = ext._filename_base[2:] else: # process the frame! ext = CompCCDExtractor(filename=path.join(ic.location, fname), plot_path=plot_path, unit='adu', **ccd_props) _pix_mask = pixel_mask_spec.get( fname, None) if pixel_mask_spec is not None else None nccd = ext.process_raw_frame( pixel_mask_spec=_pix_mask, master_bias=master_bias, master_flat=master_flat, ) nccd.write(new_fname, overwrite=overwrite) # ------------------------------------------- # Now do the 1D extraction # ------------------------------------------- fname_1d = path.join(output_path, '1d_{0}'.format(fname)) if path.exists(fname_1d) and not overwrite: logger.log(1, "\tAlready extracted! {}".format(fname_1d)) continue else: logger.debug("\tExtracting to 1D") try: tbl = ext.extract_1d() except Exception as e: logger.error('Failed! {}: {}'.format(e.__class__.__name__, str(e))) continue hdu0 = fits.PrimaryHDU(header=nccd.header) hdu1 = fits.table_to_hdu(tbl) hdulist = fits.HDUList([hdu0, hdu1]) hdulist.writeto(fname_1d, overwrite=overwrite)
def auto_identify(self): if self.line_list is None: raise ValueError("Can't auto-identify lines without a line list.") if len(self._map_dict['wavel']) < 4: msg = "Please identify at least 4 lines before trying auto-identify." logger.error(msg) self._ui['textbox'].setText("ERROR: {}".format(msg)) return None _idx = np.argsort(self._map_dict['wavel']) wvln = np.array(self._map_dict['wavel'])[_idx] pixl = np.array(self._map_dict['pixel'])[_idx] # build an approximate wavelength solution to predict where lines are spl = InterpolatedUnivariateSpline(wvln, pixl, k=1) # use linear interp. predicted_pixels = spl(self.line_list) new_wavels = [] new_pixels = [] # from Wikipedia: https://en.wikipedia.org/wiki/Voigt_profile fG = 2*self._line_std_G*np.sqrt(2*np.log(2)) fL = 2*self._line_hwhm_L lw = 0.5346*fL + np.sqrt(0.2166*fL**2 + fG**2) for pix_ctr,xmin,xmax,wave_idx,wave in zip(predicted_pixels, predicted_pixels-5*lw, predicted_pixels+5*lw, range(len(self.line_list)), self.line_list): if pix_ctr < 200 or pix_ctr > 1600: # skip if outside good rows continue elif wave_idx in self._done_wavel_idx: # skip if already fit continue logger.debug("Fitting line at predicted pix={:.2f}, λ={:.2f}" .format(pix_ctr, wave)) try: lp,gp = self.get_line_props(xmin, xmax, std_G0=self._line_std_G, hwhm_L0=self._line_hwhm_L) except Exception as e: logger.error("Failed to auto-fit line at {} ({msg})" .format(wave, msg=str(e))) continue print(lp['amp'], lp['x0']) if lp is None or lp['amp'] < 100.: # HACK continue # figure out closest line # _all_pix = np.concatenate((self._map_dict['pixel'], new_pixels)) # _all_wav = np.concatenate((self._map_dict['wavel'], new_wavels)) # _diff = np.abs(lp['x0'] - np.array(_all_pix)) # min_diff_idx = np.argmin(_diff) # min_diff_pix = _all_pix[min_diff_idx] # min_diff_wav = _all_wav[min_diff_idx] # if _diff[min_diff_idx] < 3.: # logger.error("Fit line is too close to another at pix={:.2f}, λ={:.2f}" # .format(min_diff_pix, min_diff_wav)) # continue self.draw_line_marker(lp, wave, xmin, xmax, gp=gp) new_wavels.append(wave) new_pixels.append(pix_ctr) self._done_wavel_idx.append(wave_idx) self.fig.canvas.draw() _idx = np.argsort(new_wavels) self._map_dict['wavel'] = np.array(new_wavels)[_idx] self._map_dict['pixel'] = np.array(new_pixels)[_idx]