Skip to content

Commit

Permalink
AD: small corrections
Browse files Browse the repository at this point in the history
  • Loading branch information
Allan Denis committed Jan 8, 2025
1 parent 8635da9 commit 3db40ca
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 49 deletions.
84 changes: 46 additions & 38 deletions ForMoSA/nested_sampling/forward_models.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,36 +25,39 @@ def forward_model(global_params, wav_mod_spectro, res_mod_spectro, flx_cont_obs,
'''

flx_mod *= transm_obs
flx_cont_mod = continuum_estimate(
global_params, wav_mod_spectro, flx_mod, res_mod_spectro, indobs)
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]))

weights = (1 / err_obs)**2

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)
flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, weights, 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)
flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, weights, 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)
flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, weights, 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)
flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, weights, 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):
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, weights, flx_obs, system_obs, bounds):
'''
Non linear forward model of planet and star contributions (see Landman et al. 2023)
This model is in principle more general than forward_model_estimate_speckles
because in the latter case we make the assumption that the star speckels dominate the data which is not the case here
Args:
flx_cont_obs (array): Continuum of the data
Expand All @@ -63,7 +66,7 @@ def forward_model_nonlinear_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mo
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
weights (array): Weights to apply to the data
flx_obs (array): Data
system_obs (array): Systematics
bounds (tuple): Bounds to be applied to the estimated parameters
Expand All @@ -72,35 +75,40 @@ def forward_model_nonlinear_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mo
'''

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)
res = 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))
if len(theta) > ind_star:
res += np.dot(theta[ind_star:], system_obs) - flx_obs
return weights * res


# 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

if len(system_obs) > 0:
for i in range(len(system_obs[0])):
# Arbitrary initial guesses for systematics contribution
theta0.append(1)

# Solve non linear Least Squares
res = optimize.least_squares(f, theta0, bounds=bounds)

# Full model
flx_mod_full = f(res.x)*err_obs + flx_obs
flx_mod_full = f(res.x) / weights + 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):
def forward_model_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_obs, star_flx_cont_obs, weights, 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)
Expand All @@ -111,7 +119,7 @@ def forward_model_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_fl
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
weights (array): Weights to apply to the data
flx_obs (array): Data
system_obs (array): Systematics
bounds (tuple): Bounds to be applied to the estimated parameters
Expand All @@ -129,29 +137,29 @@ def forward_model_estimate_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_fl

# Build matrix A
A = np.zeros([np.size(flx_obs), ind_system])
A[:, 0] = 1 / err_obs * (flx_mod - flx_cont_mod *
A[:, 0] = weights * (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)
A[:, star_i + 1] = weights * (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]
A[:, system_i + ind_star] = weights * system_obs[:, system_i]

# Build vector b
b = 1 / err_obs * flx_obs
b = weights * 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
flx_mod_full = np.dot(A, res.x) / weights
star_flx_obs = np.dot(A[:, 1:ind_star], res.x[1:ind_star]) / weights
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):
def forward_model_remove_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_obs_master, star_flx_cont_obs, weights, 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)
Expand All @@ -162,7 +170,7 @@ def forward_model_remove_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_
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
weights (array): Weights to apply to the data
flx_obs (array): Data
system_obs (array): Systematics
bounds (tuple): Bounds to be applied to the estimated parameters
Expand All @@ -179,28 +187,28 @@ def forward_model_remove_speckles(flx_cont_obs, flx_mod, flx_cont_mod, star_flx_
A = np.zeros([np.size(flx_obs), ind_system])

# Build matrix A
A[:, 0] = 1 / err_obs * (flx_mod - flx_cont_mod *
A[:, 0] = weights * (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]
A[:, system_i + 1] = weights * system_obs[:, system_i]

# Build vector b
b = 1 / err_obs * (flx_obs - star_flx_obs_master /
b = weights * (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
flx_mod_full = np.dot(A[:,0], res.x[0]) / weights + 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):
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, weights, flx_obs, system_obs, bounds):
'''
Linear forward model of planet and star contributions where we remove the continuums (see Wang et al. 2021)
Expand All @@ -211,7 +219,7 @@ def forward_model_estimate_speckles_remove_continuum(flx_cont_obs, flx_mod, flx_
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
weights (array): Weights to apply to the data
flx_obs (array): Data
system_obs (array): Systematics
bounds (tuple): Bounds to be applied to the estimated parameters
Expand All @@ -230,23 +238,23 @@ def forward_model_estimate_speckles_remove_continuum(flx_cont_obs, flx_mod, flx_

# 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))
A[:, 0] = weights * (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]))
A[:, star_i+1] = weights * (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]
A[:, system_i + ind_star] = weights * system_obs[:, system_i]

# Build vector b
b = 1 / err_obs * (flx_obs - flx_cont_obs + np.mean(flx_obs))
b = weights * (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
flx_mod_full = np.dot(A, res.x) / weights
flx_obs = b / weights
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:])

