Skip to content

Commit

Permalink
API
Browse files Browse the repository at this point in the history
  • Loading branch information
odunbar committed Oct 11, 2023
1 parent 53062ec commit bc057b3
Show file tree
Hide file tree
Showing 3 changed files with 84 additions and 16 deletions.
14 changes: 14 additions & 0 deletions docs/src/API/RandomFeatures.md
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,19 @@
```@meta
CurrentModule = CalibrateEmulateSample.Emulators
```
## Kernel and Covariance structure
```@docs
OneDimFactor
DiagonalFactor
CholeskyFactor
LowRankFactor
HierLowRankFactor
SeparableKernel
NonseparableKernel
calculate_n_hyperparameters
hyperparameters_from_flat
build_default_prior
```

## Scalar interface

Expand Down Expand Up @@ -36,4 +49,5 @@ get_feature_decomposition
get_optimizer_options
optimize_hyperparameters!(::ScalarRandomFeatureInterface)
optimize_hyperparameters!(::VectorRandomFeatureInterface)
shrinkage_cov
```
84 changes: 69 additions & 15 deletions src/RandomFeature.jl
Original file line number Diff line number Diff line change
Expand Up @@ -26,30 +26,64 @@ struct EnsembleThreading <: MultithreadType end


# Different types of covariance representation in the prior

abstract type CovarianceStructureType end

"""
$(DocStringExtensions.TYPEDEF)
covariance structure for a one-dimensional space
"""
struct OneDimFactor <: CovarianceStructureType end

"""
$(DocStringExtensions.TYPEDEF)
builds a diagonal covariance structure
"""
struct DiagonalFactor{FT <: AbstractFloat} <: CovarianceStructureType
"nugget term"
eps::FT
end
DiagonalFactor() = DiagonalFactor(Float64(1.0))
get_eps(df::DiagonalFactor) = df.eps

"""
$(DocStringExtensions.TYPEDEF)
builds a general positive-definite covariance structure
"""
struct CholeskyFactor{FT <: AbstractFloat} <: CovarianceStructureType
"nugget term"
eps::FT
end
CholeskyFactor() = CholeskyFactor(Float64(1.0))
get_eps(cf::CholeskyFactor) = cf.eps

"""
$(DocStringExtensions.TYPEDEF)
builds a covariance structure that deviates from the identity with a low-rank perturbation. This perturbation is diagonalized in the low-rank space
"""
struct LowRankFactor{FT <: AbstractFloat} <: CovarianceStructureType
"rank of perturbation"
rank::Int
"nugget term"
eps::FT
end
LowRankFactor(r::Int) = LowRankFactor(r, Float64(1.0))
get_eps(lrf::LowRankFactor) = lrf.eps


"""
$(DocStringExtensions.TYPEDEF)
builds a covariance structure that deviates from the identity with a more general low-rank perturbation
"""
struct HierarchicalLowRankFactor{FT <: AbstractFloat} <: CovarianceStructureType
#cst::CovarianceStructureType
"rank of perturbation"
rank::Int
"nugget term"
eps::FT
end

Expand All @@ -66,6 +100,11 @@ rank(hlrf::HierarchicalLowRankFactor) = hlrf.rank
# build_default_prior(name::SS, d::Int, t::T) where {SS <: AbstractString}
# and add a string id to cov_structure_from_string

"""
$(DocStringExtensions.TYPEDSIGNATURES)
creates some simple covariance structures from strings: ["onedim", "diagonal", "cholesky", "lowrank", hierlowrank"]. See the covariance Structures for more details.
"""
function cov_structure_from_string(s::S, d::Int) where {S <: AbstractString}
if s == "onedim"
return OneDimFactor()
Expand All @@ -91,20 +130,36 @@ cov_structure_from_string(cst::CST, args...; kwargs...) where {CST <: Covariance
# Different types of kernel
abstract type KernelStructureType end

"""
$(DocStringExtensions.TYPEDEF)
Builds a separable kernel, i.e. one that accounts for input and output covariance structure separately
"""
struct SeparableKernel{CST1 <: CovarianceStructureType, CST2 <: CovarianceStructureType} <: KernelStructureType
input_cov_structure::CST1
output_cov_structure::CST2
end
get_input_cov_structure(kernel_structure::SeparableKernel) = kernel_structure.input_cov_structure
get_output_cov_structure(kernel_structure::SeparableKernel) = kernel_structure.output_cov_structure

struct NonseparableKernel{CST <: CovarianceStructureType} <: KernelStructureType

"""
$(DocStringExtensions.TYPEDEF)
Builds a nonseparable kernel, i.e. one that accounts for a joint input and output covariance structure
"""struct NonseparableKernel{CST <: CovarianceStructureType} <: KernelStructureType
cov_structure::CST
end
get_cov_structure(kernel_structure::NonseparableKernel) = kernel_structure.cov_structure


# calculate_n_hyperparameters

"""
$(DocStringExtensions.TYPEDSIGNATURES)
calculates the number of hyperparameters generated by the choice of covariance structure
"""
calculate_n_hyperparameters(d::Int, odf::OneDimFactor) = 0
calculate_n_hyperparameters(d::Int, df::DiagonalFactor) = d
calculate_n_hyperparameters(d::Int, cf::CholeskyFactor) = Int(d * (d + 1) / 2)
Expand All @@ -113,6 +168,13 @@ calculate_n_hyperparameters(d::Int, hlrf::HierarchicalLowRankFactor) =
Int(rank(hlrf) * (rank(hlrf) + 1) / 2 + d * rank(hlrf))

# build from flat


"""
$(DocStringExtensions.TYPEDSIGNATURES)
reshapes a list of hyperparameters into a covariance matrix based on the selected structure
"""
hyperparameters_from_flat(x::V, odf::OneDimFactor) where {V <: AbstractVector} = nothing

function hyperparameters_from_flat(x::V, df::DiagonalFactor) where {V <: AbstractVector}
Expand Down Expand Up @@ -172,6 +234,11 @@ function hyperparameters_from_flat(x::V, hlrf::HierarchicalLowRankFactor) where
end


"""
$(DocStringExtensions.TYPEDSIGNATURES)
builds a prior distribution for the kernel hyperparameters to initialize optimization.
"""
build_default_prior(name::SS, n_hp::Int, odf::OneDimFactor) where {SS <: AbstractString} = nothing

function build_default_prior(name::SS, n_hp::Int, df::DiagonalFactor) where {SS <: AbstractString}
Expand Down Expand Up @@ -211,12 +278,6 @@ function build_default_prior(name::SS, n_hp::Int, hlrf::HierarchicalLowRankFacto
end

# combining input and output spaces:

"""
$(DocStringExtensions.TYPEDSIGNATURES)
Calculate number of hyperparameters required to create the default prior in the given input/output space dimensions (determined from the `CovarianceStructureType`)
"""
function calculate_n_hyperparameters(
input_dim::Int,
output_dim::Int,
Expand Down Expand Up @@ -261,13 +322,6 @@ end



"""
$(DocStringExtensions.TYPEDSIGNATURES)
Builds a prior over the hyperparameters (i.e. the low-rank/cholesky/diagaonal or individaul entries of the input/output covariances).
For example, the case where the input covariance ``U = γ_1 * (LL^T + γ_2 I)``,
we set priors for the entries of the lower triangular matrix ``L`` as normal, and constant scalings ``γ_i`` as log-normal to retain positivity.
"""
function build_default_prior(input_dim::Int, output_dim::Int, kernel_structure::SK) where {SK <: SeparableKernel}
input_cov_structure = get_input_cov_structure(kernel_structure)
output_cov_structure = get_output_cov_structure(kernel_structure)
Expand Down
2 changes: 1 addition & 1 deletion src/ScalarRandomFeature.jl
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ $(DocStringExtensions.TYPEDEF)
Structure holding the Scalar Random Feature models.
# FieldsWhen calibrated with ocean LES,
# Fields
$(DocStringExtensions.TYPEDFIELDS)
"""
Expand Down

0 comments on commit bc057b3

Please sign in to comment.