Add UniformScaling as a covariance type

CliMA · Jan 28, 2022 · 72ed703 · 72ed703
1 parent 3d67c96
commit 72ed703
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 12 deletions.
diff --git a/src/Emulator.jl b/src/Emulator.jl
@@ -51,7 +51,7 @@ struct Emulator{FT<:AbstractFloat}
     "square root of the inverse of the input covariance matrix; input_dim × input_dim"
     normalize_inputs::Bool
     "whether to fit models on normalized outputs outputs / standardize_outputs_factor"
-    sqrt_inv_input_cov::Union{AbstractMatrix{FT}, Nothing}
+    sqrt_inv_input_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing}
     " if normalizing: whether to fit models on normalized inputs ((inputs - input_mean) * sqrt_inv_input_cov)"
     standardize_outputs::Bool
     "if standardizing: Standardization factors (characteristic values of the problem)"
@@ -64,7 +64,7 @@ end
 function Emulator(
     machine_learning_tool::MachineLearningTool,
     input_output_pairs::PairedDataContainer{FT};
-    obs_noise_cov::Union{AbstractMatrix{FT}, Nothing} = nothing,
+    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing} = nothing,
     normalize_inputs::Bool = true,
     standardize_outputs::Bool = false,
     standardize_outputs_factors::Union{AbstractVector{FT}, Nothing} = nothing,
@@ -223,7 +223,7 @@ normalize with the empirical Gaussian distribution of points
 function normalize(
     inputs::AbstractVecOrMat{FT}, 
     input_mean::AbstractVector{FT}, 
-    sqrt_inv_input_cov::AbstractMatrix{FT}
+    sqrt_inv_input_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}}
 ) where {FT<:AbstractFloat}
     training_inputs = sqrt_inv_input_cov * (inputs .- input_mean)
     return training_inputs 
@@ -236,7 +236,7 @@ standardize with a vector of factors (size equal to output dimension)
 """
 function standardize(
     outputs::AbstractVecOrMat{FT}, 
-    output_covs::Vector{<:AbstractMatrix{FT}}, 
+    output_covs::Vector{<:Union{AbstractMatrix{FT}, UniformScaling{FT}}}, 
     factors::AbstractVector{FT}
 ) where {FT<:AbstractFloat}
     # Case where `output_covs` is a Vector of covariance matrices
@@ -250,7 +250,7 @@ end
 
 function standardize(
     outputs::AbstractVecOrMat{FT}, 
-    output_cov::AbstractMatrix{FT}, 
+    output_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}}, 
     factors::AbstractVector{FT}
 ) where {FT<:AbstractFloat}
     # Case where `output_cov` is a single covariance matrix
@@ -268,7 +268,7 @@ dimension). `output_cov` is a Vector of covariance matrices, such as is returned
 function reverse_standardize(
     emulator::Emulator{FT}, 
     outputs::AbstractVecOrMat{FT}, 
-    output_covs::Union{AbstractMatrix{FT}, Vector{<:AbstractMatrix{FT}}}
+    output_covs::Union{AbstractMatrix{FT}, UniformScaling{FT}, Vector{<:AbstractMatrix{FT}}}
 ) where {FT<:AbstractFloat}
     if emulator.standardize_outputs
         return standardize(outputs, output_covs, 1.0 ./ emulator.standardize_outputs_factors)
@@ -297,7 +297,7 @@ in S are sorted in descending order.
 """
 function svd_transform(
     data::AbstractMatrix{FT}, 
-    obs_noise_cov::Union{AbstractMatrix{FT}, Nothing}; 
+    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing}; 
     truncate_svd::FT=1.0
 ) where {FT<:AbstractFloat}
     if obs_noise_cov === nothing
@@ -327,7 +327,7 @@ end
 
 function svd_transform(
     data::AbstractVector{FT}, 
-    obs_noise_cov::Union{AbstractMatrix{FT}, Nothing}; 
+    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}, Nothing}; 
     truncate_svd::FT = 1.0
 ) where {FT<:AbstractFloat}
     # method for 1D data

diff --git a/src/MarkovChainMonteCarlo.jl b/src/MarkovChainMonteCarlo.jl
@@ -32,7 +32,7 @@ mutable struct MCMC{FT<:AbstractFloat, IT<:Int}
     "a single sample from the observations. Can e.g. be picked from an Obs struct using get_obs_sample"
     obs_sample::AbstractVector{FT}
     "covariance of the observational noise"
-    obs_noise_cov::AbstractMatrix{FT}
+    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}}
     "array of length N_parameters with the parameters' prior distributions"
     prior::ParameterDistribution
     "MCMC step size"
@@ -68,7 +68,7 @@ where max_iter is the number of MCMC steps to perform (e.g., 100_000)
 """
 function MCMC(
     obs_sample::AbstractVector{FT},
-    obs_noise_cov::AbstractMatrix{FT},
+    obs_noise_cov::Union{AbstractMatrix{FT}, UniformScaling{FT}},
     prior::ParameterDistribution,
     step::FT,
     param_init::AbstractVector{FT},
@@ -154,7 +154,7 @@ end
 function mcmc_sample!(
     mcmc::MCMC{FT,IT}, 
     g::AbstractVector{FT}, 
-    gcov::AbstractMatrix{FT}
+    gcov::Union{AbstractMatrix{FT}, UniformScaling{FT}}
 ) where {FT<:AbstractFloat, IT<:Int}
     if mcmc.algtype == "rwm"
         log_posterior = log_likelihood(mcmc, g, gcov) + log_prior(mcmc)
@@ -196,7 +196,7 @@ end
 function log_likelihood(
     mcmc::MCMC{FT,IT},
     g::AbstractVector{FT},
-    gcov::AbstractMatrix{FT}
+    gcov::Union{AbstractMatrix{FT}, UniformScaling{FT}}
 ) where {FT<:AbstractFloat, IT<:Int}
     log_rho = 0.0
     #if gcov == nothing