from functools import partial
import numpy.ma as ma
import autograd.numpy as np
from autograd import grad
from autograd.extend import defvjp, primitive
import proxmin
import logging
from .bbox import overlapped_slices
from .component import CombinedComponent
from .model import UpdateException
logger = logging.getLogger("scarlet.blend")
@primitive
def _add_models(*models, full_model, slices):
"""Insert the models into the full model
`slices` is a tuple `(full_model_slice, model_slices)` used
to insert a model into the full_model in the region where the
two models overlap.
"""
for i in range(len(models)):
full_model[slices[i][0]] += models[i][slices[i][1]]
return full_model
def _grad_add_models(upstream_grad, *models, full_model, slices, index):
"""Gradient for a single model
The full model is just the sum of the models,
so the gradient is 1 for each model,
we just have to slice it appropriately.
"""
model = models[index]
full_model_slices = slices[index][0]
model_slices = slices[index][1]
def result(upstream_grad):
_result = np.zeros(model.shape, dtype=model.dtype)
_result[model_slices] = upstream_grad[full_model_slices]
return _result
return result
[docs]class Blend(CombinedComponent):
"""The blended scene
The class represents a scene as collection of and provides the functions to
fit it to data.
"""
def __init__(self, sources, observations):
"""Constructor
Form a blended scene from a collection of `~scarlet.component.Component`s
Parameters
----------
sources: list of `~scarlet.component.Component` or `~scarlet.component.ComponentTree`
Intitialized components or sources to fit to the observations
observations: a `scarlet.Observation` instance or a list thereof
Data package(s) to fit
"""
if hasattr(sources, "__iter__"):
self.sources = sources
else:
self.sources = (sources,)
if hasattr(observations, "__iter__"):
self.observations = observations
else:
self.observations = (observations,)
super().__init__(self.sources)
# only for backward compatibility, use log_likelihood instead
self.loss = []
[docs] def fit(self, max_iter=200, e_rel=1e-3, min_iter=1, noise_factor=0, **alg_kwargs):
"""Fit the model for each source to the data
Parameters
----------
max_iter: int
Maximum number of iterations if the algorithm doesn't converge
e_rel: float
Relative error for convergence of the loss function
min_iter: int
Maximum number of iterations if the algorithm doesn't converge
alg_kwargs: dict
Keywords for the `proxmin.adaprox` optimizer
"""
it = 0
self._noise_factor = noise_factor
while it < max_iter:
try:
X = self.parameters + tuple(
p for obs in self.observations for p in obs.parameters
)
# compute the backward gradients
# but only for non-fixed parameters
require_grad = tuple(k for k, x in enumerate(X) if not x.fixed)
def expand_grads(*X, func=None):
G = func(*X)
expanded = [0.0] * len(X)
for k, j in enumerate(require_grad):
expanded[j] = G[k]
return expanded
grad_logL_func = grad(self._loss_func, require_grad)
grad_logL = lambda *X: expand_grads(*X, func=grad_logL_func)
# same for prior. easier her bc we call them independently
grad_logP = lambda *X: tuple(
x.prior(x.view(np.ndarray))
if x.prior is not None and not x.fixed
else 0
for x in X
)
# combine for log posterior
_grad = lambda *X: tuple(
l + p for l, p in zip(grad_logL(*X), grad_logP(*X))
)
# step sizes
_step = lambda *X, it: tuple(
x.step(x, it=it) if hasattr(x.step, "__call__") else x.step
for x in X
)
# proxes
_prox = tuple(x.constraint for x in X)
# good defaults for adaprox
scheme = alg_kwargs.pop("scheme", "amsgrad")
prox_max_iter = alg_kwargs.pop("prox_max_iter", 10)
callback = partial(
self._callback,
e_rel=e_rel,
callback=alg_kwargs.pop("callback", None),
min_iter=min_iter,
)
# do we have a current state of the optimizer to warm start?
for x in X:
if x.m is None:
x.m = np.zeros(x.shape)
if x.v is None:
x.v = np.zeros(x.shape)
if x.vhat is None:
x.vhat = np.zeros(x.shape)
M = tuple(x.m for x in X)
V = tuple(x.v for x in X)
Vhat = tuple(x.vhat for x in X)
proxmin.adaprox(
X,
_grad,
_step,
prox=_prox,
max_iter=max_iter - it,
e_rel=e_rel,
check_convergence=False,
scheme=scheme,
prox_max_iter=prox_max_iter,
callback=callback,
M=M,
V=V,
Vhat=Vhat,
**alg_kwargs
)
logger.info(
"scarlet ran for {0} iterations to logL = {1}".format(
len(self.log_likelihood), self.log_likelihood[-1]
)
)
# set convergence and standard deviation from optimizer
for p, m, v, vhat in zip(X, M, V, Vhat):
p.std = 1 / np.sqrt(
ma.masked_equal(v, 0)
) # this is rough estimate!
return len(self.log_likelihood), self.log_likelihood[-1]
# model update forces restart
except UpdateException:
it = len(self.log_likelihood)
[docs] def get_model(self, *parameters, frame=None):
"""Get the model of the entire blend
Parameters
----------
parameters: tuple of optimization parameters
frame: `scarlet.Frame`
Alternative Frame to project the model into
Returns
-------
model: array
(Bands, Height, Width) data cube
"""
# boxed models of every source
models = self.get_models_of_children(*parameters, frame=None)
if frame is None:
frame = self.frame
# if this is the model frame then the slices are already cached
if frame == self.frame:
slices = tuple(
(src._model_frame_slices, src._model_slices) for src in self.sources
)
else:
slices = tuple(
overlapped_slices(frame.bbox, src.bbox) for src in self.sources
)
# We have to declare the function that inserts sources
# into the blend with autograd.
# This has to be done each time we fit a blend,
# since the number of components => the number of arguments,
# which must be linked to the autograd primitive function.
defvjp(
_add_models,
*([partial(_grad_add_models, index=k) for k in range(len(self.sources))])
)
full_model = np.zeros(frame.shape, dtype=frame.dtype)
full_model = _add_models(*models, full_model=full_model, slices=slices)
return full_model
@property
def log_likelihood(self):
"""Log likelihood at each iteration
The fitting method computes and sums the negative log-likelihood for the each
observation given the current model as loss function for the optimization.
Returns
-------
log_likelihood: array of (positive) log-likelihood
"""
return -np.array(self.loss)
def _loss_func(self, *parameters):
n_params = len(self.parameters)
model = self.get_model(*parameters[:n_params], frame=self.frame)
# Caculate the total loss function from all of the observations
total_loss = 0
for observation in self.observations:
n_obs_params = len(observation.parameters)
obs_params = parameters[n_params : n_params + n_obs_params]
total_loss = total_loss - observation.get_log_likelihood(
model, *obs_params, noise_factor=self._noise_factor
)
n_params += n_obs_params
self.loss.append(total_loss._value)
return total_loss
def _callback(self, *parameters, it=None, e_rel=1e-3, callback=None, min_iter=1):
# raises ArithmeticError if some of the parameters have become inf/nan
for src in self.sources:
src.check_parameters()
# raises UpdateException if model updates require optimimzation interruption
throw_exception = False
if it > 0 and it % 10 == 0:
for src in self.sources:
try:
src.update()
except UpdateException:
throw_exception = True
pass
if throw_exception:
raise UpdateException
if it > min_iter and abs(self.loss[-1] - self.loss[-2]) < e_rel * np.abs(
self.loss[-1]
):
raise StopIteration(
"scarlet.Blend.fit() converged"
) # clean return from proxmin
if callback is not None:
callback(*parameters, it=it)
@property
def bbox(self):
"""Bounding box of the blend
"""
return self.frame.bbox
