Source code for pacman.s30_run

import time
import shutil
from pathlib import Path

import numpy as np
from astropy.io import ascii

from .lib import manageevent as me
from .lib import nice_fit_par
from .lib import plots
from .lib import util
from .lib.formatter import ReturnParams
from .lib.least_squares import lsq_fit
from .lib.mcmc import mcmc_fit
from .lib.model import Model
from .lib.nested import nested_sample
from .lib.options import OPTIONS
from .lib.read_data import Data


[docs]def run30(eventlabel: str, workdir: Path, meta=None): """This functions reads in the spectroscopic or white light curve(s) and fits a model to them.""" print('Starting s30\n') if meta is None: meta = me.loadevent(workdir / f'WFC3_{eventlabel}_Meta_Save') # Create directories for Stage 3 processing datetime = time.strftime('%Y-%m-%d_%H-%M-%S') # Create a directory for the white or the spectroscopic fit if meta.s30_fit_white: meta.fitdir = Path('fit_white') / f'fit_{datetime}_{meta.eventlabel}' elif meta.s30_fit_spec: meta.fitdir = Path('fit_spec') / f'fit_{datetime}_{meta.eventlabel}' fit_dir = meta.workdir / meta.fitdir if not fit_dir.exists(): fit_dir.mkdir(parents=True, exist_ok=True) # Make fit_par nicer by lining up the columns nice_fit_par.nice_fit_par(meta.workdir / "fit_par.txt") # Copy pcf and fit_par files into the new fitdirectory shutil.copy(meta.workdir / "obs_par.pcf", fit_dir) shutil.copy(meta.workdir / "fit_par.txt", fit_dir) # Reads in fit parameters from the fit_par file #TODO: Check that fit_par is configured correctly. Eg initial value has to be within boundaries! fit_par = ascii.read(meta.workdir / "fit_par.txt", Reader=ascii.CommentedHeader) # Read in the user wanted fit functions myfuncs = meta.s30_myfuncs # Indentify path of light curve files files, meta = util.read_fitfiles(meta) # Prepare some lists for diagnostics # TODO: Maybe make dict out of this? if meta.run_verbose: vals = [] errs = [] idxs = [] if meta.run_mcmc: vals_mcmc = [] errs_lower_mcmc = [] errs_upper_mcmc = [] if meta.run_nested: vals_nested = [] errs_lower_nested = [] errs_upper_nested = [] meta = util.log_run_setup(meta) # Loop through light curves for counter, f in enumerate(files): print('\n****** File: {0}/{1}'.format(counter+1, len(files))) meta.s30_file_counter = counter time.sleep(1.01) #sleep to prevent overwriting of data if saved in the same second meta.run_file = f meta.fittime = time.strftime('%Y-%m-%d_%H-%M-%S') ### RUN LEAST SQUARES ### ## IF NO CLIPPING IS WISHED if meta.run_lsq: if meta.run_clipiters == 0: print('\n') data = Data(f, meta, fit_par) model = Model(data, myfuncs) print('\n*STARTS LEAST SQUARED*') data, model, params, m = lsq_fit(fit_par, data, meta, model, myfuncs, noclip=True) #not clipping ## IF THE USER WANTS CLIPPING else: clip_idxs = [] for iii in range(meta.run_clipiters+1): print('\n') print('Sigma Iters: ', iii, 'of', meta.run_clipiters) if iii == 0: data = Data(f, meta, fit_par) else: data = Data(f, meta, fit_par, clip_idx) model = Model(data, myfuncs) if iii == meta.run_clipiters: data, model, params, m = lsq_fit(fit_par, data, meta, model, myfuncs, noclip=True) else: data, model, params, clip_idx, m = lsq_fit(fit_par, data, meta, model, myfuncs) print("rms, chi2red = ", model.rms, model.chi2red) print(clip_idx == []) if clip_idx == []: break clip_idxs.append(clip_idx) print(clip_idxs) print('length: ', len(clip_idxs) ) if len(clip_idxs)>1: clip_idx = update_clips(clip_idxs) print(clip_idx) clip_idxs = update_clips(clip_idxs) print(clip_idxs) if meta.run_verbose: print("rms, chi2red = ", model.rms, model.chi2red) # Save white systematics file if it was a white fit if meta.s30_fit_white: with (fit_dir / 'white_systematics.txt')\ .open("w", encoding="ascii") as outfile: for i in range(len(model.all_sys)): print(model.all_sys[i], file=outfile) print('Saved white_systematics.txt file') meta.labels = data.free_parnames if meta.run_mcmc: print('\n*STARTS MCMC*') time.sleep(1.01) if meta.rescale_uncert: # NOTE: Rescale error bars so reduced chi-squared is one print(f'rescale_uncert in the pcf was set to {meta.rescale_uncert}') if model.chi2red < 1: print('After the first fit, you got chi2_red < 1') print(f'rescale_uncert_smaller_1 was set to {meta.rescale_uncert_smaller_1}') if meta.rescale_uncert_smaller_1: # also rescale if chi2red < 1 print('We will therefore rescale the errorbars, so that chi2_red = 1') data.err *= np.sqrt(model.chi2red) else: print('chi2red < 1 and rescale_uncert_smaller_1 = False --> no rescaling') elif model.chi2red >= 1: print('After the first fit, you got chi2_red >= 1') print('Errorbars are being rescaled so that chi2_red = 1') data.err *= np.sqrt(model.chi2red) # if a least square hasnt been run yet or the uncertainies should be rescaled, to the lsq again if not meta.run_lsq or meta.rescale_uncert: data, model, params, m = lsq_fit(fit_par, data, meta, model, myfuncs, noclip=True) if meta.run_verbose == True: print("rms, chi2red = ", model.rms, model.chi2red) val_mcmc, err_lower_mcmc, err_upper_mcmc, fit = mcmc_fit(data, model, params, f, meta, fit_par) if meta.run_nested: print('\n*STARTS NESTED SAMPLING*') time.sleep(1.01) if meta.rescale_uncert: ##rescale error bars so reduced chi-squared is one print('Errorbars are being rescaled so that chi2_red = 1') if model.chi2red > 1: data.err *= np.sqrt(model.chi2red) if not meta.run_lsq or meta.rescale_uncert: data, model, params, m = lsq_fit(fit_par, data, meta, model, myfuncs, noclip=True) if meta.run_verbose == True: print("rms, chi2red = ", model.rms, model.chi2red) val_nested, err_lower_nested, err_upper_nested, fit = nested_sample(data, model, params, f, meta, fit_par) if meta.run_verbose: val, err, idx = ReturnParams(m, data) vals.append(val) errs.append(err) idxs.append(idx) if meta.run_mcmc: vals_mcmc.append(val_mcmc) errs_lower_mcmc.append(err_lower_mcmc) errs_upper_mcmc.append(err_upper_mcmc) if meta.run_nested: vals_nested.append(val_nested) errs_lower_nested.append(err_lower_nested) errs_upper_nested.append(err_upper_nested) if meta.run_verbose: plots.params_vs_wvl(vals, errs, idxs, meta) if meta.run_mcmc: plots.params_vs_wvl_mcmc(vals_mcmc, errs_lower_mcmc, errs_upper_mcmc, meta) if meta.run_nested: plots.params_vs_wvl_nested(vals_nested, errs_lower_nested, errs_upper_nested, meta) if not meta.s30_fit_white and ('rp' in meta.labels): # Saves rprs and wvl as a txt file util.make_lsq_rprs_txt(vals, errs, idxs, meta) # Saves rprs vs wvl as a plot plots.lsq_rprs(vals, errs, idxs, meta) if not meta.s30_fit_white and ('rp' in meta.labels) and meta.run_mcmc: plots.mcmc_rprs(vals_mcmc, errs_lower_mcmc, errs_upper_mcmc, meta) util.make_rprs_txt(vals_mcmc, errs_lower_mcmc, errs_upper_mcmc, meta, fitter='mcmc') if not meta.s30_fit_white and ('rp' in meta.labels) and meta.run_nested: plots.nested_rprs(vals_nested, errs_lower_nested, errs_upper_nested, meta) util.make_rprs_txt(vals_nested, errs_lower_nested, errs_upper_nested, meta, fitter='nested') if meta.run_mcmc or meta.run_nested: util.save_fit_output(fit, data, meta) print('Finished s30') return meta
[docs]def update_clips(clips_array): clips_old = clips_array[0] clips_new = clips_array[1] clips_new[0] + sum([i <= clips_new[0] for i in clips_old]) clips_new_updated = [clips_new[ii] + sum([i <= clips_new[ii] for i in clips_old]) for ii in range(len(clips_new))] return np.concatenate((clips_old, clips_new_updated))