AD : Multiple modifications, including : adaptation of the model on a…

… slightly larger wavelength grid, cleaning of the plotting_class function, removing of the lsq function replaced by a forward_model script, possibility of the user to chose (and custom) its own lsq_function to disentangle between planet data and star speckles data

ForMoSA/adapt/adapt_grid.py

@@ -124,6 +124,7 @@ def adapt_grid(global_params, res_mod_obs, wav_obs_spectro, res_obs_spectro, wav
     wav_mod_nativ = ds["wavelength"].values
     grid = ds['grid']
     attr = ds.attrs
+    attr['res'] = res_obs_spectro
     grid_np = grid.to_numpy()
 
     # create arrays without any assumptions on the number of parameters

ForMoSA/adapt/adapt_obs_mod.py

@@ -32,7 +32,7 @@ def launch_adapt(global_params, justobs='no'):
     res_mod_nativ = attr['res']
     ds.close()
 
-    # Check if the grid is Nyquist-sampled, else set the resolution to R = wav / 2 Deltawav to make sure we are adding any info
+    # Check if the grid is Nyquist-sampled, else set the resolution to R = wav / 2 Deltawav to make sure we are adding any info
     dwav = np.abs(wav_mod_nativ - np.roll(wav_mod_nativ, 1))
     dwav[0] = dwav[1]
     res_Nyquist = wav_mod_nativ / (2 * dwav)
@@ -61,11 +61,28 @@ def launch_adapt(global_params, justobs='no'):
 
         # Interpolate the resolution onto the wavelength of the data
         if len(obs_spectro[0]) != 0:
-            mask_mod_obs = (wav_mod_nativ <= obs_spectro[0][-1]) & (wav_mod_nativ > obs_spectro[0][0])
+            if (global_params.rv[indobs*3] == 'NA') or (global_params.rv == 'NA'):
+                mask_mod_obs = (wav_mod_nativ <= obs_spectro[0][-1]) & (wav_mod_nativ > obs_spectro[0][0])
+            else:
+                # If the user defined an RV prior, we slightly modify the strategy for the adaptation of the model
+                # Instead of adapting the model to the wavelength of the observations, we adapt the model to an extended version of the wavelength range of the observation
+                # so that we do not lose data on the edges of the wavelength of the observations when applying the RV correction
+                mask_mod_obs = (wav_mod_nativ <= 1.02 * obs_spectro[0][-1]) & (wav_mod_nativ >= 0.98 * obs_spectro[0][0])
+                wav_obs_spectro = obs_spectro[0]
+                res_obs_spectro = obs_spectro[3]
+                while wav_obs_spectro[0] > 0.98 * obs_spectro[0][0]:
+                    wav_obs_spectro = np.insert(wav_obs_spectro, 0, 2*wav_obs_spectro[0]-wav_obs_spectro[1])
+                    wav_obs_spectro = np.append(wav_obs_spectro, 2*wav_obs_spectro[-1]-wav_obs_spectro[-2])
+                    res_obs_spectro = np.insert(res_obs_spectro, 0, res_obs_spectro[0])
+                    res_obs_spectro = np.append(res_obs_spectro, res_obs_spectro[-1])
+                while wav_obs_spectro[-1] < 1.02 * obs_spectro[0][-1]:
+                    wav_obs_spectro = np.append(wav_obs_spectro, 2 * wav_obs_spectro[-1]-wav_obs_spectro[-2])
+                    res_obs_spectro = np.append(res_obs_spectro, res_obs_spectro[-1])
+
             wav_mod_cut = wav_mod_nativ[mask_mod_obs]
             res_mod_cut = res_mod_nativ[mask_mod_obs]
             interp_mod_to_obs = interp1d(wav_mod_cut, res_mod_cut, fill_value='extrapolate')
-            res_mod_obs = interp_mod_to_obs(obs_spectro[0])
+            res_mod_obs = interp_mod_to_obs(wav_obs_spectro)
         else:
             res_mod_obs = np.asarray([])
 
@@ -103,7 +120,7 @@ def launch_adapt(global_params, justobs='no'):
             print('- - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -')
             print(f"-> Sarting the adaptation of {obs_name}")
 
-            adapt_grid(global_params, res_mod_obs, obs_spectro[0], obs_spectro[3], obs_photo[0], obs_photo_ins, obs_name=obs_name, indobs=indobs)
+            adapt_grid(global_params, res_mod_obs, wav_obs_spectro, res_obs_spectro, obs_photo[0], obs_photo_ins, obs_name=obs_name, indobs=indobs)
 
 # ----------------------------------------------------------------------------------------------------------------------
 

ForMoSA/main_utilities.py

@@ -51,30 +51,51 @@ def __init__(self, config_file_path):
         model_name = model_name[0]
         self.model_name = model_name
 
-        if type(config['config_adapt']['wav_for_adapt']) != list: # Create lists if only one obs in the loop
+        if type(config['config_adapt']['wav_for_adapt']) != list: # Create lists if only one obs in the loop
             # [config_adapt] (5)
             self.wav_for_adapt = [config['config_adapt']['wav_for_adapt']]
             self.adapt_method = [config['config_adapt']['adapt_method']]
             self.custom_reso = [config['config_adapt']['custom_reso']]
             self.continuum_sub = [config['config_adapt']['continuum_sub']]
             self.wav_for_continuum = [config['config_adapt']['wav_for_continuum']]
-            self.use_lsqr = [config['config_adapt']['use_lsqr']]
 
             # [config_inversion] (4)
             self.logL_type = [config['config_inversion']['logL_type']]
             self.wav_fit = [config['config_inversion']['wav_fit']]
+
+            # [config_forward_models] (3)
+            try:
+                self.fm_type = [config['config_forward_models']['fm_type']]
+                self.fm_continuum_res = [config['config_forward_models']['fm_continuum_res']]
+                self.bounds_lsq = [config['config_forward_models']['bounds_lsq']]
+            except KeyError:
+                self.fm_type = 'NA'
+                self.fm_continuum_res = 'NA'
+                self.bounds_lsq = 'NA'
+
         else:
             # [config_adapt] (5)
             self.wav_for_adapt = config['config_adapt']['wav_for_adapt']
             self.adapt_method = config['config_adapt']['adapt_method']
             self.custom_reso = config['config_adapt']['custom_reso']
             self.continuum_sub = config['config_adapt']['continuum_sub']
             self.wav_for_continuum = config['config_adapt']['wav_for_continuum']
-            self.use_lsqr = config['config_adapt']['use_lsqr']
 
             # [config_inversion] (4)
             self.logL_type = config['config_inversion']['logL_type']
             self.wav_fit = config['config_inversion']['wav_fit']
+
+            # [config_forwarf_models] (3)
+            try:
+                self.fm_type =  config['config_forward_models']['fm_type']
+                self.fm_continuum_res = config['config_forward_models']['fm_continuum_res']
+                self.bounds_lsq = config['config_forward_models']['bounds_lsq']
+            except KeyError:
+                self.fm_type = 'NA'
+                self.fm_continuum_res = 'NA'
+                self.bounds_lsq = 'NA'
+
+            # 
 
         # parallelisation of adapt
         try:
@@ -129,7 +150,7 @@ def __init__(self, config_file_path):
         # self.p_context = eval(config['config_pymultinest']['context'])
         # self.p_write_output = config['config_pymultinest']['write_output']
         # self.p_log_zero = eval(config['config_pymultinest']['log_zero'])
-        # self.p_max_iter = eval(config['config_pymultinest']['max_iter'])
+        # self.p_max_iter = eval(config['config_pymultinest']['max_iter'])
         # self.p_init_MPI = config['config_pymultinest']['init_MPI']
         # self.p_dump_callback = config['config_pymultinest']['dump_callback']
         # self.p_use_MPI = config['config_pymultinest']['use_MPI']
@@ -158,7 +179,7 @@ def __init__(self, config_file_path):
         # print()
         #
         # config_current = self.result_path + '/past_config.ini'
-        # config.filename = ' '
+        # config.filename = ' '
         # config['config_path']['stock_interp_grid'] = stock_interp_grid
         # config['config_path']['stock_result'] = stock_result_subsub_dir
         # config.write()

ForMoSA/nested_sampling/forward_models.py

@@ -0,0 +1,257 @@
+import numpy as np
+import scipy.optimize as optimize
+from ForMoSA.adapt.extraction_functions import continuum_estimate
+import ForMoSA.utils as utils
+
+
+def forward_model(global_params, wav_mod_spectro, res_mod_spectro, flx_cont_obs, flx_mod, star_flx_obs, star_flx_cont_obs, err_obs, transm_obs, flx_obs, system_obs, indobs):
+    '''
+    For high-contrast companions, where the star speckles signal contaminate the data
+
+    Args:
+    global_params (object): Class containing each parameter used in ForMoSA
+    wav_mod_spectro (array): Wavelength grid of the model
+    res_mod_spectro (array): Resolution grid of the model
+    flx_cont_obs (array): Continuum of the data
+    flx_mod (array): Model of the companion
+    star_flx_obs (array): Star data
+    star_flx_cont_obs (array): Continuum of star data
+    err_obs (array): Noise of the data
+    transm_obs (array): Transmission
+    flx_obs (array): Data
+    system_obs (array): Systematics
+
+    Authors: Allan Denis
+    '''
+
+    flx_mod *= transm_obs
+    flx_cont_mod = continuum_estimate(
+        global_params, wav_mod_spectro, flx_mod, res_mod_spectro, indobs)
+
+    star_flx_obs_master = star_flx_obs[:, len(star_flx_obs[0]) // 2]
+
+    bounds = (float(global_params.bounds_lsq[indobs][1:-1].split(',')[0]), 
+              float(global_params.bounds_lsq[indobs][1:-1].split(',')[1]))
+
+    if global_params.fm_type[indobs] == 'nonlinear_fit_spec':
+        res, flx_mod, flx_obs, star_flx_obs, system_obs = forward_model_nonlinear_estimate_speckles(
+            flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds)
+    elif global_params.fm_type[indobs] == 'fit_spec':
+        res, flx_mod, flx_obs, star_flx_obs, system_obs = forward_model_estimate_speckles(
+            flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds)
+    elif global_params.fm_type[indobs] == 'rm_spec':
+        res, flx_mod, flx_obs, star_flx_obs, system_obs = forward_model_remove_speckles(
+            flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds)
+    elif global_params.fm_type[indobs] == 'fit_spec_rm_cont':
+        res, flx_mod, flx_obs, star_flx_obs, system_obs = forward_model_estimate_speckles_remove_continuum(
+            flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds)
+    elif global_params.fm_type[indobs] == 'fit_spec_fit_cont':
+        raise Exception(
+            'Continuum fitting-based forward model no implement yet ! Please use another function')
+
+    return res, flx_mod, flx_obs, star_flx_obs, system_obs
+
+
+def forward_model_nonlinear_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds):
+    '''
+    Non linear forward model of planet and star contributions (see Landman et al. 2023)
+
+    Args:
+    flx_cont_obs: Continuum of the data
+    flx_mod: Model of the companion
+    flx_cont_mod: Continuum of the model of the companion
+    star_flx_obs_master: Master star data
+    star_flx_obs: Star data
+    star_flx_cont_obs: Continuum of star data
+    err_obs: Noise of the data
+    flx_obs: Data
+
+    Authors: Allan Denis
+    '''
+
+    ind_star = 1 + len(star_flx_obs[0])
+
+    # # # # # # # Solve non linear Least Squares full_model(theta) = flx_obs
+
+    # Definition of f
+    def f(theta):
+        return 1 / err_obs * (theta[0] * flx_mod + np.dot(theta[1:ind_star], star_flx_obs / star_flx_cont_obs * (flx_cont_obs - theta[0] * flx_cont_mod)) + np.dot(theta[ind_star:], system_obs) - flx_obs)
+
+    # Solve non linear Least Squares
+    # Initial guess for the planetary contribution
+    theta0 = [0]
+    for i in range(len(star_flx_obs[0])):
+        # Arbitrary initial guesses for star speckles contribution
+        theta0.append((i / len(star_flx_obs[0]))**2)
+    for i in range(len(system_obs[0])):
+        # Arbitrary initial guesses for systematics contribution
+        theta0.append(1)
+
+    # Solve non linear Least Squaresr the assumtion that the star speckles dominate the data
+    res = optimize.least_squares(f, theta0)
+
+    # Full model
+    flx_mod_full = f(res.x)*err_obs + flx_obs
+    star_flx_obs = np.dot(res.x[1:ind_star], star_flx_obs / star_flx_cont_obs * flx_cont_obs)
+    system_obs = np.dot(res.x[ind_star:], system_obs)
+
+    return res.x, flx_mod_full, flx_obs, star_flx_obs, system_obs
+
+
+def forward_model_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds):
+    '''
+    Linear forward model of planet and star contributions under the assumtion that the star speckles dominate the data  (see Landman et al. 2023)
+
+    Args:
+    flx_cont_obs (array): Continuum of the data
+    flx_mod (array): Model of the companion
+    flx_cont_mod (array): Continuum of the model of the companion
+    star_flx_obs_master (array): Master star data
+    star_flx_obs (array): Star data
+    star_flx_cont_obs (array): Continuum of star data
+    err_obs (array): Noise of the data
+    flx_obs (array): Data
+    system_obs (array): Systematics
+
+    Authors: Allan Denis
+    '''
+
+    ind_star = 1 + len(star_flx_obs[0])
+    if len(system_obs) > 0:
+        ind_system = ind_star + len(system_obs[0])
+    else:
+        ind_system = ind_star
+
+    # # # # # # Solve linear Least Squares A.x = b
+
+    # Build matrix A
+    A = np.zeros([np.size(flx_obs), ind_system])
+    A[:, 0] = 1 / err_obs * (flx_mod - flx_cont_mod *
+                             star_flx_obs_master / star_flx_cont_obs)
+
+    for star_i in range(len(star_flx_obs[0])):
+        A[:, star_i + 1] = 1 / err_obs * (star_flx_obs[:, star_i] / star_flx_cont_obs * flx_cont_obs)
+
+    for system_i in range(ind_system - ind_star):
+        A[:, system_i + ind_star] = 1 / err_obs * system_obs[:, system_i]
+
+    # Build vector b
+    b = 1 / err_obs * flx_obs
+    # Solve linear Least Squares
+    res = optimize.lsq_linear(A, b, bounds=bounds)
+
+    # Full model
+    flx_mod_full = np.dot(A, res.x) * err_obs
+    star_flx_obs = np.dot(A[:, 1:ind_star], res.x[1:ind_star]) * err_obs
+    system_obs = np.dot(A[:, ind_star:], res.x[ind_star:])
+
+    return res.x, flx_mod_full, flx_obs, star_flx_obs, system_obs
+
+
+def forward_model_remove_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds):
+    '''
+    Linear forward model of planet contribution only where the speckles are filtered out from the data (see Landman et al. 2023)
+
+    Args:
+    flx_cont_obs: Continuum of the data
+    flx_mod: Model of the companion
+    flx_cont_mod: Continuum of the model of the companion
+    star_flx_obs_master: Master star data
+    star_flx_cont_obs: Continuum of star data
+    err_obs: Noise of the data
+    flx_obs: Data
+    system_obs (array): Systematics
+
+    Authors: Allan Denis
+    '''
+
+    if len(system_obs) > 0:
+        ind_system = 1 + len(system_obs[0])
+    else:
+        ind_system = 1
+
+    # # # # # # # Solve linear Least Squared A.x = b
+    A = np.zeros([np.size(flx_obs), ind_system])
+
+    # Build matrix A
+    A[:, 0] = 1 / err_obs * (flx_mod - flx_cont_mod *
+                             star_flx_obs_master / star_flx_cont_obs)
+
+    for system_i in range(ind_system-1):
+        A[:, system_i + 1] = 1 / err_obs * system_obs[:, system_i]
+
+    # Build vector b
+    b = 1 / err_obs * (flx_obs - star_flx_obs_master /
+                       star_flx_cont_obs * flx_cont_obs)
+
+    # Solve linear Least Squared
+    res = optimize.lsq_linear(A, b, bounds=bounds)
+
+    # Full model
+    star_flx_obs = star_flx_obs_master / star_flx_cont_obs + flx_cont_obs
+    flx_mod_full = np.dot(A[:,0], res.x[0])*err_obs + star_flx_obs
+    system_obs = np.dot(A[:, 1:], res.x[1:])
+
+    return res.x, flx_mod_full, flx_obs, star_flx_obs, system_obs
+
+
+def forward_model_estimate_speckles_remove_continuum(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, err_obs, flx_obs, system_obs, bounds):
+    '''
+    Linear forward model of planet and star contributions where we remove the continuums (see Wang et al. 2021)
+
+    Args:
+    flx_cont_obs: Continuum of the data
+    flx_mod: Model of the companion
+    flx_cont_mod: Continuum of the mdoel of the companion
+    star_flx_obs_master: Master star data
+    star_flx_obs: Star data
+    star_flx_cont_obs: Continuum of star data
+    err_obs: Noise of the data
+    flx_obs : Data
+    system_obs (array): Systematics
+
+    Authors: Allan Denis
+    '''
+    ind_star = 1 + len(star_flx_obs[0])
+    if len(system_obs) > 0:
+        ind_system = ind_star + len(system_obs[0])
+    else:
+        ind_system = ind_star
+
+
+    # # # # # # Solve linear Least Squares A.x = b
+
+    # Build matrix A
+    A = np.zeros([np.size(flx_obs), ind_system])
+    A[:, 0] = 1 / err_obs * (flx_mod - flx_cont_mod + np.mean(flx_mod))
+
+    for star_i in range(len(star_flx_obs[0])):
+        A[:, star_i+1] = 1 / err_obs * (star_flx_obs[:, star_i] - star_flx_cont_obs + np.mean(star_flx_obs[:, star_i]))
+
+    for system_i in range(ind_system-ind_star):
+        A[:, system_i + ind_star] = 1 / err_obs * system_obs[:, system_i]
+
+    # Build vector b
+    b = 1 / err_obs * (flx_obs - flx_cont_obs + np.mean(flx_obs))
+
+    # Solve linear Least Squares
+    res = optimize.lsq_linear(A, b, bounds=bounds)
+
+    # Full model
+    flx_mod_full = np.dot(A, res.x) * err_obs
+    flx_obs = b * err_obs
+    star_flx_obs = np.dot(A[:, 1:ind_star], res.x[1:ind_star])
+    system_obs = np.dot(A[:, ind_star:], res.x[ind_star:])
+
+    return res.x, flx_mod_full, flx_obs, star_flx_obs, system_obs
+
+
+def forward_model_estimate_speckles_estimate_continuum():
+    '''
+    Linear forward model of planet and star contributions where we fit the continuum
+    To Be Defined
+
+    Authors: Allan Denis
+    '''
+
+    return

ForMoSA/nested_sampling/nested_logL_functions.py

@@ -57,7 +57,7 @@ def logL_chi2_extended(delta_flx, err):
     N = len(delta_flx)
     chi2 = np.nansum((delta_flx / err) ** 2)
     s2 = 1/N * chi2
-    logL = -(chi2 / (2*s2) + N/2 * np.log(2*np.pi*s2) + 1/2 * np.log(np.dot(err,err)))
+    logL = -(chi2 / (2*s2) + N/2 * np.log(2*np.pi*s2))
 
     return logL