From bc057b37917e6a99f25a284ed8d7a6be7670908f Mon Sep 17 00:00:00 2001
From: odunbar <odunbar@caltech.edu>
Date: Tue, 10 Oct 2023 18:05:18 -0700
Subject: [PATCH] API

---
 docs/src/API/RandomFeatures.md | 14 ++++++
 src/RandomFeature.jl           | 84 ++++++++++++++++++++++++++++------
 src/ScalarRandomFeature.jl     |  2 +-
 3 files changed, 84 insertions(+), 16 deletions(-)

diff --git a/docs/src/API/RandomFeatures.md b/docs/src/API/RandomFeatures.md
index 2baffd992..abef7d191 100644
--- a/docs/src/API/RandomFeatures.md
+++ b/docs/src/API/RandomFeatures.md
@@ -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
 
@@ -36,4 +49,5 @@ get_feature_decomposition
 get_optimizer_options
 optimize_hyperparameters!(::ScalarRandomFeatureInterface) 
 optimize_hyperparameters!(::VectorRandomFeatureInterface) 
+shrinkage_cov
 ```
diff --git a/src/RandomFeature.jl b/src/RandomFeature.jl
index f6f793ad9..c01cf8a02 100644
--- a/src/RandomFeature.jl
+++ b/src/RandomFeature.jl
@@ -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
 
@@ -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()
@@ -91,6 +130,11 @@ 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
@@ -98,13 +142,24 @@ 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)
@@ -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}
@@ -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}
@@ -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,
@@ -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)
diff --git a/src/ScalarRandomFeature.jl b/src/ScalarRandomFeature.jl
index a23586127..2c9ccc736 100644
--- a/src/ScalarRandomFeature.jl
+++ b/src/ScalarRandomFeature.jl
@@ -6,7 +6,7 @@ $(DocStringExtensions.TYPEDEF)
 
 Structure holding the Scalar Random Feature models. 
 
-# FieldsWhen calibrated with ocean LES,
+# Fields
 $(DocStringExtensions.TYPEDFIELDS)
 
 """