diff --git a/Project.toml b/Project.toml index c67c274..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" @@ -9,6 +9,7 @@ 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" KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392" LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" 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/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 8c01c49..1879c4d 100644 --- a/src/model_functions.jl +++ b/src/model_functions.jl @@ -575,6 +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_μ, 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), @@ -668,6 +676,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 +704,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 +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=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( @@ -739,8 +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=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...) @@ -750,8 +784,15 @@ 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), (: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_λ) @@ -760,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), 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 @@ -768,17 +809,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(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), 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 +840,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 +956,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 +988,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,15 +1289,18 @@ 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`. +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_μ) @@ -1264,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) @@ -1295,22 +1350,101 @@ function model_s_prior(s, reg::Dict) return 0 end - """ 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) + + 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) + + if haskey(om.reg_model, :L1_μ) + val += L1(om.metadata[:tel_prior_spec] - lm[min(3, length(lm))]) * om.reg_model[:L1_μ] + 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 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 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 + +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` """ +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 0ff6b19..513b580 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...)) @@ -699,7 +700,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) @@ -710,7 +711,8 @@ 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 + # KM: ignore this setting if telluric priors are used + if shift_scores && !mws.om.metadata[:tel_prior] if !(typeof(mws) <: FrozenTelWorkspace) remove_lm_score_means!(mws.om.tel.lm; prop=0.2) end @@ -960,7 +962,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 +1019,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,24 +1213,41 @@ 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 +- `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; 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, - log_λ_gp_star::Real=1/SOAP_gp_params.λ, log_λ_gp_tel::Real=1/LSF_gp_params.λ, kwargs...) + 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, 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(), + ignore_μ_tel::Bool=false, 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 @@ -1247,8 +1266,18 @@ 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, reg_star=reg_star, + reg_tel=reg_tel, reg_model=reg_model, fixed_vectors=fixed_vectors, kwargs...) + # Add telluric priors to specific vectors + if use_tel_prior_h2o + om.metadata[:h2o_vector] = 1 + 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 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 @@ -1337,16 +1366,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 @@ -1358,13 +1389,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=max_iter) + else + # improve the model + 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 - 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) @@ -1390,6 +1425,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 @@ -1397,10 +1435,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.μ) @@ -1412,7 +1462,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) @@ -1474,9 +1524,70 @@ 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 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...]) + 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 @@ -1502,21 +1613,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.μ @@ -1544,23 +1656,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...") diff --git a/src/regularization_functions.jl b/src/regularization_functions.jl index 3c6aa95..d4a9f89 100644 --- a/src/regularization_functions.jl +++ b/src/regularization_functions.jl @@ -164,8 +164,10 @@ 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] """ @@ -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 """ @@ -188,18 +192,27 @@ max_reg = 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_ℓ") @@ -215,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 @@ -223,13 +236,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 +276,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 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