Expand Down
22 changes: 11 additions & 11 deletions ForMoSA/nested_sampling/nested_sampling.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,7 @@ def import_obsmod(global_params):
min_ns_u = float(ns_u.split(',')[0])
max_ns_u = float(ns_u.split(',')[1])
# Indices of each model and data
if global_params.rv[indobs] == 'NA' or global_params.rv == 'NA': # If the user didn't define a prior in RV, we juste adapt the model to the values defined in 'wav_fit'
if global_params.rv[indobs] == 'NA' or global_params.rv == 'NA': # If the user didn't define a prior in RV, we just adapt the model to the values defined in 'wav_fit'
mask_mod_spectro += (grid_spectro['wavelength'] >= min_ns_u) & (grid_spectro['wavelength'] <= max_ns_u)
else: # Otherwise we chose a slightly larger wavelength range to avoid loosing data onf the edges when applying RV correction
mask_mod_spectro += (grid_spectro['wavelength'] >= 0.98 * min_ns_u) & (grid_spectro['wavelength'] <= 1.02 * max_ns_u)
Expand All @@ -116,19 +116,19 @@ def import_obsmod(global_params):
else:
inv_cov_obs_ns_u = np.asarray([])
if len(transm_obs_spectro) != 0: # Add the transmission (if necessary)
transm_obs_ns_u = transm_obs_spectro[mask_obs_spectro]
transm_obs_spectro_ns_u = transm_obs_spectro[mask_obs_spectro]
else:
transm_obs_ns_u = np.asarray([])
transm_obs_spectro_ns_u = np.asarray([])
if len(star_flx_obs_spectro) != 0: # Add star flux (if necessary)
star_flx_obs_ns_u = star_flx_obs_spectro[mask_obs_spectro]
star_flx_cont_obs_ns_u = star_flx_cont_obs_spectro[mask_obs_spectro]
star_flx_obs_spectro_ns_u = star_flx_obs_spectro[mask_obs_spectro]
star_flx_cont_obs_spectro_ns_u = star_flx_cont_obs_spectro[mask_obs_spectro]
else:
star_flx_obs_ns_u = np.asarray([])
star_flx_cont_obs_ns_u = np.array([])
star_flx_obs_spectro_ns_u = np.asarray([])
star_flx_cont_obs_spectro_ns_u = np.array([])
if len(system_obs_spectro) != 0: # Add systematics model (if necessary)
system_obs_ns_u = system_obs_spectro[mask_obs_spectro]
system_obs_spectro_ns_u = system_obs_spectro[mask_obs_spectro]
else:
system_obs_ns_u = np.asarray([])
system_obs_spectro_ns_u = np.asarray([])
wav_obs_photo_ns_u = wav_obs_photo[mask_obs_photo]
flx_obs_photo_ns_u = flx_obs_photo[mask_obs_photo]
err_obs_photo_ns_u = err_obs_photo[mask_obs_photo]
Expand All @@ -140,7 +140,7 @@ def import_obsmod(global_params):
main_file.append([[wav_obs_spectro_ns_u, wav_obs_photo_ns_u],
[flx_obs_spectro_ns_u, flx_cont_obs_spectro_ns_u, flx_obs_photo_ns_u],
[err_obs_spectro_ns_u, err_obs_photo_ns_u],
[transm_obs_ns_u, star_flx_obs_ns_u, star_flx_cont_obs_ns_u, system_obs_ns_u],
[transm_obs_spectro_ns_u, star_flx_obs_spectro_ns_u, star_flx_cont_obs_spectro_ns_u, system_obs_spectro_ns_u],
[res_obs_spectro_ns_u, res_mod_spectro_ns_u],
inv_cov_obs_ns_u,
grid_spectro_ns_u,
Expand Down Expand Up @@ -562,7 +562,7 @@ def launch_nested_sampling(global_params):
print('WARNING. You cannot use CCF mappings without substracting the continuum')
print()
exit()
elif global_params.logL_type[indobs] == 'CCF_Zucker' and global_params.continuum_sub[indobs] == 'NA' and global_params.star_contaminated[indobs] != 'Remove':
elif global_params.logL_type[indobs] == 'CCF_Zucker' and global_params.continuum_sub[indobs] == 'NA':
print('WARNING. You cannot use CCF mappings without substracting the continuum')
print()
exit()
Expand Down

0 comments on commit 3db40ca

Please sign in to comment.