From 0138dd806b3a83252c0f42618f25dfc4c372361f Mon Sep 17 00:00:00 2001 From: kment Date: Tue, 17 Dec 2024 11:47:29 -0500 Subject: [PATCH 1/4] Changes before 12/17/2024 --- examples/_lsf.jl | 5 +-- src/model_functions.jl | 68 ++++++++++++++++++++++++++--------- src/optimization_functions.jl | 36 +++++++++++++++---- 3 files changed, 83 insertions(+), 26 deletions(-) diff --git a/examples/_lsf.jl b/examples/_lsf.jl index e1f73b1..9e8bb5e 100644 --- a/examples/_lsf.jl +++ b/examples/_lsf.jl @@ -61,7 +61,8 @@ module NEIDLSF return ans end - function neid_lsf(order::Int, log_λ_neid_order::AbstractVector, log_λ_obs::AbstractVector) + # KM: Adding the option to extrapolate + function neid_lsf(order::Int, log_λ_neid_order::AbstractVector, log_λ_obs::AbstractVector; extrapolate::Bool=false) @assert 1 <= order <= length(no_lsf_orders) if no_lsf_orders[order]; return nothing end n = length(log_λ_obs) @@ -72,7 +73,7 @@ module NEIDLSF pixel_separation_ratio = SSOF.simple_derivative(log_λ_neid_order) ./ pixel_separation_log_λ_obs.(log_λ_neid_order) # make the linear_interpolation object and evaluate it # converter(vals) = linear_interpolation(log_λ_neid_order, pixel_separation_ratio .* vals; extrapolation_bc=Line())(log_λ_obs) - converter(vals) = (DataInterpolations.LinearInterpolation(pixel_separation_ratio .* vals, log_λ_neid_order)).(log_λ_obs) + converter(vals) = (DataInterpolations.LinearInterpolation(pixel_separation_ratio .* vals, log_λ_neid_order; extrapolate=extrapolate)).(log_λ_obs) σs_converted = converter(σs[:, order]) bhws_converted = converter(bhws[:, order]) threeish_sigma_converted = converter(threeish_sigma.(σs[:, order], bhws[:, order])) diff --git a/src/model_functions.jl b/src/model_functions.jl index 8c01c49..5918531 100644 --- a/src/model_functions.jl +++ b/src/model_functions.jl @@ -575,6 +575,9 @@ default_reg_tel = Dict([(:GP_μ, 1e6), (:L2_μ, 1e6), (:L1_μ, 1e5), (:L1_μ₊_ (:GP_M, 1e7), (:L1_M, 1e7)]) default_reg_star = Dict([(:GP_μ, 1e2), (:L2_μ, 1e-2), (:L1_μ, 1e2), (:L1_μ₊_factor, 6.), (:GP_M, 1e4), (:L1_M, 1e7)]) +default_reg_model = Dict([(:L2_μ, 1e4), (:L2_h2o, 1e4), (:L2_oxy, 1e4)]) +#default_reg_model = Dict([(:L2_μ, 1e4),]) +#default_reg_model = Dict([(:L2_h2o, 1e4), (:L2_oxy, 1e4)]) default_reg_tel_full = Dict([(:GP_μ, 1e6), (:L2_μ, 1e6), (:L1_μ, 1e5), (:L1_μ₊_factor, 6.), (:GP_M, 1e7), (:L2_M, 1e4), (:L1_M, 1e7)]) default_reg_star_full = Dict([(:GP_μ, 1e2), (:L2_μ, 1e-2), (:L1_μ, 1e1), @@ -668,6 +671,8 @@ struct OrderModelDPCA{T<:Number} <: OrderModel reg_tel::Dict{Symbol, T} "Stellar regularization coefficients" reg_star::Dict{Symbol, T} + "Telluric model regularization coefficients" + reg_model::Dict{Symbol, T} "Matrices to interpolate stellar model to observed wavelengths (barycenter 2 observed)" b2o::AbstractVector{<:SparseMatrixCSC} "Matrices to interpolate telluric model to observed wavelengths (telluric 2 observed)" @@ -694,6 +699,8 @@ struct OrderModelWobble{T<:Number} <: OrderModel reg_tel::Dict{Symbol, T} "Stellar regularization coefficients" reg_star::Dict{Symbol, T} + "Telluric model regularization coefficients" + reg_model::Dict{Symbol, T} "Linear interpolation helper objects to interpolate stellar model to observed wavelengths (barycenter 2 observed)" b2o::StellarInterpolationHelper "Approximate barycentric correction \"RVs\" (using `(λ1-λ0)/λ0 = λ1/λ0 - 1 = e^D - 1 ≈ β = v / c`)" @@ -724,6 +731,7 @@ Constructor for the `OrderModel`-type objects (SSOF model for a set of 1D spectr - `log_λ_gp_tel::Real=1/LSF_gp_params.λ`: The log λ lengthscale of the telluric regularization GP - `tel_log_λ::Union{Nothing,AbstractRange}=nothing`: The log - `star_log_λ::Union{Nothing,AbstractRange}=nothing`: The log λ lengthscale of the telluric regularization GP +- `use_tel_prior`::Bool=true`: Whether to include a modeled prior in the telluric submodel - `kwargs...`: kwargs passed to `Submodel` constructors """ function OrderModel( @@ -739,6 +747,7 @@ function OrderModel( log_λ_gp_tel::Real=1/LSF_gp_params.λ, tel_log_λ::Union{Nothing,AbstractRange}=nothing, star_log_λ::Union{Nothing,AbstractRange}=nothing, + use_tel_prior::Bool=true, kwargs...) # Creating models @@ -750,8 +759,8 @@ function OrderModel( rv = zeros(n_obs) bary_rvs = D_to_rv.([median(d.log_λ_star[:, i] - d.log_λ_obs[:, i]) for i in 1:n_obs]) - todo = Dict([(:initialized, false), (:reg_improved, false), (:err_estimated, false)]) - metadata = Dict([(:todo, todo), (:instrument, instrument), (:order, order), (:star, star_str)]) + todo = Dict([(:initialized, false), (:reg_improved, false), (:err_estimated, false), (:h2o_tried, false), (:oxy_tried, false)]) + metadata = Dict([(:todo, todo), (:tel_prior, use_tel_prior), (:instrument, instrument), (:order, order), (:star, star_str), (:h2o_vector, 0), (:oxy_vector, 0)]) if dpca if oversamp b2o = oversamp_interp_helper(d.log_λ_star_bounds, star.log_λ) @@ -760,7 +769,7 @@ function OrderModel( b2o = undersamp_interp_helper(d.log_λ_star, star.log_λ) t2o = undersamp_interp_helper(d.log_λ_obs, tel.log_λ) end - return OrderModelDPCA(tel, star, rv, copy(default_reg_tel), copy(default_reg_star), b2o, t2o, metadata, n_obs) + return OrderModelDPCA(tel, star, rv, copy(default_reg_tel), copy(default_reg_star), copy(default_reg_model), b2o, t2o, metadata, n_obs) else b2o = StellarInterpolationHelper(star.log_λ, bary_rvs, d.log_λ_obs) if oversamp @@ -768,17 +777,19 @@ function OrderModel( else t2o = undersamp_interp_helper(d.log_λ_obs, tel.log_λ) end - return OrderModelWobble(tel, star, rv, copy(default_reg_tel), copy(default_reg_star), b2o, bary_rvs, t2o, metadata, n_obs) + return OrderModelWobble(tel, star, rv, copy(default_reg_tel), copy(default_reg_star), copy(default_reg_model), b2o, bary_rvs, t2o, metadata, n_obs) end end -Base.copy(om::OrderModelDPCA) = OrderModelDPCA(copy(om.tel), copy(om.star), copy(om.rv), copy(om.reg_tel), copy(om.reg_star), om.b2o, om.t2o, copy(om.metadata), om.n) +Base.copy(om::OrderModelDPCA) = OrderModelDPCA(copy(om.tel), copy(om.star), copy(om.rv), copy(om.reg_tel), copy(om.reg_star), copy(om.reg_model), + om.b2o, om.t2o, deepcopy(om.metadata), om.n) (om::OrderModelDPCA)(inds::AbstractVecOrMat) = OrderModelDPCA(om.tel(inds), om.star(inds), om.rv(inds), copy(om.reg_tel), - copy(om.reg_star), view(om.b2o, inds), view(om.t2o, inds), copy(om.metadata), length(inds)) -Base.copy(om::OrderModelWobble) = OrderModelWobble(copy(om.tel), copy(om.star), copy(om.rv), copy(om.reg_tel), copy(om.reg_star), copy(om.b2o), om.bary_rvs, om.t2o, copy(om.metadata), om.n) + copy(om.reg_star), copy(om.reg_model), view(om.b2o, inds), view(om.t2o, inds), deepcopy(om.metadata), length(inds)) +Base.copy(om::OrderModelWobble) = OrderModelWobble(copy(om.tel), copy(om.star), copy(om.rv), copy(om.reg_tel), copy(om.reg_star), copy(om.reg_model), + copy(om.b2o), om.bary_rvs, om.t2o, deepcopy(om.metadata), om.n) (om::OrderModelWobble)(inds::AbstractVecOrMat) = - OrderModelWobble(om.tel(inds), om.star(inds), view(om.rv, inds), copy(om.reg_tel), - copy(om.reg_star), om.b2o(inds, length(om.star.lm.μ)), view(om.bary_rvs, inds), view(om.t2o, inds), copy(om.metadata), length(inds)) + OrderModelWobble(om.tel(inds), om.star(inds), view(om.rv, inds), copy(om.reg_tel), copy(om.reg_star), copy(om.reg_model), + om.b2o(inds, length(om.star.lm.μ)), view(om.bary_rvs, inds), view(om.t2o, inds), deepcopy(om.metadata), length(inds)) """ @@ -797,6 +808,7 @@ end rm_regularization!(om) Remove all of the keys in the regularization Dicts +Currently does not affect om.reg_model """ function rm_regularization!(om::OrderModel) rm_dict!(om.reg_tel) @@ -912,12 +924,12 @@ downsize(m::OrderModelDPCA, n_comp_tel::Int, n_comp_star::Int) = OrderModelDPCA( downsize(m.tel, n_comp_tel), downsize(m.star, n_comp_star), - copy(m.rv), copy(m.reg_tel), copy(m.reg_star), m.b2o, m.t2o, copy(m.metadata), m.n) + copy(m.rv), copy(m.reg_tel), copy(m.reg_star), copy(m.reg_model), m.b2o, m.t2o, deepcopy(m.metadata), m.n) downsize(m::OrderModelWobble, n_comp_tel::Int, n_comp_star::Int) = OrderModelWobble( downsize(m.tel, n_comp_tel), downsize(m.star, n_comp_star), - copy(m.rv), copy(m.reg_tel), copy(m.reg_star), copy(m.b2o), m.bary_rvs, m.t2o, copy(m.metadata), m.n) + copy(m.rv), copy(m.reg_tel), copy(m.reg_star), copy(m.reg_model), copy(m.b2o), m.bary_rvs, m.t2o, deepcopy(m.metadata), m.n) """ @@ -944,12 +956,12 @@ downsize_view(m::OrderModelDPCA, n_comp_tel::Int, n_comp_star::Int) = OrderModelDPCA( downsize_view(m.tel, n_comp_tel), downsize_view(m.star, n_comp_star), - m.rv, copy(m.reg_tel), copy(m.reg_star), m.b2o, m.t2o, copy(m.metadata), m.n) + m.rv, copy(m.reg_tel), copy(m.reg_star), copy(m.reg_model), m.b2o, m.t2o, deepcopy(m.metadata), m.n) downsize_view(m::OrderModelWobble, n_comp_tel::Int, n_comp_star::Int) = OrderModelWobble( downsize_view(m.tel, n_comp_tel), downsize_view(m.star, n_comp_star), - m.rv, copy(m.reg_tel), copy(m.reg_star), m.b2o, m.bary_rvs, m.t2o, copy(m.metadata), m.n) + m.rv, copy(m.reg_tel), copy(m.reg_star), copy(m.reg_model), m.b2o, m.bary_rvs, m.t2o, deepcopy(m.metadata), m.n) """ spectra_interp(model, interp_helper) @@ -1245,7 +1257,7 @@ end """ model_prior(lm, om, key) -Calulate the model prior on `lm` with the regularization terms in `om.reg_` * `key` +Calculate the model prior on `lm` with the regularization terms in `om.reg_` * `key` """ function model_prior(lm, om::OrderModel, key::Symbol) reg = getfield(om, Symbol(:reg_, key)) @@ -1296,18 +1308,40 @@ function model_s_prior(s, reg::Dict) end +## TEMPORARY - TODO: INTEGRATE FUNCTIONS INTO SSOF +include("/home/kment/RV/telluricfunc.jl") """ tel_prior(om) -Calulate the telluric model prior on `om.tel.lm` with the regularization terms in `om.reg_tel` +Calculate the telluric model prior on `om.tel.lm` with the regularization terms in `om.reg_model` """ +function tel_prior(lm, om::OrderModel) + + val = model_prior(lm, om, :tel) + + if om.metadata[:tel_prior] && haskey(om.reg_model, :L2_μ) + val += L2(get_telluric_prior_spectrum(om.tel.λ) - lm[min(3, length(lm))]) * om.reg_model[:L2_μ] + end + + if om.metadata[:tel_prior] && length(lm) >= 3 + if om.metadata[:h2o_vector] >= 1 && size(lm[1], 2) >= om.metadata[:h2o_vector] && haskey(om.reg_model, :L2_h2o) + val += L2(get_telluric_prior_spectrum(om.tel.λ, :h2o) - lm[1][:,om.metadata[:h2o_vector]] .- 1) * om.reg_model[:L2_h2o] + end + if om.metadata[:oxy_vector] >= 1 && size(lm[1], 2) >= om.metadata[:oxy_vector] && haskey(om.reg_model, :L2_oxy) + val += L2(get_telluric_prior_spectrum(om.tel.λ, :oxy) - lm[1][:,om.metadata[:oxy_vector]] .- 1) * om.reg_model[:L2_oxy] + end + end + + return val + +end +tel_prior(lm::Union{FullLinearModel, TemplateModel}, om::OrderModel) = tel_prior(vec(lm), om) tel_prior(om::OrderModel) = tel_prior(om.tel.lm, om) -tel_prior(lm, om::OrderModel) = model_prior(lm, om, :tel) """ star_prior(om) -Calulate the stellar model prior on `om.star.lm` with the regularization terms in `om.star_tel` +Calculate the stellar model prior on `om.star.lm` with the regularization terms in `om.star_tel` """ star_prior(om::OrderModel) = star_prior(om.star.lm, om) star_prior(lm, om::OrderModel) = model_prior(lm, om, :star) diff --git a/src/optimization_functions.jl b/src/optimization_functions.jl index 0ff6b19..7b623c4 100644 --- a/src/optimization_functions.jl +++ b/src/optimization_functions.jl @@ -49,6 +49,7 @@ end _loss_diagnostic(mws::ModelWorkspace; kwargs...) = _loss_diagnostic(mws.o, mws.om, mws.d; kwargs...) # summed χ² loss functions +#_loss(model, data, variance) = sum(_χ²_loss(model, data, variance)) _loss(tel, star, rv, d::Data; kwargs...) = sum(__loss_diagnostic(tel, star, rv, d; kwargs...)) _loss(tel, star, d::Data; kwargs...) = sum(__loss_diagnostic(tel, star, d; kwargs...)) _loss(o::Output, om::OrderModel, d::Data; kwargs...) = sum(_loss_diagnostic(o, om, d; kwargs...)) @@ -73,6 +74,7 @@ _loss_recalc_rv_basis(mws::ModelWorkspace; kwargs...) = _loss_recalc_rv_basis(mw Create a loss function for the model and data in `mws` """ function loss_func(mws::ModelWorkspace; include_priors::Bool=false) + println("loss_func() called with include_priors=", include_priors) # DEBUG if include_priors return (; kwargs...) -> _loss(mws; kwargs...) + tel_prior(mws.om) + star_prior(mws.om) else @@ -92,6 +94,7 @@ Create loss functions for changing Used to fit scores efficiently with L-BFGS """ function loss_funcs_telstar(o::Output, om::OrderModel, d::Data) + println("loss_funcs_telstar() called") # DEBUG l_telstar(telstar; kwargs...) = _loss(o, om, d; tel=telstar[1], star=telstar[2], kwargs...) + tel_prior(telstar[1], om) + star_prior(telstar[2], om) @@ -135,6 +138,7 @@ Create loss functions for changing Used to fit models with ADAM """ function loss_funcs_total(o::Output, om::OrderModelDPCA, d::Data) + println("loss_funcs_total() called") # DEBUG l_total(total) = _loss_recalc_rv_basis(o, om, d; tel=total[1], star=total[2], rv=total[3]) + tel_prior(total[1], om) + star_prior(total[2], om) @@ -165,9 +169,17 @@ function loss_funcs_total(o::Output, om::OrderModelDPCA, d::Data) return l_total, l_total_s end function loss_funcs_total(o::Output, om::OrderModelWobble, d::Data) - l_total(total) = + println("loss_funcs_total() wobble called") # DEBUG + """l_total(total) = _loss(o, om, d; tel=total[1], star=total[2], rv=total[3]) + - tel_prior(total[1], om) + star_prior(total[2], om) + tel_prior(total[1], om) + star_prior(total[2], om)""" + function l_total(total) # DEBUG + p_loss = _loss(o, om, d; tel=total[1], star=total[2], rv=total[3]) + p_tel = tel_prior(total[1], om) + p_star = star_prior(total[2], om) + #println("LOSS: ", p_loss, " PTEL: ", p_tel, " PSTAR: ", p_star) + return p_loss + p_tel + p_star + end is_tel_time_variable = is_time_variable(om.tel) is_star_time_variable = is_time_variable(om.star) function l_total_s(total_s) @@ -699,7 +711,7 @@ function train_OrderModel!(mws::AdamWorkspace; ignore_regularization::Bool=false update_interpolation_locations!(mws) - # optionally ignore the regularization in `mws.om` + # optionally ignore the regularization in `mws.om` (except om.reg_model) if ignore_regularization reg_tel_holder = copy(mws.om.reg_tel) reg_star_holder = copy(mws.om.reg_star) @@ -960,7 +972,7 @@ function train_OrderModel!(ow::OptimTelStarWorkspace; verbose::Bool=_verbose_def optim_cb = optim_cb_f(; verbose=verbose) - # optionally ignore the regularization in `ow.om` + # optionally ignore the regularization in `ow.om` (except ow.om.reg_model) if ignore_regularization reg_tel_holder = copy(ow.om.reg_tel) reg_star_holder = copy(ow.om.reg_star) @@ -1017,7 +1029,7 @@ function train_OrderModel!(ow::OptimTotalWorkspace; verbose::Bool=_verbose_def, optim_cb = optim_cb_f(; verbose=verbose) - # optionally ignore the regularization in `ow.om` + # optionally ignore the regularization in `ow.om` (except ow.om.reg_model) if ignore_regularization reg_tel_holder = copy(ow.om.reg_tel) reg_star_holder = copy(ow.om.reg_star) @@ -1211,6 +1223,7 @@ Defaults to returning the AIC-minimum model - `max_n_tel::Int=5`: The maximum amount of telluric feature vectors to look for - `max_n_star::Int=5`: The maximum amount of stellar feature vectors to look for - `use_all_comps::Bool=false`: Whether to use all feature vectors, regardless of AIC +- `use_tel_prior::Bool=false`: Whether to use a modeled prior on telluric feature vectors - `careful_first_step::Bool=true`: Whether to shrink the learning rates until the loss improves on the first iteration - `speed_up::Bool=false`: Whether to inflate the learning rates until the loss is no longer improving throughout the optimization - `log_λ_gp_star::Real=1/SOAP_gp_params.λ`: The log λ lengthscale of the stellar regularization GP @@ -1221,7 +1234,8 @@ function calculate_initial_model(data::Data; instrument::String="None", desired_order::Int=0, star::String="None", times::AbstractVector=1:size(data.flux, 2), μ_min::Real=0, μ_max::Real=Inf, use_mean::Bool=true, stop_early::Bool=false, remove_reciprocal_continuum::Bool=false, return_full_path::Bool=false, - max_n_tel::Int=5, max_n_star::Int=5, use_all_comps::Bool=false, careful_first_step::Bool=true, speed_up::Bool=false, + max_n_tel::Int=5, max_n_star::Int=5, use_all_comps::Bool=false, use_tel_prior::Bool=false, + careful_first_step::Bool=true, speed_up::Bool=false, log_λ_gp_star::Real=1/SOAP_gp_params.λ, log_λ_gp_tel::Real=1/LSF_gp_params.λ, kwargs...) # TODO: Make this work for OrderModelDPCA @@ -1247,8 +1261,14 @@ function calculate_initial_model(data::Data; comp2ind(n_tel::Int, n_star::Int) = (n_tel+2, n_star+1) # converts number of components to storage matrix index n_obs = size(d.flux, 2) - om = OrderModel(d; instrument=instrument, order=desired_order, star_str=star, n_comp_tel=max_n_tel, n_comp_star=max_n_star, log_λ_gp_star=log_λ_gp_star, log_λ_gp_tel=log_λ_gp_tel, kwargs...) + om = OrderModel(d; instrument=instrument, order=desired_order, star_str=star, n_comp_tel=max_n_tel, n_comp_star=max_n_star, log_λ_gp_star=log_λ_gp_star, log_λ_gp_tel=log_λ_gp_tel, use_tel_prior=use_tel_prior, kwargs...) + # TEMPORARY - FIX TELLURIC PRIOR TO BASIS VECTORS 1 AND 2 + if use_tel_prior + om.metadata[:h2o_vector] = 1 + om.metadata[:oxy_vector] = 2 + end + # get the stellar model wavelengths in observed frame as a function of time star_log_λ_tel = _shift_log_λ_model(d.log_λ_obs, d.log_λ_star, om.star.log_λ) # get the telluric model wavelengths in stellar frame as a function of time @@ -1595,6 +1615,8 @@ function calculate_initial_model(data::Data; # flux_star .= _eval_lm(oms[i...].star.lm) interp_to_tel!(oms[i...])# .+ rv_to_D(oms[i...].rv)') # the rv is a small effect that we could just be getting wrong if n_tel_cur + 1 > 0 # if we are trying to add a feature vector + println(oms[i...].metadata) # DEBUG + println(oms[i...].reg_model) EMPCA.EMPCA!(oms[i...].tel.lm, flux_tel, 1 ./ vars_tel; inds=(n_tel_cur+1):(n_tel_cur+1)) else # if we are trying to add a template oms[i...].tel.lm.μ .= make_template(flux_tel, vars_tel; min=μ_min, max=μ_max, use_mean=use_mean) From 548d139440789912b7437c98a435e52b20431fef Mon Sep 17 00:00:00 2001 From: Kristo Ment Date: Mon, 14 Jul 2025 13:38:01 -0400 Subject: [PATCH 2/4] July 2025 update --- Project.toml | 2 + src/model_functions.jl | 157 ++++++++++++++++++++----- src/optimization_functions.jl | 199 +++++++++++++++++++++++--------- src/regularization_functions.jl | 39 ++++++- 4 files changed, 312 insertions(+), 85 deletions(-) diff --git a/Project.toml b/Project.toml index c67c274..df3ba6c 100644 --- a/Project.toml +++ b/Project.toml @@ -9,6 +9,8 @@ DSP = "717857b8-e6f2-59f4-9121-6e50c889abd2" DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ExpectationMaximizationPCA = "20521544-8895-46ad-ab44-5135df82bd8f" +FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb" +Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/src/model_functions.jl b/src/model_functions.jl index 5918531..00bc99e 100644 --- a/src/model_functions.jl +++ b/src/model_functions.jl @@ -575,9 +575,14 @@ default_reg_tel = Dict([(:GP_μ, 1e6), (:L2_μ, 1e6), (:L1_μ, 1e5), (:L1_μ₊_ (:GP_M, 1e7), (:L1_M, 1e7)]) default_reg_star = Dict([(:GP_μ, 1e2), (:L2_μ, 1e-2), (:L1_μ, 1e2), (:L1_μ₊_factor, 6.), (:GP_M, 1e4), (:L1_M, 1e7)]) -default_reg_model = Dict([(:L2_μ, 1e4), (:L2_h2o, 1e4), (:L2_oxy, 1e4)]) -#default_reg_model = Dict([(:L2_μ, 1e4),]) -#default_reg_model = Dict([(:L2_h2o, 1e4), (:L2_oxy, 1e4)]) +default_reg_model = Dict([(:L2_μ, 1e6), (:L2_h2o, 1e6), (:L2_oxy, 1e6)]) +#default_reg_model = Dict([(:L2_μ, 1e10), (:L2_h2o, 1e10), (:L2_oxy, 1e10), (:GP_μ_star, 1e10)]) +#default_reg_model = Dict([(:L2_μ, 1e10), (:GP_μ, 1e10)]) +#default_reg_model = Dict([(:L2_μ, 1e10),]) +#default_reg_model = Dict([(:L2_μ, 1e10), (:GP_μ, 1e10), (:GP_h2o, 1e10), (:GP_oxy, 1e10)]) +#default_reg_model = Dict([(:L2_μ, 1e10), (:L2_h2o, 1e10), (:L2_oxy, 1e10), (:GP_μ, 1e10)]) +#default_reg_model = Dict([(:L2_μ, 1e10), (:L2_h2o, 1e10), (:L2_oxy, 1e10), (:GP_μ, 1e10), (:GP_h2o, 1e10), (:GP_oxy, 1e10)]) +#default_reg_model = Dict([(:L2_h2o, 1e10), (:L2_oxy, 1e10)]) default_reg_tel_full = Dict([(:GP_μ, 1e6), (:L2_μ, 1e6), (:L1_μ, 1e5), (:L1_μ₊_factor, 6.), (:GP_M, 1e7), (:L2_M, 1e4), (:L1_M, 1e7)]) default_reg_star_full = Dict([(:GP_μ, 1e2), (:L2_μ, 1e-2), (:L1_μ, 1e1), @@ -731,7 +736,12 @@ Constructor for the `OrderModel`-type objects (SSOF model for a set of 1D spectr - `log_λ_gp_tel::Real=1/LSF_gp_params.λ`: The log λ lengthscale of the telluric regularization GP - `tel_log_λ::Union{Nothing,AbstractRange}=nothing`: The log - `star_log_λ::Union{Nothing,AbstractRange}=nothing`: The log λ lengthscale of the telluric regularization GP -- `use_tel_prior`::Bool=true`: Whether to include a modeled prior in the telluric submodel +- `use_tel_prior::Bool=false`: Whether to include a modeled prior in the telluric submodel +- `reg_star::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Overrides default_reg_star values if given +- `reg_tel::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Overrides default_reg_tel values if given +- `reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Overrides default_reg_model values if given +- `tel_below_one::Bool=false`: Whether to impose a heavy penalty for any telluric vector flux going above 1 +- `fixed_vectors::Dict=Dict()`: Fix any stellar or telluric vectors. Must contain 'log_λ' - `kwargs...`: kwargs passed to `Submodel` constructors """ function OrderModel( @@ -747,9 +757,24 @@ function OrderModel( log_λ_gp_tel::Real=1/LSF_gp_params.λ, tel_log_λ::Union{Nothing,AbstractRange}=nothing, star_log_λ::Union{Nothing,AbstractRange}=nothing, - use_tel_prior::Bool=true, + use_tel_prior::Bool=false, + reg_star::Union{Nothing,Dict{Symbol,Float64}}=nothing, + reg_tel::Union{Nothing,Dict{Symbol,Float64}}=nothing, + reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing, + tel_below_one::Bool=false, + fixed_vectors::Dict=Dict(), kwargs...) + if isnothing(reg_star) + reg_star = default_reg_star + end + if isnothing(reg_tel) + reg_tel = default_reg_tel + end + if isnothing(reg_model) + reg_model = default_reg_model + end + # Creating models tel = Submodel(d.log_λ_obs, n_comp_tel, log_λ_gp_tel; log_λ=tel_log_λ, kwargs...) star = Submodel(d.log_λ_star, n_comp_star, log_λ_gp_star; log_λ=star_log_λ, kwargs...) @@ -760,7 +785,14 @@ function OrderModel( bary_rvs = D_to_rv.([median(d.log_λ_star[:, i] - d.log_λ_obs[:, i]) for i in 1:n_obs]) todo = Dict([(:initialized, false), (:reg_improved, false), (:err_estimated, false), (:h2o_tried, false), (:oxy_tried, false)]) - metadata = Dict([(:todo, todo), (:tel_prior, use_tel_prior), (:instrument, instrument), (:order, order), (:star, star_str), (:h2o_vector, 0), (:oxy_vector, 0)]) + metadata = Dict([(:todo, todo), (:tel_prior, use_tel_prior), (:tel_below_one, tel_below_one), (:instrument, instrument), (:order, order), + (:star, star_str), (:h2o_vector, 0), (:oxy_vector, 0), (:fixed_vectors, fixed_vectors)]) + # Pre-compute prior vectors + if use_tel_prior + metadata[:tel_prior_spec] = get_telluric_prior_spectrum(tel.λ, :all) + metadata[:h2o_prior_spec] = get_telluric_prior_spectrum(tel.λ, :h2o) + metadata[:oxy_prior_spec] = get_telluric_prior_spectrum(tel.λ, :oxy) + end if dpca if oversamp b2o = oversamp_interp_helper(d.log_λ_star_bounds, star.log_λ) @@ -769,7 +801,7 @@ function OrderModel( b2o = undersamp_interp_helper(d.log_λ_star, star.log_λ) t2o = undersamp_interp_helper(d.log_λ_obs, tel.log_λ) end - return OrderModelDPCA(tel, star, rv, copy(default_reg_tel), copy(default_reg_star), copy(default_reg_model), b2o, t2o, metadata, n_obs) + return OrderModelDPCA(tel, star, rv, copy(reg_tel), copy(reg_star), copy(reg_model), b2o, t2o, metadata, n_obs) else b2o = StellarInterpolationHelper(star.log_λ, bary_rvs, d.log_λ_obs) if oversamp @@ -777,7 +809,7 @@ function OrderModel( else t2o = undersamp_interp_helper(d.log_λ_obs, tel.log_λ) end - return OrderModelWobble(tel, star, rv, copy(default_reg_tel), copy(default_reg_star), copy(default_reg_model), b2o, bary_rvs, t2o, metadata, n_obs) + return OrderModelWobble(tel, star, rv, copy(reg_tel), copy(reg_star), copy(reg_model), b2o, bary_rvs, t2o, metadata, n_obs) end end Base.copy(om::OrderModelDPCA) = OrderModelDPCA(copy(om.tel), copy(om.star), copy(om.rv), copy(om.reg_tel), copy(om.reg_star), copy(om.reg_model), @@ -1257,15 +1289,18 @@ end """ model_prior(lm, om, key) -Calculate the model prior on `lm` with the regularization terms in `om.reg_` * `key` +Calculate the model prior on `lm` with the regularization terms in `om.reg_` * `key`. +Can optionally exclude a number of the first basis vectors. """ -function model_prior(lm, om::OrderModel, key::Symbol) +function model_prior(lm, om::OrderModel, key::Symbol; exclude_vectors::Int64 = 0) + reg = getfield(om, Symbol(:reg_, key)) sm = getfield(om, key) isFullLinearModel = length(lm) > 2 val = 0. if haskey(reg, :GP_μ) || haskey(reg, :L2_μ) || haskey(reg, :L1_μ) || haskey(reg, :L1_μ₊_factor) + μ_mod = lm[1+2*isFullLinearModel] .- 1 if haskey(reg, :L2_μ); val += L2(μ_mod) * reg[:L2_μ] end if haskey(reg, :L1_μ) @@ -1276,19 +1311,27 @@ function model_prior(lm, om::OrderModel, key::Symbol) # if haskey(reg, :GP_μ); val -= logpdf(SOAP_gp(getfield(om, key).log_λ), μ_mod) * reg[:GP_μ] end # if haskey(reg, :GP_μ); val -= gp_ℓ_nabla(μ_mod, sm.A_sde, sm.Σ_sde) * reg[:GP_μ] end if haskey(reg, :GP_μ); val -= gp_ℓ_precalc(sm.Δℓ_coeff, μ_mod, sm.A_sde, sm.Σ_sde) * reg[:GP_μ] end + end - if isFullLinearModel - if haskey(reg, :shared_M); val += shared_attention(lm[1]) * reg[:shared_M] end - if haskey(reg, :L2_M); val += L2(lm[1]) * reg[:L2_M] end - if haskey(reg, :L1_M); val += L1(lm[1]) * reg[:L1_M] end + + if isFullLinearModel && (size(lm[1], 2) > exclude_vectors) + + lm_M = (exclude_vectors > 0 ? lm[1][:, (1+exclude_vectors):size(lm[1], 2)] : lm[1]) + lm_s = (exclude_vectors > 0 ? lm[2][(1+exclude_vectors):size(lm[2], 1), :] : lm[2]) + + if haskey(reg, :shared_M); val += shared_attention(lm_M) * reg[:shared_M] end + if haskey(reg, :L2_M); val += L2(lm_M) * reg[:L2_M] end + if haskey(reg, :L1_M); val += L1(lm_M) * reg[:L1_M] end # if haskey(reg, :GP_μ); val -= gp_ℓ_precalc(sm.Δℓ_coeff, view(lm[1], :, 1), sm.A_sde, sm.Σ_sde) * reg[:GP_μ] end if haskey(reg, :GP_M) - for i in 1:size(lm[1], 2) - val -= gp_ℓ_precalc(sm.Δℓ_coeff, lm[1][:, i], sm.A_sde, sm.Σ_sde) * reg[:GP_M] + for i in 1:size(lm_M, 2) + val -= gp_ℓ_precalc(sm.Δℓ_coeff, lm_M[:, i], sm.A_sde, sm.Σ_sde) * reg[:GP_M] end end - val += model_s_prior(lm[2], reg) + val += model_s_prior(lm_s, reg) + end + return val end model_prior(lm::Union{FullLinearModel, TemplateModel}, om::OrderModel, key::Symbol) = model_prior(vec(lm), om, key) @@ -1316,20 +1359,60 @@ include("/home/kment/RV/telluricfunc.jl") Calculate the telluric model prior on `om.tel.lm` with the regularization terms in `om.reg_model` """ function tel_prior(lm, om::OrderModel) + + if om.metadata[:tel_prior] + + exclude_vectors = (om.metadata[:h2o_vector] > 0) + (om.metadata[:oxy_vector] > 0) + val = model_prior(lm, om, :tel; exclude_vectors=exclude_vectors) - val = model_prior(lm, om, :tel) + if haskey(om.reg_model, :L1_μ) + val += L1(om.metadata[:tel_prior_spec] - lm[min(3, length(lm))]) * om.reg_model[:L1_μ] + end - if om.metadata[:tel_prior] && haskey(om.reg_model, :L2_μ) - val += L2(get_telluric_prior_spectrum(om.tel.λ) - lm[min(3, length(lm))]) * om.reg_model[:L2_μ] - end + if haskey(om.reg_model, :L2_μ) + val += L2(om.metadata[:tel_prior_spec] - lm[min(3, length(lm))]) * om.reg_model[:L2_μ] + end - if om.metadata[:tel_prior] && length(lm) >= 3 - if om.metadata[:h2o_vector] >= 1 && size(lm[1], 2) >= om.metadata[:h2o_vector] && haskey(om.reg_model, :L2_h2o) - val += L2(get_telluric_prior_spectrum(om.tel.λ, :h2o) - lm[1][:,om.metadata[:h2o_vector]] .- 1) * om.reg_model[:L2_h2o] + if haskey(om.reg_model, :GP_μ) + μ_mod = lm[min(3, length(lm))] - om.metadata[:tel_prior_spec] + val -= gp_ℓ_precalc(om.tel.Δℓ_coeff, μ_mod, om.tel.A_sde, om.tel.Σ_sde) * om.reg_model[:GP_μ] end - if om.metadata[:oxy_vector] >= 1 && size(lm[1], 2) >= om.metadata[:oxy_vector] && haskey(om.reg_model, :L2_oxy) - val += L2(get_telluric_prior_spectrum(om.tel.λ, :oxy) - lm[1][:,om.metadata[:oxy_vector]] .- 1) * om.reg_model[:L2_oxy] + + if length(lm) >= 3 + if om.metadata[:h2o_vector] >= 1 && size(lm[1], 2) >= om.metadata[:h2o_vector] + if haskey(om.reg_model, :L1_h2o) + val += L1(log.(om.metadata[:h2o_prior_spec]) - lm[1][:,om.metadata[:h2o_vector]]) * om.reg_model[:L1_h2o] + end + if haskey(om.reg_model, :L2_h2o) + val += L2(log.(om.metadata[:h2o_prior_spec]) - lm[1][:,om.metadata[:h2o_vector]]) * om.reg_model[:L2_h2o] + end + if haskey(om.reg_model, :GP_h2o) + M_mod = lm[1][:,om.metadata[:h2o_vector]] - log.(om.metadata[:h2o_prior_spec]) + val -= gp_ℓ_precalc(om.tel.Δℓ_coeff, M_mod, om.tel.A_sde, om.tel.Σ_sde) * om.reg_model[:GP_h2o] + end + end + if om.metadata[:oxy_vector] >= 1 && size(lm[1], 2) >= om.metadata[:oxy_vector] + if haskey(om.reg_model, :L1_oxy) + val += L1(log.(om.metadata[:oxy_prior_spec]) - lm[1][:,om.metadata[:oxy_vector]]) * om.reg_model[:L1_oxy] + end + if haskey(om.reg_model, :L2_oxy) + val += L2(log.(om.metadata[:oxy_prior_spec]) - lm[1][:,om.metadata[:oxy_vector]]) * om.reg_model[:L2_oxy] + end + if haskey(om.reg_model, :GP_oxy) + M_mod = lm[1][:,om.metadata[:oxy_vector]] - log.(om.metadata[:oxy_prior_spec]) + val -= gp_ℓ_precalc(om.tel.Δℓ_coeff, M_mod, om.tel.A_sde, om.tel.Σ_sde) * om.reg_model[:GP_oxy] + end + end end + + else + + val = model_prior(lm, om, :tel) + + end + + if om.metadata[:tel_below_one] + val += sum(exp.((lm[min(3, length(lm))] .- 1) .* (lm[min(3, length(lm))] .> 1)) .- 1) * 1e10 end return val @@ -1343,8 +1426,28 @@ tel_prior(om::OrderModel) = tel_prior(om.tel.lm, om) Calculate the stellar model prior on `om.star.lm` with the regularization terms in `om.star_tel` """ +function star_prior(lm, om::OrderModel) + + val = model_prior(lm, om, :star) + + if om.metadata[:tel_prior] && haskey(om.reg_model, :L2_μ_star) && haskey(om.metadata[:fixed_vectors], "log_λ") && haskey(om.metadata[:fixed_vectors], "star_μ") + _flux_star_interp = CubicSpline(om.metadata[:fixed_vectors]["star_μ"], om.metadata[:fixed_vectors]["log_λ"]; extrapolate=true) + flux_star_interp = _flux_star_interp(om.star.log_λ) + flux_star_interp[om.star.log_λ .< minimum(om.metadata[:fixed_vectors]["log_λ"])] .= 1 + flux_star_interp[om.star.log_λ .> maximum(om.metadata[:fixed_vectors]["log_λ"])] .= 1 + val += L2(flux_star_interp - lm[min(3, length(lm))]) * om.reg_model[:L2_μ_star] + end + + if om.metadata[:tel_prior] && haskey(om.reg_model, :GP_μ_star) + μ_mod = lm[min(3, length(lm))] .- 1 + val -= gp_ℓ_precalc(om.star.Δℓ_coeff, μ_mod, om.star.A_sde, om.star.Σ_sde) * om.reg_model[:GP_μ_star] + end + + return val + +end star_prior(om::OrderModel) = star_prior(om.star.lm, om) -star_prior(lm, om::OrderModel) = model_prior(lm, om, :star) +#star_prior(lm, om::OrderModel) = model_prior(lm, om, :star) """ diff --git a/src/optimization_functions.jl b/src/optimization_functions.jl index 7b623c4..e49c9e1 100644 --- a/src/optimization_functions.jl +++ b/src/optimization_functions.jl @@ -74,7 +74,6 @@ _loss_recalc_rv_basis(mws::ModelWorkspace; kwargs...) = _loss_recalc_rv_basis(mw Create a loss function for the model and data in `mws` """ function loss_func(mws::ModelWorkspace; include_priors::Bool=false) - println("loss_func() called with include_priors=", include_priors) # DEBUG if include_priors return (; kwargs...) -> _loss(mws; kwargs...) + tel_prior(mws.om) + star_prior(mws.om) else @@ -94,7 +93,6 @@ Create loss functions for changing Used to fit scores efficiently with L-BFGS """ function loss_funcs_telstar(o::Output, om::OrderModel, d::Data) - println("loss_funcs_telstar() called") # DEBUG l_telstar(telstar; kwargs...) = _loss(o, om, d; tel=telstar[1], star=telstar[2], kwargs...) + tel_prior(telstar[1], om) + star_prior(telstar[2], om) @@ -138,7 +136,6 @@ Create loss functions for changing Used to fit models with ADAM """ function loss_funcs_total(o::Output, om::OrderModelDPCA, d::Data) - println("loss_funcs_total() called") # DEBUG l_total(total) = _loss_recalc_rv_basis(o, om, d; tel=total[1], star=total[2], rv=total[3]) + tel_prior(total[1], om) + star_prior(total[2], om) @@ -169,17 +166,9 @@ function loss_funcs_total(o::Output, om::OrderModelDPCA, d::Data) return l_total, l_total_s end function loss_funcs_total(o::Output, om::OrderModelWobble, d::Data) - println("loss_funcs_total() wobble called") # DEBUG - """l_total(total) = + l_total(total) = _loss(o, om, d; tel=total[1], star=total[2], rv=total[3]) + - tel_prior(total[1], om) + star_prior(total[2], om)""" - function l_total(total) # DEBUG - p_loss = _loss(o, om, d; tel=total[1], star=total[2], rv=total[3]) - p_tel = tel_prior(total[1], om) - p_star = star_prior(total[2], om) - #println("LOSS: ", p_loss, " PTEL: ", p_tel, " PSTAR: ", p_star) - return p_loss + p_tel + p_star - end + tel_prior(total[1], om) + star_prior(total[2], om) is_tel_time_variable = is_time_variable(om.tel) is_star_time_variable = is_time_variable(om.star) function l_total_s(total_s) @@ -722,14 +711,15 @@ function train_OrderModel!(mws::AdamWorkspace; ignore_regularization::Bool=false function cb(as::AdamState) # optionally shift the score means to be near 0 - if shift_scores - if !(typeof(mws) <: FrozenTelWorkspace) - remove_lm_score_means!(mws.om.tel.lm; prop=0.2) - end - if typeof(mws.om) <: OrderModelWobble - remove_lm_score_means!(mws.om.star.lm; prop=0.2) - end - end + # DEBUG this has been turned off to test simulated telluric models + #if shift_scores + # if !(typeof(mws) <: FrozenTelWorkspace) + # remove_lm_score_means!(mws.om.tel.lm; prop=0.2) + # end + # if typeof(mws.om) <: OrderModelWobble + # remove_lm_score_means!(mws.om.star.lm; prop=0.2) + # end + #end # optionally make the templates always positive if μ_positive @@ -1228,6 +1218,10 @@ Defaults to returning the AIC-minimum model - `speed_up::Bool=false`: Whether to inflate the learning rates until the loss is no longer improving throughout the optimization - `log_λ_gp_star::Real=1/SOAP_gp_params.λ`: The log λ lengthscale of the stellar regularization GP - `log_λ_gp_tel::Real=1/LSF_gp_params.λ`: The log λ lengthscale of the telluric regularization GP +- `reg_star::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Override default_reg_star if given +- `reg_tel::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Override default_reg_tel if given +- `reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Override default_reg_model if given +- `fixed_vectors::Dict=Dict()`: Fix any stellar or telluric vectors. Must contain 'log_λ' - `kwargs...`: kwargs passed to `OrderModel` constructor """ function calculate_initial_model(data::Data; @@ -1236,13 +1230,21 @@ function calculate_initial_model(data::Data; remove_reciprocal_continuum::Bool=false, return_full_path::Bool=false, max_n_tel::Int=5, max_n_star::Int=5, use_all_comps::Bool=false, use_tel_prior::Bool=false, careful_first_step::Bool=true, speed_up::Bool=false, - log_λ_gp_star::Real=1/SOAP_gp_params.λ, log_λ_gp_tel::Real=1/LSF_gp_params.λ, kwargs...) + log_λ_gp_star::Real=1/SOAP_gp_params.λ, log_λ_gp_tel::Real=1/LSF_gp_params.λ, + reg_star::Union{Nothing,Dict{Symbol,Float64}}=nothing, reg_tel::Union{Nothing,Dict{Symbol,Float64}}=nothing, + reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing, fixed_vectors::Dict=Dict(), kwargs...) # TODO: Make this work for OrderModelDPCA # Get non-LSF version of `data` d = GenericData(data) - @assert max_n_tel >= -1 + # Determine the types of telluric priors used (if any) + use_tel_prior_μ = use_tel_prior && !isnothing(reg_model) && haskey(reg_model, :L2_μ) + use_tel_prior_h2o = use_tel_prior && !isnothing(reg_model) && haskey(reg_model, :L2_h2o) + use_tel_prior_oxy = use_tel_prior && !isnothing(reg_model) && haskey(reg_model, :L2_oxy) + #use_tel_full = !isnothing(reg_model) && haskey(reg_model, :L2_μ) && haskey(reg_model, :L2_h2o) && haskey(reg_model, :L2_oxy) + + @assert max_n_tel >= -1 + use_tel_prior_μ + use_tel_prior_h2o + use_tel_prior_oxy @assert max_n_star >= 0 # which amounts of feature vectors to test @@ -1261,12 +1263,16 @@ function calculate_initial_model(data::Data; comp2ind(n_tel::Int, n_star::Int) = (n_tel+2, n_star+1) # converts number of components to storage matrix index n_obs = size(d.flux, 2) - om = OrderModel(d; instrument=instrument, order=desired_order, star_str=star, n_comp_tel=max_n_tel, n_comp_star=max_n_star, log_λ_gp_star=log_λ_gp_star, log_λ_gp_tel=log_λ_gp_tel, use_tel_prior=use_tel_prior, kwargs...) + om = OrderModel(d; instrument=instrument, order=desired_order, star_str=star, n_comp_tel=max_n_tel, n_comp_star=max_n_star, + log_λ_gp_star=log_λ_gp_star, log_λ_gp_tel=log_λ_gp_tel, use_tel_prior=use_tel_prior, reg_star=reg_star, + reg_tel=reg_tel, reg_model=reg_model, fixed_vectors=fixed_vectors, kwargs...) - # TEMPORARY - FIX TELLURIC PRIOR TO BASIS VECTORS 1 AND 2 - if use_tel_prior + # Add telluric priors to specific vectors + if use_tel_prior_h2o om.metadata[:h2o_vector] = 1 - om.metadata[:oxy_vector] = 2 + end + if use_tel_prior_oxy + om.metadata[:oxy_vector] = (use_tel_prior_h2o ? 2 : 1) end # get the stellar model wavelengths in observed frame as a function of time @@ -1357,16 +1363,18 @@ function calculate_initial_model(data::Data; # remove the score means and flip the feature vectors function nicer_model!(mws::ModelWorkspace) - remove_lm_score_means!(mws.om) - flip_feature_vectors!(mws.om) + if !use_tel_prior + remove_lm_score_means!(mws.om) + flip_feature_vectors!(mws.om) + end mws.om.metadata[:todo][:initialized] = true mws.om.metadata[:todo][:downsized] = true # copy_dict!(mws.om.reg_tel, default_reg_tel) # copy_dict!(mws.om.reg_star, default_reg_star) end - # - function get_metrics!(mws::ModelWorkspace, i::Int, j::Int) + # If only_s then optimizes scores only + function get_metrics!(mws::ModelWorkspace, i::Int, j::Int; only_s::Bool=false) # # could set very small regularizations beforehand if we wanted # for (k, v) in mws.om.reg_tel @@ -1378,13 +1386,17 @@ function calculate_initial_model(data::Data; try - # improve the model - improve_initial_model!(mws; careful_first_step=careful_first_step, speed_up=speed_up, iter=50) + if only_s + finalize_scores!(mws; iter=50) + else + # improve the model + improve_initial_model!(mws; careful_first_step=careful_first_step, speed_up=speed_up, iter=50) - # if there is an LSF, do some more fitting - if mws.d != data - mws = typeof(mws)(copy(mws.om), data) - improve_initial_model!(mws; careful_first_step=careful_first_step, speed_up=speed_up, iter=30) + # if there is an LSF, do some more fitting + if mws.d != data + mws = typeof(mws)(copy(mws.om), data) + improve_initial_model!(mws; careful_first_step=careful_first_step, speed_up=speed_up, iter=30) + end end nicer_model!(mws) @@ -1410,6 +1422,9 @@ function calculate_initial_model(data::Data; end end + # Whether the stellar mean vector stays fixed + μ_star_fixed = haskey(fixed_vectors, "log_λ") && haskey(fixed_vectors, "star_μ") + # stellar template model assuming no tellurics n_tel_cur = -1 n_star_cur = 0 @@ -1417,10 +1432,22 @@ function calculate_initial_model(data::Data; search_new_star = n_star_cur+1 <= max_n_star oms[1,1] = downsize(om, 0, 0) oms[1,1].tel.lm.μ .= 1 - _spectra_interp_gp!(flux_star, vars_star, oms[1,1].star.log_λ, d.flux, d.var .+ SOAP_gp_var, d.log_λ_star; gp_mean=1., λ_kernel=1/log_λ_gp_star) + if μ_star_fixed + println("Using fixed stellar mean vector.") + _flux_star_interp = CubicSpline(fixed_vectors["star_μ"], fixed_vectors["log_λ"]; extrapolate=true) + flux_star_interp = _flux_star_interp(oms[1,1].star.log_λ) + flux_star_interp[oms[1,1].star.log_λ .< minimum(fixed_vectors["log_λ"])] .= 1 + flux_star_interp[oms[1,1].star.log_λ .> maximum(fixed_vectors["log_λ"])] .= 1 + oms[1,1].star.lm.μ[:] = flux_star_interp + flux_star .*= flux_star_interp + #oms[1,1].star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) + println(sum(oms[1,1].star.lm.μ)) + else + _spectra_interp_gp!(flux_star, vars_star, oms[1,1].star.log_λ, d.flux, d.var .+ SOAP_gp_var, d.log_λ_star; gp_mean=1., λ_kernel=1/log_λ_gp_star) + oms[1,1].star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) + end flux_star_no_tel = copy(flux_star) vars_star_no_tel = copy(vars_star) - oms[1,1].star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) # how good is the stellar template at modeling each pixel dop_comp = doppler_component(oms[1,1].star.λ, oms[1,1].star.lm.μ) @@ -1432,7 +1459,7 @@ function calculate_initial_model(data::Data; # get aic for base, only stellar template model mws = FrozenTelWorkspace(oms[1,1], d) - om_cur = get_metrics!(mws, 1, 1) + om_cur = get_metrics!(mws, 1, 1; only_s=μ_star_fixed) # telluric template assuming no stellar (will be overwritten later) _om = downsize(om, 0, 0) @@ -1494,9 +1521,68 @@ function calculate_initial_model(data::Data; # interp_to_tel!(; kwargs...) = interp_helper!(flux_tel, vars_tel, om.tel.log_λ, # flux_star, vars_star, star_log_λ_tel, # d.log_λ_obs; kwargs...) + + # Make a telluric template after dividing out a partial stellar template + function make_telluric_template!(om::OrderModel) + + use_tel = χ²_star .> χ²_tel # which pixels are telluric dominated + + # mask out the telluric dominated pixels for partial stellar template estimation + _var = copy(d.var) + _var[use_tel, :] .= Inf + + # get stellar template in portions of spectra where it is dominant + _spectra_interp_gp!(flux_star, vars_star, om.star.log_λ, d.flux, _var .+ SOAP_gp_var, d.log_λ_star; gp_mean=1., λ_kernel=1/log_λ_gp_star) + om.star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) + + # get telluric template after dividing out the partial stellar template + # flux_star .= oms[2,1].star.lm.μ + interp_to_tel!(om; mask_extrema=false) + om.tel.lm.μ[:] = make_template(flux_tel, vars_tel; min=μ_min, max=μ_max, use_mean=use_mean) + + end + + # If the telluric prior is turned on, add the respective mean and/or vectors (all added simultaneously) + if use_tel_prior + + n_tel_next = use_tel_prior_h2o + use_tel_prior_oxy + i = comp2ind(n_tel_next, n_star_cur) + oms[i...] = downsize(om, n_tel_next, n_star_cur) + + if use_tel_prior_μ + oms[i...].tel.lm.μ[:] = get_telluric_prior_spectrum(oms[i...].tel.λ) + else + make_telluric_template!(oms[i...]) + end + + if use_tel_prior_h2o + oms[i...].tel.lm.M[:,1] = log.(get_telluric_prior_spectrum(oms[i...].tel.λ, :h2o)) + end + + if use_tel_prior_oxy + oms[i...].tel.lm.M[:,1+use_tel_prior_h2o] = log.(get_telluric_prior_spectrum(oms[i...].tel.λ, :oxy)) + end + + # Optimize scores first (with no stellar spectrum) if priors are on any vectors + if n_tel_next > 0 + oms[i...].star.lm.μ .= 1 + _mws = TotalWorkspace(oms[i...], d; only_s=true) + finalize_scores!(_mws; iter=50) + end + + # get stellar template after dividing out full telluric template + interp_to_star!(oms[i...]; mask_extrema=false) + oms[i...].star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) + + mws = TotalWorkspace(oms[i...], d) + # flux_tel .= oms[2,1].tel.lm.μ + # interp_to_star!(; mask_extrema=false) + # mws.om.rv .= vec(project_doppler_comp(flux_star .- mws.om.star.lm.μ, doppler_component(mws.om.star.λ, mws.om.star.lm.μ), 1 ./ vars_star)) + om_add_tel = get_metrics!(mws, i...) + added_tel_better = true # if one wants to search for a telluric template - if search_new_tel + elseif search_new_tel oms[2,1] = downsize(om, 0, 0) oms[2,1].star.lm.μ .= 1 @@ -1522,21 +1608,22 @@ function calculate_initial_model(data::Data; # end # end - - if sum(use_tel) > 0 # if any pixels are telluric dominated + if sum(use_tel) > 0 && !μ_star_fixed # if any pixels are telluric dominated (unless μ_star is fixed) # mask out the telluric dominated pixels for partial stellar template estimation - _var = copy(d.var) - _var[use_tel, :] .= Inf + #_var = copy(d.var) + #_var[use_tel, :] .= Inf # get stellar template in portions of spectra where it is dominant - _spectra_interp_gp!(flux_star, vars_star, oms[2,1].star.log_λ, d.flux, _var .+ SOAP_gp_var, d.log_λ_star; gp_mean=1., λ_kernel=1/log_λ_gp_star) - oms[2,1].star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) + #_spectra_interp_gp!(flux_star, vars_star, oms[2,1].star.log_λ, d.flux, _var .+ SOAP_gp_var, d.log_λ_star; gp_mean=1., λ_kernel=1/log_λ_gp_star) + #oms[2,1].star.lm.μ[:] = make_template(flux_star, vars_star; min=μ_min, max=μ_max, use_mean=use_mean) # get telluric template after dividing out the partial stellar template # flux_star .= oms[2,1].star.lm.μ - interp_to_tel!(oms[2,1]; mask_extrema=false) - oms[2,1].tel.lm.μ[:] = make_template(flux_tel, vars_tel; min=μ_min, max=μ_max, use_mean=use_mean) + #interp_to_tel!(oms[2,1]; mask_extrema=false) + #oms[2,1].tel.lm.μ[:] = make_template(flux_tel, vars_tel; min=μ_min, max=μ_max, use_mean=use_mean) + + make_telluric_template!(oms[2,1]) # get stellar template after dividing out full telluric template # flux_tel .= oms[2,1].tel.lm.μ @@ -1564,23 +1651,27 @@ function calculate_initial_model(data::Data; # mws.om.rv .= vec(project_doppler_comp(flux_star .- mws.om.star.lm.μ, doppler_component(mws.om.star.λ, mws.om.star.lm.μ), 1 ./ vars_star)) om_add_tel = get_metrics!(mws, 2, 1) + # Was the model with a telluric template better than the current model (just a stellar template)? + added_tel_better = aics[comp2ind(n_tel_cur+1, n_star_cur)...] < aics[comp2ind(n_tel_cur, n_star_cur)...] + else om_add_tel = om_cur + added_tel_better = false end j = comp2ind(n_tel_cur, n_star_cur) # if we looked for a telluric template, get the aic - search_new_tel ? aic_tel = aics[comp2ind(n_tel_cur+1, n_star_cur)...] : aic_tel = Inf + #search_new_tel ? aic_tel = aics[comp2ind(n_tel_cur+1, n_star_cur)...] : aic_tel = Inf # was the model with a telluric template better than the current model (just a stellar template)? - added_tel_better = aic_tel < aics[j...] + #added_tel_better = aic_tel < aics[j...] # if including the telluric template model helped, use it going forward if added_tel_better; oms[j...] = om_cur end n_star_next = n_star_cur - n_tel_next = n_tel_cur+added_tel_better - added_tel_better ? aic_next = aic_tel : aic_next = aics[j...] + n_tel_next = n_tel_cur + added_tel_better + use_tel_prior_h2o + use_tel_prior_oxy + #added_tel_better ? aic_next = aic_tel : aic_next = aics[j...] add_comp = true println("looking for time variability...") @@ -1615,8 +1706,6 @@ function calculate_initial_model(data::Data; # flux_star .= _eval_lm(oms[i...].star.lm) interp_to_tel!(oms[i...])# .+ rv_to_D(oms[i...].rv)') # the rv is a small effect that we could just be getting wrong if n_tel_cur + 1 > 0 # if we are trying to add a feature vector - println(oms[i...].metadata) # DEBUG - println(oms[i...].reg_model) EMPCA.EMPCA!(oms[i...].tel.lm, flux_tel, 1 ./ vars_tel; inds=(n_tel_cur+1):(n_tel_cur+1)) else # if we are trying to add a template oms[i...].tel.lm.μ .= make_template(flux_tel, vars_tel; min=μ_min, max=μ_max, use_mean=use_mean) diff --git a/src/regularization_functions.jl b/src/regularization_functions.jl index 3c6aa95..1761abc 100644 --- a/src/regularization_functions.jl +++ b/src/regularization_functions.jl @@ -39,6 +39,7 @@ function fit_regularization_helper!(reg_fields::Vector{Symbol}, reg_key::Symbol, # only do anything if if haskey(getfield(mws.om, reg_fields[1]), reg_key) + println("Regularization of $reg_key in $(reg_fields[1])") # DEBUG om = mws.om @assert 0 < reg_min < reg_max < Inf ℓs = Array{Float64}(undef, 2) @@ -166,6 +167,7 @@ end _key_list = [:GP_μ, :L2_μ, :L1_μ, :L1_μ₊_factor, :GP_M, :L2_M, :L1_M, :shared_M] _key_list_fit = [:GP_μ, :L2_μ, :L1_μ, :GP_M, :L2_M, :L1_M] _key_list_bases = [:GP_M, :L2_M, :L1_M, :shared_M] +_key_list_model = [:GP_μ, :L2_μ, :L1_μ, :L2_h2o, :L1_h2o, :L2_oxy, :L1_oxy] """ @@ -179,8 +181,10 @@ function check_for_valid_regularization(reg::Dict{Symbol, <:Real}) end end -min_reg = 1e-3 -max_reg = 1e12 +min_reg = 1e-1 #1e-3 +max_reg = 1e7 #1e12 +min_reg_model = 1e0 +max_reg_model = 1e12 """ @@ -223,13 +227,39 @@ function fit_regularization!(mws::ModelWorkspace, testing_inds::AbstractVecOrMat end +""" + fit_regularization_model!(mws, testing_inds; key_list=_key_list_model, kwargs...) + +Fit all of the regularization values in `key_list` for `mws.om.reg_model` +""" +function fit_regularization_model!(mws::ModelWorkspace, testing_inds::AbstractVecOrMat; key_list::Vector{Symbol}=_key_list_model, kwargs...) + om = mws.om + n_obs = size(mws.d.flux, 2) + training_inds = [i for i in 1:n_obs if !(i in testing_inds)] + hold_reg = copy(om.reg_model) + for (key, value) in om.reg_model + om.reg_model[key] = 0 + end + println("starting regularization searches for mws.om.reg_model") + before_ℓ = _eval_regularization(copy(mws.om), mws, training_inds, testing_inds) + println("initial training χ²: $before_ℓ") + test_factor, reg_min, reg_max = 10, min_reg_model, max_reg_model + for key in key_list + if haskey(hold_reg, key) + before_ℓ = fit_regularization_helper!([:reg_model], key, before_ℓ, mws, training_inds, testing_inds, test_factor, reg_min, reg_max; start=hold_reg[key], kwargs...) + end + end +end + + """ fit_regularization!(mws; verbose=true, testing_ratio=0.33, careful_first_step=true, speed_up=false, kwargs...) Find the best fit model without regularization then fit all of the regularization values in `key_list` for the model in `mws` """ -function fit_regularization!(mws::ModelWorkspace; verbose::Bool=true, testing_ratio::Real=0.33, careful_first_step::Bool=true, speed_up::Bool=false, kwargs...) +function fit_regularization!(mws::ModelWorkspace; verbose::Bool=true, testing_ratio::Real=0.33, careful_first_step::Bool=true, speed_up::Bool=false, + fit_reg_model::Bool=true, kwargs...) # if mws.om.metadata[:todo][:reg_improved] n_obs = size(mws.d.flux, 2) train_OrderModel!(mws; verbose=verbose, ignore_regularization=true, careful_first_step=careful_first_step, speed_up=speed_up) @@ -237,6 +267,9 @@ function fit_regularization!(mws::ModelWorkspace; verbose::Bool=true, testing_ra test_start_ind = max(1, Int(round(rand() * (n_obs - n_obs_test)))) testing_inds = test_start_ind:test_start_ind+n_obs_test-1 fit_regularization!(mws, testing_inds; kwargs...) + if mws.om.metadata[:tel_prior] && fit_reg_model + fit_regularization_model!(mws, testing_inds; kwargs...) + end mws.om.metadata[:todo][:reg_improved] = true # end end \ No newline at end of file From 4256312670aab09ddc9e915b889cbcea8ef62468 Mon Sep 17 00:00:00 2001 From: Kristo Ment Date: Fri, 15 Aug 2025 16:38:08 -0400 Subject: [PATCH 3/4] Telluric prior implementation --- Project.toml | 5 +- src/StellarSpectraObservationFitting.jl | 1 + src/error_estimation.jl | 4 +- src/model_functions.jl | 3 - src/optimization_functions.jl | 33 ++++--- src/regularization_functions.jl | 15 ++- src/telluric_functions.jl | 122 ++++++++++++++++++++++++ 7 files changed, 158 insertions(+), 25 deletions(-) create mode 100644 src/telluric_functions.jl diff --git a/Project.toml b/Project.toml index df3ba6c..21996eb 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "StellarSpectraObservationFitting" uuid = "6b1189bd-9150-44f8-8055-f6ca7d16451e" -authors = ["Chistian Gilbertson "] -version = "0.1.3" +authors = ["Christian Gilbertson ", "Kristo Ment "] +version = "0.2.0" [deps] AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918" @@ -10,7 +10,6 @@ DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ExpectationMaximizationPCA = "20521544-8895-46ad-ab44-5135df82bd8f" FITSIO = "525bcba6-941b-5504-bd06-fd0dc1a4d2eb" -Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" diff --git a/src/StellarSpectraObservationFitting.jl b/src/StellarSpectraObservationFitting.jl index a4a1296..9f8a906 100644 --- a/src/StellarSpectraObservationFitting.jl +++ b/src/StellarSpectraObservationFitting.jl @@ -14,5 +14,6 @@ include("Nabla_extension.jl") include("prior_gp_functions.jl") include("rassine.jl") include("error_estimation.jl") +include("telluric_functions.jl") end # module diff --git a/src/error_estimation.jl b/src/error_estimation.jl index 8c8697d..161af1d 100644 --- a/src/error_estimation.jl +++ b/src/error_estimation.jl @@ -79,12 +79,12 @@ function estimate_σ_curvature_helper_finalizer!(σs::AbstractVecOrMat, _ℓs::A # fit a parabola (or line if using gradient) to `_ℓs` and convert to uncertainties if use_gradient poly_f = ordinary_lst_sq_f(_ℓs, 1; x=x_test) - σs[i] = sqrt(1 / poly_f.w[2]) + σs[i] = poly_f.w[2] > 0 ? sqrt(1 / poly_f.w[2]) : 0 # KM: added check for negative values max_dif = maximum(abs.((poly_f.(x_test)./_ℓs) .- 1)) if verbose; println("∇_$i: $(poly_f.w[1] + poly_f.w[2] * x[i])") end else poly_f = ordinary_lst_sq_f(_ℓs, 2; x=x_test) - σs[i] = sqrt(1 / (2 * poly_f.w[3])) + σs[i] = poly_f.w[3] > 0 ? sqrt(1 / (2 * poly_f.w[3])) : 0 # KM: added check for negative values max_dif = maximum(abs.((poly_f.(x_test)./_ℓs) .- 1)) if verbose; println("∇_$i: $(poly_f.w[2] + 2 * poly_f.w[3] * x[i])") end end diff --git a/src/model_functions.jl b/src/model_functions.jl index 00bc99e..1879c4d 100644 --- a/src/model_functions.jl +++ b/src/model_functions.jl @@ -1350,9 +1350,6 @@ function model_s_prior(s, reg::Dict) return 0 end - -## TEMPORARY - TODO: INTEGRATE FUNCTIONS INTO SSOF -include("/home/kment/RV/telluricfunc.jl") """ tel_prior(om) diff --git a/src/optimization_functions.jl b/src/optimization_functions.jl index e49c9e1..0bb0166 100644 --- a/src/optimization_functions.jl +++ b/src/optimization_functions.jl @@ -711,15 +711,15 @@ function train_OrderModel!(mws::AdamWorkspace; ignore_regularization::Bool=false function cb(as::AdamState) # optionally shift the score means to be near 0 - # DEBUG this has been turned off to test simulated telluric models - #if shift_scores - # if !(typeof(mws) <: FrozenTelWorkspace) - # remove_lm_score_means!(mws.om.tel.lm; prop=0.2) - # end - # if typeof(mws.om) <: OrderModelWobble - # remove_lm_score_means!(mws.om.star.lm; prop=0.2) - # end - #end + # KM: ignore this setting if telluric priors are used + if shift_scores && !om.metadata[:tel_prior] + if !(typeof(mws) <: FrozenTelWorkspace) + remove_lm_score_means!(mws.om.tel.lm; prop=0.2) + end + if typeof(mws.om) <: OrderModelWobble + remove_lm_score_means!(mws.om.star.lm; prop=0.2) + end + end # optionally make the templates always positive if μ_positive @@ -1216,12 +1216,14 @@ Defaults to returning the AIC-minimum model - `use_tel_prior::Bool=false`: Whether to use a modeled prior on telluric feature vectors - `careful_first_step::Bool=true`: Whether to shrink the learning rates until the loss improves on the first iteration - `speed_up::Bool=false`: Whether to inflate the learning rates until the loss is no longer improving throughout the optimization +- `max_iter::Int=50`: Maximum number of iterations for initial model optimization - `log_λ_gp_star::Real=1/SOAP_gp_params.λ`: The log λ lengthscale of the stellar regularization GP - `log_λ_gp_tel::Real=1/LSF_gp_params.λ`: The log λ lengthscale of the telluric regularization GP - `reg_star::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Override default_reg_star if given - `reg_tel::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Override default_reg_tel if given - `reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing`: Override default_reg_model if given - `fixed_vectors::Dict=Dict()`: Fix any stellar or telluric vectors. Must contain 'log_λ' +- `ignore_μ_tel::Bool=false`: Set the telluric mean to 1 if use_tel_prior - `kwargs...`: kwargs passed to `OrderModel` constructor """ function calculate_initial_model(data::Data; @@ -1229,10 +1231,11 @@ function calculate_initial_model(data::Data; μ_min::Real=0, μ_max::Real=Inf, use_mean::Bool=true, stop_early::Bool=false, remove_reciprocal_continuum::Bool=false, return_full_path::Bool=false, max_n_tel::Int=5, max_n_star::Int=5, use_all_comps::Bool=false, use_tel_prior::Bool=false, - careful_first_step::Bool=true, speed_up::Bool=false, + careful_first_step::Bool=true, speed_up::Bool=false, max_iter::Int=50, log_λ_gp_star::Real=1/SOAP_gp_params.λ, log_λ_gp_tel::Real=1/LSF_gp_params.λ, reg_star::Union{Nothing,Dict{Symbol,Float64}}=nothing, reg_tel::Union{Nothing,Dict{Symbol,Float64}}=nothing, - reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing, fixed_vectors::Dict=Dict(), kwargs...) + reg_model::Union{Nothing,Dict{Symbol,Float64}}=nothing, fixed_vectors::Dict=Dict(), + ignore_μ_tel::Bool=false, kwargs...) # TODO: Make this work for OrderModelDPCA # Get non-LSF version of `data` @@ -1387,10 +1390,10 @@ function calculate_initial_model(data::Data; try if only_s - finalize_scores!(mws; iter=50) + finalize_scores!(mws; iter=max_iter) else # improve the model - improve_initial_model!(mws; careful_first_step=careful_first_step, speed_up=speed_up, iter=50) + improve_initial_model!(mws; careful_first_step=careful_first_step, speed_up=speed_up, iter=max_iter) # if there is an LSF, do some more fitting if mws.d != data @@ -1549,7 +1552,9 @@ function calculate_initial_model(data::Data; i = comp2ind(n_tel_next, n_star_cur) oms[i...] = downsize(om, n_tel_next, n_star_cur) - if use_tel_prior_μ + if ignore_μ_tel + oms[i...].tel.lm.μ[:] .= 1 + elseif use_tel_prior_μ oms[i...].tel.lm.μ[:] = get_telluric_prior_spectrum(oms[i...].tel.λ) else make_telluric_template!(oms[i...]) diff --git a/src/regularization_functions.jl b/src/regularization_functions.jl index 1761abc..d4a9f89 100644 --- a/src/regularization_functions.jl +++ b/src/regularization_functions.jl @@ -39,7 +39,6 @@ function fit_regularization_helper!(reg_fields::Vector{Symbol}, reg_key::Symbol, # only do anything if if haskey(getfield(mws.om, reg_fields[1]), reg_key) - println("Regularization of $reg_key in $(reg_fields[1])") # DEBUG om = mws.om @assert 0 < reg_min < reg_max < Inf ℓs = Array{Float64}(undef, 2) @@ -165,6 +164,7 @@ end _key_list = [:GP_μ, :L2_μ, :L1_μ, :L1_μ₊_factor, :GP_M, :L2_M, :L1_M, :shared_M] +_key_list_μ = [:GP_μ, :L2_μ, :L1_μ, :L1_μ₊_factor] _key_list_fit = [:GP_μ, :L2_μ, :L1_μ, :GP_M, :L2_M, :L1_M] _key_list_bases = [:GP_M, :L2_M, :L1_M, :shared_M] _key_list_model = [:GP_μ, :L2_μ, :L1_μ, :L2_h2o, :L1_h2o, :L2_oxy, :L1_oxy] @@ -192,18 +192,27 @@ max_reg_model = 1e12 Fit all of the regularization values in `key_list` for the model in `mws` """ -function fit_regularization!(mws::ModelWorkspace, testing_inds::AbstractVecOrMat; key_list::Vector{Symbol}=_key_list_fit, share_regs::Bool=false, kwargs...) +function fit_regularization!(mws::ModelWorkspace, testing_inds::AbstractVecOrMat; key_list::Vector{Symbol}=_key_list_fit, share_regs::Bool=false, ignore_μ_tel::Bool=false, kwargs...) om = mws.om n_obs = size(mws.d.flux, 2) training_inds = [i for i in 1:n_obs if !(i in testing_inds)] check_for_valid_regularization(om.reg_tel) check_for_valid_regularization(om.reg_star) if share_regs; @assert keys(om.reg_tel) == keys(om.reg_star) end + if share_regs; @assert !ignore_μ_tel end hold_tel = copy(default_reg_star_full) hold_star = copy(default_reg_tel_full) + initial_tel = copy(om.reg_tel) copy_dict!(hold_tel, om.reg_tel) copy_dict!(hold_star, om.reg_star) zero_regularization(om) + if ignore_μ_tel + for key in _key_list_μ + if key in keys(initial_tel) + om.reg_tel[key] = initial_tel[key] + end + end + end println("starting regularization searches") before_ℓ = _eval_regularization(copy(mws.om), mws, training_inds, testing_inds) println("initial training χ²: $before_ℓ") @@ -219,7 +228,7 @@ function fit_regularization!(mws::ModelWorkspace, testing_inds::AbstractVecOrMat if (!(key in _key_list_bases)) || is_time_variable(om.star) before_ℓ = fit_regularization_helper!([:reg_star], key, before_ℓ, mws, training_inds, testing_inds, test_factor, reg_min, reg_max; start=hold_star[key], kwargs...) end - if (!(key in _key_list_bases)) || is_time_variable(om.tel) + if ((!(key in _key_list_bases)) || is_time_variable(om.tel)) && !((key in _key_list_μ) && ignore_μ_tel) before_ℓ = fit_regularization_helper!([:reg_tel], key, before_ℓ, mws, training_inds, testing_inds, test_factor, reg_min, reg_max; start=hold_tel[key], kwargs...) end end diff --git a/src/telluric_functions.jl b/src/telluric_functions.jl new file mode 100644 index 0000000..464807b --- /dev/null +++ b/src/telluric_functions.jl @@ -0,0 +1,122 @@ +using FITSIO +using DataInterpolations + +teldir = "/home/kment/RV/telluric" # Directory containing telluric model spectra + +get_h2o_spectrum(λ_values) = interp_telluric_file("$teldir/kpno_telluric_4.0_0.0_h2o_only.fits", λ_values) +get_oxy_spectrum(λ_values) = interp_telluric_file("$teldir/kpno_telluric_4.0_0.0_oxy_only.fits", λ_values) +get_tel_spectrum(λ_values) = interp_telluric_file("$teldir/kpno_telluric_4.0_0.0.fits", λ_values) + +""" + get_telluric_model(λ_values, pwv, z_angle; verbose=true) + +Estimate the telluric spectrum at a range of wavelenths (λ_values) for a specific zenith angle +(z_angle, in degrees) and precipitable water vapor (pwv) value. +""" +function get_telluric_model(λ_values, pwv, z_angle; verbose::Bool=true) + + @assert pwv <= 16 "Precipitable water vapor cannot exceed 16." + @assert z_angle < 60 "Zenith angle must be below 60." + + pwv_files = [1, 2, 4, 8, 16] + z_files = [0, 34, 44, 51, 60] + i_pwv = 1 # Index of PWV file + i_z = 1 # Index of zenith angle file + + for (i, x) in enumerate(pwv_files) + if pwv > x + i_pwv = i + 1 + end + end + for (i, x) in enumerate(z_files) + if z_angle >= x + i_z = i + end + end + + fn_model1 = "$teldir/kpno_telluric_$(pwv_files[i_pwv-1]).0_$(z_files[i_z]).0.fits" + fn_model2 = "$teldir/kpno_telluric_$(pwv_files[i_pwv]).0_$(z_files[i_z]).0.fits" + fn_model3 = "$teldir/kpno_telluric_$(pwv_files[i_pwv-1]).0_$(z_files[i_z+1]).0.fits" + fn_model4 = "$teldir/kpno_telluric_$(pwv_files[i_pwv]).0_$(z_files[i_z+1]).0.fits" + if verbose + println("Using models: $fn_model1 $fn_model2 $fn_model3 $fn_model4") + end + + t_model1 = interp_telluric_file(fn_model1, λ_values) + t_model2 = interp_telluric_file(fn_model2, λ_values) + t_model3 = interp_telluric_file(fn_model3, λ_values) + t_model4 = interp_telluric_file(fn_model4, λ_values) + + X = sec(z_angle * π / 180) # Airmass + X_model12 = sec(z_files[i_z] * π / 180) + X_model34 = sec(z_files[i_z+1] * π / 180) + X_power12 = X / X_model12 + X_power34 = X / X_model34 + pwv_power12 = (pwv - pwv_files[i_pwv-1]) / (pwv_files[i_pwv] - pwv_files[i_pwv-1]) * X_power12 + t_scaled12 = t_model1 .^ (X_power12 - pwv_power12) .* t_model2 .^ pwv_power12 + pwv_power34 = (pwv - pwv_files[i_pwv-1]) / (pwv_files[i_pwv] - pwv_files[i_pwv-1]) * X_power34 + t_scaled34 = t_model3 .^ (X_power34 - pwv_power34) .* t_model4 .^ pwv_power34 + + weight = (X - X_model12) / (X_model34 - X_model12) + t_scaled = t_scaled12 .* weight .+ t_scaled34 .* (1 - weight) + + return t_scaled + +end + +""" + get_telluric_prior_spectrum(λ_values, species::Symbol=:all) + +Evaluate the spectrum used as a telluric prior at a range of wavelenths (λ_values). +""" +function get_telluric_prior_spectrum(λ_values, species::Symbol=:all) + + if species == :all + return get_tel_spectrum(λ_values) + elseif species == :h2o + return get_h2o_spectrum(λ_values) + elseif species == :oxy + return get_oxy_spectrum(λ_values) + end + +end + +""" + interp_telluric_file(filename, λ_values) + +Estimate the telluric spectrum at a range of wavelenths (λ_values) by interpolating a model spectrum +in a given FITS file. +""" +function interp_telluric_file(filename, λ_values) + + f = FITS(filename) + λ_all = reverse(read(f[2], "wavelength")) .* 10 + t_all = reverse(read(f[2], "transmittance")) + + interp = LinearInterpolation(t_all, λ_all) + t_model = interp.(λ_values) + t_model[t_model .< 0] .= 0 + + return t_model + +end + +""" + simulate_telluric_data(λ_values, pwv_values, z_values) + +Simulate the telluric spectrum at a range of wavelenths (λ_values) for a list of precipitable water +vapor values (pwv_values) and zenith angle values (z_values, in degrees). +""" +function simulate_telluric_data(λ_values, pwv_values, z_values) + + @assert length(pwv_values) == length(z_values) "Must provide an equal number of PWV and zenith angle values." + + t_all = Matrix{Float64}(undef, length(λ_values), length(pwv_values)) + + for i in range(1, length(pwv_values)) + t_all[:,i] = get_telluric_model(λ_values, pwv_values[i], z_values[i]; verbose=false) + end + + return t_all + +end \ No newline at end of file From 17fd8ca5deb75d5338ad854b4fa34d761087a931 Mon Sep 17 00:00:00 2001 From: Kristo Ment Date: Thu, 21 Aug 2025 17:34:54 +0300 Subject: [PATCH 4/4] Update optimization_functions.jl (quick fix) --- src/optimization_functions.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/optimization_functions.jl b/src/optimization_functions.jl index 0bb0166..513b580 100644 --- a/src/optimization_functions.jl +++ b/src/optimization_functions.jl @@ -712,7 +712,7 @@ function train_OrderModel!(mws::AdamWorkspace; ignore_regularization::Bool=false # optionally shift the score means to be near 0 # KM: ignore this setting if telluric priors are used - if shift_scores && !om.metadata[:tel_prior] + if shift_scores && !mws.om.metadata[:tel_prior] if !(typeof(mws) <: FrozenTelWorkspace) remove_lm_score_means!(mws.om.tel.lm; prop=0.2) end