def __init__(
self,
factor_approx: FactorApproximation,
fixed_kws=None,
param_bounds=None,
opt_only=None,
**kwargs,
):
self.factor_approx = factor_approx
self.opt_params = {**self._opt_params, **kwargs}
param_shapes = {}
param_bounds = {} if param_bounds is None else param_bounds
fixed_kws = {} if fixed_kws is None else fixed_kws
for v in factor_approx.factor.variables:
dist = factor_approx.model_dist[v]
if isinstance(dist, FixedMessage):
fixed_kws[v] = dist.mean
else:
param_shapes[v] = dist.shape
param_bounds[v] = dist._support
self.fixed_kws = fixed_kws
self.param_shapes = FlattenArrays(param_shapes)
if opt_only is None:
opt_only = tuple(v for v, d in factor_approx.cavity_dist.items()
if not isinstance(d, FixedMessage))
self.opt_only = opt_only
self.resid_means = {
k: factor_approx.cavity_dist[k].mean
for k in self.opt_only
}
self.resid_scales = {
k: factor_approx.cavity_dist[k].scale
for k in self.opt_only
}
self.resid_shapes = FlattenArrays(
{k: np.shape(m)
for k, m in self.resid_means.items()})
self.bounds = tuple(
np.array(
list(
zip(*[
b for k, s in param_shapes.items()
for bound in param_bounds[k]
for b in repeat(bound, np.prod(s, dtype=int))
]))))
def to_full(self):
full_ops = [op.to_full() for op in self.operators]
full = block_diag(*(op.operator.to_dense() for op in full_ops))
param_shapes = FlattenArrays(
(v, s) for op in full_ops for v, s in op.param_shapes.items())
return VariableFullOperator(full_ops[0].operator.from_dense(full),
param_shapes)
def update(self, *args):
(u, v), *next_args = args
param_shapes = FlattenArrays.from_arrays(u)
uv = param_shapes.flatten(u)[:,
None] * param_shapes.flatten(v)[None, :]
out = VariableFullOperator.from_dense(uv, param_shapes)
return out.update(*next_args)
def from_approx(
cls,
factor_approx: FactorApproximation,
transform: Optional[AbstractLinearTransform] = None,
) -> 'OptFactor':
value_shapes = {}
fixed_kws = {}
bounds = {}
for v in factor_approx.variables:
dist = factor_approx.model_dist[v]
if isinstance(dist, FixedMessage):
fixed_kws[v] = dist.mean
else:
value_shapes[v] = dist.shape
bounds[v] = dist._support
return cls(
factor_approx,
FlattenArrays(value_shapes),
fixed_kws=fixed_kws,
model_dist=factor_approx.model_dist,
transform=transform,
bounds=bounds,
)
class LeastSquaresOpt:
_opt_params = dict(
jac='2-point', method='trf', ftol=1e-08,
xtol=1e-08, gtol=1e-08, x_scale=1.0, loss='linear',
f_scale=1.0, diff_step=None, tr_solver=None,
tr_options={}, jac_sparsity=None, max_nfev=None,
verbose=0)
def __init__(
self,
factor_approx: FactorApproximation,
fixed_kws=None,
param_bounds=None,
opt_only=None,
**kwargs):
self.factor_approx = factor_approx
self.opt_params = {**self._opt_params, **kwargs}
param_shapes = {}
param_bounds = {} if param_bounds is None else param_bounds
fixed_kws = {} if fixed_kws is None else fixed_kws
for v in factor_approx.factor.variables:
dist = factor_approx.model_dist[v]
if isinstance(dist, FixedMessage):
fixed_kws[v] = dist.mean
else:
param_shapes[v] = dist.shape
param_bounds[v] = dist._support
self.fixed_kws = fixed_kws
self.param_shapes = FlattenArrays(param_shapes)
if opt_only is None:
opt_only = tuple(
v for v, d in factor_approx.cavity_dist.items()
if not isinstance(d, FixedMessage)
)
self.opt_only = opt_only
self.resid_means = {
k: factor_approx.cavity_dist[k].mean for k in self.opt_only}
self.resid_scales = {
k: factor_approx.cavity_dist[k].scale for k in self.opt_only}
self.resid_shapes = FlattenArrays({
k: np.shape(m) for k, m in self.resid_means.items()})
self.bounds = tuple(np.array(list(zip(*[
b for k, s in param_shapes.items()
for bound in param_bounds[k]
for b in repeat(bound, np.prod(s, dtype=int))]))))
def __call__(self, arr):
p0 = self.param_shapes.unflatten(arr)
values = {**p0, **self.fixed_kws}
fvals = self.factor_approx.factor(values)
values.update(fvals.deterministic_values)
residuals = {
v: (values[v] - mean) / self.resid_scales[v]
for v, mean in self.resid_means.items()
}
return self.resid_shapes.flatten(residuals)
def least_squares(self, values=None):
values = values or {}
model_dist = self.factor_approx.model_dist
p0 = {
v: values[v] if v in values else model_dist[v].sample()
for v in self.param_shapes.keys()}
arr = self.param_shapes.flatten(p0)
res = least_squares(
self, arr, bounds=self.bounds, **self.opt_params)
sol = self.param_shapes.unflatten(res.x)
fval = self.factor_approx.factor(
{**sol, **self.fixed_kws}
)
det_vars = fval.deterministic_values
jac = {
(d, k): b
for k, a in self.param_shapes.unflatten(
res.jac.T, ndim=1).items()
for d, b in self.resid_shapes.unflatten(
a.T, ndim=1).items()}
hess = self.param_shapes.unflatten(
res.jac.T.dot(res.jac))
def inv(a):
shape = np.shape(a)
ndim = len(shape)
if ndim:
a = np.asanyarray(a)
s = shape[:ndim // 2]
n = np.prod(s, dtype=int)
return np.linalg.inv(
a.reshape(n, n)).reshape(s + s)
else:
return 1 / a
invhess = {
k: inv(h) for k, h in hess.items()}
for det in det_vars:
invhess[det] = 0.
for v in sol:
invhess[det] += propagate_uncertainty(
invhess[v], jac[det, v])
mode = {**sol, **det_vars}
return mode, invhess, res
def lparam_shapes(self):
return FlattenArrays(
{v: op.lshape
for v, op in self.operators.items()})
def from_blocks(cls, Ms: VariableData) -> "VariableFullOperator":
param_shapes = FlattenArrays(
{v: np.shape(M)[:np.ndim(M) // 2]
for v, M in Ms.items()})
return cls.from_dense(param_shapes.flatten2d(Ms), param_shapes)
def from_diagonal(cls, diag: VariableData) -> "VariableFullOperator":
param_shapes = FlattenArrays.from_arrays(diag)
operator = DiagonalMatrix(param_shapes.flatten(diag))
return cls(operator, param_shapes)
def from_state(cls, state):
param_shapes = FlattenArrays.from_arrays(state.parameters)
return cls(state, param_shapes)