Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
ift
NIFTy
Commits
1c4f555b
Commit
1c4f555b
authored
Feb 25, 2021
by
Jakob Knollmüller
Committed by
Philipp Arras
May 30, 2021
Browse files
Add parametric variational inference
parent
0acc9ff7
Changes
7
Hide whitespace changes
Inline
Side-by-side
demos/meanfield_demo.py
0 → 100644
View file @
1c4f555b
import
nifty7
as
ift
from
matplotlib
import
pyplot
as
plt
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
# Copyright(C) 2013-2020 Max-Planck-Society
#
# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
###############################################################################
# Log-normal field reconstruction from Poissonian data with inhomogenous
# exposure (in case for 2D mode)
# 1D (set mode=0), 2D (mode=1), or on the sphere (mode=2)
###############################################################################
import
sys
import
numpy
as
np
# def exposure_2d(domain):
# # Structured exposure for 2D mode
# x_shape, y_shape = domain.shape
# exposure = np.ones(domain.shape)
# exposure[x_shape//3:x_shape//2, :] *= 2.
# exposure[x_shape*4//5:x_shape, :] *= .1
# exposure[x_shape//2:x_shape*3//2, :] *= 3.
# exposure[:, x_shape//3:x_shape//2] *= 2.
# exposure[:, x_shape*4//5:x_shape] *= .1
# exposure[:, x_shape//2:x_shape*3//2] *= 3.
# return ift.Field.from_raw(domain, exposure)
if
__name__
==
'__main__'
:
# Two-dimensional regular grid with inhomogeneous exposure
position_space
=
ift
.
RGSpace
([
100
])
# Define harmonic space and harmonic transform
harmonic_space
=
position_space
.
get_default_codomain
()
HT
=
ift
.
HarmonicTransformOperator
(
harmonic_space
,
position_space
)
# Domain on which the field's degrees of freedom are defined
domain
=
ift
.
DomainTuple
.
make
(
harmonic_space
)
# Define amplitude (square root of power spectrum)
def
sqrtpspec
(
k
):
return
1.
/
(
1.
+
k
**
2
)
p_space
=
ift
.
PowerSpace
(
harmonic_space
)
pd
=
ift
.
PowerDistributor
(
harmonic_space
,
p_space
)
a
=
ift
.
PS_field
(
p_space
,
sqrtpspec
)
A
=
pd
(
a
)
# Define sky operator
sky
=
10
*
ift
.
exp
(
HT
(
ift
.
makeOp
(
A
))).
ducktape
(
'xi'
)
# M = ift.DiagonalOperator(exposure)
GR
=
ift
.
GeometryRemover
(
position_space
)
# Define instrumental response
# R = GR(M)
R
=
GR
# Generate mock data and define likelihood operator
d_space
=
R
.
target
[
0
]
lamb
=
R
(
sky
)
mock_position
=
ift
.
from_random
(
sky
.
domain
,
'normal'
)
data
=
lamb
(
mock_position
)
data
=
ift
.
random
.
current_rng
().
poisson
(
data
.
val
.
astype
(
np
.
float64
))
data
=
ift
.
Field
.
from_raw
(
d_space
,
data
)
likelihood
=
ift
.
PoissonianEnergy
(
data
)
@
lamb
# Settings for minimization
ic_newton
=
ift
.
DeltaEnergyController
(
name
=
'Newton'
,
iteration_limit
=
1
,
tol_rel_deltaE
=
1e-8
)
# minimizer = ift.L_BFGS(ic_newton)
minimizer_fc
=
ift
.
ADVIOptimizer
(
steps
=
10
)
minimizer_mf
=
ift
.
ADVIOptimizer
(
steps
=
10
)
# Compute MAP solution by minimizing the information Hamiltonian
H
=
ift
.
StandardHamiltonian
(
likelihood
)
# initial_position = ift.from_random(domain, 'normal')
# meanfield_model = ift.MeanfieldModel(H.domain)
fullcov_model
=
ift
.
FullCovarianceModel
(
H
.
domain
)
meanfield_model
=
ift
.
MeanfieldModel
(
H
.
domain
)
position_fc
=
fullcov_model
.
get_initial_pos
(
initial_sig
=
0.01
)
position_mf
=
meanfield_model
.
get_initial_pos
(
initial_sig
=
0.01
)
KL_fc
=
ift
.
ParametricGaussianKL
.
make
(
position_fc
,
H
,
fullcov_model
,
3
,
True
)
KL_mf
=
ift
.
ParametricGaussianKL
.
make
(
position_mf
,
H
,
meanfield_model
,
3
,
True
)
# plt.figure('data')
# plt.imshow(sky(mock_position).val)
plt
.
pause
(
0.001
)
for
i
in
range
(
3000
):
# KL = ParametricGaussianKL.make(position,H,meanfield_model,3,True)
KL_fc
,
_
=
minimizer_fc
(
KL_fc
)
KL_mf
,
_
=
minimizer_mf
(
KL_mf
)
plt
.
figure
(
'result'
)
plt
.
cla
()
plt
.
plot
(
sky
(
fullcov_model
.
generator
(
KL_fc
.
position
)).
val
,
'b-'
,
label
=
'fc'
)
plt
.
plot
(
sky
(
meanfield_model
.
generator
(
KL_mf
.
position
)).
val
,
'r-'
,
label
=
'mf'
)
for
samp
in
KL_fc
.
samples
:
plt
.
plot
(
sky
(
fullcov_model
.
generator
(
KL_fc
.
position
+
samp
)).
val
,
'b-'
,
alpha
=
0.3
)
for
samp
in
KL_mf
.
samples
:
plt
.
plot
(
sky
(
meanfield_model
.
generator
(
KL_mf
.
position
+
samp
)).
val
,
'r-'
,
alpha
=
0.3
)
plt
.
plot
(
data
.
val
,
'kx'
)
plt
.
plot
(
sky
(
mock_position
).
val
,
'k-'
,
label
=
'true'
)
plt
.
legend
()
plt
.
ylim
(
0
,
data
.
val
.
max
()
+
10
)
plt
.
pause
(
0.001
)
\ No newline at end of file
demos/mgvi_visualized.py
View file @
1c4f555b
...
...
@@ -61,26 +61,27 @@ def main():
return
lh
+
prior
z
=
np
.
exp
(
-
1
*
np_ham
(
xx
,
yy
))
plt
.
plot
(
y
,
np
.
sum
(
z
,
axis
=
0
))
plt
.
xlabel
(
'y'
)
plt
.
ylabel
(
'unnormalized pdf'
)
plt
.
title
(
'Marginal density'
)
plt
.
pause
(
2.0
)
plt
.
close
()
plt
.
plot
(
x
*
scale
,
np
.
sum
(
z
,
axis
=
1
))
plt
.
xlabel
(
'x'
)
plt
.
ylabel
(
'unnormalized pdf'
)
plt
.
title
(
'Marginal density'
)
plt
.
pause
(
2.0
)
plt
.
close
()
#
plt.plot(y, np.sum(z, axis=0))
#
plt.xlabel('y')
#
plt.ylabel('unnormalized pdf')
#
plt.title('Marginal density')
#
plt.pause(2.0)
#
plt.close()
#
plt.plot(x*scale, np.sum(z, axis=1))
#
plt.xlabel('x')
#
plt.ylabel('unnormalized pdf')
#
plt.title('Marginal density')
#
plt.pause(2.0)
#
plt.close()
pos
=
ift
.
from_random
(
ham
.
domain
,
'normal'
)
MAP
=
ift
.
EnergyAdapter
(
pos
,
ham
,
want_metric
=
True
)
minimizer
=
ift
.
NewtonCG
(
ift
.
GradientNormController
(
iteration_limit
=
20
,
name
=
'Mini'
))
minimizer_mf
=
ift
.
ADVIOptimizer
(
10
)
MAP
,
_
=
minimizer
(
MAP
)
map_xs
,
map_ys
=
[],
[]
for
ii
in
range
(
1
0
):
for
ii
in
range
(
2
0
):
samp
=
(
MAP
.
metric
.
draw_sample
(
from_inverse
=
True
)
+
MAP
.
position
).
val
map_xs
.
append
(
samp
[
'a'
])
map_ys
.
append
(
samp
[
'b'
])
...
...
@@ -88,11 +89,14 @@ def main():
minimizer
=
ift
.
NewtonCG
(
ift
.
GradientNormController
(
iteration_limit
=
2
,
name
=
'Mini'
))
pos
=
ift
.
from_random
(
ham
.
domain
,
'normal'
)
plt
.
figure
(
figsize
=
[
12
,
8
])
for
ii
in
range
(
15
):
if
ii
%
3
==
0
:
mgkl
=
ift
.
MetricGaussianKL
.
make
(
pos
,
ham
,
40
,
False
)
mf_model
=
ift
.
MeanfieldModel
(
ham
.
domain
)
pos_mf
=
mf_model
.
get_initial_pos
(
initial_mean
=
pos
)
mfkl
=
ift
.
ParametricGaussianKL
.
make
(
pos_mf
,
ham
,
mf_model
,
20
,
True
)
plt
.
figure
(
figsize
=
[
12
,
8
])
for
ii
in
range
(
300
):
if
ii
%
1
==
0
:
mgkl
=
ift
.
MetricGaussianKL
.
make
(
pos
,
ham
,
20
,
True
)
plt
.
cla
()
plt
.
imshow
(
z
.
T
,
origin
=
'lower'
,
norm
=
LogNorm
(),
vmin
=
1e-3
,
vmax
=
np
.
max
(
z
),
cmap
=
'gist_earth_r'
,
...
...
@@ -105,17 +109,27 @@ def main():
samp
=
(
samp
+
pos
).
val
xs
.
append
(
samp
[
'a'
])
ys
.
append
(
samp
[
'b'
])
xxs
,
yys
=
[],
[]
for
samp
in
mfkl
.
samples
:
samp
=
mf_model
.
generator
((
samp
+
mfkl
.
position
)).
val
xxs
.
append
(
samp
[
'a'
])
yys
.
append
(
samp
[
'b'
])
plt
.
scatter
(
np
.
array
(
xs
)
*
scale
,
np
.
array
(
ys
),
label
=
'MGVI samples'
)
plt
.
scatter
(
np
.
array
(
xxs
)
*
scale
,
np
.
array
(
yys
),
label
=
'MFVI samples'
)
plt
.
scatter
(
pos
.
val
[
'a'
]
*
scale
,
pos
.
val
[
'b'
],
label
=
'MGVI latent mean'
)
plt
.
scatter
(
np
.
array
(
map_xs
)
*
scale
,
np
.
array
(
map_ys
),
label
=
'Laplace samples'
)
plt
.
scatter
(
MAP
.
position
.
val
[
'a'
]
*
scale
,
MAP
.
position
.
val
[
'b'
],
label
=
'Maximum a posterior solution'
)
plt
.
legend
()
plt
.
ylim
(
-
4
,
4
)
plt
.
xlim
(
-
8
,
8
)
plt
.
draw
()
plt
.
pause
(
1
.0
)
plt
.
pause
(
0
.0
1
)
mgkl
,
_
=
minimizer
(
mgkl
)
mfkl
,
_
=
minimizer_mf
(
mfkl
)
pos
=
mgkl
.
position
ift
.
logger
.
info
(
'Finished'
)
# Uncomment the following line in order to leave the plots open
...
...
src/__init__.py
View file @
1c4f555b
...
...
@@ -53,7 +53,7 @@ from .operators.energy_operators import (
Squared2NormOperator
,
StudentTEnergy
,
VariableCovarianceGaussianEnergy
)
from
.operators.convolution_operators
import
FuncConvolutionOperator
from
.operators.normal_operators
import
NormalTransform
,
LognormalTransform
from
.operators.multifield_flattener
import
MultifieldFlattener
from
.probing
import
probe_with_posterior_samples
,
probe_diagonal
,
\
StatCalculator
,
approximation2endo
...
...
@@ -67,11 +67,12 @@ from .minimization.nonlinear_cg import NonlinearCG
from
.minimization.descent_minimizers
import
(
DescentMinimizer
,
SteepestDescent
,
VL_BFGS
,
L_BFGS
,
RelaxedNewton
,
NewtonCG
)
from
.minimization.stochastic_minimizer
import
ADVIOptimizer
from
.minimization.scipy_minimizer
import
L_BFGS_B
from
.minimization.energy
import
Energy
from
.minimization.quadratic_energy
import
QuadraticEnergy
from
.minimization.energy_adapter
import
EnergyAdapter
from
.minimization.
metric_
gaussian_kl
import
MetricGaussianKL
from
.minimization.gaussian_kl
import
MetricGaussianKL
,
ParametricGaussianKL
from
.sugar
import
*
...
...
@@ -90,6 +91,7 @@ from .library.adjust_variances import (make_adjust_variances_hamiltonian,
from
.library.nft
import
Gridder
,
FinuFFT
from
.library.correlated_fields
import
CorrelatedFieldMaker
from
.library.correlated_fields_simple
import
SimpleCorrelatedField
from
.library.variational_models
import
MeanfieldModel
,
FullCovarianceModel
from
.
import
extra
...
...
src/library/variational_models.py
0 → 100644
View file @
1c4f555b
import
numpy
as
np
from
..operators.multifield_flattener
import
MultifieldFlattener
from
..operators.simple_linear_operators
import
FieldAdapter
,
PartialExtractor
from
..operators.energy_operators
import
EnergyOperator
from
..operators.sandwich_operator
import
SandwichOperator
from
..operators.linear_operator
import
LinearOperator
from
..operators.einsum
import
MultiLinearEinsum
from
..sugar
import
full
,
from_random
,
makeField
,
domain_union
from
..linearization
import
Linearization
from
..field
import
Field
from
..multi_field
import
MultiField
from
..domain_tuple
import
DomainTuple
from
..domains.unstructured_domain
import
UnstructuredDomain
class
MeanfieldModel
():
def
__init__
(
self
,
domain
):
self
.
domain
=
domain
self
.
Flat
=
MultifieldFlattener
(
self
.
domain
)
self
.
std
=
FieldAdapter
(
self
.
Flat
.
target
,
'var'
).
absolute
()
self
.
latent
=
FieldAdapter
(
self
.
Flat
.
target
,
'latent'
)
self
.
mean
=
FieldAdapter
(
self
.
Flat
.
target
,
'mean'
)
self
.
generator
=
self
.
Flat
.
adjoint
(
self
.
mean
+
self
.
std
*
self
.
latent
)
self
.
entropy
=
GaussianEntropy
(
self
.
std
.
target
)
@
self
.
std
def
get_initial_pos
(
self
,
initial_mean
=
None
,
initial_sig
=
1
):
initial_pos
=
from_random
(
self
.
generator
.
domain
).
to_dict
()
initial_pos
[
'latent'
]
=
full
(
self
.
generator
.
domain
[
'latent'
],
0.
)
initial_pos
[
'var'
]
=
full
(
self
.
generator
.
domain
[
'var'
],
initial_sig
)
if
initial_mean
is
None
:
initial_mean
=
0.1
*
from_random
(
self
.
generator
.
target
)
initial_pos
[
'mean'
]
=
self
.
Flat
(
initial_mean
)
return
MultiField
.
from_dict
(
initial_pos
)
class
FullCovarianceModel
():
def
__init__
(
self
,
domain
):
self
.
domain
=
domain
self
.
Flat
=
MultifieldFlattener
(
self
.
domain
)
one_space
=
UnstructuredDomain
(
1
)
self
.
flat_domain
=
self
.
Flat
.
target
[
0
]
N_tri
=
self
.
flat_domain
.
shape
[
0
]
*
(
self
.
flat_domain
.
shape
[
0
]
+
1
)
//
2
triangular_space
=
DomainTuple
.
make
(
UnstructuredDomain
(
N_tri
))
tri
=
FieldAdapter
(
triangular_space
,
'cov'
)
mat_space
=
DomainTuple
.
make
((
self
.
flat_domain
,
self
.
flat_domain
))
lat_mat_space
=
DomainTuple
.
make
((
one_space
,
self
.
flat_domain
))
lat
=
FieldAdapter
(
lat_mat_space
,
'latent'
)
LT
=
LowerTriangularProjector
(
triangular_space
,
mat_space
)
mean
=
FieldAdapter
(
self
.
flat_domain
,
'mean'
)
cov
=
LT
@
tri
co
=
FieldAdapter
(
cov
.
target
,
'co'
)
matmul_setup_dom
=
domain_union
((
co
.
domain
,
lat
.
domain
))
co_part
=
PartialExtractor
(
matmul_setup_dom
,
co
.
domain
)
lat_part
=
PartialExtractor
(
matmul_setup_dom
,
lat
.
domain
)
matmul_setup
=
lat_part
.
adjoint
@
lat
.
adjoint
@
lat
+
co_part
.
adjoint
@
co
.
adjoint
@
cov
MatMult
=
MultiLinearEinsum
(
matmul_setup
.
target
,
'ij,ki->jk'
,
key_order
=
(
'co'
,
'latent'
))
Resp
=
Respacer
(
MatMult
.
target
,
mean
.
target
)
self
.
generator
=
self
.
Flat
.
adjoint
@
(
mean
+
Resp
@
MatMult
@
matmul_setup
)
Diag
=
DiagonalSelector
(
cov
.
target
,
self
.
Flat
.
target
)
diag_cov
=
Diag
(
cov
).
absolute
()
self
.
entropy
=
GaussianEntropy
(
diag_cov
.
target
)
@
diag_cov
def
get_initial_pos
(
self
,
initial_mean
=
None
,
initial_sig
=
1
):
initial_pos
=
from_random
(
self
.
generator
.
domain
).
to_dict
()
initial_pos
[
'latent'
]
=
full
(
self
.
generator
.
domain
[
'latent'
],
0.
)
diag_tri
=
np
.
diag
(
np
.
full
(
self
.
flat_domain
.
shape
[
0
],
initial_sig
))[
np
.
tril_indices
(
self
.
flat_domain
.
shape
[
0
])]
initial_pos
[
'cov'
]
=
makeField
(
self
.
generator
.
domain
[
'cov'
],
diag_tri
)
if
initial_mean
is
None
:
initial_mean
=
0.1
*
from_random
(
self
.
generator
.
target
)
initial_pos
[
'mean'
]
=
self
.
Flat
(
initial_mean
)
return
MultiField
.
from_dict
(
initial_pos
)
class
GaussianEntropy
(
EnergyOperator
):
def
__init__
(
self
,
domain
):
self
.
_domain
=
domain
def
apply
(
self
,
x
):
self
.
_check_input
(
x
)
res
=
-
0.5
*
(
2
*
np
.
pi
*
np
.
e
*
x
**
2
).
log
().
sum
()
if
not
isinstance
(
x
,
Linearization
):
return
Field
.
scalar
(
res
)
if
not
x
.
want_metric
:
return
res
return
res
.
add_metric
(
SandwichOperator
.
make
(
res
.
jac
))
#FIXME not sure about metric
class
LowerTriangularProjector
(
LinearOperator
):
def
__init__
(
self
,
domain
,
target
):
self
.
_domain
=
domain
self
.
_target
=
target
self
.
_indices
=
np
.
tril_indices
(
target
.
shape
[
0
])
self
.
_capability
=
self
.
TIMES
|
self
.
ADJOINT_TIMES
def
apply
(
self
,
x
,
mode
):
self
.
_check_mode
(
mode
)
if
mode
==
self
.
TIMES
:
mat
=
np
.
zeros
(
self
.
_target
.
shape
)
mat
[
self
.
_indices
]
=
x
.
val
return
makeField
(
self
.
_target
,
mat
)
return
makeField
(
self
.
_domain
,
x
.
val
[
self
.
_indices
].
reshape
(
self
.
_domain
.
shape
))
class
DiagonalSelector
(
LinearOperator
):
def
__init__
(
self
,
domain
,
target
):
self
.
_domain
=
domain
self
.
_target
=
target
self
.
_capability
=
self
.
TIMES
|
self
.
ADJOINT_TIMES
def
apply
(
self
,
x
,
mode
):
self
.
_check_mode
(
mode
)
if
mode
==
self
.
TIMES
:
result
=
np
.
diag
(
x
.
val
)
return
makeField
(
self
.
_target
,
result
)
return
makeField
(
self
.
_domain
,
np
.
diag
(
x
.
val
).
reshape
(
self
.
_domain
.
shape
))
class
Respacer
(
LinearOperator
):
def
__init__
(
self
,
domain
,
target
):
self
.
_domain
=
domain
self
.
_target
=
target
self
.
_capability
=
self
.
TIMES
|
self
.
ADJOINT_TIMES
def
apply
(
self
,
x
,
mode
):
self
.
_check_mode
(
mode
)
if
mode
==
self
.
TIMES
:
return
makeField
(
self
.
_target
,
x
.
val
.
reshape
(
self
.
_target
.
shape
))
return
makeField
(
self
.
_domain
,
x
.
val
.
reshape
(
self
.
_domain
.
shape
))
src/minimization/
metric_
gaussian_kl.py
→
src/minimization/gaussian_kl.py
View file @
1c4f555b
...
...
@@ -23,12 +23,14 @@ from ..logger import logger
from
..multi_field
import
MultiField
from
..operators.endomorphic_operator
import
EndomorphicOperator
from
..operators.energy_operators
import
StandardHamiltonian
from
..operators.multifield_flattener
import
MultifieldFlattener
from
..probing
import
approximation2endo
from
..sugar
import
makeOp
from
..sugar
import
makeOp
,
full
,
from_random
from
..utilities
import
myassert
from
.energy
import
Energy
class
_KLMetric
(
EndomorphicOperator
):
def
__init__
(
self
,
KL
):
self
.
_KL
=
KL
...
...
@@ -258,3 +260,157 @@ class MetricGaussianKL(Energy):
yield
s
if
self
.
_mirror_samples
:
yield
-
s
class
ParametricGaussianKL
(
Energy
):
"""Provides the sampled Kullback-Leibler divergence between a distribution
and a Parametric Gaussian.
Notes
-----
See also
"""
def
__init__
(
self
,
variational_parameters
,
hamiltonian
,
variational_model
,
n_samples
,
mirror_samples
,
comm
,
local_samples
,
nanisinf
,
_callingfrommake
=
False
):
if
not
_callingfrommake
:
raise
NotImplementedError
super
(
ParametricGaussianKL
,
self
).
__init__
(
variational_parameters
)
assert
variational_model
.
generator
.
target
is
hamiltonian
.
domain
self
.
_hamiltonian
=
hamiltonian
self
.
_variational_model
=
variational_model
self
.
_full_model
=
hamiltonian
(
variational_model
.
generator
)
+
variational_model
.
entropy
self
.
_n_samples
=
int
(
n_samples
)
self
.
_mirror_samples
=
bool
(
mirror_samples
)
self
.
_comm
=
comm
self
.
_local_samples
=
local_samples
self
.
_nanisinf
=
bool
(
nanisinf
)
lin
=
Linearization
.
make_partial_var
(
variational_parameters
,
[
'latent'
])
v
,
g
=
[],
[]
for
s
in
self
.
_local_samples
:
# s = _modify_sample_domain(s, variational_parameters.domain)
tmp
=
self
.
_full_model
(
lin
+
s
)
tv
=
tmp
.
val
.
val
tg
=
tmp
.
gradient
if
mirror_samples
:
tmp
=
self
.
_full_model
(
lin
-
s
)
tv
=
tv
+
tmp
.
val
.
val
tg
=
tg
+
tmp
.
gradient
v
.
append
(
tv
)
g
.
append
(
tg
)
self
.
_val
=
utilities
.
allreduce_sum
(
v
,
self
.
_comm
)[()]
/
self
.
n_eff_samples
if
np
.
isnan
(
self
.
_val
)
and
self
.
_nanisinf
:
self
.
_val
=
np
.
inf
self
.
_grad
=
utilities
.
allreduce_sum
(
g
,
self
.
_comm
)
/
self
.
n_eff_samples
@
staticmethod
def
make
(
variational_parameters
,
hamiltonian
,
variational_model
,
n_samples
,
mirror_samples
,
comm
=
None
,
nanisinf
=
False
):
"""Return instance of :class:`MetricGaussianKL`.
Parameters
----------
mean : Field
Mean of the Gaussian probability distribution.
hamiltonian : StandardHamiltonian
Hamiltonian of the approximated probability distribution.
n_samples : integer
Number of samples used to stochastically estimate the KL.
mirror_samples : boolean
Whether the negative of the drawn samples are also used, as they are
equally legitimate samples. If true, the number of used samples
doubles. Mirroring samples stabilizes the KL estimate as extreme
sample variation is counterbalanced. Since it improves stability in
many cases, it is recommended to set `mirror_samples` to `True`.
constants : list
List of parameter keys that are kept constant during optimization.
Default is no constants.
napprox : int
Number of samples for computing preconditioner for sampling. No
preconditioning is done by default.
comm : MPI communicator or None
If not None, samples will be distributed as evenly as possible
across this communicator. If `mirror_samples` is set, then a sample and
its mirror image will always reside on the same task.
nanisinf : bool
If true, nan energies which can happen due to overflows in the forward
model are interpreted as inf. Thereby, the code does not crash on
these occaisions but rather the minimizer is told that the position it
has tried is not sensible.
Note
----
The two lists `constants` and `point_estimates` are independent from each
other. It is possible to sample along domains which are kept constant
during minimization and vice versa.
"""
if
not
isinstance
(
hamiltonian
,
StandardHamiltonian
):
raise
TypeError
if
hamiltonian
.
domain
is
not
variational_model
.
generator
.
target
:
raise
ValueError
if
not
isinstance
(
n_samples
,
int
):
raise
TypeError
if
not
isinstance
(
mirror_samples
,
bool
):
raise
TypeError
n_samples
=
int
(
n_samples
)
mirror_samples
=
bool
(
mirror_samples
)
local_samples
=
[]
sseq
=
random
.
spawn_sseq
(
n_samples
)
#FIXME dirty trick, many multiplications with zero
DirtyMaskDict
=
full
(
variational_model
.
generator
.
domain
,
0.
).
to_dict
()
DirtyMaskDict
[
'latent'
]
=
full
(
variational_model
.
generator
.
domain
[
'latent'
],
1.
)
DirtyMask
=
MultiField
.
from_dict
(
DirtyMaskDict
)
for
i
in
range
(
*
_get_lo_hi
(
comm
,
n_samples
)):
with
random
.
Context
(
sseq
[
i
]):
local_samples
.
append
(
DirtyMask
*
from_random
(
variational_model
.
generator
.
domain
))
local_samples
=
tuple
(
local_samples
)
return
ParametricGaussianKL
(
variational_parameters
,
hamiltonian
,
variational_model
,
n_samples
,
mirror_samples
,
comm
,
local_samples
,
nanisinf
,
_callingfrommake
=
True
)
@
property
def
n_eff_samples
(
self
):
if
self
.
_mirror_samples
:
return
2
*
self
.
_n_samples
return
self
.
_n_samples
def
at
(
self
,
position
):
return
ParametricGaussianKL
(
position
,
self
.
_hamiltonian
,
self
.
_variational_model
,
self
.
_n_samples
,
self
.
_mirror_samples
,
self
.
_comm
,
self
.
_local_samples
,
self
.
_nanisinf
,
True
)
@
property
def
value
(
self
):
return
self
.
_val
@
property
def
gradient
(
self
):
return
self
.
_grad
@
property
def
samples
(
self
):
ntask
,
rank
,
_
=
utilities
.
get_MPI_params_from_comm
(
self
.
_comm
)
if
ntask
==
1
:
for
s
in
self
.
_local_samples
:
yield
s
if
self
.
_mirror_samples
:
yield
-
s
else
:
rank_lo_hi
=
[
utilities
.
shareRange
(
self
.
_n_samples
,
ntask
,
i
)
for
i
in
range
(
ntask
)]
lo
,
_
=
_get_lo_hi
(
self
.
_comm
,
self
.
_n_samples
)
for
itask
,
(
l
,
h
)
in
enumerate
(
rank_lo_hi
):
for
i
in
range
(
l
,
h
):
data
=
self
.
_local_samples
[
i
-
lo
]
if
rank
==
itask
else
None
s
=
self
.
_comm
.
bcast
(
data
,
root
=
itask
)
yield
s
if
self
.
_mirror_samples
:
yield
-
s
src/minimization/stochastic_minimizer.py
0 → 100644
View file @
1c4f555b
from
..minimization.gaussian_kl
import
ParametricGaussianKL
from
.minimizer
import
Minimizer
import
numpy
as
np
class
ADVIOptimizer
(
Minimizer
):
def
__init__
(
self
,
steps
,
eta
=
1
,
alpha
=
1
,
tau
=
1
,
epsilon
=
1e-16
):
self
.
alpha
=
alpha
self
.
eta
=
eta
self
.
tau
=
tau
self
.
epsilon
=
epsilon
self
.
counter
=
1
self
.
steps
=
steps
self
.
s
=
None
def
_step
(
self
,
position
,
gradient
):
self
.
s
=
self
.
alpha
*
gradient
**
2
+
(
1
-
self
.
alpha
)
*
self
.
s
self
.
rho
=
self
.
eta
*
self
.
counter
**
(
-
0.5
+
self
.
epsilon
)
/
(
self
.
tau
+
(
self
.
s
).
sqrt
())
new_position
=
position
-
self
.
rho
*
gradient
self
.
counter
+=
1
return
new_position
def
__call__
(
self
,
E
):
if
self
.
s
is
None
:
self
.
s
=
E
.
gradient
**
2
convergence
=
0
# come up with somthing how to determine convergence
for
i
in
range
(
self
.
steps
):
x
=
self
.
_step
(
E
.
position
,
E
.
gradient
)
# maybe some KL function for resample? Should make it more generic.
E
=
ParametricGaussianKL
.
make
(
x
,
E
.
_hamiltonian
,
E
.
_variational_model
,
E
.
_n_samples
,
E
.
_mirror_samples
)
return
E
,
convergence
def
reset
(
self
):
self
.
counter
=
1
self
.
s
=
None
\ No newline at end of file
src/operators/multifield_flattener.py
0 → 100644
View file @
1c4f555b
import
numpy
as
np
from
.linear_operator
import
LinearOperator