Skip to content

Commit cde5b82

Browse files
ArneBouillonArneBouillon
committed
Add first draft of likelihood-informed data processor
1 parent 8152241 commit cde5b82

3 files changed

Lines changed: 202 additions & 2 deletions

File tree

‎Project.toml‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ GaussianProcesses = "891a1506-143c-57d2-908e-e1f8e92e6de9"
1717
KernelFunctions = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
1818
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
1919
MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
20+
Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
21+
Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
2022
Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
2123
Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
2224
ProgressBars = "49802e3a-d2f1-5c88-81d8-b72133a6f568"
@@ -42,6 +44,8 @@ ForwardDiff = "0.10.38, 1"
4244
GaussianProcesses = "0.12"
4345
KernelFunctions = "0.10.64"
4446
MCMCChains = "4.14, 5, 6, 7"
47+
Manifolds = "0.10.23"
48+
Manopt = "0.5.20"
4549
Printf = "1"
4650
ProgressBars = "1"
4751
PyCall = "1.93"

‎src/Utilities.jl‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,5 @@
11
module Utilities
22

3-
4-
53
using DocStringExtensions
64
using LinearAlgebra
75
using Statistics

‎src/Utilities/likelihood_informed.jl‎

Lines changed: 198 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,198 @@
1+
# included in Utilities.jl
2+
3+
using Manifolds, Manopt
4+
5+
mutable struct LikelihoodInformed{FT <: Real} <: PairedDataContainerProcessor
6+
encoder_mat::Union{Nothing, AbstractMatrix}
7+
decoder_mat::Union{Nothing, AbstractMatrix}
8+
apply_to::Union{Nothing, AbstractString}
9+
dim_criterion::Tuple{Symbol, <:Number}
10+
α::FT
11+
grad_type::Symbol
12+
use_prior_samples::Bool
13+
end
14+
15+
function likelihood_informed(retain_KL; alpha = 0.0, grad_type = :localsl, use_prior_samples = true)
16+
if grad_type ∉ [:linreg, :localsl]
17+
@error "Unknown grad_type=$grad_type"
18+
end
19+
20+
LikelihoodInformed(nothing, nothing, nothing, (:retain_KL, retain_KL), alpha, grad_type, use_prior_samples)
21+
end
22+
23+
get_encoder_mat(li::LikelihoodInformed) = li.encoder_mat
24+
get_decoder_mat(li::LikelihoodInformed) = li.decoder_mat
25+
26+
function initialize_processor!(
27+
li::LikelihoodInformed,
28+
in_data::MM,
29+
out_data::MM,
30+
::Dict{Symbol, <:StructureMatrix},
31+
output_structure_matrices::Dict{Symbol, <:StructureMatrix},
32+
input_structure_vectors::Dict{Symbol, <:StructureVector},
33+
output_structure_vectors::Dict{Symbol, <:StructureVector},
34+
apply_to::AbstractString,
35+
) where {MM <: AbstractMatrix}
36+
output_dim = size(out_data, 2)
37+
38+
if isnothing(get_encoder_mat(li))
39+
α = li.α
40+
y = if α ≈ 0.0
41+
# For α=0, it doesn't matter what this value is, so we avoid requiring its presence
42+
zeros(size(out_data, 1))
43+
else
44+
get_structure_vec(output_structure_vectors, :observation)
45+
end
46+
samples_in, samples_out = if li.use_prior_samples
47+
@assert α ≈ 0.0
48+
(
49+
get_structure_vec(input_structure_vectors, :prior_samples_in),
50+
get_structure_vec(output_structure_vectors, :prior_samples_out),
51+
)
52+
else
53+
(in_data, out_data)
54+
end
55+
obs_noise_cov = get_structure_mat(output_structure_matrices, :obs_noise_cov)
56+
noise_cov_inv = inv(obs_noise_cov)
57+
58+
li.apply_to = apply_to
59+
60+
grads = if li.grad_type == :linreg
61+
grad = (samples_out .- mean(samples_out; dims = 2)) / (samples_in .- mean(samples_in; dims = 2))
62+
fill(grad, size(samples_in, 2))
63+
else
64+
@assert li.grad_type == :localsl
65+
66+
map(eachcol(samples_in)) do u
67+
# TODO: It might be interesting to introduce a parameter to weight this distance with.
68+
# This can be a scalar or a matrix; in the latter case, we can even use the covariance
69+
# of the samples (or the prior covariance).
70+
weights = exp.(-1/2 * norm.(eachcol(u .- samples_in)).^2)
71+
D = Diagonal(sqrt.(weights))
72+
uw = (samples_in .- mean(samples_in * Diagonal(weights); dims = 2)) * D
73+
gw = (samples_out .- mean(samples_out * Diagonal(weights); dims = 2)) * D
74+
gw / uw
75+
end
76+
end
77+
78+
li.encoder_mat = if apply_to == "in" || α ≈ 0
79+
decomp = if apply_to == "in"
80+
eigen(mean(grad' * noise_cov_inv * ((1-α)obs_noise_cov + α^2 * (y - g) * (y - g)') * noise_cov_inv * grad for (g, grad) in zip(eachcol(samples_out), grads)), sortby = (-))
81+
else
82+
@assert apply_to == "out"
83+
eigen(mean(grad * grad' for grad in grads), obs_noise_cov, sortby = (-))
84+
end
85+
86+
if li.dim_criterion[1] == :retain_KL
87+
retain_KL = li.dim_criterion[2]
88+
sv_cumsum = cumsum(decomp.values) / sum(decomp.values)
89+
trunc_val = findfirst(x -> (x ≥ retain_KL), sv_cumsum)
90+
else
91+
@assert li.dim_criterion[1] == :dimension
92+
trunc_val = li.dim_criterion[2]
93+
end
94+
li.encoder_mat = decomp.vectors[:, 1:trunc_val]'
95+
else
96+
@assert apply_to == "out"
97+
@warn "Using LikelihoodInformed on output data with α≠0 triggers a manifold optimization process that may take some time."
98+
99+
k = if li.dim_criterion[1] == :retain_KL
100+
1
101+
else
102+
@assert li.dim_criterion[1] == :dimension
103+
li.dim_criterion[2]
104+
end
105+
Vs = nothing
106+
while true
107+
M = Grassmann(output_dim, k)
108+
109+
f = (_, Vs) -> begin
110+
prec = noise_cov_inv - Vs * inv(Vs' * obs_noise_cov * Vs) * Vs'
111+
tr(mean(
112+
grad' * prec * ((1-α)I + α^2 * (y - g)*(y - g)') * prec * grad
113+
for (g, grad) in zip(eachcol(out_data), grads)
114+
))
115+
end
116+
egrad = (_, Vs) -> begin
117+
B = Vs * inv(Vs' * obs_noise_cov * Vs) * Vs'
118+
prec = noise_cov_inv - B
119+
120+
121+
-2mean(begin
122+
A = ((1-α)I + α^2 * (y - g)*(y - g)')
123+
S = grad * grad'
124+
(I - obs_noise_cov * B) * (S * prec * A + A * prec * S)
125+
end for (g, grad) in zip(eachcol(out_data), grads)) * B * Vs
126+
end
127+
rgrad = (M, Vs) -> begin
128+
(I - Vs*Vs') * egrad(M, Vs)
129+
end
130+
131+
Vs = Matrix(qr(randn(output_dim, k)))
132+
quasi_Newton!(M, f, rgrad, Vs; stopping_criterion = StopWhenGradientNormLess(3.0))
133+
134+
if li.dim_criterion[1] == :retain_KL
135+
retain_KL = li.dim_criterion[2]
136+
ref = f(M, zeros(output_dim, 0))
137+
if f(M, Vs) / ref ≤ 1 - retain_KL
138+
break # TODO: Start bisecting?
139+
else
140+
k *= 2
141+
end
142+
else
143+
@assert li.dim_criterion[1] == :dimension
144+
break
145+
end
146+
end
147+
148+
Vs'
149+
end
150+
li.decoder_mat = li.encoder_mat'
151+
end
152+
end
153+
154+
"""
155+
$(TYPEDSIGNATURES)
156+
157+
Apply the `LikelihoodInformed` encoder, on a columns-are-data matrix
158+
"""
159+
function encode_data(li::LikelihoodInformed, data::MM) where {MM <: AbstractMatrix}
160+
encoder_mat = get_encoder_mat(li)
161+
return encoder_mat * data
162+
end
163+
164+
"""
165+
$(TYPEDSIGNATURES)
166+
167+
Apply the `LikelihoodInformed` decoder, on a columns-are-data matrix
168+
"""
169+
function decode_data(li::LikelihoodInformed, data::MM) where {MM <: AbstractMatrix}
170+
decoder_mat = get_decoder_mat(li)
171+
return decoder_mat * data
172+
end
173+
174+
"""
175+
$(TYPEDSIGNATURES)
176+
177+
Apply the `LikelihoodInformed` encoder to a provided structure matrix
178+
"""
179+
function encode_structure_matrix(
180+
li::LikelihoodInformed,
181+
structure_matrix::SM,
182+
) where {SM <: StructureMatrix}
183+
encoder_mat = get_encoder_mat(li)
184+
return encoder_mat * structure_matrix * encoder_mat'
185+
end
186+
187+
"""
188+
$(TYPEDSIGNATURES)
189+
190+
Apply the `LikelihoodInformed` decoder to a provided structure matrix
191+
"""
192+
function decode_structure_matrix(
193+
li::LikelihoodInformed,
194+
structure_matrix::SM,
195+
) where {SM <: StructureMatrix}
196+
decoder_mat = get_decoder_mat(li)
197+
return decoder_mat * structure_matrix * decoder_mat'
198+
end

0 commit comments

Comments
 (0)