diff --git a/src/Utilities.jl b/src/Utilities.jl index 1d3e34a20..7a85a811d 100644 --- a/src/Utilities.jl +++ b/src/Utilities.jl @@ -1,5 +1,7 @@ module Utilities + + using DocStringExtensions using LinearAlgebra using Statistics @@ -12,6 +14,23 @@ using ..DataContainers export get_training_points export orig2zscore export zscore2orig + +export PairedDataContainerProcessor, DataContainerProcessor +export UnivariateAffineScaling, ElementwiseScaler, QuartileScaling, MinMaxScaling, ZScoreScaling +export quartile_scale, minmax_scale, zscore_scale +export Decorrelater, decorrelate_sample_cov, decorrelate_structure_mat, decorrelate +export get_type, + get_shift, get_scale, get_data_mean, get_encoder_mat, get_decoder_mat, get_retain_var, get_decorrelate_with +export create_encoder_schedule, encode_with_schedule, decode_with_schedule +export initialize_processor!, + initialize_and_encode_data!, encode_data, decode_data, encode_structure_matrix, decode_structure_matrix + +export CanonicalCorrelation +export canonical_correlation +export get_apply_to + + + """ $(DocStringExtensions.TYPEDSIGNATURES) @@ -90,4 +109,992 @@ function zscore2orig(Z::AbstractMatrix{FT}, mean::AbstractVector{FT}, std::Abstr return X end +# Data processing tooling: + +abstract type PairedDataContainerProcessor end # tools that operate on inputs and outputs +abstract type DataContainerProcessor end # tools that operate on only inputs or outputs + +abstract type UnivariateAffineScaling end +abstract type QuartileScaling <: UnivariateAffineScaling end +abstract type MinMaxScaling <: UnivariateAffineScaling end +abstract type ZScoreScaling <: UnivariateAffineScaling end + +# Processors +""" +$(TYPEDEF) + +The ElementwiseScaler{T} will create an encoding of the data_container via elementwise affine transformations. + +Different methods `T` will build different transformations: +- [`quartile_scale`](@ref) : creates `QuartileScaling`, +- [`minmax_scale`](@ref) : creates `MinMaxScaling` +- [`zscore_scale`](@ref) : creates `ZScoreScaling` + +and are accessed with [`get_type`](@ref) +""" +struct ElementwiseScaler{T, VV <: AbstractVector} <: DataContainerProcessor + "storage for the shift applied to data" + shift::VV + "storage for the scaling" + scale::VV +end + +""" +$(TYPEDSIGNATURES) + +Constructs `ElementwiseScaler{QuartileScaling}` processor. +As part of an encoder schedule, it will apply the transform ``\\frac{x - Q2(x)}{Q3(x) - Q1(x)}`` to each data dimension. +Also known as "robust scaling" +""" +quartile_scale() = ElementwiseScaler(QuartileScaling) + +""" +$(TYPEDSIGNATURES) + +Constructs `ElementwiseScaler{MinMaxScaling}` processor. +As part of an encoder schedule, this will apply the transform ``\\frac{x - \\min(x)}{\\max(x) - \\min(x)}`` to each data dimension. +""" +minmax_scale() = ElementwiseScaler(MinMaxScaling) + +""" +$(TYPEDSIGNATURES) + +Constructs `ElementwiseScaler{ZScoreScaling}` processor. +As part of an encoder schedule, this will apply the transform ``\\frac{x-\\mu}{\\sigma}``, (where ``x\\sim N(\\mu,\\sigma)``), to each data dimension. +For multivariate standardization, see [`Decorrelater`](@ref) +""" +zscore_scale() = ElementwiseScaler(ZScoreScaling) + +ElementwiseScaler(::Type{UAS}) where {UAS <: UnivariateAffineScaling} = + ElementwiseScaler{UAS, Vector{Float64}}(Float64[], Float64[]) + +""" +$(TYPEDSIGNATURES) + +Gets the UnivariateAffineScaling type `T` +""" +get_type(es::ElementwiseScaler{T}) where {T} = T + +""" +$(TYPEDSIGNATURES) + +Gets the `shift` field of the `ElementwiseScaler` +""" +get_shift(es::ElementwiseScaler) = es.shift + +""" +$(TYPEDSIGNATURES) + +Gets the `scale` field of the `ElementwiseScaler` +""" +get_scale(es::ElementwiseScaler) = es.scale + +function Base.show(io::IO, es::ElementwiseScaler) + out = "ElementwiseScaler: $(get_type(es))" + print(io, out) +end + +function initialize_processor!( + es::ElementwiseScaler, + data::MM, + T::Type{QS}, +) where {MM <: AbstractMatrix, QS <: QuartileScaling} + quartiles_vec = [quantile(dd, [0.25, 0.5, 0.75]) for dd in eachrow(data)] + quartiles_mat = reduce(hcat, quartiles_vec) # 3 rows: Q1, Q2, and Q3 + append!(get_shift(es), quartiles_mat[2, :]) + append!(get_scale(es), (quartiles_mat[3, :] - quartiles_mat[1, :])) +end + +function initialize_processor!( + es::ElementwiseScaler, + data::MM, + T::Type{MMS}, +) where {MM <: AbstractMatrix, MMS <: MinMaxScaling} + minmax_vec = [[minimum(dd), maximum(dd)] for dd in eachrow(data)] + minmax_mat = reduce(hcat, minmax_vec) # 2 rows: min max + append!(get_shift(es), minmax_mat[1, :]) + append!(get_scale(es), (minmax_mat[2, :] - minmax_mat[1, :])) +end + +function initialize_processor!( + es::ElementwiseScaler, + data::MM, + T::Type{ZSS}, +) where {MM <: AbstractMatrix, ZSS <: ZScoreScaling} + stat_vec = [[mean(dd), std(dd)] for dd in eachrow(data)] + stat_mat = reduce(hcat, stat_vec) # 2 rows: mean, std + append!(get_shift(es), stat_mat[1, :]) + append!(get_scale(es), stat_mat[2, :]) +end + +function initialize_processor!(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} + if length(get_shift(es)) == 0 + T = get_type(es) + initialize_processor!(es, data, T) + end +end + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` encoder, on a columns-are-data matrix +""" +function encode_data(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} + out = deepcopy(data) + for i in 1:size(out, 1) + out[i, :] .-= get_shift(es)[i] + out[i, :] /= get_scale(es)[i] + end + return out +end + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` decoder, on a columns-are-data matrix +""" +function decode_data(es::ElementwiseScaler, data::MM) where {MM <: AbstractMatrix} + out = deepcopy(data) + for i in 1:size(out, 1) + out[i, :] *= get_scale(es)[i] + out[i, :] .+= get_shift(es)[i] + end + return out +end + +""" +$(TYPEDSIGNATURES) + +Computes and populates the `shift` and `scale` fields for the `ElementwiseScaler` +""" +initialize_processor!(es::ElementwiseScaler, data::MM, structure_matrix) where {MM <: AbstractMatrix} = + initialize_processor!(es, data) + + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` encoder to a provided structure matrix +""" +function encode_structure_matrix(es::ElementwiseScaler, structure_matrix::MM) where {MM <: AbstractMatrix} + return Diagonal(1 ./ get_scale(es)) * structure_matrix * Diagonal(1 ./ get_scale(es)) +end + +""" +$(TYPEDSIGNATURES) + +Apply the `ElementwiseScaler` decoder to a provided structure matrix +""" +function decode_structure_matrix(es::ElementwiseScaler, enc_structure_matrix::MM) where {MM <: AbstractMatrix} + return Diagonal(get_scale(es)) * enc_structure_matrix * Diagonal(get_scale(es)) +end + + +""" +$(TYPEDEF) + +Decorrelate the data via taking an SVD decomposition and projecting onto the singular-vectors. + +Preferred construction is with the methods +- [`decorrelate_structure_mat`](@ref) +- [`decorrelate_sample_cov`](@ref) +- [`decorrelate`](@ref) + +For `decorrelate_structure_mat`: +The SVD is taken over a structure matrix (e.g., `prior_cov` for inputs, `obs_noise_cov` for outputs). The structure matrix will become exactly `I` after processing. + +For `decorrelate_sample_cov`: +The SVD is taken over the estimated covariance of the data. The data samples will have a `Normal(0,I)` distribution after processing, + +For `decorrelate(;decorrelate_with="combined")` (default): +The SVD is taken to be the sum of structure matrix and estimated covariance. This may be more robust to ill-specification of structure matrix, or poor estimation of the sample covariance. + +# Fields +$(TYPEDFIELDS) +""" +struct Decorrelater{VV1, VV2, VV3, FT, AS <: AbstractString} <: DataContainerProcessor + "storage for the data mean" + data_mean::VV1 + "the matrix used to perform encoding" + encoder_mat::VV2 + "the inverse of the the matrix used to perform encoding" + decoder_mat::VV3 + "the fraction of variance to be retained after truncating singular values (1 implies no truncation)" + retain_var::FT + "Switch to choose what form of matrix to use to decorrelate the data" + decorrelate_with::AS +end + +""" +$(TYPEDSIGNATURES) + +Constructs the `Decorrelater` struct. Users can add optional keyword arguments: +- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained +- `decorrelate_with` [=`"combined"`]: from which matrix do we provide subspace directions, options are + - `"structure_mat"`, see [`decorrelate_structure_mat`](@ref) + - `"sample_cov"`, see [`decorrelate_sample_cov`](@ref) + - `"combined"`, sums the `"sample_cov"` and `"structure_mat"` matrices +""" +decorrelate(; retain_var::FT = Float64(1.0), decorrelate_with = "combined") where {FT} = + Decorrelater([], [], [], clamp(retain_var, FT(0), FT(1)), decorrelate_with) + +""" +$(TYPEDSIGNATURES) + +Constructs the `Decorrelater` struct, setting decorrelate_with = "sample_cov". Encoding data with this will ensure that the distribution of data samples after encoding will be `Normal(0,I)`. One can additionally add keywords: +- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained +""" +decorrelate_sample_cov(; retain_var::FT = Float64(1.0)) where {FT} = + Decorrelater([], [], [], clamp(retain_var, FT(0), FT(1)), "sample_cov") + +""" +$(TYPEDSIGNATURES) + +Constructs the `Decorrelater` struct, setting decorrelate_with = "structure_mat". This encoding will transform a provided structure matrix into `I`. One can additionally add keywords: +- `retain_var`[=`1.0`]: to project onto the leading singular vectors such that `retain_var` variance is retained +""" +decorrelate_structure_mat(; retain_var::FT = Float64(1.0)) where {FT} = + Decorrelater([], [], [], clamp(retain_var, FT(0), FT(1)), "structure_mat") + +""" +$(TYPEDSIGNATURES) + +returns the `data_mean` field of the `Decorrelater`. +""" +get_data_mean(dd::Decorrelater) = dd.data_mean + +""" +$(TYPEDSIGNATURES) + +returns the `encoder_mat` field of the `Decorrelater`. +""" +get_encoder_mat(dd::Decorrelater) = dd.encoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `decoder_mat` field of the `Decorrelater`. +""" +get_decoder_mat(dd::Decorrelater) = dd.decoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `retain_var` field of the `Decorrelater`. +""" +get_retain_var(dd::Decorrelater) = dd.retain_var + +""" +$(TYPEDSIGNATURES) + +returns the `decorrelate_with` field of the `Decorrelater`. +""" +get_decorrelate_with(dd::Decorrelater) = dd.decorrelate_with + +function Base.show(io::IO, dd::Decorrelater) + out = "Decorrelater" + out *= ": decorrelate_with=$(get_decorrelate_with(dd))" + if get_retain_var(dd) < 1.0 + out *= ", retain_var=$(get_retain_var(dd))" + end + print(io, out) +end + +""" +$(TYPEDSIGNATURES) + +Computes and populates the `data_mean` and `encoder_mat` and `decoder_mat` fields for the `Decorrelater` +""" +function initialize_processor!( + dd::Decorrelater, + data::MM, + structure_matrix::USorM, +) where {MM <: AbstractMatrix, USorM <: Union{UniformScaling, AbstractMatrix}} + if length(get_data_mean(dd)) == 0 + push!(get_data_mean(dd), vec(mean(data, dims = 2))) + end + + if length(get_encoder_mat(dd)) == 0 + + # Can do tsvd here for large matrices + decorrelate_with = get_decorrelate_with(dd) + if decorrelate_with == "structure_mat" + svdA = svd(structure_matrix) + rk = rank(structure_matrix) + elseif decorrelate_with == "sample_cov" + cd = cov(data, dims = 2) + svdA = svd(cd) + rk = rank(cd) + elseif decorrelate_with == "combined" + spluscd = structure_matrix + cov(data, dims = 2) + svdA = svd(spluscd) + rk = rank(spluscd) + else + throw( + ArgumentError( + "Keyword `decorrelate_with` must be taken from [\"sample_cov\", \"structure_mat\", \"combined\"]. Received $(decorrelate_with)", + ), + ) + end + ret_var = get_retain_var(dd) + if ret_var < 1.0 + sv_cumsum = cumsum(svdA.S .^ 2) / sum(svdA.S .^ 2) # variance contributions are (sing_val)^2 + trunc = minimum(findall(x -> (x > ret_var), sv_cumsum)) + @info " truncating at $(trunc)/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc])% of the variance of the structure matrix" + else + trunc = rk + end + sqrt_inv_sv = Diagonal(1.0 ./ sqrt.(svdA.S[1:trunc])) + sqrt_sv = Diagonal(sqrt.(svdA.S[1:trunc])) + + # as we have svd of cov-matrix we can use U or Vt + encoder_mat = sqrt_inv_sv * svdA.Vt[1:trunc, :] + decoder_mat = svdA.Vt[1:trunc, :]' * sqrt_sv + + push!(get_encoder_mat(dd), encoder_mat) + push!(get_decoder_mat(dd), decoder_mat) + end +end + + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelater` encoder, on a columns-are-data matrix +""" +function encode_data(dd::Decorrelater, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(dd)[1] + encoder_mat = get_encoder_mat(dd)[1] + return encoder_mat * (data .- data_mean) +end + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelater` decoder, on a columns-are-data matrix +""" +function decode_data(dd::Decorrelater, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(dd)[1] + decoder_mat = get_decoder_mat(dd)[1] + return decoder_mat * data .+ data_mean +end + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelater` encoder to a provided structure matrix +""" +function encode_structure_matrix(dd::Decorrelater, structure_matrix::MM) where {MM <: AbstractMatrix} + encoder_mat = get_encoder_mat(dd)[1] + return encoder_mat * structure_matrix * encoder_mat' +end + +""" +$(TYPEDSIGNATURES) + +Apply the `Decorrelater` decoder to a provided structure matrix +""" +function decode_structure_matrix(dd::Decorrelater, enc_structure_matrix::MM) where {MM <: AbstractMatrix} + decoder_mat = get_decoder_mat(dd)[1] + return decoder_mat * enc_structure_matrix * decoder_mat' +end + +# ...struct VariationalAutoEncoder <: DataContainerProcessor end + +# PDCProcessors +# struct InverseProblemInformed <: PairedDataContainerProcessor end +# struct LikelihoodInformed <: PairedDataContainerProcessor end + + +""" +$(TYPEDEF) + +Uses both input and output data to learn a subspace of maximal correlation between inputs and outputs. The subspace for a pair (X,Y) will be of size minimum(rank(X),rank(Y)), computed using SVD-based method +e.g. See e.g., https://numerical.recipes/whp/notes/CanonCorrBySVD.pdf + +Preferred construction is with the [`canonical_correlation`](@ref) method + +# Fields +$(TYPEDFIELDS) +""" +struct CanonicalCorrelation{VV1, VV2, VV3, FT, VV4} <: PairedDataContainerProcessor + "storage for the input or output data mean" + data_mean::VV1 + "the encoding matrix of input or output canonical correlations" + encoder_mat::VV2 + "the decoding matrix of input or output canonical correlations" + decoder_mat::VV3 + "the fraction of variance to be retained after truncating singular values (1 implies no truncation)" + retain_var::FT + "Stores whether this is an input or output encoder (vector with string \"in\" or \"out\")" + apply_to::VV4 +end + +""" +$(TYPEDSIGNATURES) + +Constructs the `CanonicalCorrelation` struct. Can optionally provide the keyword +- `retain_var`[=1.0]: to project onto the leading singular vectors (of the input-output product) such that `retain_var` variance is retained. +""" +canonical_correlation(; retain_var::FT = Float64(1.0)) where {FT} = + CanonicalCorrelation(Any[], Any[], Any[], clamp(retain_var, FT(0), FT(1)), AbstractString[]) + +""" +$(TYPEDSIGNATURES) + +returns the `data_mean` field of the `CanonicalCorrelation`. +""" +get_data_mean(cc::CanonicalCorrelation) = cc.data_mean + +""" +$(TYPEDSIGNATURES) + +returns the `encoder_mat` field of the `CanonicalCorrelation`. +""" +get_encoder_mat(cc::CanonicalCorrelation) = cc.encoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `decoder_mat` field of the `CanonicalCorrelation`. +""" +get_decoder_mat(cc::CanonicalCorrelation) = cc.decoder_mat + +""" +$(TYPEDSIGNATURES) + +returns the `retain_var` field of the `CanonicalCorrelation`. +""" +get_retain_var(cc::CanonicalCorrelation) = cc.retain_var + +""" +$(TYPEDSIGNATURES) + +returns the `apply_to` field of the `CanonicalCorrelation`. +""" +get_apply_to(cc::CanonicalCorrelation) = cc.apply_to + +function Base.show(io::IO, cc::CanonicalCorrelation) + + out = "CanonicalCorrelation:" + if length(get_apply_to(cc)) > 0 + out *= " apply_to=$(get_apply_to(cc)[1])" + end + if get_retain_var(cc) < 1.0 + out *= " retain_var=$(get_retain_var(cc))" + end + print(io, out) +end + + +function initialize_processor!( + cc::CanonicalCorrelation, + in_data::MM, + out_data::MM, + apply_to::AS, +) where {MM <: AbstractMatrix, AS <: AbstractString} + + if apply_to ∉ ["in", "out"] + bad_apply_to(apply_to) + end + + if length(get_apply_to(cc)) == 0 + push!(get_apply_to(cc), apply_to) + end + + if length(get_data_mean(cc)) == 0 + if apply_to == "in" + push!(get_data_mean(cc), vec(mean(in_data, dims = 2))) + elseif apply_to == "out" + push!(get_data_mean(cc), vec(mean(out_data, dims = 2))) + end + end + + if length(get_encoder_mat(cc)) == 0 + + if size(in_data, 2) < size(in_data, 1) || size(out_data, 2) < size(out_data, 1) + throw( + ArgumentError( + "CanonicalCorrelation implementation not defined for # data samples < dimensions, please obtain more samples, or perform prior dimension reduction approaches until this is satisfied", + ), + ) + end + + # Individually decompose in and out + svdi = svd(in_data .- mean(in_data, dims = 2)) + svdo = svd(out_data .- mean(out_data, dims = 2)) + + # ensure correct shaping (in_mat = (in_dim x n_samples), out_mat = (out_dim x n_samples)) + in_mat_sq, in_mat_nonsq = (size(svdi.U, 1) == size(svdi.U, 2)) ? (svdi.U, svdi.Vt) : (svdi.Vt, svdi.U) + out_mat_sq, out_mat_nonsq = (size(svdo.U, 1) == size(svdo.U, 2)) ? (svdo.U, svdo.Vt) : (svdo.Vt, svdo.U) + + svdio = svd(out_mat_nonsq * in_mat_nonsq') + + # retain variance + ret_var = get_retain_var(cc) + if ret_var < 1.0 + sv_cumsum = cumsum(svdio.S .^ 2) / sum(svdio.S .^ 2) # variance contributions are (sing_val)^2 + trunc = minimum(findall(x -> (x > ret_var), sv_cumsum)) + @info " truncating at $(trunc)/$(length(sv_cumsum)) retaining $(100.0*sv_cumsum[trunc])% of the variance in the joint space" + else + trunc = min(rank(in_data), rank(out_data)) + end + + if apply_to == "in" + in_dim = size(in_data, 1) + svdio_mat = (size(svdio.U, 1) == in_dim) ? svdio.U : svdio.V + # mat' * Sx⁻¹ * Uxt + encoder_mat = svdio_mat[:, 1:trunc]' * Diagonal(1 ./ svdi.S) * in_mat_sq' + decoder_mat = in_mat_sq * Diagonal(svdi.S) * svdio_mat[:, 1:trunc] + elseif apply_to == "out" + out_dim = size(out_data, 1) + svdio_mat = (size(svdio.U, 1) == out_dim) ? svdio.U : svdio.V + # Vt * Sy⁻¹ * Uyt + encoder_mat = svdio_mat[:, 1:trunc]' * Diagonal(1 ./ svdo.S) * out_mat_sq' + decoder_mat = out_mat_sq * Diagonal(svdo.S) * svdio_mat[:, 1:trunc] + end + + push!(get_encoder_mat(cc), encoder_mat) + push!(get_decoder_mat(cc), decoder_mat) + + # Note: To check CCA: + # u = in_encoder * (in - mean(in)) + # v = out_encoder * (out - mean(out)) + # u * u' = v * v' = I, + # v * u' = u * v' = Diagonal(svdio.S[1:trunc]) + end +end + +""" +$(TYPEDSIGNATURES) + +Computes and populates the `data_mean`, `encoder_mat`, `decoder_mat` and `apply_to` fields for the `CanonicalCorrelation` +""" +initialize_processor!( + cc::CanonicalCorrelation, + in_data::MM, + out_data::MM, + structure_matrix, + apply_to::AS, +) where {MM <: AbstractMatrix, AS <: AbstractString} = initialize_processor!(cc, in_data, out_data, apply_to) + + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` encoder, on a columns-are-data matrix +""" +function encode_data(cc::CanonicalCorrelation, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(cc)[1] + encoder_mat = get_encoder_mat(cc)[1] + return encoder_mat * (data .- data_mean) +end + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` decoder, on a columns-are-data matrix +""" +function decode_data(cc::CanonicalCorrelation, data::MM) where {MM <: AbstractMatrix} + data_mean = get_data_mean(cc)[1] + decoder_mat = get_decoder_mat(cc)[1] + return decoder_mat * data .+ data_mean +end + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` encoder to a provided structure matrix +""" +function encode_structure_matrix(cc::CanonicalCorrelation, structure_matrix::MM) where {MM <: AbstractMatrix} + encoder_mat = get_encoder_mat(cc)[1] + return encoder_mat * structure_matrix * encoder_mat' +end + +""" +$(TYPEDSIGNATURES) + +Apply the `CanonicalCorrelation` decoder to a provided structure matrix +""" +function decode_structure_matrix(cc::CanonicalCorrelation, enc_structure_matrix::MM) where {MM <: AbstractMatrix} + decoder_mat = get_decoder_mat(cc)[1] + return decoder_mat * enc_structure_matrix * decoder_mat' +end + + + + + +#### + + +# generic functions to initialize, encode, and decode data. With Data processors or PairedData processors +function initialize_and_encode_data!( + dcp::DCP, + data::MM, + structure_mat::USorM, +) where {DCP <: DataContainerProcessor, USorM <: Union{UniformScaling, AbstractMatrix}, MM <: AbstractMatrix} + initialize_processor!(dcp, data, structure_mat) + return encode_data(dcp, data) +end + +""" +$(TYPEDSIGNATURES) + +Initializes the `DataContainerProcessor` encoder (often requires data, and structure matrices), then encodes the provided columns-are-data matrix +""" +initialize_and_encode_data!( + dcp::DCP, + data::MM, + structure_mat::USorM, + apply_to::AS, +) where { + DCP <: DataContainerProcessor, + MM <: AbstractMatrix, + USorM <: Union{UniformScaling, AbstractMatrix}, + AS <: AbstractString, +} = initialize_and_encode_data!(dcp, data, structure_mat) + + +""" +$(TYPEDSIGNATURES) + +decodes the columns-are-data matrix with the processor. +""" +decode_data( + dcp::DCP, + data::MM, + apply_to::AS, +) where {DCP <: DataContainerProcessor, MM <: AbstractMatrix, AS <: AbstractString} = decode_data(dcp, data) + +""" +$(TYPEDSIGNATURES) + +Initializes the `PairedDataContainerProcesser` encoder (often requires input & output data, and structure matrices), then encodes either the input or output data (pair of columns-are-data matrices) based on `apply_to`. +""" +function initialize_and_encode_data!( + dcp::PDCP, + data, + structure_mat::USorM, + apply_to::AS, +) where {PDCP <: PairedDataContainerProcessor, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString} + input_data, output_data = data + initialize_processor!(dcp, input_data, output_data, structure_mat, apply_to) + + if apply_to == "in" + return encode_data(dcp, input_data) + elseif apply_to == "out" + return encode_data(dcp, output_data) + else + bad_apply_to(apply_to) + end +end + +""" +$(TYPEDSIGNATURES) + +decodes the input or output dat (pair of columns-are-data matrices) a with the processor, based on `apply_to`. +""" +function decode_data(dcp::PDCP, data, apply_to::AS) where {PDCP <: PairedDataContainerProcessor, AS <: AbstractString} + input_data, output_data = data + if apply_to == "in" + return decode_data(dcp, input_data) + elseif apply_to == "out" + return decode_data(dcp, output_data) + else + bad_apply_to(apply_to) + end +end + +""" +$TYPEDSIGNATURES + +Create a flatter encoder schedule for the +from the user's proposed schedule of the form: +```julia +enc_schedule = [ + (DataProcessor1(...), "in"), + (DataProcessor2(...), "out"), + (PairedDataProcessor3(...), "in"), + (DataProcessor4(...), "in_and_out"), +] +``` +This function creates the encoder scheduler that is also machine readable +```julia +enc_schedule = [ + (DataProcessor1(...), x -> get_inputs(x), "in"), + (DataProcessor2(...), x -> get_outputs(x), "out"), + (DataProcessor2(...), x -> get_outputs(x), "out"), + (PairedDataProcessor3(...), x -> (get_outputs(x), get_outputs(x)), "in"), + (DataProcessor4(...), x -> get_inputs(x), "in"), + (DataProcessor4(...), x -> get_outputs(x), "out"), +] +``` +and the decoder schedule is a copy of the encoder schedule reversed (and processors copied) +""" +function create_encoder_schedule(schedule_in::VV) where {VV <: AbstractVector} + + encoder_schedule = [] + for (processor, apply_to) in schedule_in + # converts the string into the extraction of data + if isa(processor, DataContainerProcessor) + if apply_to == "in" + func = x -> get_inputs(x) + push!(encoder_schedule, (processor, func, apply_to)) + + elseif apply_to == "out" + func = x -> get_outputs(x) + push!(encoder_schedule, (processor, func, apply_to)) + + elseif apply_to == "in_and_out" + func1 = x -> get_inputs(x) + func2 = x -> get_outputs(x) + push!(encoder_schedule, (processor, func1, "in")) + push!(encoder_schedule, (deepcopy(processor), func2, "out")) + else + @warn( + "Expected schedule keywords ∈ {\"in\",\"out\",\"in_and_out\"}. Received $(apply_to), ignoring processor $(processor)..." + ) + end + + # extract all the data (needed for the paired processor, but still pass through what you apply to) + elseif isa(processor, PairedDataContainerProcessor) + if apply_to ∈ ["in", "out"] + func = x -> (get_inputs(x), get_outputs(x)) + push!(encoder_schedule, (processor, func, apply_to)) + elseif apply_to == "in_and_out" + func = x -> (get_inputs(x), get_outputs(x)) + push!(encoder_schedule, (processor, func, "in")) + push!(encoder_schedule, (deepcopy(processor), func, "out")) + else + @warn( + "Expected schedule keywords ∈ {\"in\",\"out\",\"in_and_out\"}. Received $(apply_to), ignoring processor $(processor)..." + ) + end + end + end + + return encoder_schedule +end + +""" +Create size-1 encoder schedule with a tuple of `(DataProcessor1(...), apply_to)` with `apply_to = "in"`, `"out"` or `"in_and_out"`. +""" +create_encoder_schedule(schedule_in::TT) where {TT <: Tuple} = create_encoder_schedule([schedule_in]) + + +function bad_apply_to(apply_to::AS) where {AS <: AbstractString} + throw( + ArgumentError( + "processer can only be applied to inputs (\"in\") or outputs (\"out\"). received $(apply_to). \n Please use `create_encoder_schedule` prior to encoding/decoding to ensure correct schedule format", + ), + ) +end + +# Functions to encode/decode with uninitialized schedule (require structure matrices as input) +""" +$TYPEDSIGNATURES + +Takes in the created encoder schedule (See [`create_encoder_schedule`](@ref)), and initializes it, and encodes the paired data container, and structure matrices with it. +""" +function encode_with_schedule( + encoder_schedule::VV, + io_pairs::PDC, + input_structure_mat::USorM1, + output_structure_mat::USorM2, +) where { + VV <: AbstractVector, + PDC <: PairedDataContainer, + USorM1 <: Union{UniformScaling, AbstractMatrix}, + USorM2 <: Union{UniformScaling, AbstractMatrix}, +} + processed_io_pairs = deepcopy(io_pairs) + processed_input_structure_mat = deepcopy(input_structure_mat) + processed_output_structure_mat = deepcopy(output_structure_mat) + + # apply_to is the string "in", "out" etc. + for (processor, extract_data, apply_to) in encoder_schedule + @info "Initialize encoding of data: \"$(apply_to)\" with $(processor)" + if apply_to == "in" + structure_matrix = processed_input_structure_mat + elseif apply_to == "out" + structure_matrix = processed_output_structure_mat + else + bad_apply_to(apply_to) + end + processed = initialize_and_encode_data!(processor, extract_data(processed_io_pairs), structure_matrix, apply_to) + if apply_to == "in" + processed_input_structure_mat = encode_structure_matrix(processor, structure_matrix) + processed_io_pairs = PairedDataContainer(processed, get_outputs(processed_io_pairs)) + elseif apply_to == "out" + processed_output_structure_mat = encode_structure_matrix(processor, structure_matrix) + processed_io_pairs = PairedDataContainer(get_inputs(processed_io_pairs), processed) + end + end + + return processed_io_pairs, processed_input_structure_mat, processed_output_structure_mat +end + +# Functions to encode/decode with initialized schedule +""" +$TYPEDSIGNATURES + +Takes in an already initialized encoder schedule, and encodes a `DataContainer`, the `in_or_out` string indicates if the data is input `"in"` or output `"out"` data (and thus encoded differently) +""" +function encode_with_schedule( + encoder_schedule::VV, + data_container::DC, + in_or_out::AS, +) where {VV <: AbstractVector, DC <: DataContainer, AS <: AbstractString} + + if in_or_out ∉ ["in", "out"] + bad_in_or_out(in_or_out) + end + processed_container = deepcopy(data_container) + + # apply_to is the string "in", "out" etc. + for (processor, extract_data, apply_to) in encoder_schedule + if apply_to == in_or_out + processed = encode_data(processor, get_data(processed_container)) + processed_container = DataContainer(processed) + + end + end + + return processed_container +end + +function bad_in_or_out(in_or_out::AS) where {AS <: AbstractString} + throw( + ArgumentError( + "`in_or_out` must be either \"in\" (data is an input) or \"out\" (data is an output). Received $(in_or_out)", + ), + ) +end + + +""" +$TYPEDSIGNATURES + +Takes in an already initialized encoder schedule, and encodes a structure matrix, the `in_or_out` string indicates if the structure matrix is for input `"in"` or output `"out"` space (and thus encoded differently) +""" +function encode_with_schedule( + encoder_schedule::VV, + structure_matrix::USorM, + in_or_out::AS, +) where {VV <: AbstractVector, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString} + + if !(in_or_out ∈ ["in", "out"]) + bad_in_or_out(in_or_out) + end + processed_structure_matrix = deepcopy(structure_matrix) + + # apply_to is the string "in", "out" etc. + for (processor, extract_data, apply_to) in encoder_schedule + if apply_to == in_or_out + processed_structure_matrix = encode_structure_matrix(processor, processed_structure_matrix) + end + end + + return processed_structure_matrix +end + +""" +$TYPEDSIGNATURES + +Takes in an already initialized encoder schedule, and decodes a `DataContainer`, and structure matrices with it, the `in_or_out` string indicates if the data is input `"in"` or output `"out"` data (and thus decoded differently) +""" +function decode_with_schedule( + encoder_schedule::VV, + io_pairs::PDC, + input_structure_mat::USorM1, + output_structure_mat::USorM2, +) where { + VV <: AbstractVector, + USorM1 <: Union{UniformScaling, AbstractMatrix}, + USorM2 <: Union{UniformScaling, AbstractMatrix}, + PDC <: PairedDataContainer, +} + + processed_io_pairs = deepcopy(io_pairs) + processed_input_structure_mat = deepcopy(input_structure_mat) + processed_output_structure_mat = deepcopy(output_structure_mat) + + # apply_to is the string "in", "out" etc. + for idx in reverse(eachindex(encoder_schedule)) + (processor, extract_data, apply_to) = encoder_schedule[idx] + processed = decode_data(processor, extract_data(processed_io_pairs), apply_to) + + if apply_to == "in" + processed_input_structure_mat = decode_structure_matrix(processor, processed_input_structure_mat) + processed_io_pairs = PairedDataContainer(processed, get_outputs(processed_io_pairs)) + elseif apply_to == "out" + processed_output_structure_mat = decode_structure_matrix(processor, processed_output_structure_mat) + processed_io_pairs = PairedDataContainer(get_inputs(processed_io_pairs), processed) + else + bad_apply_to(apply_to) + end + end + + return processed_io_pairs, processed_input_structure_mat, processed_output_structure_mat +end + +""" +$TYPEDSIGNATURES + +Takes in an already initialized encoder schedule, and decodes a `DataContainer`, the `in_or_out` string indicates if the data is input `"in"` or output `"out"` data (and thus decoded differently) +""" +function decode_with_schedule( + encoder_schedule::VV, + data_container::DC, + in_or_out::AS, +) where {VV <: AbstractVector, DC <: DataContainer, AS <: AbstractString} + + if in_or_out ∉ ["in", "out"] + bad_in_or_out(in_or_out) + end + processed_container = deepcopy(data_container) + + # apply_to is the string "in", "out" etc. + for idx in reverse(eachindex(encoder_schedule)) + (processor, extract_data, apply_to) = encoder_schedule[idx] + if apply_to == in_or_out + processed = decode_data(processor, get_data(processed_container)) + processed_container = DataContainer(processed) + + end + end + + return processed_container +end + +""" +$TYPEDSIGNATURES + +Takes in an already initialized encoder schedule, and decodes a structure matrix, the `in_or_out` string indicates if the structure matrix is for input `"in"` or output `"out"` space (and thus decoded differently) +""" +function decode_with_schedule( + encoder_schedule::VV, + structure_matrix::USorM, + in_or_out::AS, +) where {VV <: AbstractVector, USorM <: Union{UniformScaling, AbstractMatrix}, AS <: AbstractString} + + if in_or_out ∉ ["in", "out"] + bad_in_or_out(in_or_out) + end + processed_structure_matrix = deepcopy(structure_matrix) + + # apply_to is the string "in", "out" etc. + for idx in reverse(eachindex(encoder_schedule)) + (processor, extract_data, apply_to) = encoder_schedule[idx] + if apply_to == in_or_out + processed_structure_matrix = decode_structure_matrix(processor, processed_structure_matrix) + end + end + + return processed_structure_matrix +end + + + end # module diff --git a/test/Utilities/runtests.jl b/test/Utilities/runtests.jl index 7fe46f522..687532b22 100644 --- a/test/Utilities/runtests.jl +++ b/test/Utilities/runtests.jl @@ -1,6 +1,7 @@ using Test using Random using Statistics +using Distributions using LinearAlgebra using CalibrateEmulateSample.Utilities @@ -68,4 +69,291 @@ using CalibrateEmulateSample.DataContainers @test isposdef(pdmat2) @test minimum(eigvals(pdmat2)) >= (1 - 1e-4) * tol +end + + +@testset "Data Preprocessing" begin + + # Seed for pseudo-random number generator + rng = Random.MersenneTwister(4154) + + # quick build tests and test getters + zs = zscore_scale() + mm = minmax_scale() + qq = quartile_scale() + QQ = ElementwiseScaler{QuartileScaling, Vector{Int}}([1], [2]) + @test isa(zs, ElementwiseScaler) + @test get_type(zs) == ZScoreScaling + @test isa(mm, ElementwiseScaler) + @test get_type(mm) == MinMaxScaling + @test isa(qq, ElementwiseScaler) + @test get_type(qq) == QuartileScaling + @test get_shift(QQ) == [1] + @test get_scale(QQ) == [2] + + dd = decorrelate() + @test get_retain_var(dd) == 1.0 + @test get_decorrelate_with(dd) == "combined" + dd2 = decorrelate_sample_cov(retain_var = 0.7) + @test get_retain_var(dd2) == 0.7 + @test get_decorrelate_with(dd2) == "sample_cov" + dd3 = decorrelate_structure_mat(retain_var = 0.7) + @test get_retain_var(dd3) == 0.7 + @test get_decorrelate_with(dd3) == "structure_mat" + DD = Decorrelater([1], [2], [3], 1.0, "test") + @test get_data_mean(DD) == [1] + @test get_encoder_mat(DD) == [2] + @test get_decoder_mat(DD) == [3] + # get some data as IO pairs for functional tests + + in_dim = 10 + out_dim = 50 + samples = 120 + + x = randn(rng, in_dim, in_dim) + prior_cov = x * x' + in_data = rand(rng, MvNormal(zeros(in_dim), prior_cov), samples) + obs_noise_cov = [max(5.0 - abs(i - j), 0.0) for i in 1:out_dim, j in 1:out_dim] # [5 4 3 2 1 0 0 ...] off diagonal + out_data = rand(rng, MvNormal(-10 * ones(out_dim), obs_noise_cov), samples) + + io_pairs = PairedDataContainer(in_data, out_data) + test_names = [ # order as in schedules below + "zscore", + "quartile", + "minmax", + "decorrelate-sample-cov", + "decorrelate-structure-mat", + "decorrelate-combined", + "canonical-correlation", + "decorrelate-structure-mat-retain-0.95-var", + "canonical-correlation-0.95-var", + ] + + # Test encodings-decodings individually + schedules = [ + (zscore_scale(), "in_and_out"), + (quartile_scale(), "in_and_out"), + (minmax_scale(), "in_and_out"), + (decorrelate_sample_cov(), "in_and_out"), + (decorrelate_structure_mat(), "in_and_out"), + (decorrelate(), "in_and_out"), # combined + (canonical_correlation(), "in_and_out"), + (decorrelate_structure_mat(retain_var = 0.95), "in_and_out"), + (canonical_correlation(retain_var = 0.95), "in_and_out"), + ] + + lossless = [fill(true, 6); fill(false, 4)] # are these lossy approximations? + + # functional test pipeline + tol = 1e-12 + + for (name, sch, ll_flag) in zip(test_names, schedules, lossless) + encoder_schedule = create_encoder_schedule(sch) + (encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) = + encode_with_schedule(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + + (decoded_io_pairs, decoded_prior_cov, decoded_obs_noise_cov) = + decode_with_schedule(encoder_schedule, encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) + for (enc_dat, dec_dat, test_dat, enc_covv, dec_covv, test_covv, dim) in zip( + (get_inputs(encoded_io_pairs), get_outputs(encoded_io_pairs)), + (get_inputs(decoded_io_pairs), get_outputs(decoded_io_pairs)), + (get_inputs(io_pairs), get_outputs(io_pairs)), + (encoded_prior_cov, encoded_obs_noise_cov), + (decoded_prior_cov, decoded_obs_noise_cov), + (prior_cov, obs_noise_cov), + (in_dim, out_dim), + ) + # univariate "rescaling" tests + if name == "zscore" + stat_vec = [[mean(dd), std(dd)] for dd in eachrow(enc_dat)] + stat_mat = reduce(hcat, stat_vec) + @test all(isapprox.(stat_mat[1, :], zeros(dim), atol = tol)) + @test all(isapprox.(stat_mat[2, :], ones(dim), atol = tol)) + + test_vec = [[mean(dd), std(dd)] for dd in eachrow(test_dat)] + test_mat = reduce(hcat, test_vec) + @test isapprox( + norm(enc_covv - Diagonal(1 ./ test_mat[2, :]) * test_covv * Diagonal(1 ./ test_mat[2, :])), + 0.0, + atol = tol * dim^2, + ) + elseif name == "quartile" + quartiles_vec = [quantile(dd, [0.25, 0.5, 0.75]) for dd in eachrow(enc_dat)] + quartiles_mat = reduce(hcat, quartiles_vec) # 3 rows: Q1, Q2, and Q3 + @test all(isapprox.(quartiles_mat[2, :], zeros(dim), atol = tol)) + @test all(isapprox.(quartiles_mat[3, :] - quartiles_mat[1, :], ones(dim), atol = tol)) + test_vec = [quantile(dd, [0.25, 0.5, 0.75]) for dd in eachrow(test_dat)] + test_mat = reduce(hcat, test_vec) + @test isapprox( + norm( + enc_covv - + Diagonal(1 ./ (test_mat[3, :] - test_mat[1, :])) * + test_covv * + Diagonal(1 ./ (test_mat[3, :] - test_mat[1, :])), + ), + 0.0, + atol = tol * dim^2, + ) + elseif name == "minmax" + minmax_vec = [[minimum(dd), maximum(dd)] for dd in eachrow(enc_dat)] + minmax_mat = reduce(hcat, minmax_vec) # 2 rows: min, max + @test all(isapprox.(minmax_mat[1, :], zeros(dim), atol = tol)) + @test all(isapprox.(minmax_mat[2, :], ones(dim), atol = tol)) + test_vec = [[minimum(dd), maximum(dd)] for dd in eachrow(test_dat)] + test_mat = reduce(hcat, test_vec) + @test isapprox( + norm( + enc_covv - + Diagonal(1 ./ (test_mat[2, :] - test_mat[1, :])) * + test_covv * + Diagonal(1 ./ (test_mat[2, :] - test_mat[1, :])), + ), + 0.0, + atol = tol * dim^2, + ) + end + + # Multivariate lossless tests + pop_mean = mean(enc_dat, dims = 2) + pop_cov = cov(enc_dat, dims = 2) + dimm = size(pop_cov, 1) + big_tol = 0.1 + if name == "decorrelate-structure-mat" + @test all(isapprox.(pop_mean, zeros(dimm), atol = tol)) + @test isapprox(norm(pop_cov - I), 0.0, atol = big_tol * dimm^2) # expect poorly accurate + @test isapprox(norm(enc_covv - I), 0.0, atol = tol * dimm^2) # expect very accurate + + elseif name == "decorrelate-sample-cov" + @test all(isapprox.(pop_mean, zeros(dimm), atol = tol)) + @test isapprox(norm(pop_cov - I), 0.0, atol = tol * dimm^2) # expect very accurate + @test isapprox(norm(enc_covv - I), 0.0, atol = big_tol * dimm^2) # expect poorly accurate, particularly if dimm < dim + + elseif name == "decorrelate-combined" + @test all(isapprox.(pop_mean, zeros(dimm), atol = tol)) + @test isapprox(norm(pop_cov - I), 0.0, atol = big_tol * dimm^2) # expect poorly accurate + @test isapprox(norm(enc_covv - I), 0.0, atol = big_tol * dimm^2) # expect poorly accurate + end + + # Multivariate lossy dim-reduction tests + if name == "decorrelate-structure-mat-retain-0.95-var" + svdc = svd(test_covv) + var_cumsum = cumsum(svdc.S .^ 2) ./ sum(svdc.S .^ 2) + @test var_cumsum[dimm] > 0.95 + @test var_cumsum[dimm - 1] < 0.95 + @test all(isapprox.(pop_mean, zeros(dimm), atol = tol)) + @test isapprox(norm(pop_cov - I), 0.0, atol = big_tol * dimm^2) # expect poorly accurate + @test isapprox(norm(enc_covv - I), 0.0, atol = tol * dimm^2) # expect very accurate + end + + # Paired data processor reduction: + if name == "canonical-correlation" + @test dimm == min(rank(get_inputs(io_pairs)), rank(get_outputs(io_pairs))) + @test isapprox(norm(enc_dat * enc_dat' - I), 0.0, atol = tol * dimm^2) # test in or out orthogonality + + # check cross-orthogonality is diagonal (nb this test will be duplicate) + enc_in = get_inputs(encoded_io_pairs) + enc_out = get_outputs(encoded_io_pairs) + @test isapprox(norm(enc_in * enc_out' - Diagonal(diag(enc_in * enc_out'))), 0.0, atol = tol * dimm^2) # test cross - orthogonality + + @test isapprox(norm(enc_out * enc_in' - Diagonal(diag(enc_out * enc_in'))), 0.0, atol = tol * dimm^2) # test cross - orthogonality + + # decoder test is lossless only for smaller dimension + if dim == min(in_dim, out_dim) + @test isapprox(norm(dec_dat - test_dat), 0.0, atol = tol * dim * samples) + @test isapprox(norm(dec_covv - test_covv), 0.0, atol = tol * dim^2) + end + + end + + # test decode approximation of lossless options + if ll_flag + # when dimm < dim, loss can occur in some tests + tol1 = (name == "decorrelate-structure-mat" && dimm < dim) ? big_tol : tol + tol2 = (name == "decorrelate-sample-cov" && dimm < dim) ? big_tol : tol + @test isapprox(norm(dec_dat - test_dat), 0.0, atol = tol1 * dim * samples) + @test isapprox(norm(dec_covv - test_covv), 0.0, atol = tol2 * dim^2) + + end + + end + + end + + + + # combine a few lossless encoding schedules (lossless requires samples>dims) + samples = 150 # for full test coverage have samples in_dim < samples < out_dim + in_data = rand(MvNormal(zeros(in_dim), prior_cov), samples) + out_data = rand(MvNormal(-10 * ones(out_dim), obs_noise_cov), samples) + io_pairs = PairedDataContainer(in_data, out_data) + + schedule_builder = [ + (zscore_scale(), "in_and_out"), + (quartile_scale(), "in"), + (decorrelate_sample_cov(), "in_and_out"), + (minmax_scale(), "out"), + (decorrelate_structure_mat(), "in_and_out"), + (canonical_correlation(), "in"), + ] + + # make schedule more parsable + + @test_logs (:warn,) create_encoder_schedule((canonical_correlation(), "bad")) + @test_logs (:warn,) create_encoder_schedule((zscore_scale(), "bad")) + func = x -> (get_inputs(x), get_outputs(x)) + bad_encoder_schedule = [(canonical_correlation(), func, "bad")] + @test_throws ArgumentError encode_with_schedule(bad_encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + @test_throws ArgumentError decode_with_schedule(bad_encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + @test_throws ArgumentError initialize_and_encode_data!(canonical_correlation(), func(io_pairs), prior_cov, "bad") + @test_throws ArgumentError decode_data(canonical_correlation(), func(io_pairs), "bad") + + + encoder_schedule = create_encoder_schedule(schedule_builder) + + # encode the data using the schedule + (encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) = + encode_with_schedule(encoder_schedule, io_pairs, prior_cov, obs_noise_cov) + + # decode the data using the schedule + (decoded_io_pairs, decoded_prior_cov, decoded_obs_noise_cov) = + decode_with_schedule(encoder_schedule, encoded_io_pairs, encoded_prior_cov, encoded_obs_noise_cov) + + tol = 1e-12 + @test all(isapprox.(get_inputs(io_pairs), get_inputs(decoded_io_pairs), atol = tol)) + @test all(isapprox.(get_outputs(io_pairs), get_outputs(decoded_io_pairs), atol = tol)) + @test isapprox(norm(prior_cov - decoded_prior_cov), 0.0, atol = tol * in_dim^2) + @test isapprox(norm(obs_noise_cov - decoded_obs_noise_cov), 0.0, atol = tol * out_dim^2) + + # enc/dec just data + samples = 1 # try one sample + new_in_data = rand(rng, MvNormal(zeros(in_dim), prior_cov), samples) + new_out_data = rand(rng, MvNormal(-10 * ones(out_dim), obs_noise_cov), samples) + id = DataContainer(new_in_data) + od = DataContainer(new_out_data) + + encoded_id = encode_with_schedule(encoder_schedule, id, "in") + encoded_od = encode_with_schedule(encoder_schedule, od, "out") + decoded_id = decode_with_schedule(encoder_schedule, encoded_id, "in") + decoded_od = decode_with_schedule(encoder_schedule, encoded_od, "out") + # tests in latent space? + @test isapprox(norm(get_data(decoded_id) - get_data(id)), 0.0, atol = tol * in_dim * samples) + @test isapprox(norm(get_data(decoded_od) - get_data(od)), 0.0, atol = tol * out_dim * samples) + @test_throws ArgumentError encode_with_schedule(encoder_schedule, id, "bad") + @test_throws ArgumentError decode_with_schedule(encoder_schedule, id, "bad") + + # enc/dec just mats + encoded_pc = encode_with_schedule(encoder_schedule, prior_cov, "in") + encoded_oc = encode_with_schedule(encoder_schedule, obs_noise_cov, "out") + decoded_pc = decode_with_schedule(encoder_schedule, encoded_pc, "in") + decoded_oc = decode_with_schedule(encoder_schedule, encoded_oc, "out") + @test isapprox(norm(encoded_pc - encoded_prior_cov), 0.0, atol = tol * in_dim^2) + @test isapprox(norm(encoded_oc - encoded_obs_noise_cov), 0.0, atol = tol * out_dim^2) + @test isapprox(norm(decoded_pc - decoded_prior_cov), 0.0, atol = tol * in_dim^2) + @test isapprox(norm(decoded_oc - decoded_obs_noise_cov), 0.0, atol = tol * out_dim^2) + @test_throws ArgumentError encode_with_schedule(encoder_schedule, prior_cov, "bad") + @test_throws ArgumentError decode_with_schedule(encoder_schedule, encoded_pc, "bad") + + + + end