def exact_fit(
self, factor: Factor, model_approx: EPMeanField, status: Status = Status()
) -> Tuple[EPMeanField, Status]:
with LogWarnings(logger=self.logger, action='always') as caught_warnings:
if factor._calc_exact_update:
factor_approx = model_approx.factor_approximation(factor)
new_approx = model_approx if self.inplace else model_approx.copy()
new_approx.update_factor_mean_field(
factor, factor.calc_exact_update(factor_approx.cavity_dist)
)
elif factor._calc_exact_projection:
factor_approx = model_approx.factor_approximation(factor)
new_model_dist = factor.calc_exact_projection(factor_approx.cavity_dist)
new_approx, status = self.update_model_approx(
new_model_dist, factor_approx, model_approx, status
)
else:
raise NotImplementedError(
"Factor does not have exact updates methods"
)
status_kws = status._asdict()
status_kws['messages'] = status.messages + tuple(caught_warnings.messages)
status = Status(**status_kws)
return new_approx, status
def _parse_result(
self,
result: OptimizeResult,
status: Status = Status()) -> OptResult:
success, messages = status
success = result.success
message = result.message.decode()
messages += (
"optimise.find_factor_mode: "
f"nfev={result.nfev}, nit={result.nit}, "
f"status={result.status}, message={message}",)
full_hess_inv = result.hess_inv
if not isinstance(full_hess_inv, np.ndarray):
# if optimiser is L-BFGS-B then convert
# implicit hess_inv into dense matrix
full_hess_inv = full_hess_inv.todense()
# make inverse transform back
M = self.transform
x = M.ldiv(result.x)
full_hess_inv = M.ldiv(M.ldiv(full_hess_inv).T)
mode = {**self.param_shapes.unflatten(x), **self.fixed_kws}
hess_inv = self.param_shapes.unflatten(full_hess_inv)
return OptResult(
mode,
hess_inv,
self.sign * result.fun, # minimized negative logpdf of factor approximation
full_hess_inv, # full inverse hessian of optimisation
result,
Status(success, messages))
def factor_step(
self,
factor,
subset_approx,
optimiser=None
):
factor_logger = self.factor_loggers.get(factor.name, logging.getLogger(factor.name))
factor_logger.debug("Optimising...")
optimiser = optimiser or self.factor_optimisers[factor]
subset_factor = subset_approx._factor_subset_factor[factor]
try:
with LogWarnings(logger=factor_logger.debug, action='always') as caught_warnings:
subset_approx, status = optimiser.optimise(
subset_factor,
subset_approx,
)
messages = status.messages + tuple(caught_warnings.messages)
status = Status(status.success, messages, status.flag)
except (ValueError, ArithmeticError, RuntimeError) as e:
logger.exception(e)
status = Status(
False,
status.messages + (f"Factor: {factor} experienced error {e}",),
StatusFlag.FAILURE,
)
factor_logger.debug(status)
return subset_approx, status
def lstsq_laplace_factor_approx(
model_approx: EPMeanField,
factor: Factor,
delta: float = 0.5,
opt_kws: Optional[Dict[str, Any]] = None):
"""
"""
factor_approx = model_approx.factor_approximation(factor)
opt = LeastSquaresOpt(
factor_approx, **({} if opt_kws is None else opt_kws))
mode, covar, result = opt.least_squares()
message = (
"optimise.lsq_sq_laplace_factor_approx: "
f"nfev={result.nfev}, njev={result.njev}, "
f"optimality={result.optimality}, "
f"cost={result.cost}, "
f"status={result.status}, message={result.message}",)
status = Status(result.success, message)
model_dist = MeanField({
v: factor_approx.factor_dist[v].from_mode(
mode[v],
covar.get(v))
for v in mode
})
projection, status = factor_approx.project(
model_dist, delta=delta, status=status)
return model_approx.project(projection, status=status)
def laplace_factor_approx(
model_approx: EPMeanField,
factor: Factor,
delta: float = 1.,
status: Status = Status(),
opt_kws: Optional[Dict[str, Any]] = None
):
opt_kws = {} if opt_kws is None else opt_kws
factor_approx = model_approx.factor_approximation(factor)
res = find_factor_mode(
factor_approx,
return_cov=True,
status=status,
**opt_kws
)
model_dist = factor_approx.model_dist.project_mode(res)
projection, status = factor_approx.project(
model_dist,
delta=delta,
status=res.status
)
new_approx, status = model_approx.project(
projection, status=status)
return new_approx, status
def find_factor_mode(
factor_approx: FactorApproximation,
return_cov: bool = True,
status: Status = Status(),
min_iter: int = 2,
opt_kws: Optional[dict] = None,
**kwargs
) -> OptResult:
"""
"""
opt_kws = {} if opt_kws is None else opt_kws
opt = OptFactor.from_approx(factor_approx, **kwargs)
res = opt.maximise(status=status, **opt_kws)
if return_cov:
# Calculate deterministic values
value = factor_approx.factor(res.mode)
res.mode.update(value.deterministic_values)
# Calculate covariance of deterministic values
jacobian = factor_approx.factor.jacobian(
res.mode, opt.free_vars)
update_det_cov(res, jacobian)
return res
def optimise(
self,
factor: Factor,
model_approx: EPMeanField,
status: Optional[Status] = Status(),
) -> Tuple[EPMeanField, Status]:
whiten = self.transforms[factor]
opt_kws = self.opt_kws[factor]
start = self.initial_values.get(factor)
factor_approx = model_approx.factor_approximation(factor)
opt = OptFactor.from_approx(factor_approx, transform=whiten)
res = opt.maximise(start, status=status, **opt_kws)
# Calculate covariance of deterministic values
# TODO: estimate this Jacobian using Broyden's method
# https://en.wikipedia.org/wiki/Broyden%27s_method
value = factor_approx.factor(res.mode)
res.mode.update(value.deterministic_values)
jacobian = factor_approx.factor.jacobian(res.mode, opt.free_vars)
update_det_cov(res, jacobian)
self.transforms[factor] = self.transform_cls.from_dense(
res.full_hess_inv)
# Project Laplace's approximation
new_model_dist = factor_approx.model_dist.from_mode_covariance(
res.mode, res.hess_inv, res.log_norm)
return self.update_model_approx(new_model_dist, factor_approx,
model_approx, status)
def minimise(
self,
arrays_dict: Optional[ArraysDict] = None,
status: Status = Status(),
**kwargs,
):
self.sign = 1
p0 = self.get_random_start(arrays_dict or {})
res = self._minimise(p0, **kwargs)
return self._parse_result(res, status=status)
def maximise(
self,
arrays_dict: Optional[Dict[Variable, np.ndarray]] = None,
status: Status = Status(),
**kwargs,
):
self.sign = -1
p0 = self.get_random_start(arrays_dict or {})
res = self._minimise(p0, **kwargs)
return self._parse_result(res, status=status)
def run(
self,
model_approx: EPMeanField,
factors: Optional[List[Factor]] = None,
status: Status = Status()
) -> EPMeanField:
new_approx = model_approx
for i in range(self.n_iter):
for factor, new_approx, status in self.step(new_approx, factors):
self.history[i, factor] = new_approx
return new_approx, status
def optimise(
self,
factor: Factor,
model_approx: EPMeanField,
status: Status = Status(),
) -> Tuple[EPMeanField, Status]:
delta = self.deltas[factor]
sample_kws = self.sample_kws[factor]
factor_approx = model_approx.factor_approximation(factor)
sample = self(factor_approx, **sample_kws)
model_dist = project_factor_approx_sample(factor_approx, sample)
projection, status = factor_approx.project(model_dist, delta=delta)
return model_approx.project(projection, status=status)
def step(
self,
model_approx,
factors: Optional[List[Factor]] = None,
status: Status = Status(),
) -> Iterator[Tuple[Factor, EPMeanField, Status]]:
new_approx = model_approx
factors = model_approx.factor_graph.factors if factors is None else factors
for factor in factors:
new_approx, status = laplace_factor_approx(new_approx,
factor,
self.delta,
status=status,
opt_kws=self.opt_kws)
yield factor, new_approx, status
def update_model_approx(
self,
new_model_dist: MeanField,
factor_approx: FactorApproximation,
model_approx: EPMeanField,
status: Optional[Status] = Status(),
delta: Optional[float] = None,
) -> Tuple[EPMeanField, Status]:
delta = delta or self.deltas[factor_approx.factor]
new_approx, status = model_approx.project_mean_field(
new_model_dist,
factor_approx,
delta=delta,
status=status,
)
return new_approx, status
def optimise(
self, factor: Factor, model_approx: EPMeanField, status: Status = Status()
) -> Tuple[EPMeanField, Status]:
pass