diff --git a/.gitignore b/.gitignore index bd74b0361..877038827 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ *.DS_Store *.xlsx .idea +randomseed Spectractor.egg-info/ build/ dist/ diff --git a/runFitter.py b/runFitter.py index a6ee79c63..37411d1aa 100644 --- a/runFitter.py +++ b/runFitter.py @@ -1,5 +1,5 @@ from spectractor import parameters -from spectractor.fit.fitter import run_minimisation, SpectrogramFitWorkspace +from spectractor.fit.fit_spectrogram import SpectrogramFitWorkspace, run_spectrogram_minimisation from spectractor.config import load_config @@ -36,6 +36,6 @@ w = SpectrogramFitWorkspace(file_name, atmgrid_file_name=atmgrid_filename, nsteps=2000, burnin=1000, nbins=10, verbose=0, plot=False, live_fit=False) - run_minimisation(w, method="newton") + run_spectrogram_minimisation(w, method="newton") # run_emcee(fit_workspace) # fit_workspace.analyze_chains() diff --git a/spectractor/extractor/background.py b/spectractor/extractor/background.py index 1601d622f..9b82e04eb 100644 --- a/spectractor/extractor/background.py +++ b/spectractor/extractor/background.py @@ -178,7 +178,7 @@ def extract_background_photutils(data, err, ws=(20, 30), mask_signal_region=True ax[0].set_title(f'Data background: mean={np.nanmean(bgd_bands):.3f}, std={np.nanstd(bgd_bands):.3f}') ax[0].set_xlabel('X [pixels]') ax[0].set_ylabel('Y [pixels]') - bkg.plot_meshes(outlines=True, color='#1f77b4', ax=ax[0]) + bkg.plot_meshes(outlines=True, color='#1f77b4', axes=ax[0]) b = bkg.background im = ax[1].imshow(b, origin='lower', aspect="auto") ax[1].set_xlabel('X [pixels]') diff --git a/spectractor/extractor/psf.py b/spectractor/extractor/psf.py index b6833c7f3..3ce726e4e 100644 --- a/spectractor/extractor/psf.py +++ b/spectractor/extractor/psf.py @@ -11,10 +11,11 @@ from astropy.modeling.models import Moffat1D from astropy.table import Table -from spectractor.tools import dichotomie, fit_poly1d, fit_moffat1d_outlier_removal +from spectractor.tools import dichotomie, fit_poly1d, fit_moffat1d_outlier_removal, plot_image_simple from spectractor.extractor.background import extract_background_photutils from spectractor import parameters from spectractor.config import set_logger +from spectractor.fit.fitter import FitWorkspace, run_minimisation class PSF1D(Fittable1DModel): @@ -29,7 +30,7 @@ class PSF1D(Fittable1DModel): stddev = Parameter('stddev', default=1) saturation = Parameter('saturation', default=1) - param_titles = ["A", "y", "\gamma", r"\alpha", r"\eta", r"\sigma", "saturation"] + param_titles = ["A", "y", r"\gamma", r"\alpha", r"\eta", r"\sigma", "saturation"] @staticmethod def evaluate(x, amplitude_moffat, x_mean, gamma, alpha, eta_gauss, stddev, saturation): @@ -632,8 +633,8 @@ def from_profile_params_to_shape_params(self, profile_params): p = profile_params[x, :] self.PSF.parameters = p fwhm = self.PSF.fwhm(x_array=pixel_y) - self.table['flux_integral'][x] = self.PSF.integrate(bounds=(-10 * fwhm + p[1], 10 * fwhm + p[1]), - x_array=pixel_y) + integral = self.PSF.integrate(bounds=(-10 * fwhm + p[1], 10 * fwhm + p[1]), x_array=pixel_y) + self.table['flux_integral'][x] = integral self.table['fwhm'][x] = fwhm self.table['Dy_mean'][x] = 0 @@ -976,6 +977,9 @@ def fit_transverse_PSF1D_profile(self, data, err, w, ws, pixel_step=1, bgd_model fit = fit_moffat1d_outlier_removal(index, signal, sigma=sigma, niter=2, guess=guess[:4], bounds=np.array(bounds[:4]).T) moffat_guess = [getattr(fit, p).value for p in fit.param_names] + signal_width_guess = moffat_guess[2] + bounds[2] = (0.1, min(Ny // 2, 5*signal_width_guess)) + bounds[5] = (0.1, min(Ny // 2, 5*signal_width_guess)) guess[:4] = moffat_guess init_guess = np.copy(guess) # Go from max to right, then from max to left @@ -1014,61 +1018,19 @@ def fit_transverse_PSF1D_profile(self, data, err, w, ws, pixel_step=1, bgd_model # if guess[0] * (1 + guess[4]) > 1.2 * maxi: # guess = [0.9 * maxi, mean, std, 2, 0, std, saturation] PSF_guess = PSF1D(*guess) - outliers = [] - # fit with outlier removal to clean background stars - # first a simple moffat to get the general shape - # moffat1d = fit_moffat1d_outlier_removal(index, signal, sigma=5, niter=2, - # guess=guess[:4], bounds=np.array(bounds[:4]).T) - # guess[:4] = [getattr(moffat1d, p).value for p in moffat1d.param_names] - # then PSF model using the result from the Moffat1D fit - # fit, outliers = fit_PSF1D_outlier_removal(index, signal, sub_errors=err[:, x], method='basinhopping', - # guess=guess, bounds=bounds, sigma=5, niter=2, - # niter_basinhopping=5, T_basinhopping=1) - # fit = fit_PSF1D_minuit(index[good_indices], signal[good_indices], guess=guess, bounds=bounds, - # data_errors=err[good_indices, x]) fit, outliers = fit_PSF1D_minuit_outlier_removal(index, signal, guess=guess, bounds=bounds, data_errors=err[:, x], sigma=sigma, niter=2, consecutive=4) - # good_indices = [i for i in index if i not in outliers] - # outliers = [i for i in index if np.abs((signal[i] - fit(i)) / err[i, x]) > sigma] - """ - This part is relevant if outliers rejection above is activated - # test if 3 consecutive pixels are in the outlier list - test = 0 - consecutive_outliers = False - for o in range(1, len(outliers)): - t = outliers[o] - outliers[o - 1] - if t == 1: - test += t - else: - test = 0 - if test > 1: - consecutive_outliers = True - break - # test if the fit has badly fitted the two highest data points or the middle points - test = np.copy(signal) - max_badfit = False - max_index = signal.argmax() - if max_index in outliers: - test[max_index] = 0 - if test.argmax() in outliers: - max_badfit = True - # if there are consecutive outliers or max is badly fitted, re-estimate the guess and refi - if consecutive_outliers: # or max_badfit: - my_logger.warning(f'\n\tRefit because of max_badfit={max_badfit} or ' - f'consecutive_outliers={consecutive_outliers}') - guess = [1.3 * maxi, middle, guess[2], guess[3], -0.3, std / 2, saturation] - fit, outliers = fit_PSF1D_outlier_removal(index, signal, sub_errors=err[:, x], method='basinhopping', - guess=guess, bounds=bounds, sigma=5, niter=2, - niter_basinhopping=20, T_basinhopping=1) - """ - # compute the flux_sum - guess = [getattr(fit, p).value for p in fit.param_names] - self.profile_params[x, 0] = guess[0] - self.profile_params[x, -6:] = guess[1:] + # It is better not to propagate the guess to further pixel columns + # otherwise fit_chromatic_psf1D is more likely to get trapped in a local minimum + # Randomness of the slice fit is better : + # guess = [getattr(fit, p).value for p in fit.param_names] + best_fit = [getattr(fit, p).value for p in fit.param_names] + self.profile_params[x, 0] = best_fit[0] + self.profile_params[x, -6:] = best_fit[1:] self.table['flux_err'][x] = np.sqrt(np.sum(err[:, x] ** 2)) self.table['flux_sum'][x] = np.sum(signal) if live_fit and parameters.DISPLAY: # pragma: no cover - plot_transverse_PSF1D_profile(x, index, bgd_index, data, err, fit, bgd_model_func, guess, + plot_transverse_PSF1D_profile(x, index, bgd_index, data, err, fit, bgd_model_func, best_fit, PSF_guess, outliers, sigma, live_fit) # interpolate the skipped pixels with splines x = np.arange(Nx) @@ -1135,7 +1097,7 @@ def generate_test_poly_params(self): def evaluate(self, poly_params): """ - Simulate a 2D spectrogram of size Nx times Ny. + Simulate a 2D spectrogram of size Nx times Ny with transverse 1D PSF profiles. Parameters ---------- @@ -1170,7 +1132,7 @@ def evaluate(self, poly_params): output[:, k] = PSF1D.evaluate(y, *profile_params[k]) return output - def fit_chromatic_PSF1D(self, data, bgd_model_func=None, data_errors=None): + def fit_chromatic_PSF1D_minuit(self, data, bgd_model_func=None, data_errors=None): """ Fit a chromatic PSF model on 2D data. @@ -1212,15 +1174,14 @@ def fit_chromatic_PSF1D(self, data, bgd_model_func=None, data_errors=None): >>> s.plot_summary(truth=s0) # Fit the data: - >>> s.fit_chromatic_PSF1D(data, bgd_model_func=bgd_model_func, data_errors=data_errors) + >>> s.fit_chromatic_PSF1D_minuit(data, bgd_model_func=bgd_model_func, data_errors=data_errors) >>> s.plot_summary(truth=s0) """ - my_logger = set_logger(__name__) Ny, Nx = data.shape if Ny != self.Ny: - my_logger.error(f"\n\tData y shape {Ny} different from ChromaticPSF1D input self.Ny {self.Ny}.") + self.my_logger.error(f"\n\tData y shape {Ny} different from ChromaticPSF1D input self.Ny {self.Ny}.") if Nx != self.Nx: - my_logger.error(f"\n\tData x shape {Nx} different from ChromaticPSF1D input self.Nx {self.Nx}.") + self.my_logger.error(f"\n\tData x shape {Nx} different from ChromaticPSF1D input self.Nx {self.Nx}.") guess = np.copy(self.poly_params) pixels = np.arange(Ny) @@ -1265,7 +1226,7 @@ def spectrogram_chisq(shape_params): else: return np.nansum(((mod - data_subtracted) / data_errors_cropped) ** 2) + penalty - my_logger.debug(f'\n\tStart chisq: {spectrogram_chisq(guess[Nx:])} with {guess[Nx:]}') + self.my_logger.debug(f'\n\tStart chisq: {spectrogram_chisq(guess[Nx:])} with {guess[Nx:]}') error = 0.01 * np.abs(guess) * np.ones_like(guess) fix = [False] * (self.n_poly_params - Nx) fix[-1] = True @@ -1288,6 +1249,220 @@ def spectrogram_chisq(shape_params): self.plot_summary() self.plot_chromatic_PSF1D_residuals(bgd, data, data_errors, guess=guess, title='Best fit') + def fit_chromatic_PSF1D(self, data, bgd_model_func=None, data_errors=None): + """ + Fit a chromatic PSF model on 2D data. + + Parameters + ---------- + data: array_like + 2D array containing the image data. + bgd_model_func: callable, optional + A 2D function to model the extracted background (default: None -> null background) + data_errors: np.array + the 2D array uncertainties. + + Examples + -------- + + # Set the parameters + >>> parameters.PIXDIST_BACKGROUND = 40 + >>> parameters.PIXWIDTH_BACKGROUND = 10 + >>> parameters.PIXWIDTH_SIGNAL = 30 + >>> parameters.VERBOSE = True + + # Build a mock spectrogram with random Poisson noise: + >>> s0 = ChromaticPSF1D(Nx=100, Ny=100, deg=4, saturation=1000) + >>> params = s0.generate_test_poly_params() + >>> s0.poly_params = params + >>> saturation = params[-1] + >>> data = s0.evaluate(params) + >>> bgd = 10*np.ones_like(data) + >>> data += bgd + >>> data = np.random.poisson(data) + >>> data_errors = np.sqrt(data+1) + + # Extract the background + >>> bgd_model_func = extract_background_photutils(data, data_errors, ws=[30,50]) + + # Estimate the first guess values + >>> s = ChromaticPSF1D(Nx=100, Ny=100, deg=4, saturation=saturation) + >>> s.fit_transverse_PSF1D_profile(data, data_errors, w=20, ws=[30,50], + ... pixel_step=1, bgd_model_func=bgd_model_func, saturation=saturation, live_fit=False) + >>> s.plot_summary(truth=s0) + + # Fit the data: + >>> s.fit_chromatic_PSF1D(data, bgd_model_func=bgd_model_func, data_errors=data_errors) + >>> s.plot_summary(truth=s0) + """ + w = ChromaticPSF1DFitWorkspace(self, data, data_errors, bgd_model_func=bgd_model_func) + + guess = np.copy(self.poly_params) + run_minimisation(w, method="newton", ftol=1e-4, xtol=1e-4) + + self.poly_params = w.poly_params + self.profile_params = self.from_poly_params_to_profile_params(self.poly_params, force_positive=True) + self.fill_table_with_profile_params(self.profile_params) + self.from_profile_params_to_shape_params(self.profile_params) + if parameters.DEBUG or True: + # Plot data, best fit model and residuals: + self.plot_summary() + self.plot_chromatic_PSF1D_residuals(w.bgd, data, data_errors, guess=guess, title='Best fit') + + +class ChromaticPSF1DFitWorkspace(FitWorkspace): + + def __init__(self, chromatic_psf, data, data_errors, bgd_model_func=None, file_name="", + nwalkers=18, nsteps=1000, burnin=100, nbins=10, + verbose=0, plot=False, live_fit=False, truth=None): + FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, + live_fit, truth=truth) + self.my_logger = set_logger(self.__class__.__name__) + self.chromatic_psf = chromatic_psf + self.data = data + self.err = data_errors + self.bgd_model_func = bgd_model_func + length = len(self.chromatic_psf.table) + self.p = np.copy(self.chromatic_psf.poly_params[length:-1]) # remove saturation (fixed parameter)) + self.ndim = self.p.size + self.poly_params = np.copy(self.chromatic_psf.poly_params) + self.input_labels = list(np.copy(self.chromatic_psf.poly_params_labels[length:-1])) + self.axis_names = list(np.copy(self.chromatic_psf.poly_params_names[length:-1])) + self.bounds = self.chromatic_psf.set_bounds(data=None)[:-1] + self.nwalkers = max(2 * self.ndim, nwalkers) + + # prepare the fit + self.Ny, self.Nx = self.data.shape + if self.Ny != self.chromatic_psf.Ny: + self.my_logger.error(f"\n\tData y shape {self.Ny} different from ChromaticPSF input Ny {self.chromatic_psf.Ny}.") + if self.Nx != self.chromatic_psf.Nx: + self.my_logger.error(f"\n\tData x shape {self.Nx} different from ChromaticPSF input Nx {self.chromatic_psf.Nx}.") + self.pixels = np.arange(self.Ny) + + # prepare the background, data and errors + self.bgd = np.zeros_like(self.data) + if self.bgd_model_func is not None: + # xx, yy = np.meshgrid(np.arange(Nx), pixels) + self.bgd = self.bgd_model_func(np.arange(self.Nx), self.pixels) + self.data = self.data - self.bgd + self.bgd_std = float(np.std(np.random.poisson(self.bgd))) + + # crop spectrogram to fit faster + self.bgd_width = parameters.PIXWIDTH_BACKGROUND + parameters.PIXDIST_BACKGROUND - parameters.PIXWIDTH_SIGNAL + self.data = self.data[self.bgd_width:-self.bgd_width, :] + self.pixels = np.arange(self.data.shape[0]) + self.err = np.copy(self.err[self.bgd_width:-self.bgd_width, :]) + + # error matrix + self.W = 1. / (self.err * self.err) + self.W = [np.diag(self.W[:, x]) for x in range(self.Nx)] + self.W_dot_data = [self.W[x] @ self.data[:, x] for x in range(self.Nx)] + + def simulate(self, *shape_params): + """ + Compute a ChromaticPSF model given PSF shape parameters and minimizing + amplitude parameters given a spectrogram data array. + + Parameters + ---------- + shape_params: array_like + PSF shape polynomial parameter array. + + Examples + -------- + + # Set the parameters + >>> parameters.PIXDIST_BACKGROUND = 40 + >>> parameters.PIXWIDTH_BACKGROUND = 10 + >>> parameters.PIXWIDTH_SIGNAL = 30 + + # Build a mock spectrogram with random Poisson noise: + >>> s0 = ChromaticPSF1D(Nx=100, Ny=100, deg=4, saturation=1000) + >>> params = s0.generate_test_poly_params() + >>> s0.poly_params = params + >>> saturation = params[-1] + >>> data = s0.evaluate(params) + >>> bgd = 10*np.ones_like(data) + >>> data += bgd + >>> data = np.random.poisson(data) + >>> data_errors = np.sqrt(data+1) + + # Extract the background + >>> bgd_model_func = extract_background_photutils(data, data_errors, ws=[30,50]) + + # Estimate the first guess values + >>> s = ChromaticPSF1D(Nx=100, Ny=100, deg=4, saturation=saturation) + >>> s.fit_transverse_PSF1D_profile(data, data_errors, w=20, ws=[30,50], + ... pixel_step=1, bgd_model_func=bgd_model_func, saturation=saturation, live_fit=False) + + # Simulate the data: + >>> w = ChromaticPSF1DFitWorkspace(s, data, data_errors, bgd_model_func=bgd_model_func, verbose=True) + >>> y, mod, mod_err = w.simulate(s.poly_params[s.Nx:-1]) + >>> assert mod is not None + >>> w.plot_fit() + """ + # linear regression for the amplitude parameters + poly_params = np.copy(self.chromatic_psf.poly_params) + poly_params[self.Nx:-1] = np.copy(shape_params) + profile_params = self.chromatic_psf.from_poly_params_to_profile_params(poly_params, force_positive=True) + profile_params[:self.Nx, 0] = 1 + profile_params[:self.Nx, 1] -= self.bgd_width + J = np.array([self.chromatic_psf.PSF.evaluate(self.pixels, *profile_params[x, :]) for x in range(self.Nx)]) + J_dot_W_dot_J = np.array([J[x].T @ self.W[x] @ J[x] for x in range(self.Nx)]) + amplitude_params = [ + J[x].T @ self.W_dot_data[x] / (J_dot_W_dot_J[x]) if J_dot_W_dot_J[x] > 0 else 0.1 * self.bgd_std + for x in range(self.Nx)] + poly_params[:self.Nx] = amplitude_params + # in_bounds, penalty, name = self.chromatic_psf.check_bounds(poly_params, noise_level=self.bgd_std) + self.model = self.chromatic_psf.evaluate(poly_params)[self.bgd_width:-self.bgd_width, :] + self.model_err = np.zeros_like(self.model) + self.poly_params = np.copy(poly_params) + return self.pixels, self.model, self.model_err + + def plot_fit(self): + gs_kw = dict(width_ratios=[3, 0.15], height_ratios=[1, 1, 1, 1]) + fig, ax = plt.subplots(nrows=4, ncols=2, figsize=(7, 7), constrained_layout=True, gridspec_kw=gs_kw) + norm = np.max(self.data) + plot_image_simple(ax[0, 0], data=self.model / norm, aspect='auto', cax=ax[0, 1], vmin=0, vmax=1, + units='1/max(data)') + ax[0, 0].set_title("Model", fontsize=10, loc='center', color='white', y=0.8) + plot_image_simple(ax[1, 0], data=self.data / norm, title='Data', aspect='auto', + cax=ax[1, 1], vmin=0, vmax=1, units='1/max(data)') + ax[1, 0].set_title('Data', fontsize=10, loc='center', color='white', y=0.8) + residuals = (self.data - self.model) + # residuals_err = self.spectrum.spectrogram_err / self.model + norm = self.err + residuals /= norm + std = float(np.std(residuals)) + plot_image_simple(ax[2, 0], data=residuals, vmin=-5 * std, vmax=5 * std, title='(Data-Model)/Err', + aspect='auto', cax=ax[2, 1], units='', cmap="bwr") + ax[2, 0].set_title('(Data-Model)/Err', fontsize=10, loc='center', color='black', y=0.8) + ax[2, 0].text(0.05, 0.05, f'mean={np.mean(residuals):.3f}\nstd={np.std(residuals):.3f}', + horizontalalignment='left', verticalalignment='bottom', + color='black', transform=ax[2, 0].transAxes) + ax[0, 0].set_xticks(ax[2, 0].get_xticks()[1:-1]) + ax[0, 1].get_yaxis().set_label_coords(3.5, 0.5) + ax[1, 1].get_yaxis().set_label_coords(3.5, 0.5) + ax[2, 1].get_yaxis().set_label_coords(3.5, 0.5) + ax[3, 1].remove() + ax[3, 0].plot(np.arange(self.Nx), self.data.sum(axis=0), label='Data') + ax[3, 0].plot(np.arange(self.Nx), self.model.sum(axis=0), label='Model') + ax[3, 0].set_ylabel('Transverse sum') + ax[3, 0].set_xlabel(r'X [pixels]') + ax[3, 0].legend(fontsize=7) + ax[3, 0].grid(True) + if self.live_fit: # pragma: no cover + plt.draw() + plt.pause(1e-8) + plt.close() + else: + if parameters.DISPLAY and self.verbose: + plt.show() + if parameters.SAVE: # pragma: no cover + figname = self.filename.replace(self.filename.split('.')[-1], "_bestfit.pdf") + self.my_logger.info(f"\n\tSave figure {figname}.") + fig.savefig(figname, dpi=100, bbox_inches='tight') + class ChromaticPSF2D(ChromaticPSF): @@ -1857,13 +2032,15 @@ def fit_PSF2D_minuit(x, y, data, guess=None, bounds=None, data_errors=None): model = PSF2D() my_logger = set_logger(__name__) + if bounds is not None: + bounds = np.array(bounds) + if bounds.shape[0] == 2 and bounds.shape[1] > 2: + bounds = bounds.T + + guess = np.array(guess) error = 0.001 * np.abs(guess) * np.ones_like(guess) z = np.where(np.isclose(error, 0.0, 1e-6)) error[z] = 0.001 - bounds = np.array(bounds) - if bounds.shape[0] == 2 and bounds.shape[1] > 2: - bounds = bounds.T - guess = np.array(guess) def chisq_PSF2D(params): return PSF2D_chisq(params, model, x, y, data, data_errors) diff --git a/spectractor/extractor/spectrum.py b/spectractor/extractor/spectrum.py index eaf79bc4b..f9769e64b 100644 --- a/spectractor/extractor/spectrum.py +++ b/spectractor/extractor/spectrum.py @@ -910,8 +910,9 @@ def shift_minimizer(params): fwhm_func=fwhm_func, ax=None) chisq += (shift * shift) / (parameters.PIXSHIFT_PRIOR / 2) ** 2 if parameters.DEBUG and parameters.DISPLAY: - spectrum.lambdas = lambdas_test - spectrum.plot_spectrum(live_fit=True, label=f'Order {spectrum.order:d} spectrum' + if parameters.LIVE_FIT: + spectrum.lambdas = lambdas_test + spectrum.plot_spectrum(live_fit=True, label=f'Order {spectrum.order:d} spectrum' f'\nD={D:.2f}mm, shift={shift:.2f}pix') return chisq @@ -1051,7 +1052,7 @@ def extract_spectrum_from_image(image, spectrum, w=10, ws=(20, 30), right_edge=p my_logger.info(f'\n\tStart PSF1D transverse fit...') s = ChromaticPSF1D(Nx=Nx, Ny=Ny, deg=parameters.PSF_POLY_ORDER, saturation=image.saturation) s.fit_transverse_PSF1D_profile(data, err, w, ws, pixel_step=1, sigma=5, bgd_model_func=bgd_model_func, - saturation=image.saturation, live_fit=parameters.DEBUG) + saturation=image.saturation, live_fit=False) # Fill spectrum object spectrum.pixels = np.arange(pixel_start, pixel_end, 1).astype(int) diff --git a/spectractor/fit/fit_spectrogram.py b/spectractor/fit/fit_spectrogram.py new file mode 100644 index 000000000..573c04b0a --- /dev/null +++ b/spectractor/fit/fit_spectrogram.py @@ -0,0 +1,404 @@ +import time +import matplotlib.pyplot as plt +import numpy as np + +from spectractor import parameters +from spectractor.config import set_logger +from spectractor.tools import plot_image_simple, from_lambda_to_colormap +from spectractor.simulation.simulator import SimulatorInit, SpectrogramModel +from spectractor.simulation.atmosphere import Atmosphere, AtmosphereGrid +from spectractor.fit.fitter import FitWorkspace, run_minimisation, run_gradient_descent, save_gradient_descent + +plot_counter = 0 + + +class SpectrogramFitWorkspace(FitWorkspace): + + def __init__(self, file_name, atmgrid_file_name="", nwalkers=18, nsteps=1000, burnin=100, nbins=10, + verbose=0, plot=False, live_fit=False, truth=None): + FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, + live_fit, truth=truth) + self.spectrum, self.telescope, self.disperser, self.target = SimulatorInit(file_name) + self.airmass = self.spectrum.header['AIRMASS'] + self.pressure = self.spectrum.header['OUTPRESS'] + self.temperature = self.spectrum.header['OUTTEMP'] + self.my_logger = set_logger(self.__class__.__name__) + if atmgrid_file_name == "": + self.atmosphere = Atmosphere(self.airmass, self.pressure, self.temperature) + else: + self.use_grid = True + self.atmosphere = AtmosphereGrid(file_name, atmgrid_file_name) + if parameters.VERBOSE: + self.my_logger.info(f'\n\tUse atmospheric grid models from file {atmgrid_file_name}. ') + self.crop_spectrogram() + self.lambdas = self.spectrum.lambdas + self.data = self.spectrum.spectrogram + self.err = self.spectrum.spectrogram_err + self.A1 = 1.0 + self.A2 = 0.01 + self.ozone = 400. + self.pwv = 3 + self.aerosols = 0.05 + self.D = self.spectrum.header['D2CCD'] + self.psf_poly_params = self.spectrum.chromatic_psf.from_table_to_poly_params() + length = len(self.spectrum.chromatic_psf.table) + self.psf_poly_params = self.psf_poly_params[length:-1] # remove saturation (fixed parameter) + self.psf_poly_params_labels = np.copy(self.spectrum.chromatic_psf.poly_params_labels[length:-1]) + self.psf_poly_params_names = np.copy(self.spectrum.chromatic_psf.poly_params_names[length:-1]) + self.psf_poly_params_bounds = self.spectrum.chromatic_psf.set_bounds(data=None) + self.shift_x = self.spectrum.header['PIXSHIFT'] + self.shift_y = 0. + self.angle = self.spectrum.rotation_angle + self.saturation = self.spectrum.spectrogram_saturation + self.p = np.array([self.A1, self.A2, self.ozone, self.pwv, self.aerosols, + self.D, self.shift_x, self.shift_y, self.angle]) + self.psf_params_start_index = self.p.size + self.p = np.concatenate([self.p, self.psf_poly_params]) + self.ndim = self.p.size + self.input_labels = ["A1", "A2", "ozone [db]", "PWV [mm]", "VAOD", r"D_CCD [mm]", + r"shift_x [pix]", r"shift_y [pix]", r"angle [deg]"] + list(self.psf_poly_params_labels) + self.axis_names = ["$A_1$", "$A_2$", "ozone [db]", "PWV [mm]", "VAOD", r"$D_{CCD}$ [mm]", + r"$\Delta_{\mathrm{x}}$ [pix]", r"$\Delta_{\mathrm{y}}$ [pix]", + r"$\theta$ [deg]"] + list(self.psf_poly_params_names) + self.bounds = np.concatenate([np.array([(0, 2), (0, 0.5), (0, 800), (1, 10), (0, 1), + (50, 60), (-3, 3), (-3, 3), (-90, 90)]), + self.psf_poly_params_bounds[:-1]]) # remove saturation + if atmgrid_file_name != "": + self.bounds[2] = (min(self.atmosphere.OZ_Points), max(self.atmosphere.OZ_Points)) + self.bounds[3] = (min(self.atmosphere.PWV_Points), max(self.atmosphere.PWV_Points)) + self.bounds[4] = (min(self.atmosphere.AER_Points), max(self.atmosphere.AER_Points)) + self.nwalkers = max(2 * self.ndim, nwalkers) + self.simulation = SpectrogramModel(self.spectrum, self.atmosphere, self.telescope, self.disperser, + with_background=True, fast_sim=False) + self.get_spectrogram_truth() + + def crop_spectrogram(self): + bgd_width = parameters.PIXWIDTH_BACKGROUND + parameters.PIXDIST_BACKGROUND - parameters.PIXWIDTH_SIGNAL + # spectrogram must have odd size in y for the fourier simulation + yeven = 0 + if (self.spectrum.spectrogram_Ny - 2 * bgd_width) % 2 == 0: + yeven = 1 + self.spectrum.spectrogram_ymax = self.spectrum.spectrogram_ymax - bgd_width + yeven + self.spectrum.spectrogram_ymin += bgd_width + self.spectrum.spectrogram_bgd = self.spectrum.spectrogram_bgd[bgd_width:-bgd_width + yeven, :] + self.spectrum.spectrogram = self.spectrum.spectrogram[bgd_width:-bgd_width + yeven, :] + self.spectrum.spectrogram_err = self.spectrum.spectrogram_err[bgd_width:-bgd_width + yeven, :] + self.spectrum.spectrogram_y0 -= bgd_width + self.spectrum.spectrogram_Ny, self.spectrum.spectrogram_Nx = self.spectrum.spectrogram.shape + self.my_logger.debug(f'\n\tSize of the spectrogram region after cropping: ' + f'({self.spectrum.spectrogram_Nx},{self.spectrum.spectrogram_Ny})') + + def get_spectrogram_truth(self): + if 'A1' in list(self.spectrum.header.keys()): + A1_truth = self.spectrum.header['A1'] + A2_truth = self.spectrum.header['A2'] + ozone_truth = self.spectrum.header['OZONE'] + pwv_truth = self.spectrum.header['PWV'] + aerosols_truth = self.spectrum.header['VAOD'] + D_truth = self.spectrum.header['D2CCD'] + # shift_x_truth = self.spectrum.header['X0SHIFT'] + # shift_y_truth = self.spectrum.header['Y0SHIFT'] + angle_truth = self.spectrum.header['ROTANGLE'] + self.truth = (A1_truth, A2_truth, ozone_truth, pwv_truth, aerosols_truth, + D_truth, angle_truth) + else: + self.truth = None + self.truth = None + + def plot_spectrogram_comparison_simple(self, ax, title='', extent=None, dispersion=False): + lambdas = self.spectrum.lambdas + sub = np.where((lambdas > parameters.LAMBDA_MIN) & (lambdas < parameters.LAMBDA_MAX))[0] + sub = np.where(sub < self.spectrum.spectrogram_Nx)[0] + if extent is not None: + sub = np.where((lambdas > extent[0]) & (lambdas < extent[1]))[0] + if len(sub) > 0: + norm = np.max(self.spectrum.spectrogram[:, sub]) + plot_image_simple(ax[0, 0], data=self.model[:, sub] / norm, aspect='auto', cax=ax[0, 1], vmin=0, vmax=1, + units='1/max(data)') + if dispersion: + x = self.spectrum.chromatic_psf.table['Dx'][sub[2:-3]] + self.spectrum.spectrogram_x0 - sub[0] + y = np.ones_like(x) + ax[0, 0].scatter(x, y, cmap=from_lambda_to_colormap(self.lambdas[sub[2:-3]]), edgecolors='None', + c=self.lambdas[sub[2:-3]], + label='', marker='o', s=10) + # p0 = ax.plot(lambdas, self.model(lambdas), label='model') + # # ax.plot(self.lambdas, self.model_noconv, label='before conv') + if title != '': + ax[0, 0].set_title(title, fontsize=10, loc='center', color='white', y=0.8) + plot_image_simple(ax[1, 0], data=self.spectrum.spectrogram[:, sub] / norm, title='Data', aspect='auto', + cax=ax[1, 1], vmin=0, vmax=1, units='1/max(data)') + ax[1, 0].set_title('Data', fontsize=10, loc='center', color='white', y=0.8) + residuals = (self.spectrum.spectrogram - self.model) + # residuals_err = self.spectrum.spectrogram_err / self.model + norm = self.spectrum.spectrogram_err + residuals /= norm + std = float(np.std(residuals[:, sub])) + plot_image_simple(ax[2, 0], data=residuals[:, sub], vmin=-5 * std, vmax=5 * std, title='(Data-Model)/Err', + aspect='auto', cax=ax[2, 1], units='', cmap="bwr") + ax[2, 0].set_title('(Data-Model)/Err', fontsize=10, loc='center', color='black', y=0.8) + ax[2, 0].text(0.05, 0.05, f'mean={np.mean(residuals[:, sub]):.3f}\nstd={np.std(residuals[:, sub]):.3f}', + horizontalalignment='left', verticalalignment='bottom', + color='black', transform=ax[2, 0].transAxes) + ax[0, 0].set_xticks(ax[2, 0].get_xticks()[1:-1]) + ax[0, 1].get_yaxis().set_label_coords(3.5, 0.5) + ax[1, 1].get_yaxis().set_label_coords(3.5, 0.5) + ax[2, 1].get_yaxis().set_label_coords(3.5, 0.5) + # ax[0, 0].get_yaxis().set_label_coords(-0.15, 0.6) + # ax[2, 0].get_yaxis().set_label_coords(-0.15, 0.5) + # remove the underlying axes + # for ax in ax[3, 1]: + ax[3, 1].remove() + ax[3, 0].plot(self.lambdas[sub], self.spectrum.spectrogram.sum(axis=0)[sub], label='Data') + ax[3, 0].plot(self.lambdas[sub], self.model.sum(axis=0)[sub], label='Model') + ax[3, 0].set_ylabel('Cross spectrum') + ax[3, 0].set_xlabel(r'$\lambda$ [nm]') + ax[3, 0].legend(fontsize=7) + ax[3, 0].grid(True) + + def simulate(self, A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, angle, *psf_poly_params): + global plot_counter + self.simulation.fix_psf_cube = False + if np.all(np.isclose(psf_poly_params, self.p[self.psf_params_start_index:], rtol=1e-6)): + self.simulation.fix_psf_cube = True + lambdas, model, model_err = \ + self.simulation.simulate(A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, angle, psf_poly_params) + self.p = np.array([A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, angle] + list(psf_poly_params)) + self.lambdas = lambdas + self.model = model + self.model_err = model_err + if self.live_fit and (plot_counter % 30) == 0: + self.plot_fit() + plot_counter += 1 + return lambdas, model, model_err + + def jacobian(self, params, epsilon, fixed_params=None): + start = time.time() + lambdas, model, model_err = self.simulate(*params) + model = model.flatten() + J = np.zeros((params.size, model.size)) + for ip, p in enumerate(params): + if fixed_params[ip]: + continue + if ip < self.psf_params_start_index: + self.simulation.fix_psf_cube = True + else: + self.simulation.fix_psf_cube = False + tmp_p = np.copy(params) + if tmp_p[ip] + epsilon[ip] < self.bounds[ip][0] or tmp_p[ip] + epsilon[ip] > self.bounds[ip][1]: + epsilon[ip] = - epsilon[ip] + tmp_p[ip] += epsilon[ip] + tmp_lambdas, tmp_model, tmp_model_err = self.simulate(*tmp_p) + J[ip] = (tmp_model.flatten() - model) / epsilon[ip] + # print(ip, self.input_labels[ip], p, tmp_p[ip] + epsilon[ip], J[ip]) + # if False: + # plt.imshow(J, origin="lower", aspect="auto") + # plt.show() + print(f"\tjacobian time computation = {time.time() - start:.1f}s") + return J + + def plot_fit(self): + """ + Examples + -------- + >>> file_name = 'outputs/reduc_20170530_130_spectrum.fits' + >>> atmgrid_filename = file_name.replace('sim', 'reduc').replace('spectrum', 'atmsim') + >>> fit_workspace = SpectrogramFitWorkspace(file_name, atmgrid_filename=atmgrid_filename, + ... nwalkers=28, nsteps=20000, burnin=10000, nbins=10, verbose=1, plot=True, live_fit=False) + >>> A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, shift_t, angle, *psf = fit_workspace.p + >>> lambdas, model, model_err = fit_workspace.simulation.simulate(A1, A2, ozone, pwv, aerosols, + ... D, shift_x, shift_y, shift_t, angle, psf) + >>> fit_workspace.lambdas = lambdas + >>> fit_workspace.model = model + >>> fit_workspace.model_err = model_err + >>> fit_workspace.plot_fit() + """ + gs_kw = dict(width_ratios=[3, 0.15, 1, 0.15, 1, 0.15], height_ratios=[1, 1, 1, 1]) + fig, ax = plt.subplots(nrows=4, ncols=6, figsize=(12, 8), constrained_layout=True, gridspec_kw=gs_kw) + + A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, shift_t, *psf = self.p + plt.suptitle(f'A1={A1:.3f}, A2={A2:.3f}, PWV={pwv:.3f}, OZ={ozone:.3g}, VAOD={aerosols:.3f}, ' + f'D={D:.2f}mm, shift_x={shift_x:.2f}pix, shift_y={shift_y:.2f}pix') + # main plot + self.plot_spectrogram_comparison_simple(ax[:, 0:2], title='Spectrogram model', dispersion=True) + # zoom O2 + self.plot_spectrogram_comparison_simple(ax[:, 2:4], extent=[730, 800], title='Zoom $O_2$', dispersion=True) + # zoom H2O + self.plot_spectrogram_comparison_simple(ax[:, 4:6], extent=[870, 1000], title='Zoom $H_2 O$', dispersion=True) + # fig.tight_layout() + if self.live_fit: + plt.draw() + plt.pause(1e-8) + plt.close() + else: + if parameters.DISPLAY and self.verbose: + plt.show() + if parameters.SAVE: + figname = self.filename.replace(self.filename.split('.')[-1], "_bestfit.pdf") + self.my_logger.info(f"\n\tSave figure {figname}.") + fig.savefig(figname, dpi=100, bbox_inches='tight') + + +def lnprob_spectrogram(p): + global fit_workspace + lp = fit_workspace.lnprior(p) + if not np.isfinite(lp): + return -1e20 + return lp + fit_workspace.lnlike_spectrogram(p) + + +def plot_psf_poly_params(psf_poly_params): + from spectractor.extractor.psf import PSF1D + psf = PSF1D() + truth_psf_poly_params = [0.11298966008548948, -0.396825836448203, 0.2060387678061209, 2.0649268678546955, + -1.3753936625491252, 0.9242067418613167, 1.6950153822467129, -0.6942452135351901, + 0.3644178350759512, -0.0028059253333737044, -0.003111527339787137, -0.00347648933169673, + 528.3594585697788, 628.4966480821147, 12.438043546369354, 499.99999999999835] + + x = np.linspace(-1, 1, 100) + for i in range(5): + plt.plot(x, np.polynomial.legendre.legval(x, truth_psf_poly_params[3 * i:3 * i + 3]), + label="truth " + psf.param_names[1 + i]) + plt.plot(x, np.polynomial.legendre.legval(x, psf_poly_params[3 * i:3 * i + 3]), + label="fit " + psf.param_names[1 + i]) + + plt.legend() + plt.show() + + +def run_spectrogram_minimisation(fit_workspace, method="newton"): + my_logger = set_logger(__name__) + bounds = fit_workspace.bounds + + nll = lambda params: -fit_workspace.lnlike(params) + + # sim_134 + # guess = fit_workspace.p + # truth sim_134 + # guess = np.array([1., 0.05, 300, 5, 0.03, 55.45, 0.0, 0.0, -1.54, 0.11298966008548948, -0.396825836448203, 0.2060387678061209, 2.0649268678546955, -1.3753936625491252, 0.9242067418613167, 1.6950153822467129, -0.6942452135351901, 0.3644178350759512, -0.0028059253333737044, -0.003111527339787137, -0.00347648933169673, 528.3594585697788, 628.4966480821147, 12.438043546369354]) + guess = np.array( + [1., 0.05, 300, 5, 0.03, 55.45, -0.275, 0.0, -1.54, -1.47570237e-01, -5.00195918e-01, 4.74296776e-01, + 2.85776501e+00, -1.86436219e+00, 1.83899390e+00, 1.89342052e+00, + -9.43239034e-01, 1.06985560e+00, 0.00000000e+00, 0.00000000e+00, + 0.00000000e+00, 1.44368271e+00, -9.95896258e-01, 1.59015965e+00]) + # 5.00000000e+02 + guess = fit_workspace.p + if method != "newton": + run_minimisation(fit_workspace, method=method) + else: + fit_workspace.simulation.fast_sim = True + costs = np.array([fit_workspace.chisq(guess)]) + if parameters.DISPLAY: + fit_workspace.plot_fit() + params_table = np.array([guess]) + start = time.time() + epsilon = 1e-4 * guess + epsilon[epsilon == 0] = 1e-3 + epsilon[0] = 1e-3 # A1 + epsilon[1] = 1e-4 # A2 + epsilon[2] = 1 # ozone + epsilon[3] = 0.01 # pwv + epsilon[4] = 0.001 # aerosols + epsilon[5] = 0.001 # DCCD + epsilon[6] = 0.0005 # shift_x + my_logger.info(f"\n\tStart guess: {guess}") + + # cancel the Gaussian part of the PSF + # TODO: solve this Gaussian PSF part issue + guess[-6:] = 0 + + # fit trace + fix = [True] * guess.size + fix[0] = False # A1 + fix[1] = False # A2 + fix[6] = True # x0 + fix[7] = True # y0 + fix[8] = True # angle + fit_workspace.simulation.fast_sim = True + fix[fit_workspace.psf_params_start_index:fit_workspace.psf_params_start_index + 3] = [False] * 3 + fit_workspace.simulation.fix_psf_cube = False + params_table, costs = run_gradient_descent(fit_workspace, guess, epsilon, params_table, costs, + fix=fix, xtol=1e-2, ftol=1e-2, niter=20) + + # fit PSF + guess = np.array(fit_workspace.p) + fix = [True] * guess.size + fix[0] = False # A1 + fix[1] = False # A2 + fit_workspace.simulation.fast_sim = True + fix[fit_workspace.psf_params_start_index:fit_workspace.psf_params_start_index + 9] = [False] * 9 + # fix[fit_workspace.psf_params_start_index:] = [False] * (guess.size - fit_workspace.psf_params_start_index) + fit_workspace.simulation.fix_psf_cube = False + params_table, costs = run_gradient_descent(fit_workspace, guess, epsilon, params_table, costs, + fix=fix, xtol=1e-2, ftol=1e-2, niter=20) + + # fit dispersion + guess = np.array(fit_workspace.p) + fix = [True] * guess.size + fix[0] = False + fix[1] = False + fix[5] = False # DCCD + fix[6] = False # x0 + fit_workspace.simulation.fix_psf_cube = True + fit_workspace.simulation.fast_sim = True + params_table, costs = run_gradient_descent(fit_workspace, guess, epsilon, params_table, costs, + fix=fix, xtol=1e-3, ftol=1e-2, niter=10) + + # fit all except Gaussian part of the PSF + # TODO: solve this Gaussian PSF part issue + guess = np.array(fit_workspace.p) + fit_workspace.simulation.fast_sim = False + fix = [False] * guess.size + fix[6] = False # x0 + fix[7] = True # y0 + fix[8] = True # angle + fix[-6:] = [True] * 6 # gaussian part + parameters.SAVE = True + fit_workspace.simulation.fix_psf_cube = False + params_table, costs = run_gradient_descent(fit_workspace, guess, epsilon, params_table, costs, + fix=fix, xtol=1e-5, ftol=1e-5, niter=40) + if fit_workspace.filename != "": + fit_workspace.save_parameters_summary(header=fit_workspace.spectrum.date_obs) + save_gradient_descent(fit_workspace, costs, params_table) + print(f"Newton: total computation time: {time.time() - start}s") + + +if __name__ == "__main__": + from argparse import ArgumentParser + from spectractor.config import load_config + + parser = ArgumentParser() + parser.add_argument(dest="input", metavar='path', default=["tests/data/reduc_20170530_134_spectrum.fits"], + help="Input fits file name. It can be a list separated by spaces, or it can use * as wildcard.", + nargs='*') + parser.add_argument("-d", "--debug", dest="debug", action="store_true", + help="Enter debug mode (more verbose and plots).", default=False) + parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", + help="Enter verbose (print more stuff).", default=False) + parser.add_argument("-o", "--output_directory", dest="output_directory", default="outputs/", + help="Write results in given output directory (default: ./outputs/).") + parser.add_argument("-l", "--logbook", dest="logbook", default="ctiofulllogbook_jun2017_v5.csv", + help="CSV logbook file. (default: ctiofulllogbook_jun2017_v5.csv).") + parser.add_argument("-c", "--config", dest="config", default="config/ctio.ini", + help="INI config file. (default: config.ctio.ini).") + args = parser.parse_args() + + parameters.VERBOSE = args.verbose + if args.debug: + parameters.DEBUG = True + parameters.VERBOSE = True + + file_names = args.input + + load_config(args.config) + + # filename = 'outputs/reduc_20170530_130_spectrum.fits' + # filename = 'outputs/sim_20170530_134_spectrum.fits' + # 062 + filename = 'CTIODataJune2017_reduced_RG715_v2_prod6/data_30may17/sim_20170530_067_spectrum.fits' + atmgrid_filename = filename.replace('sim', 'reduc').replace('spectrum', 'atmsim') + + w = SpectrogramFitWorkspace(filename, atmgrid_file_name=atmgrid_filename, nsteps=1000, + burnin=2, nbins=10, verbose=1, plot=True, live_fit=False) + run_spectrogram_minimisation(w, method="newton") + # run_emcee(w, ln=lnprob_spectrogram) + # w.analyze_chains() diff --git a/spectractor/fit/fit_spectrum.py b/spectractor/fit/fit_spectrum.py new file mode 100644 index 000000000..a0b34a16d --- /dev/null +++ b/spectractor/fit/fit_spectrum.py @@ -0,0 +1,138 @@ +import matplotlib.pyplot as plt +from mpl_toolkits.axes_grid1 import make_axes_locatable +import numpy as np + +from spectractor import parameters +from spectractor.config import set_logger +from spectractor.simulation.simulator import SimulatorInit, SpectrumSimulation +from spectractor.simulation.atmosphere import Atmosphere, AtmosphereGrid +from spectractor.fit.fitter import FitWorkspace + + +class SpectrumFitWorkspace(FitWorkspace): + + def __init__(self, file_name, atmgrid_file_name="", nwalkers=18, nsteps=1000, burnin=100, nbins=10, + verbose=0, plot=False, live_fit=False, truth=None): + FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, live_fit, truth=truth) + self.my_logger = set_logger(self.__class__.__name__) + self.spectrum, self.telescope, self.disperser, self.target = SimulatorInit(file_name) + self.airmass = self.spectrum.header['AIRMASS'] + self.pressure = self.spectrum.header['OUTPRESS'] + self.temperature = self.spectrum.header['OUTTEMP'] + if atmgrid_file_name == "": + self.atmosphere = Atmosphere(self.airmass, self.pressure, self.temperature) + else: + self.use_grid = True + self.atmosphere = AtmosphereGrid(file_name, atmgrid_file_name) + if parameters.VERBOSE: + self.my_logger.info(f'\n\tUse atmospheric grid models from file {atmgrid_file_name}. ') + self.lambdas = self.spectrum.lambdas + self.data = self.spectrum.data + self.err = self.spectrum.err + self.A1 = 1.0 + self.A2 = 0.05 + self.ozone = 300. + self.pwv = 3 + self.aerosols = 0.03 + self.reso = 1.5 + self.D = self.spectrum.header['D2CCD'] + self.shift = self.spectrum.header['PIXSHIFT'] + self.p = np.array([self.A1, self.A2, self.ozone, self.pwv, self.aerosols, self.reso, self.D, self.shift]) + self.ndim = len(self.p) + self.input_labels = ["A1", "A2", "ozone", "PWV", "VAOD", "reso [pix]", r"D_CCD [mm]", + r"alpha_pix [pix]"] + self.axis_names = ["$A_1$", "$A_2$", "ozone", "PWV", "VAOD", "reso [pix]", r"$D_{CCD}$ [mm]", + r"$\alpha_{\mathrm{pix}}$ [pix]"] + self.bounds = [(0, 2), (0, 0.5), (0, 800), (0, 10), (0, 1), (1, 10), (50, 60), (-20, 20)] + if atmgrid_filename != "": + self.bounds[2] = (min(self.atmosphere.OZ_Points), max(self.atmosphere.OZ_Points)) + self.bounds[3] = (min(self.atmosphere.PWV_Points), max(self.atmosphere.PWV_Points)) + self.bounds[4] = (min(self.atmosphere.AER_Points), max(self.atmosphere.AER_Points)) + self.nwalkers = max(2 * self.ndim, nwalkers) + self.simulation = SpectrumSimulation(self.spectrum, self.atmosphere, self.telescope, self.disperser) + self.get_truth() + + def get_truth(self): + if 'A1' in list(self.spectrum.header.keys()): + A1_truth = self.spectrum.header['A1'] + A2_truth = self.spectrum.header['A2'] + ozone_truth = self.spectrum.header['OZONE'] + pwv_truth = self.spectrum.header['PWV'] + aerosols_truth = self.spectrum.header['VAOD'] + reso_truth = self.spectrum.header['RESO'] + D_truth = self.spectrum.header['D2CCD'] + shift_truth = self.spectrum.header['X0SHIFT'] + self.truth = (A1_truth, A2_truth, ozone_truth, pwv_truth, aerosols_truth, + reso_truth, D_truth, shift_truth) + else: + self.truth = None + + def plot_spectrum_comparison_simple(self, ax, title='', extent=None, size=0.4): + lambdas = self.spectrum.lambdas + sub = np.where((lambdas > parameters.LAMBDA_MIN) & (lambdas < parameters.LAMBDA_MAX)) + if extent is not None: + sub = np.where((lambdas > extent[0]) & (lambdas < extent[1])) + self.spectrum.plot_spectrum_simple(ax, lambdas=lambdas) + p0 = ax.plot(lambdas, self.model(lambdas), label='model') + ax.fill_between(lambdas, self.model - self.model_err, + self.model(lambdas) + self.model_err, alpha=0.3, color=p0[0].get_color()) + # ax.plot(self.lambdas, self.model_noconv, label='before conv') + if title != '': + ax.set_title(title, fontsize=10) + ax.legend() + divider = make_axes_locatable(ax) + ax2 = divider.append_axes("bottom", size=size, pad=0) + ax.figure.add_axes(ax2) + residuals = (self.spectrum.data - self.model(lambdas)) / self.model(lambdas) + residuals_err = self.spectrum.err / self.model(lambdas) + ax2.errorbar(lambdas, residuals, yerr=residuals_err, fmt='ro', markersize=2) + ax2.axhline(0, color=p0[0].get_color()) + ax2.grid(True) + residuals_model = self.model_err / self.model(lambdas) + ax2.fill_between(lambdas, -residuals_model, residuals_model, alpha=0.3, color=p0[0].get_color()) + std = np.std(residuals[sub]) + ax2.set_ylim([-2. * std, 2. * std]) + ax2.set_xlabel(ax.get_xlabel()) + ax2.set_ylabel('(data-fit)/fit') + ax2.set_xlim((lambdas[sub][0], lambdas[sub][-1])) + ax.set_xlim((lambdas[sub][0], lambdas[sub][-1])) + ax.set_ylim((0.9 * np.min(self.spectrum.data[sub]), 1.1 * np.max(self.spectrum.data[sub]))) + ax.set_xticks(ax2.get_xticks()[1:-1]) + ax.get_yaxis().set_label_coords(-0.15, 0.6) + ax2.get_yaxis().set_label_coords(-0.15, 0.5) + + def simulate(self, A1, A2, ozone, pwv, aerosols, reso, D, shift): + lambdas, model, model_err = self.simulation.simulate(A1, A2, ozone, pwv, aerosols, reso, D, shift) + # if self.live_fit: + # self.plot_fit() + self.model = model + self.model_err = model_err + return model, model_err + + def plot_fit(self): + fig = plt.figure(figsize=(12, 6)) + ax1 = plt.subplot(222) + ax2 = plt.subplot(224) + ax3 = plt.subplot(121) + A1, A2, ozone, pwv, aerosols, reso, D, shift = self.p + self.title = f'A1={A1:.3f}, A2={A2:.3f}, PWV={pwv:.3f}, OZ={ozone:.3g}, VAOD={aerosols:.3f},\n ' \ + f'reso={reso:.2f}pix, D={D:.2f}mm, shift={shift:.2f}pix ' + # main plot + self.plot_spectrum_comparison_simple(ax3, title=self.title, size=0.8) + # zoom O2 + self.plot_spectrum_comparison_simple(ax2, extent=[730, 800], title='Zoom $O_2$', size=0.8) + # zoom H2O + self.plot_spectrum_comparison_simple(ax1, extent=[870, 1000], title='Zoom $H_2 O$', size=0.8) + fig.tight_layout() + if self.live_fit: + plt.draw() + plt.pause(1e-8) + plt.close() + else: + if parameters.DISPLAY and parameters.VERBOSE: + plt.show() + if parameters.SAVE: + figname = self.filename.replace('.fits', '_bestfit.pdf') + self.my_logger.info(f'Save figure {figname}.') + fig.savefig(figname, dpi=100, bbox_inches='tight') + diff --git a/spectractor/fit/fitter.py b/spectractor/fit/fitter.py index df43953f1..76290f55e 100644 --- a/spectractor/fit/fitter.py +++ b/spectractor/fit/fitter.py @@ -1,6 +1,5 @@ from iminuit import Minuit from scipy import optimize -from mpl_toolkits.axes_grid1 import make_axes_locatable from schwimmbad import MPIPool import emcee import time @@ -11,17 +10,13 @@ from spectractor import parameters from spectractor.config import set_logger -from spectractor.tools import formatting_numbers, plot_image_simple, from_lambda_to_colormap -from spectractor.simulation.simulator import SimulatorInit, SpectrumSimulation, SpectrogramModel -from spectractor.simulation.atmosphere import Atmosphere, AtmosphereGrid +from spectractor.tools import formatting_numbers from spectractor.fit.statistics import Likelihood -plot_counter = 0 - class FitWorkspace: - def __init__(self, file_name, nwalkers=18, nsteps=1000, burnin=100, nbins=10, + def __init__(self, file_name="", nwalkers=18, nsteps=1000, burnin=100, nbins=10, verbose=0, plot=False, live_fit=False, truth=None): self.my_logger = set_logger(self.__class__.__name__) self.filename = file_name @@ -60,10 +55,13 @@ def __init__(self, file_name, nwalkers=18, nsteps=1000, burnin=100, nbins=10, self.global_std = None self.title = "" self.use_grid = False - if "." in self.filename: - self.emcee_filename = self.filename.split('.')[0] + "_emcee.h5" + if self.filename != "": + if "." in self.filename: + self.emcee_filename = self.filename.split('.')[0] + "_emcee.h5" + else: + self.my_logger.warning("\n\tFile name must have an extension.") else: - self.my_logger.warning("\n\tFile name must have an extension.") + self.emcee_filename = "emcee.h5" def set_start(self): self.start = np.array( @@ -115,7 +113,7 @@ def simulate(self, *p): Examples -------- - >>> w = FitWorkspace("filename.txt") + >>> w = FitWorkspace() >>> p = np.zeros(3) >>> x, model, model_err = w.simulate(*p) >>> assert x is not None @@ -153,7 +151,7 @@ def plot_fit(self): title += f"{label} = {self.p[i]:.3g}" if self.cov.size > 0: title += rf" $\pm$ {np.sqrt(self.cov[i, i]):.3g}" - if i < len(self.input_labels)-1: + if i < len(self.input_labels) - 1: title += ", " plt.title(title) plt.legend() @@ -178,7 +176,9 @@ def chain2likelihood(self, pdfonly=False, walker_index=-1): for j in rangedim: if i != j: likelihood.contours[i][j].fill_histogram(chains[:, i], chains[:, j], weights=None) - output_file = self.filename.replace('.fits', '_bestfit.txt') + output_file = "" + if self.filename != "": + output_file = self.filename.replace('.fits', '_bestfit.txt') likelihood.stats(output=output_file) else: for i in rangedim: @@ -373,7 +373,7 @@ def plot_correlation_matrix(self, ipar=None): cbar = fig.colorbar(im) cbar.ax.tick_params(labelsize=9) fig.tight_layout() - if parameters.SAVE: + if parameters.SAVE and self.filename != "": figname = self.filename.replace(self.filename.split('.')[-1], "_correlation.pdf") self.my_logger.info(f"Save figure {figname}.") fig.savefig(figname, dpi=100, bbox_inches='tight') @@ -409,7 +409,7 @@ def lnprior(self, p): return -np.inf def jacobian(self, params, epsilon, fixed_params=None): - lambdas, model, model_err = self.simulate(*params) + x, model, model_err = self.simulate(*params) model = model.flatten() J = np.zeros((params.size, model.size)) for ip, p in enumerate(params): @@ -419,7 +419,7 @@ def jacobian(self, params, epsilon, fixed_params=None): if tmp_p[ip] + epsilon[ip] < self.bounds[ip][0] or tmp_p[ip] + epsilon[ip] > self.bounds[ip][1]: epsilon[ip] = - epsilon[ip] tmp_p[ip] += epsilon[ip] - tmp_lambdas, tmp_model, tmp_model_err = self.simulate(*tmp_p) + tmp_x, tmp_model, tmp_model_err = self.simulate(*tmp_p) J[ip] = (tmp_model.flatten() - model) / epsilon[ip] return J @@ -433,368 +433,6 @@ def hessian(J, W): return inv_JT_W_J -class SpectrumFitWorkspace(FitWorkspace): - - def __init__(self, file_name, atmgrid_file_name="", nwalkers=18, nsteps=1000, burnin=100, nbins=10, - verbose=0, plot=False, live_fit=False, truth=None): - FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, live_fit, truth=truth) - self.my_logger = set_logger(self.__class__.__name__) - self.spectrum, self.telescope, self.disperser, self.target = SimulatorInit(file_name) - self.airmass = self.spectrum.header['AIRMASS'] - self.pressure = self.spectrum.header['OUTPRESS'] - self.temperature = self.spectrum.header['OUTTEMP'] - if atmgrid_file_name == "": - self.atmosphere = Atmosphere(self.airmass, self.pressure, self.temperature) - else: - self.use_grid = True - self.atmosphere = AtmosphereGrid(file_name, atmgrid_file_name) - if parameters.VERBOSE: - self.my_logger.info(f'\n\tUse atmospheric grid models from file {atmgrid_file_name}. ') - self.lambdas = self.spectrum.lambdas - self.data = self.spectrum.data - self.err = self.spectrum.err - self.A1 = 1.0 - self.A2 = 0.05 - self.ozone = 300. - self.pwv = 3 - self.aerosols = 0.03 - self.reso = 1.5 - self.D = self.spectrum.header['D2CCD'] - self.shift = self.spectrum.header['PIXSHIFT'] - self.p = np.array([self.A1, self.A2, self.ozone, self.pwv, self.aerosols, self.reso, self.D, self.shift]) - self.ndim = len(self.p) - self.input_labels = ["A1", "A2", "ozone", "PWV", "VAOD", "reso [pix]", r"D_CCD [mm]", - r"alpha_pix [pix]"] - self.axis_names = ["$A_1$", "$A_2$", "ozone", "PWV", "VAOD", "reso [pix]", r"$D_{CCD}$ [mm]", - r"$\alpha_{\mathrm{pix}}$ [pix]"] - self.bounds = [(0, 2), (0, 0.5), (0, 800), (0, 10), (0, 1), (1, 10), (50, 60), (-20, 20)] - if atmgrid_filename != "": - self.bounds[2] = (min(self.atmosphere.OZ_Points), max(self.atmosphere.OZ_Points)) - self.bounds[3] = (min(self.atmosphere.PWV_Points), max(self.atmosphere.PWV_Points)) - self.bounds[4] = (min(self.atmosphere.AER_Points), max(self.atmosphere.AER_Points)) - self.nwalkers = max(2 * self.ndim, nwalkers) - self.simulation = SpectrumSimulation(self.spectrum, self.atmosphere, self.telescope, self.disperser) - self.get_truth() - - def get_truth(self): - if 'A1' in list(self.spectrum.header.keys()): - A1_truth = self.spectrum.header['A1'] - A2_truth = self.spectrum.header['A2'] - ozone_truth = self.spectrum.header['OZONE'] - pwv_truth = self.spectrum.header['PWV'] - aerosols_truth = self.spectrum.header['VAOD'] - reso_truth = self.spectrum.header['RESO'] - D_truth = self.spectrum.header['D2CCD'] - shift_truth = self.spectrum.header['X0SHIFT'] - self.truth = (A1_truth, A2_truth, ozone_truth, pwv_truth, aerosols_truth, - reso_truth, D_truth, shift_truth) - else: - self.truth = None - - def plot_spectrum_comparison_simple(self, ax, title='', extent=None, size=0.4): - lambdas = self.spectrum.lambdas - sub = np.where((lambdas > parameters.LAMBDA_MIN) & (lambdas < parameters.LAMBDA_MAX)) - if extent is not None: - sub = np.where((lambdas > extent[0]) & (lambdas < extent[1])) - self.spectrum.plot_spectrum_simple(ax, lambdas=lambdas) - p0 = ax.plot(lambdas, self.model(lambdas), label='model') - ax.fill_between(lambdas, self.model - self.model_err, - self.model(lambdas) + self.model_err, alpha=0.3, color=p0[0].get_color()) - # ax.plot(self.lambdas, self.model_noconv, label='before conv') - if title != '': - ax.set_title(title, fontsize=10) - ax.legend() - divider = make_axes_locatable(ax) - ax2 = divider.append_axes("bottom", size=size, pad=0) - ax.figure.add_axes(ax2) - residuals = (self.spectrum.data - self.model(lambdas)) / self.model(lambdas) - residuals_err = self.spectrum.err / self.model(lambdas) - ax2.errorbar(lambdas, residuals, yerr=residuals_err, fmt='ro', markersize=2) - ax2.axhline(0, color=p0[0].get_color()) - ax2.grid(True) - residuals_model = self.model_err / self.model(lambdas) - ax2.fill_between(lambdas, -residuals_model, residuals_model, alpha=0.3, color=p0[0].get_color()) - std = np.std(residuals[sub]) - ax2.set_ylim([-2. * std, 2. * std]) - ax2.set_xlabel(ax.get_xlabel()) - ax2.set_ylabel('(data-fit)/fit') - ax2.set_xlim((lambdas[sub][0], lambdas[sub][-1])) - ax.set_xlim((lambdas[sub][0], lambdas[sub][-1])) - ax.set_ylim((0.9 * np.min(self.spectrum.data[sub]), 1.1 * np.max(self.spectrum.data[sub]))) - ax.set_xticks(ax2.get_xticks()[1:-1]) - ax.get_yaxis().set_label_coords(-0.15, 0.6) - ax2.get_yaxis().set_label_coords(-0.15, 0.5) - - def simulate(self, A1, A2, ozone, pwv, aerosols, reso, D, shift): - lambdas, model, model_err = self.simulation.simulate(A1, A2, ozone, pwv, aerosols, reso, D, shift) - # if self.live_fit: - # self.plot_fit() - self.model = model - self.model_err = model_err - return model, model_err - - def plot_fit(self): - fig = plt.figure(figsize=(12, 6)) - ax1 = plt.subplot(222) - ax2 = plt.subplot(224) - ax3 = plt.subplot(121) - A1, A2, ozone, pwv, aerosols, reso, D, shift = self.p - self.title = f'A1={A1:.3f}, A2={A2:.3f}, PWV={pwv:.3f}, OZ={ozone:.3g}, VAOD={aerosols:.3f},\n ' \ - f'reso={reso:.2f}pix, D={D:.2f}mm, shift={shift:.2f}pix ' - # main plot - self.plot_spectrum_comparison_simple(ax3, title=self.title, size=0.8) - # zoom O2 - self.plot_spectrum_comparison_simple(ax2, extent=[730, 800], title='Zoom $O_2$', size=0.8) - # zoom H2O - self.plot_spectrum_comparison_simple(ax1, extent=[870, 1000], title='Zoom $H_2 O$', size=0.8) - fig.tight_layout() - if self.live_fit: - plt.draw() - plt.pause(1e-8) - plt.close() - else: - if parameters.DISPLAY and parameters.VERBOSE: - plt.show() - if parameters.SAVE: - figname = self.filename.replace('.fits', '_bestfit.pdf') - self.my_logger.info(f'Save figure {figname}.') - fig.savefig(figname, dpi=100, bbox_inches='tight') - - -class SpectrogramFitWorkspace(FitWorkspace): - - def __init__(self, file_name, atmgrid_file_name="", nwalkers=18, nsteps=1000, burnin=100, nbins=10, - verbose=0, plot=False, live_fit=False, truth=None): - FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, - live_fit, truth=truth) - self.spectrum, self.telescope, self.disperser, self.target = SimulatorInit(file_name) - self.airmass = self.spectrum.header['AIRMASS'] - self.pressure = self.spectrum.header['OUTPRESS'] - self.temperature = self.spectrum.header['OUTTEMP'] - self.my_logger = set_logger(self.__class__.__name__) - if atmgrid_file_name == "": - self.atmosphere = Atmosphere(self.airmass, self.pressure, self.temperature) - else: - self.use_grid = True - self.atmosphere = AtmosphereGrid(file_name, atmgrid_file_name) - if parameters.VERBOSE: - self.my_logger.info(f'\n\tUse atmospheric grid models from file {atmgrid_file_name}. ') - self.crop_spectrogram() - self.lambdas = self.spectrum.lambdas - self.data = self.spectrum.spectrogram - self.err = self.spectrum.spectrogram_err - self.A1 = 1.0 - self.A2 = 0.01 - self.ozone = 400. - self.pwv = 3 - self.aerosols = 0.05 - self.D = self.spectrum.header['D2CCD'] - self.psf_poly_params = self.spectrum.chromatic_psf.from_table_to_poly_params() - length = len(self.spectrum.chromatic_psf.table) - self.psf_poly_params = self.psf_poly_params[length:-1] # remove saturation (fixed parameter) - self.psf_poly_params_labels = np.copy(self.spectrum.chromatic_psf.poly_params_labels[length:-1]) - self.psf_poly_params_names = np.copy(self.spectrum.chromatic_psf.poly_params_names[length:-1]) - self.psf_poly_params_bounds = self.spectrum.chromatic_psf.set_bounds(data=None) - self.shift_x = self.spectrum.header['PIXSHIFT'] - self.shift_y = 0. - self.angle = self.spectrum.rotation_angle - self.saturation = self.spectrum.spectrogram_saturation - self.p = np.array([self.A1, self.A2, self.ozone, self.pwv, self.aerosols, - self.D, self.shift_x, self.shift_y, self.angle]) - self.psf_params_start_index = self.p.size - self.p = np.concatenate([self.p, self.psf_poly_params]) - self.ndim = self.p.size - self.input_labels = ["A1", "A2", "ozone [db]", "PWV [mm]", "VAOD", r"D_CCD [mm]", - r"shift_x [pix]", r"shift_y [pix]", r"angle [deg]"] + list(self.psf_poly_params_labels) - self.axis_names = ["$A_1$", "$A_2$", "ozone [db]", "PWV [mm]", "VAOD", r"$D_{CCD}$ [mm]", - r"$\Delta_{\mathrm{x}}$ [pix]", r"$\Delta_{\mathrm{y}}$ [pix]", - r"$\theta$ [deg]"] + list(self.psf_poly_params_names) - self.bounds = np.concatenate([np.array([(0, 2), (0, 0.5), (0, 800), (1, 10), (0, 1), - (50, 60), (-3, 3), (-3, 3), (-90, 90)]), - self.psf_poly_params_bounds[:-1]]) # remove saturation - if atmgrid_file_name != "": - self.bounds[2] = (min(self.atmosphere.OZ_Points), max(self.atmosphere.OZ_Points)) - self.bounds[3] = (min(self.atmosphere.PWV_Points), max(self.atmosphere.PWV_Points)) - self.bounds[4] = (min(self.atmosphere.AER_Points), max(self.atmosphere.AER_Points)) - self.nwalkers = max(2 * self.ndim, nwalkers) - self.simulation = SpectrogramModel(self.spectrum, self.atmosphere, self.telescope, self.disperser, - with_background=True, fast_sim=False) - self.get_spectrogram_truth() - - def crop_spectrogram(self): - bgd_width = parameters.PIXWIDTH_BACKGROUND + parameters.PIXDIST_BACKGROUND - parameters.PIXWIDTH_SIGNAL - # spectrogram must have odd size in y for the fourier simulation - yeven = 0 - if (self.spectrum.spectrogram_Ny - 2 * bgd_width) % 2 == 0: - yeven = 1 - self.spectrum.spectrogram_ymax = self.spectrum.spectrogram_ymax - bgd_width + yeven - self.spectrum.spectrogram_ymin += bgd_width - self.spectrum.spectrogram_bgd = self.spectrum.spectrogram_bgd[bgd_width:-bgd_width + yeven, :] - self.spectrum.spectrogram = self.spectrum.spectrogram[bgd_width:-bgd_width + yeven, :] - self.spectrum.spectrogram_err = self.spectrum.spectrogram_err[bgd_width:-bgd_width + yeven, :] - self.spectrum.spectrogram_y0 -= bgd_width - self.spectrum.spectrogram_Ny, self.spectrum.spectrogram_Nx = self.spectrum.spectrogram.shape - self.my_logger.debug(f'\n\tSize of the spectrogram region after cropping: ' - f'({self.spectrum.spectrogram_Nx},{self.spectrum.spectrogram_Ny})') - - def get_spectrogram_truth(self): - if 'A1' in list(self.spectrum.header.keys()): - A1_truth = self.spectrum.header['A1'] - A2_truth = self.spectrum.header['A2'] - ozone_truth = self.spectrum.header['OZONE'] - pwv_truth = self.spectrum.header['PWV'] - aerosols_truth = self.spectrum.header['VAOD'] - D_truth = self.spectrum.header['D2CCD'] - # shift_x_truth = self.spectrum.header['X0SHIFT'] - # shift_y_truth = self.spectrum.header['Y0SHIFT'] - angle_truth = self.spectrum.header['ROTANGLE'] - self.truth = (A1_truth, A2_truth, ozone_truth, pwv_truth, aerosols_truth, - D_truth, angle_truth) - else: - self.truth = None - self.truth = None - - def plot_spectrogram_comparison_simple(self, ax, title='', extent=None, dispersion=False): - lambdas = self.spectrum.lambdas - sub = np.where((lambdas > parameters.LAMBDA_MIN) & (lambdas < parameters.LAMBDA_MAX))[0] - sub = np.where(sub < self.spectrum.spectrogram_Nx)[0] - if extent is not None: - sub = np.where((lambdas > extent[0]) & (lambdas < extent[1]))[0] - if len(sub) > 0: - norm = np.max(self.spectrum.spectrogram[:, sub]) - plot_image_simple(ax[0, 0], data=self.model[:, sub] / norm, aspect='auto', cax=ax[0, 1], vmin=0, vmax=1, - units='1/max(data)') - if dispersion: - x = self.spectrum.chromatic_psf.table['Dx'][sub[2:-3]] + self.spectrum.spectrogram_x0 - sub[0] - y = np.ones_like(x) - ax[0, 0].scatter(x, y, cmap=from_lambda_to_colormap(self.lambdas[sub[2:-3]]), edgecolors='None', - c=self.lambdas[sub[2:-3]], - label='', marker='o', s=10) - # p0 = ax.plot(lambdas, self.model(lambdas), label='model') - # # ax.plot(self.lambdas, self.model_noconv, label='before conv') - if title != '': - ax[0, 0].set_title(title, fontsize=10, loc='center', color='white', y=0.8) - plot_image_simple(ax[1, 0], data=self.spectrum.spectrogram[:, sub] / norm, title='Data', aspect='auto', - cax=ax[1, 1], vmin=0, vmax=1, units='1/max(data)') - ax[1, 0].set_title('Data', fontsize=10, loc='center', color='white', y=0.8) - residuals = (self.spectrum.spectrogram - self.model) - # residuals_err = self.spectrum.spectrogram_err / self.model - norm = self.spectrum.spectrogram_err - residuals /= norm - std = float(np.std(residuals[:, sub])) - plot_image_simple(ax[2, 0], data=residuals[:, sub], vmin=-5 * std, vmax=5 * std, title='(Data-Model)/Err', - aspect='auto', cax=ax[2, 1], units='', cmap="bwr") - ax[2, 0].set_title('(Data-Model)/Err', fontsize=10, loc='center', color='black', y=0.8) - ax[2, 0].text(0.05, 0.05, f'mean={np.mean(residuals[:, sub]):.3f}\nstd={np.std(residuals[:, sub]):.3f}', - horizontalalignment='left', verticalalignment='bottom', - color='black', transform=ax[2, 0].transAxes) - ax[0, 0].set_xticks(ax[2, 0].get_xticks()[1:-1]) - ax[0, 1].get_yaxis().set_label_coords(3.5, 0.5) - ax[1, 1].get_yaxis().set_label_coords(3.5, 0.5) - ax[2, 1].get_yaxis().set_label_coords(3.5, 0.5) - # ax[0, 0].get_yaxis().set_label_coords(-0.15, 0.6) - # ax[2, 0].get_yaxis().set_label_coords(-0.15, 0.5) - # remove the underlying axes - # for ax in ax[3, 1]: - ax[3, 1].remove() - ax[3, 0].plot(self.lambdas[sub], self.spectrum.spectrogram.sum(axis=0)[sub], label='Data') - ax[3, 0].plot(self.lambdas[sub], self.model.sum(axis=0)[sub], label='Model') - ax[3, 0].set_ylabel('Cross spectrum') - ax[3, 0].set_xlabel(r'$\lambda$ [nm]') - ax[3, 0].legend(fontsize=7) - ax[3, 0].grid(True) - - def simulate(self, A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, angle, *psf_poly_params): - global plot_counter - self.simulation.fix_psf_cube = False - if np.all(np.isclose(psf_poly_params, self.p[self.psf_params_start_index:], rtol=1e-6)): - self.simulation.fix_psf_cube = True - lambdas, model, model_err = \ - self.simulation.simulate(A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, angle, psf_poly_params) - self.p = np.array([A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, angle] + list(psf_poly_params)) - self.lambdas = lambdas - self.model = model - self.model_err = model_err - if self.live_fit and (plot_counter % 30) == 0: - self.plot_fit() - plot_counter += 1 - return lambdas, model, model_err - - def jacobian(self, params, epsilon, fixed_params=None): - start = time.time() - lambdas, model, model_err = self.simulate(*params) - model = model.flatten() - J = np.zeros((params.size, model.size)) - for ip, p in enumerate(params): - if fixed_params[ip]: - continue - if ip < self.psf_params_start_index: - self.simulation.fix_psf_cube = True - else: - self.simulation.fix_psf_cube = False - tmp_p = np.copy(params) - if tmp_p[ip] + epsilon[ip] < self.bounds[ip][0] or tmp_p[ip] + epsilon[ip] > self.bounds[ip][1]: - epsilon[ip] = - epsilon[ip] - tmp_p[ip] += epsilon[ip] - tmp_lambdas, tmp_model, tmp_model_err = self.simulate(*tmp_p) - J[ip] = (tmp_model.flatten() - model) / epsilon[ip] - # print(ip, self.input_labels[ip], p, tmp_p[ip] + epsilon[ip], J[ip]) - # if False: - # plt.imshow(J, origin="lower", aspect="auto") - # plt.show() - print(f"\tjacobian time computation = {time.time() - start:.1f}s") - return J - - def plot_fit(self): - """ - Examples - -------- - >>> file_name = 'outputs/reduc_20170530_130_spectrum.fits' - >>> atmgrid_filename = file_name.replace('sim', 'reduc').replace('spectrum', 'atmsim') - >>> fit_workspace = SpectrogramFitWorkspace(file_name, atmgrid_filename=atmgrid_filename, - ... nwalkers=28, nsteps=20000, burnin=10000, nbins=10, verbose=1, plot=True, live_fit=False) - >>> A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, shift_t, angle, *psf = fit_workspace.p - >>> lambdas, model, model_err = fit_workspace.simulation.simulate(A1, A2, ozone, pwv, aerosols, - ... D, shift_x, shift_y, shift_t, angle, psf) - >>> fit_workspace.lambdas = lambdas - >>> fit_workspace.model = model - >>> fit_workspace.model_err = model_err - >>> fit_workspace.plot_fit() - """ - gs_kw = dict(width_ratios=[3, 0.15, 1, 0.15, 1, 0.15], height_ratios=[1, 1, 1, 1]) - fig, ax = plt.subplots(nrows=4, ncols=6, figsize=(12, 8), constrained_layout=True, gridspec_kw=gs_kw) - - A1, A2, ozone, pwv, aerosols, D, shift_x, shift_y, shift_t, *psf = self.p - plt.suptitle(f'A1={A1:.3f}, A2={A2:.3f}, PWV={pwv:.3f}, OZ={ozone:.3g}, VAOD={aerosols:.3f}, ' - f'D={D:.2f}mm, shift_x={shift_x:.2f}pix, shift_y={shift_y:.2f}pix') - # main plot - self.plot_spectrogram_comparison_simple(ax[:, 0:2], title='Spectrogram model', dispersion=True) - # zoom O2 - self.plot_spectrogram_comparison_simple(ax[:, 2:4], extent=[730, 800], title='Zoom $O_2$', dispersion=True) - # zoom H2O - self.plot_spectrogram_comparison_simple(ax[:, 4:6], extent=[870, 1000], title='Zoom $H_2 O$', dispersion=True) - # fig.tight_layout() - if self.live_fit: - plt.draw() - plt.pause(1e-8) - plt.close() - else: - if parameters.DISPLAY and self.verbose: - plt.show() - if parameters.SAVE: - figname = self.filename.replace(self.filename.split('.')[-1], "_bestfit.pdf") - self.my_logger.info(f"\n\tSave figure {figname}.") - fig.savefig(figname, dpi=100, bbox_inches='tight') - - -def lnprob_spectrogram(p): - global fit_workspace - lp = fit_workspace.lnprior(p) - if not np.isfinite(lp): - return -1e20 - return lp + fit_workspace.lnlike_spectrogram(p) - - def lnprob(p): global fit_workspace lp = fit_workspace.lnprior(p) @@ -804,6 +442,7 @@ def lnprob(p): def gradient_descent(fit_workspace, params, epsilon, niter=10, fixed_params=None, xtol=1e-3, ftol=1e-3): + my_logger = set_logger(__name__) tmp_params = np.copy(params) W = 1 / fit_workspace.err.flatten() ** 2 ipar = np.arange(params.size) @@ -813,14 +452,12 @@ def gradient_descent(fit_workspace, params, epsilon, niter=10, fixed_params=None params_table = [] inv_JT_W_J = np.zeros((len(ipar), len(ipar))) for i in range(niter): - print(f"start iteration={i}", end=' ') # , tmp_params) start = time.time() tmp_lambdas, tmp_model, tmp_model_err = fit_workspace.simulate(*tmp_params) # if fit_workspace.live_fit: # fit_workspace.plot_fit() residuals = (tmp_model - fit_workspace.data).flatten() cost = np.sum((residuals ** 2) * W) - print(f"cost={cost:.3f} chisq_red={cost / tmp_model.size:.3f}") J = fit_workspace.jacobian(tmp_params, epsilon, fixed_params=fixed_params) # remove parameters with unexpected null Jacobian vectors for ip in range(J.shape[0]): @@ -829,7 +466,7 @@ def gradient_descent(fit_workspace, params, epsilon, niter=10, fixed_params=None if np.all(J[ip] == np.zeros(J.shape[1])): ipar = np.delete(ipar, list(ipar).index(ip)) tmp_params[ip] = 0 - print(f"Step {i}: {fit_workspace.input_labels[ip]} has a null Jacobian; parameter is fixed at 0 " + my_logger.warning(f"\n\tStep {i}: {fit_workspace.input_labels[ip]} has a null Jacobian; parameter is fixed at 0 " f"in the following instead of its current value ({tmp_params[ip]}).") # remove fixed parameters J = J[ipar].T @@ -852,10 +489,7 @@ def line_search(alpha): # tol parameter acts on alpha (not func) alpha_min, fval, iter, funcalls = optimize.brent(line_search, full_output=True, tol=1e-2) - print(f"\talpha_min={alpha_min:.3g} iter={iter} funcalls={funcalls}") tmp_params[ipar] += alpha_min * dparams - print(f"shift: {alpha_min * dparams}") - print(f"new params: {tmp_params[ipar]}") # check bounds for ip, p in enumerate(tmp_params): if p < fit_workspace.bounds[ip][0]: @@ -869,40 +503,30 @@ def line_search(alpha): # prepare outputs costs.append(fval) params_table.append(np.copy(tmp_params)) - print(f"end iteration={i} in {time.time() - start:.2f}s cost={fval:.3f}") - # if np.sum(np.abs(alpha_min * dparams)) / np.sum(np.abs(tmp_params[ipar])) < xtol \ - # or (len(costs) > 1 and np.abs(costs[-2] - fval) / np.max([np.abs(fval), np.abs(costs[-2]), 1]) < ftol): - # break + my_logger.info(f"\n\tIteration={i}: initial cost={cost:.3f} initial chisq_red={cost / tmp_model.size:.3f}" + f"\n\t\t Line search: alpha_min={alpha_min:.3g} iter={iter} funcalls={funcalls}" + f"\n\tParameter shifts: {alpha_min * dparams}" + f"\n\tNew parameters: {tmp_params[ipar]}" + f"\n\tFinal cost={fval:.3f} final chisq_red={fval / tmp_model.size:.3f} computed in {time.time() - start:.2f}s") + if np.sum(np.abs(alpha_min * dparams)) / np.sum(np.abs(tmp_params[ipar])) < xtol : + my_logger.info(f"\n\tGradient descent terminated in {i} iterations because the sum of parameter shift " + f"relative to the sum of the parameters is below xtol={xtol}.") + break if len(costs) > 1 and np.abs(costs[-2] - fval) / np.max([np.abs(fval), np.abs(costs[-2]), 1]) < ftol: + my_logger.info(f"\n\tGradient descent terminated in {i} iterations because the " + f"relative change of cost is below ftol={ftol}.") break plt.close() return tmp_params, inv_JT_W_J, np.array(costs), np.array(params_table) -def plot_psf_poly_params(psf_poly_params): - from spectractor.extractor.psf import PSF1D - psf = PSF1D() - truth_psf_poly_params = [0.11298966008548948, -0.396825836448203, 0.2060387678061209, 2.0649268678546955, - -1.3753936625491252, 0.9242067418613167, 1.6950153822467129, -0.6942452135351901, - 0.3644178350759512, -0.0028059253333737044, -0.003111527339787137, -0.00347648933169673, - 528.3594585697788, 628.4966480821147, 12.438043546369354, 499.99999999999835] - - x = np.linspace(-1, 1, 100) - for i in range(5): - plt.plot(x, np.polynomial.legendre.legval(x, truth_psf_poly_params[3 * i:3 * i + 3]), - label="truth " + psf.param_names[1 + i]) - plt.plot(x, np.polynomial.legendre.legval(x, psf_poly_params[3 * i:3 * i + 3]), - label="fit " + psf.param_names[1 + i]) - - plt.legend() - plt.show() - - def print_parameter_summary(params, cov, labels): + my_logger = set_logger(__name__) + txt = "" for ip in np.arange(0, cov.shape[0]).astype(int): - txt = "%s: %s +%s -%s" % formatting_numbers(params[ip], np.sqrt(cov[ip, ip]), np.sqrt(cov[ip, ip]), + txt += "%s: %s +%s -%s\n\t" % formatting_numbers(params[ip], np.sqrt(cov[ip, ip]), np.sqrt(cov[ip, ip]), label=labels[ip]) - print(txt) + my_logger.info(f"\n\t{txt}") def plot_gradient_descent(fit_workspace, costs, params_table): @@ -921,7 +545,7 @@ def plot_gradient_descent(fit_workspace, costs, params_table): ax[1].set_xlabel("Iterations") fig.tight_layout() plt.subplots_adjust(wspace=0, hspace=0) - if parameters.SAVE: + if parameters.SAVE and fit_workspace.filename != "": figname = fit_workspace.filename.replace(fit_workspace.filename.split('.')[-1], "_fitting.pdf") fit_workspace.my_logger.info(f"\n\tSave figure {figname}.") fig.savefig(figname, dpi=100, bbox_inches='tight') @@ -945,23 +569,31 @@ def save_gradient_descent(fit_workspace, costs, params_table): fit_workspace.my_logger.info(f"\n\tSave gradient descent log {output_filename}.") -def run_minimisation(fit_workspace, method="newton"): +def run_gradient_descent(fit_workspace, guess, epsilon, params_table, costs, fix, xtol, ftol, niter): + fit_workspace.p, fit_workspace.cov, tmp_costs, tmp_params_table = gradient_descent(fit_workspace, guess, + epsilon, niter=niter, + fixed_params=fix, + xtol=xtol, ftol=ftol) + params_table = np.concatenate([params_table, tmp_params_table]) + costs = np.concatenate([costs, tmp_costs]) + ipar = np.array(np.where(np.array(fix).astype(int) == 0)[0]) + print_parameter_summary(fit_workspace.p[ipar], fit_workspace.cov, + [fit_workspace.input_labels[ip] for ip in ipar]) + if parameters.DEBUG: + # plot_psf_poly_params(fit_workspace.p[fit_workspace.psf_params_start_index:]) + plot_gradient_descent(fit_workspace, costs, params_table) + fit_workspace.plot_correlation_matrix(ipar=ipar) + return params_table, costs + + +def run_minimisation(fit_workspace, method="newton", epsilon=None, fix=None, xtol=1e-4, ftol=1e-4, niter=50): my_logger = set_logger(__name__) bounds = fit_workspace.bounds nll = lambda params: -fit_workspace.lnlike(params) - # sim_134 - # guess = fit_workspace.p - # truth sim_134 - # guess = np.array([1., 0.05, 300, 5, 0.03, 55.45, 0.0, 0.0, -1.54, 0.11298966008548948, -0.396825836448203, 0.2060387678061209, 2.0649268678546955, -1.3753936625491252, 0.9242067418613167, 1.6950153822467129, -0.6942452135351901, 0.3644178350759512, -0.0028059253333737044, -0.003111527339787137, -0.00347648933169673, 528.3594585697788, 628.4966480821147, 12.438043546369354]) - guess = np.array( - [1., 0.05, 300, 5, 0.03, 55.45, -0.275, 0.0, -1.54, -1.47570237e-01, -5.00195918e-01, 4.74296776e-01, - 2.85776501e+00, -1.86436219e+00, 1.83899390e+00, 1.89342052e+00, - -9.43239034e-01, 1.06985560e+00, 0.00000000e+00, 0.00000000e+00, - 0.00000000e+00, 1.44368271e+00, -9.95896258e-01, 1.59015965e+00]) - # 5.00000000e+02 - guess = fit_workspace.p + guess = fit_workspace.p.astype('float64') + if method == "minimize": start = time.time() result = optimize.minimize(nll, fit_workspace.p, method='L-BFGS-B', @@ -970,7 +602,7 @@ def run_minimisation(fit_workspace, method="newton"): 'maxls': 50, 'maxcor': 30}, bounds=bounds) fit_workspace.p = result['x'] - print(f"Minimize: total computation time: {time.time() - start}s") + my_logger.info(f"\n\tMinimize: total computation time: {time.time() - start}s") fit_workspace.simulate(*fit_workspace.p) fit_workspace.live_fit = False fit_workspace.plot_fit() @@ -982,7 +614,7 @@ def run_minimisation(fit_workspace, method="newton"): p = optimize.least_squares(fit_workspace.weighted_residuals, guess, verbose=2, ftol=1e-6, x_scale=x_scale, diff_step=0.001, bounds=bounds.T) fit_workspace.p = p.x # m.np_values() - print(f"Least_squares: total computation time: {time.time() - start}s") + my_logger.info(f"\n\tLeast_squares: total computation time: {time.time() - start}s") fit_workspace.simulate(*fit_workspace.p) fit_workspace.live_fit = False fit_workspace.plot_fit() @@ -1000,102 +632,29 @@ def run_minimisation(fit_workspace, method="newton"): m.tol = 10 m.migrad() fit_workspace.p = m.np_values() - print(f"Minuit: total computation time: {time.time() - start}s") + my_logger.info(f"\n\tMinuit: total computation time: {time.time() - start}s") fit_workspace.simulate(*fit_workspace.p) fit_workspace.live_fit = False fit_workspace.plot_fit() elif method == "newton": - - def bloc_gradient_descent(guess, epsilon, params_table, costs, fix, xtol, ftol, niter): - fit_workspace.p, fit_workspace.cov, tmp_costs, tmp_params_table = gradient_descent(fit_workspace, guess, - epsilon, niter=niter, - fixed_params=fix, - xtol=xtol, ftol=ftol) - params_table = np.concatenate([params_table, tmp_params_table]) - costs = np.concatenate([costs, tmp_costs]) - ipar = np.array(np.where(np.array(fix).astype(int) == 0)[0]) - print_parameter_summary(fit_workspace.p[ipar], fit_workspace.cov, - [fit_workspace.input_labels[ip] for ip in ipar]) - if True: - # plot_psf_poly_params(fit_workspace.p[fit_workspace.psf_params_start_index:]) - plot_gradient_descent(fit_workspace, costs, params_table) - fit_workspace.plot_correlation_matrix(ipar=ipar) - return params_table, costs - - fit_workspace.simulation.fast_sim = True - costs = np.array([fit_workspace.chisq_spectrogram(guess)]) - if parameters.DISPLAY: - fit_workspace.plot_fit() - params_table = np.array([guess]) - start = time.time() - epsilon = 1e-4 * guess - epsilon[epsilon == 0] = 1e-3 - epsilon[0] = 1e-3 # A1 - epsilon[1] = 1e-4 # A2 - epsilon[2] = 1 # ozone - epsilon[3] = 0.01 # pwv - epsilon[4] = 0.001 # aerosols - epsilon[5] = 0.001 # DCCD - epsilon[6] = 0.0005 # shift_x my_logger.info(f"\n\tStart guess: {guess}") + costs = np.array([fit_workspace.chisq(guess)]) - # cancel the Gaussian part of the PSF - # TODO: solve this Gaussian PSF part issue - guess[-6:] = 0 - - # fit trace - fix = [True] * guess.size - fix[0] = False # A1 - fix[1] = False # A2 - fix[6] = True # x0 - fix[7] = True # y0 - fix[8] = True # angle - fit_workspace.simulation.fast_sim = True - fix[fit_workspace.psf_params_start_index:fit_workspace.psf_params_start_index + 3] = [False] * 3 - fit_workspace.simulation.fix_psf_cube = False - params_table, costs = bloc_gradient_descent(guess, epsilon, params_table, costs, - fix=fix, xtol=1e-2, ftol=1e-2, niter=20) - - # fit PSF - guess = np.array(fit_workspace.p) - fix = [True] * guess.size - fix[0] = False # A1 - fix[1] = False # A2 - fit_workspace.simulation.fast_sim = True - fix[fit_workspace.psf_params_start_index:fit_workspace.psf_params_start_index + 9] = [False] * 9 - # fix[fit_workspace.psf_params_start_index:] = [False] * (guess.size - fit_workspace.psf_params_start_index) - fit_workspace.simulation.fix_psf_cube = False - params_table, costs = bloc_gradient_descent(guess, epsilon, params_table, costs, - fix=fix, xtol=1e-2, ftol=1e-2, niter=20) - - # fit dispersion - guess = np.array(fit_workspace.p) - fix = [True] * guess.size - fix[0] = False - fix[1] = False - fix[5] = False # DCCD - fix[6] = False # x0 - fit_workspace.simulation.fix_psf_cube = True - fit_workspace.simulation.fast_sim = True - params_table, costs = bloc_gradient_descent(guess, epsilon, params_table, costs, - fix=fix, xtol=1e-3, ftol=1e-2, niter=10) + params_table = np.array([guess]) + if epsilon is None: + epsilon = 1e-4 * guess + epsilon[epsilon == 0] = 1e-4 + if fix is None: + fix = [False] * guess.size - # fit all except Gaussian part of the PSF - # TODO: solve this Gaussian PSF part issue - guess = np.array(fit_workspace.p) - fit_workspace.simulation.fast_sim = False - fix = [False] * guess.size - fix[6] = False # x0 - fix[7] = True # y0 - fix[8] = True # angle - fix[-6:] = [True] * 6 # gaussian part - parameters.SAVE = True - fit_workspace.simulation.fix_psf_cube = False - params_table, costs = bloc_gradient_descent(guess, epsilon, params_table, costs, - fix=fix, xtol=1e-5, ftol=1e-5, niter=40) - fit_workspace.save_parameters_summary(header=fit_workspace.spectrum.date_obs) - save_gradient_descent(fit_workspace, costs, params_table) - print(f"Newton: total computation time: {time.time() - start}s") + start = time.time() + params_table, costs = run_gradient_descent(fit_workspace, guess, epsilon, params_table, costs, + fix=fix, xtol=xtol, ftol=ftol, niter=niter) + fit_workspace.costs = costs + my_logger.info(f"\n\tNewton: total computation time: {time.time() - start}s") + if fit_workspace.filename != "": + fit_workspace.save_parameters_summary() + save_gradient_descent(fit_workspace, costs, params_table) def run_emcee(fit_workspace, ln=lnprob): @@ -1112,7 +671,7 @@ def run_emcee(fit_workspace, ln=lnprob): sys.exit(0) sampler = emcee.EnsembleSampler(fit_workspace.nwalkers, fit_workspace.ndim, ln, args=(), pool=pool, backend=backend) - print(f"Initial size: {backend.iteration}") + my_logger.info(f"\n\tInitial size: {backend.iteration}") if backend.iteration > 0: p0 = backend.get_last_sample() if nsamples - backend.iteration > 0: @@ -1121,52 +680,10 @@ def run_emcee(fit_workspace, ln=lnprob): except ValueError: sampler = emcee.EnsembleSampler(fit_workspace.nwalkers, fit_workspace.ndim, ln, args=(), threads=multiprocessing.cpu_count(), backend=backend) - print(f"Initial size: {backend.iteration}") + my_logger.info(f"\n\tInitial size: {backend.iteration}") if backend.iteration > 0: p0 = sampler.get_last_sample() for _ in sampler.sample(p0, iterations=max(0, nsamples - backend.iteration), progress=True, store=True): continue fit_workspace.chains = sampler.chain fit_workspace.lnprobs = sampler.lnprobability - - -if __name__ == "__main__": - from argparse import ArgumentParser - from spectractor.config import load_config - - parser = ArgumentParser() - parser.add_argument(dest="input", metavar='path', default=["tests/data/reduc_20170530_134_spectrum.fits"], - help="Input fits file name. It can be a list separated by spaces, or it can use * as wildcard.", - nargs='*') - parser.add_argument("-d", "--debug", dest="debug", action="store_true", - help="Enter debug mode (more verbose and plots).", default=False) - parser.add_argument("-v", "--verbose", dest="verbose", action="store_true", - help="Enter verbose (print more stuff).", default=False) - parser.add_argument("-o", "--output_directory", dest="output_directory", default="outputs/", - help="Write results in given output directory (default: ./outputs/).") - parser.add_argument("-l", "--logbook", dest="logbook", default="ctiofulllogbook_jun2017_v5.csv", - help="CSV logbook file. (default: ctiofulllogbook_jun2017_v5.csv).") - parser.add_argument("-c", "--config", dest="config", default="config/ctio.ini", - help="INI config file. (default: config.ctio.ini).") - args = parser.parse_args() - - parameters.VERBOSE = args.verbose - if args.debug: - parameters.DEBUG = True - parameters.VERBOSE = True - - file_names = args.input - - load_config(args.config) - - # filename = 'outputs/reduc_20170530_130_spectrum.fits' - # filename = 'outputs/sim_20170530_134_spectrum.fits' - # 062 - filename = 'CTIODataJune2017_reduced_RG715_v2_prod6/data_30may17/sim_20170530_067_spectrum.fits' - atmgrid_filename = filename.replace('sim', 'reduc').replace('spectrum', 'atmsim') - - w = SpectrogramFitWorkspace(filename, atmgrid_file_name=atmgrid_filename, nsteps=1000, - burnin=2, nbins=10, verbose=1, plot=True, live_fit=False) - run_minimisation(w, method="newton") - # run_emcee(w, ln=lnprob_spectrogram) - # w.analyze_chains() diff --git a/spectractor/simulation/atmosphere.py b/spectractor/simulation/atmosphere.py index d47e8d90a..c5d7cffab 100644 --- a/spectractor/simulation/atmosphere.py +++ b/spectractor/simulation/atmosphere.py @@ -127,10 +127,10 @@ def plot_transmission(self): >>> a.plot_transmission() """ - plt.figure() + fig = plt.figure() plot_transmission_simple(plt.gca(), parameters.LAMBDAS, self.transmission(parameters.LAMBDAS), title=self.title, label=self.label) - if parameters.DISPLAY: + if parameters.DISPLAY: # pragma: no cover plt.show() @@ -261,6 +261,13 @@ def compute(self): >>> a = AtmosphereGrid(image_filename='tests/data/reduc_20170605_028.fits', ... pwv_grid=[5, 5, 1], ozone_grid=[400, 400, 1], aerosol_grid=[0.0, 0.1, 2]) >>> atmospheric_grid = a.compute() + >>> atmospheric_grid + array([[0.000000e+00, 0.000000e+00, 0.000000e+00, ..., 1.099400e+03, + 1.099600e+03, 1.099800e+03], + [1.000000e+00, 0.000000e+00, 5.000000e+00, ..., 9.520733e-01, + 9.520733e-01, 9.520733e-01], + [2.000000e+00, 1.000000e-01, 5.000000e+00, ..., 9.213718e-01, + 9.213718e-01, 9.213718e-01]]) >>> assert np.all(np.isclose(a.atmgrid[0, a.index_atm_data:], parameters.LAMBDAS)) >>> assert not np.any(np.isclose(a.atmgrid[1, a.index_atm_data:], np.zeros_like(parameters.LAMBDAS), rtol=1e-6)) >>> assert a.atmgrid.shape == (3, a.index_atm_data+len(parameters.LAMBDAS)) @@ -304,7 +311,7 @@ def plot_transmission(self): f'VAOD={self.atmgrid[int(count), self.index_atm_aer]}' plot_transmission_simple(plt.gca(), self.lambdas, self.atmgrid[int(count), self.index_atm_data:], title="Atmospheric grid", label=label) - if parameters.DISPLAY: + if parameters.DISPLAY: # pragma: no cover plt.show() def plot_transmission_image(self): @@ -333,6 +340,14 @@ def save_file(self, filename=""): ---------- filename: str The output file name. + + Examples + -------- + >>> a = AtmosphereGrid(image_filename='tests/data/reduc_20170605_028.fits', + ... pwv_grid=[5, 5, 1], ozone_grid=[400, 400, 1], aerosol_grid=[0.0, 0.1, 2]) + >>> atmospheric_grid = a.compute() + >>> a.save_file(a.image_filename.replace('.fits', '_atmsim.fits')) + >>> assert os.path.isfile('tests/data/reduc_20170605_028_atmsim.fits') """ hdr = fits.Header() @@ -390,6 +405,21 @@ def load_file(self, filename): ---------- filename: str The input file name. + + Examples + -------- + >>> a = AtmosphereGrid(image_filename='tests/data/reduc_20170605_028.fits', + ... pwv_grid=[5, 5, 1], ozone_grid=[400, 400, 1], aerosol_grid=[0.0, 0.1, 2]) + >>> atmospheric_grid = a.compute() + >>> a.save_file(a.image_filename.replace('.fits', '_atmsim.fits')) + >>> assert os.path.isfile('tests/data/reduc_20170605_028_atmsim.fits') + >>> a.load_file(a.image_filename.replace('.fits', '_atmsim.fits')) + >>> a.AER_Points + array([0. , 0.1]) + >>> a.PWV_Points + array([5.]) + >>> a.OZ_Points + array([400.]) """ if filename != "": diff --git a/spectractor/simulation/image_simulation.py b/spectractor/simulation/image_simulation.py index f9f9cb4d6..240852ebc 100644 --- a/spectractor/simulation/image_simulation.py +++ b/spectractor/simulation/image_simulation.py @@ -20,9 +20,55 @@ class StarModel: + """Class to model a star in the image simulation process. + + Attributes + ---------- + x0: float + X position of the star centroid in pixels. + y0: float + Y position of the star centroid in pixels. + amplitude: amplitude + The amplitude of the star in image units. + target: Target + The associated Target instance (default: None). + """ def __init__(self, pixcoords, model, amplitude, target=None): - """ [x0, y0] coords in pixels, sigma width in pixels, A height in image units""" + """Create a StarModel instance. + + The model is based on an Astropy Fittable2DModel. The centroid and amplitude + parameters of the given model are updated by the dedicated arguments. + + Parameters + ---------- + pixcoords: array_like + Tuple of (x,y) coordinates of the desired star centroid in pixels. + model: Fittable2DModel + Astropy fittable 2D model + amplitude: float + The desired amplitude of the star in image units. + target: Target + The associated Target instance (default: None). + + Examples + -------- + >>> from spectractor.extractor.psf import PSF2D + >>> from spectractor.extractor.images import fit_PSF2D_minuit + >>> p = (100, 50, 50, 3, 2, -0.1, 1, 200) + >>> psf = PSF2D(*p) + >>> yy, xx = np.mgrid[:100,:100] + >>> data = psf.evaluate(xx, yy, *p) + >>> model = fit_PSF2D_minuit(xx, yy, data, guess=p) + >>> s = StarModel((20, 20), model, 200, target=None) + >>> s.plot_model() + >>> s.x0 + 20 + >>> s.y0 + 20 + >>> s.model.amplitude + 200 + """ self.my_logger = set_logger(self.__class__.__name__) self.x0 = pixcoords[0] self.y0 = pixcoords[1] @@ -37,6 +83,9 @@ def __init__(self, pixcoords, model, amplitude, target=None): self.sigma = self.model.stddev / 2 def plot_model(self): + """ + Plot the star model. + """ x = np.linspace(self.x0 - 10 * self.fwhm, self.x0 + 10 * self.fwhm, 50) y = np.linspace(self.y0 - 10 * self.fwhm, self.y0 + 10 * self.fwhm, 50) yy, xx = np.meshgrid(x, y) @@ -147,38 +196,81 @@ def plot_model(self): class BackgroundModel: + """Class to model the background of the simulated image. + + The background model size is set with the parameters.CCD_IMSIZE global keyword. + + Attributes + ---------- + level: float + The mean level of the background in image units. + frame: array_like + (x, y, smooth) right and upper limits in pixels of a vignetting frame, + and the smoothing gaussian width (default: None). + """ def __init__(self, level, frame=None): + """Create a BackgroundModel instance. + + The background model size is set with the parameters.CCD_IMSIZE global keyword. + + Parameters + ---------- + level: float + The mean level of the background in image units. + frame: array_like, None + (x, y, smooth) right and upper limits in pixels of a vignetting frame, + and the smoothing gaussian width (default: None). + + Examples + -------- + >>> from spectractor import parameters + >>> parameters.CCD_IMSIZE = 200 + >>> bgd = BackgroundModel(10) + >>> model = bgd.model() + >>> np.all(model==10) + True + >>> model.shape + (200, 200) + >>> bgd = BackgroundModel(10, frame=(160, 180, 3)) + >>> bgd.plot_model() + """ self.my_logger = set_logger(self.__class__.__name__) self.level = level if self.level <= 0: self.my_logger.warning('\n\tBackground level must be strictly positive.') self.frame = frame - def model(self, x, y): - """Background model for the image simulation. - Args: - x: - y: + def model(self): + """Compute the background model for the image simulation in image units. - Returns: + A shadowing vignetting frame is roughly simulated if self.frame is set. + The background model size is set with the parameters.CCD_IMSIZE global keyword. + + Returns + ------- + bkgd: array_like + The array of the background model. """ - bkgd = self.level * np.ones_like(x) + yy, xx = np.mgrid[0:parameters.CCD_IMSIZE:1, 0:parameters.CCD_IMSIZE:1] + bkgd = self.level * np.ones_like(xx) if self.frame is None: return bkgd else: - xlim, ylim = self.frame + xlim, ylim, width = self.frame bkgd[ylim:, :] = self.level / 100 bkgd[:, xlim:] = self.level / 100 - kernel = np.outer(gaussian(parameters.CCD_IMSIZE, 50), gaussian(parameters.CCD_IMSIZE, 50)) + kernel = np.outer(gaussian(parameters.CCD_IMSIZE, width), gaussian(parameters.CCD_IMSIZE, width)) bkgd = fftconvolve(bkgd, kernel, mode='same') bkgd *= self.level / bkgd[parameters.CCD_IMSIZE // 2, parameters.CCD_IMSIZE // 2] return bkgd def plot_model(self): - yy, xx = np.mgrid[0:parameters.CCD_IMSIZE:1, 0:parameters.CCD_IMSIZE:1] - bkgd = self.model(xx, yy) + """Plot the background model. + + """ + bkgd = self.model() fig, ax = plt.subplots(1, 1) im = plt.imshow(bkgd, origin='lower', cmap='jet') ax.grid(color='white', ls='solid') @@ -195,41 +287,6 @@ def plot_model(self): plt.show() -class SpectrumModel: - - def __init__(self, base_image, spectrumsim, sigma, A1=1, A2=0, reso=None, rotation=False): - self.my_logger = set_logger(self.__class__.__name__) - self.base_image = base_image - self.spectrumsim = spectrumsim - self.disperser = base_image.disperser - self.sigma = sigma - self.A1 = A1 - self.A2 = A2 - self.reso = reso - self.rotation = rotation - self.yprofile = models.Gaussian1D(1, 0, sigma) - self.true_lambdas = None - self.true_spectrum = None - - def model(self, x, y): - self.true_lambdas = np.arange(parameters.LAMBDA_MIN, parameters.LAMBDA_MAX) - self.true_spectrum = np.copy(self.spectrumsim.model(self.true_lambdas)) - x0, y0 = self.base_image.target_pixcoords - if self.rotation: - theta = self.disperser.theta(self.base_image.target_pixcoords) * np.pi / 180. - u = np.cos(theta) * (x - x0) + np.sin(theta) * (y - y0) - v = -np.sin(theta) * (x - x0) + np.cos(theta) * (y - y0) - x = u + x0 - y = v + y0 - l = self.disperser.grating_pixel_to_lambda(x - x0, x0=self.base_image.target_pixcoords, order=1) - amp = self.A1 * self.spectrumsim.model(l) * self.yprofile(y - y0) + self.A1 * self.A2 * self.spectrumsim.model( - l / 2) * self.yprofile(y - y0) - amp = amp * parameters.FLAM_TO_ADURATE * l * np.gradient(l, axis=1) - if self.reso is not None: - amp = fftconvolve_gaussian(amp, self.reso) - return amp - - class ImageModel(Image): def __init__(self, filename, target=None): @@ -240,7 +297,7 @@ def __init__(self, filename, target=None): def compute(self, star, background, spectrogram, starfield=None): yy, xx = np.mgrid[0:parameters.CCD_IMSIZE:1, 0:parameters.CCD_IMSIZE:1] - self.data = star.model(xx, yy) + background.model(xx, yy) + self.data = star.model(xx, yy) + background.model() self.data[spectrogram.spectrogram_ymin:spectrogram.spectrogram_ymax, spectrogram.spectrogram_xmin:spectrogram.spectrogram_xmax] += spectrogram.data # - spectrogram.spectrogram_bgd) self.true_lambdas = spectrogram.lambdas @@ -313,7 +370,7 @@ def ImageSim(image_filename, spectrum_filename, outputdir, pwv=5, ozone=300, aer # Background model my_logger.info('\n\tBackground model...') yy, xx = np.mgrid[:parameters.XWINDOW, :parameters.YWINDOW] - bgd_level = np.mean(image.target_bkgd2D(xx, yy)) + bgd_level = float(np.mean(image.target_bkgd2D(xx, yy))) background = BackgroundModel(level=bgd_level, frame=None) # frame=(1600, 1650)) if parameters.DEBUG: background.plot_model() @@ -375,7 +432,7 @@ def ImageSim(image_filename, spectrum_filename, outputdir, pwv=5, ozone=300, aer image.data = np.around(image.data) # Plot - if parameters.VERBOSE and parameters.DISPLAY: + if parameters.VERBOSE and parameters.DISPLAY: # pragma: no cover image.plot_image(scale="log", title="Image simulation", target_pixcoords=target_pixcoords, units=image.units) # Set output path diff --git a/spectractor/tools.py b/spectractor/tools.py index 39a88411e..bdbd4a493 100644 --- a/spectractor/tools.py +++ b/spectractor/tools.py @@ -887,7 +887,7 @@ def fit_moffat1d_outlier_removal(x, y, sigma=3.0, niter=3, guess=None, bounds=No or_fit = fitting.FittingWithOutlierRemoval(fit, sigma_clip, niter=niter, sigma=sigma) # get fitted model and filtered data or_fitted_model, filtered_data = or_fit(gg_init, x, y) - my_logger.info(f'\n\t{or_fitted_model}') + my_logger.debug(f'\n\t{or_fitted_model}') # my_logger.debug(f'\n\t{fit.fit_info}') return or_fitted_model diff --git a/tests/test_extractor.py b/tests/test_extractor.py index 936965d46..5cf66f5c0 100644 --- a/tests/test_extractor.py +++ b/tests/test_extractor.py @@ -4,6 +4,7 @@ from spectractor.extractor.extractor import Spectractor from spectractor.logbook import LogBook import os +import numpy as np def test_logbook(): @@ -30,10 +31,16 @@ def test_extractor(): if target is None or xpos is None or ypos is None: continue spectrum = Spectractor(file_name, './outputs/', [xpos, ypos], target, disperser_label, - config='./config/ctio.ini', - line_detection=True, atmospheric_lines=True) + config='./config/ctio.ini', line_detection=True, atmospheric_lines=True) assert spectrum.data is not None + assert np.sum(spectrum.data) > 1e-10 + assert np.isclose(spectrum.lambdas[0], 296.56935941, atol=0.2) + assert np.isclose(spectrum.lambdas[-1], 1083.9470213, atol=0.2) + assert np.all(np.isclose(spectrum.x0 , [743.6651370068676, 683.0577836601408], atol=0.2)) + assert np.isclose(spectrum.spectrogram_x0, -239.3348629931324, atol=0.2) + assert 2 < np.mean(spectrum.chromatic_psf.table['gamma']) < 3 assert os.path.isfile('./outputs/' + tag.replace('.fits', '_spectrum.fits')) is True + assert os.path.isfile('./outputs/' + tag.replace('.fits', '_spectrogram.fits')) is True if __name__ == "__main__": diff --git a/tests/test_fitter.py b/tests/test_fitter.py index 3c5245bd2..b993bba33 100644 --- a/tests/test_fitter.py +++ b/tests/test_fitter.py @@ -2,16 +2,16 @@ import numpy as np from spectractor.fit.fitter import FitWorkspace, run_minimisation, run_emcee from spectractor.config import set_logger +from spectractor import parameters import os class LineFitWorkspace(FitWorkspace): - def __init__(self, file_name, x, y, yerr, nwalkers=18, nsteps=1000, burnin=100, nbins=10, + def __init__(self, x, y, yerr, file_name="", nwalkers=18, nsteps=1000, burnin=100, nbins=10, verbose=0, plot=False, live_fit=False, truth=None): - FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, - live_fit, truth=truth) + FitWorkspace.__init__(self, file_name, nwalkers, nsteps, burnin, nbins, verbose, plot, live_fit, truth=truth) self.my_logger = set_logger(self.__class__.__name__) self.x = x self.data = y @@ -42,16 +42,22 @@ def test_fitworkspace(): y = a * x + b yerr = sigma * np.ones_like(y) y += np.random.normal(scale=sigma, size=N) + parameters.VERBOSE = True # Do the fits file_name = "test_linefitworkspace.txt" - w = LineFitWorkspace(file_name, x, y, yerr, truth=truth, nwalkers=20, nsteps=3000, burnin=500, nbins=20) + w = LineFitWorkspace(x, y, yerr, file_name, truth=truth, nwalkers=20, nsteps=5000, burnin=1000, nbins=20) run_minimisation(w, method="minimize") assert np.all([np.abs(w.p[i] - truth[i]) / sigma < 1 for i in range(w.ndim)]) + w.p = np.array([1, 1]) run_minimisation(w, method="least_squares") assert np.all([np.abs(w.p[i] - truth[i]) / sigma < 1 for i in range(w.ndim)]) + w.p = np.array([1, 1]) run_minimisation(w, method="minuit") assert np.all([np.abs(w.p[i] - truth[i]) / sigma < 1 for i in range(w.ndim)]) + w.p = np.array([1, 1]) + run_minimisation(w, method="newton") + assert np.all([np.abs(w.p[i] - truth[i]) / sigma < 1 for i in range(w.ndim)]) def lnprob(p): lp = w.lnprior(p) @@ -62,7 +68,7 @@ def lnprob(p): run_emcee(w, ln=lnprob) w.analyze_chains() - assert w.chains.shape == (3000, 20, 2) + assert w.chains.shape == (5000, 20, 2) assert np.all(w.gelmans < 0.03) assert os.path.exists("test_linefitworkspace.txt") assert os.path.exists(file_name.replace(".txt", "_emcee.h5")) diff --git a/tests/test_simulator.py b/tests/test_simulator.py index 3aeef1472..90353ae6f 100644 --- a/tests/test_simulator.py +++ b/tests/test_simulator.py @@ -1,4 +1,5 @@ from numpy.testing import run_module_suite +import numpy as np from spectractor import parameters from spectractor.simulation.simulator import (SpectrumSimulator, SpectrumSimulatorSimGrid, @@ -8,10 +9,43 @@ import os +def test_atmosphere(): + a = Atmosphere(airmass=1.2, pressure=800, temperature=5) + transmission = a.simulate(ozone=400, pwv=5, aerosols=0.05) + assert transmission is not None + assert a.transmission(500) > 0 + assert a.ozone == 400 + assert a.pwv == 5 + assert a.aerosols == 0.05 + a.plot_transmission() + + a = AtmosphereGrid(image_filename='tests/data/reduc_20170605_028.fits', pwv_grid = [5, 5, 1], ozone_grid = [400, 400, 1], aerosol_grid = [0.0, 0.1, 2]) + atmospheric_grid = a.compute() + assert np.sum(atmospheric_grid) > 0 + assert np.all(np.isclose(a.atmgrid[0, a.index_atm_data:], parameters.LAMBDAS)) + assert not np.any(np.isclose(a.atmgrid[1, a.index_atm_data:], np.zeros_like(parameters.LAMBDAS), rtol=1e-6)) + assert a.atmgrid.shape == (3, a.index_atm_data + len(parameters.LAMBDAS)) + a.save_file(a.image_filename.replace('.fits', '_atmsim.fits')) + assert os.path.isfile('tests/data/reduc_20170605_028_atmsim.fits') + a.load_file(a.image_filename.replace('.fits', '_atmsim.fits')) + assert np.all(a.AER_Points == np.array([0., 0.1])) + assert np.all(a.PWV_Points == np.array([5.])) + assert np.all(a.OZ_Points == np.array([400.])) + + a.plot_transmission() + a.plot_transmission_image() + + a = AtmosphereGrid(filename='tests/data/reduc_20170530_134_atmsim.fits') + lambdas = np.arange(200, 1200) + transmission = a.simulate(ozone=400, pwv=5, aerosols=0.05) + assert np.max(transmission(lambdas)) < 1 and np.min(transmission(lambdas)) >= 0 + + def test_simulator(): file_names = ['tests/data/reduc_20170530_134_spectrum.fits'] parameters.VERBOSE = True + parameters.DEBUG = True load_config('config/ctio.ini') for file_name in file_names: @@ -27,6 +61,7 @@ def test_simulator(): aerosol_grid=[0, 0.1, 2]) atmgrid = AtmosphereGrid(file_name, file_name.replace('spectrum', 'atmsim')) atm = Atmosphere(atmgrid.airmass, atmgrid.pressure, atmgrid.temperature) + assert os.path.isfile('./tests/data/' + tag.replace('_spectrum.fits', '_atmsim.fits')) is True assert image_simulation.data is not None # assert spectrum_simulation.data is not None assert spectrogram_simulation.data is not None