(Gradually) add advanced interface to solvers

[kohr-h]

Copied from the discussion in #1252:

I am not sure that this change in API makes a big difference. It means that if one wants to have more specialised callbacks, then one can recycle some code by using the update function. However, in some applications one also might want to choose which values to compute and which ones to store. Thus, for these one needs to change the update function. Thus, I am not convinced one gains that much with the proposed changes.

If you want to write callbacks that do more interesting stuff, then yes, agreed. But the whole point of writing a class that exposes elements of the optimization method is to not have to use callbacks to do interesting stuff.
See, the purpose of callbacks is basically to get information out of a black box. You don't know if or when things will happen, but you say to the black box "Whenever you do X, please call this function." But if you have a "grey box" rather than a black box, you have all kinds of possibilities, depending on how much the box exposes.

So in our example, think of something like this:

Basic version

class Landweber(object):
    def __init__(self, op, x, rhs, omega=1):
        self.op = op
        self.x = x
        self.rhs = rhs
        self.omega = omega
        
    def step(self):
        self.x.assign(self.x - self.omega * self.op.adjoint(self.op(self.x) - self.rhs))
        
    def current(self):
        return self.x

More advanced version

class Landweber(object):
    def __init__(self, op, x, rhs, omega=1):
        self.op = op
        self.x = x
        self.rhs = rhs
        self.omega = omega
        
    def step(self):
        self.x.assign(self.current() - self.omega * self.current_l2sq_grad())

    def current_l2sq_grad(self):
        return self.op.adjoint(self.current_residual())

    def current_residual(self):
        return self.op(self.x) - self.rhs
        
    def current(self):
        return self.x

With intermediate storage

class Landweber(object):
    def __init__(self, op, x, rhs, omega=1, use_tmp_dom=True, use_tmp_ran=True):
        self.op = op
        self.x = x
        self.rhs = rhs
        self.omega = omega
        self.tmp_dom = self.op.domain.element() if use_tmp_dom else None
        self.tmp_ran = self.op.range.element() if use_tmp_ran else None
        
    def step(self):
        self.x.lincomb(1, self.x, -self.omega, self.current_l2sq_grad())

    def current_l2sq_grad(self):
        return self.op.adjoint(self.current_residual(), out=self.tmp_dom)

    def current_residual(self):
        tmp = self.op(self.x, out=self.tmp_ran)
        tmp -= self.rhs
        return tmp
        
    def current(self):
        return self.x

This now gives you all kinds of possibilities to alter the iteration, inspect intermediate results etc. without having to resort to callbacks.

Also note that these changes can be added gradually without changing the interface (step).

[kohr-h]

I've been thinking a bit more about this issue lately and tried to come up with a useful implementation of this kind of functionality.

On the surface, it's pretty simple: You just break an optimization method into steps, and then make a class that exposes the individual steps as methods. Users can then pull the triggers one by one, or do a whole step, whatever they like. In any case, there's gonna be some code sitting around this new object that drives the computation.

For instance, let's take the Landweber class from above (slightly adapted):

class Landweber(object):
    def __init__(self, op, x, rhs, omega):
        self.op = op
        self.x = x
        self.rhs = rhs
        self.omega = omega
        
    def step(self):
        self.x.assign(self.x - self.omega * self.l2sq_grad())

    def l2sq_grad(self):
        return self.op.adjoint(self.residual())

    def residual(self):
        return self.op(self.x) - self.rhs

This looks nice, but has bad performance characteristics: If I run residual() and then l2sq_grad(), the residual will be re-computed. Thus, the broken-down version is slower than the wholesale step method that only computes stuff once (albeit not using out parameters and such), and it only gets worse with more complex methods.

Caching to the rescue:

class Landweber(object):
    def __init__(self, op, x, rhs, omega):
        ...
        self.cache = {'residual': None, 'l2sq_grad': None}

    def l2sq_grad(self):
        if self.cache['l2sq_grad'] is None:
            self.cache['l2sq_grad'] = self.op.adjoint(self.residual())
        return self.cache['l2sq_grad']

    def residual(self):
        if self.cache['residual'] is None:
            self.cache['residual'] = self.op(self.x) - self.rhs
        return self.cache['residual']

Of course, this code is broken since it will just return the same cached values over and over again after the first round. The issue here is that we have no notion of cache validity, so we need to add that. However, that brings us to the next interesting question: When should the cache be invalidated?

A simple answer would be "at the end of each step":

class Landweber(object):
    def __init__(self, op, x, rhs, omega):
        ...
        self.cache = {'residual': {'is_valid': False, 'value': None},
                      'l2sq_grad': {'is_valid': False, 'value': None}}

    def step(self):
        self.x.assign(self.x - self.omega * self.l2sq_grad())
        self._invalidate_cache()

    def l2sq_grad(self):
        if not self.cache['l2sq_grad']['is_valid']:
            self.cache['l2sq_grad']['value'] = self.op.adjoint(self.residual())
        return self.cache['l2sq_grad']['value']

    def residual(self):
        if not self.cache['residual']['is_valid']:
            self.cache['residual']['value'] = self.op(self.x) - self.rhs
        return self.cache['residual']['value']

    def _invalidate_cache(self):
        for k in self.cache:
            self.cache['k']['is_valid'] = False

This is somewhat workable, but not really great, since it allows the system to be in inconsistent state. Suppose someone first computes residual() and then mutates x. The next call to residual(), and by transitivity also l2sq_grad(), will reflect the old value of x, and that's a bummer.

What we could do to fix this is to define (direct) dependencies between cache entries within one iteration. For instance:

class Landweber(object):
    def __init__(self, op, x, rhs, omega):
        ...
        self._x = x
        self.cache = {
            'x': {'is_valid': True, 'value': x, 'rdeps': ['residual']},
            'residual': {'is_valid': False, 'value': None, 'rdeps': ['l2sq_grad']},
            'l2sq_grad': {'is_valid': False, 'value': None, 'rdeps': []}}

    @property
    def x(self):
        return self.cache['x']['value']

    @x.setter
    def x(self, new_x):
        self._invalidate_cache('x')
        self.cache['x']['value'] = new_x
        self.cache['x']['is_valid'] = True

    def step(self):
        self.x.assign(self.x - self.omega * self.l2sq_grad())
        self._invalidate_cache('x')

    def l2sq_grad(self):
        if not self.cache['l2sq_grad']['is_valid']:
            self.cache['l2sq_grad']['value'] = self.op.adjoint(self.residual())
        return self.cache['l2sq_grad']['value']

    def residual(self):
        if not self.cache['residual']['is_valid']:
            self.cache['residual']['value'] = self.op(self.x) - self.rhs
        return self.cache['residual']['value']

    def _invalidate_cache(self, start):
        self.cache[start]['is_valid'] = False
        for rdep in self.cache[start]['rdeps']:
            self._invalidate_cache(rdep)

So this starts to look a bit complicated, but it should work in principle: Now, whenx gets set, the cache of dependent quantities is cleared, and if they're requested subsequently, they will be recomputed. And if you squint hard enough, you'll see that we effectively built a computation graph of our optimizer (a rather coarse-grained one, but still) to keep track of all dependencies.

However, there's one huge caveat: x is still a mutable reference, and we've got no chance to detect an in-place modification and do the cache invalidation thing. Even some hacky monkey-patching of x would only get us so far. Users could do something like np.put(lw.x.asarray(), indices, values), and then it's game over for us.

So here I don't know exactly where to go next. The only solution I can think of is to make x read-only by setting the Numpy array flag writable to False. Internally, I guess we could still sail around that obstacle by making x writable again when needed, but it's a bit of a hack.

Anyway, this became way too long, but I find the challenge very interesting and see a great potential in having a well thought-out functionality like this. I'd be curious to hear other opinions, @adler-j, @aringh, @mehrhardt, @sbanert for instance.

[sbanert]

Would it be feasible to use Tensorflow (or another learning framework) as a backend to build the computation graph and executing? Another useful aspect of this would be that you wouldn’t have to re-implement these if you want to use the solvers within Tensorflow.

[kohr-h]

Would it be feasible to use Tensorflow (or another learning framework) as a backend to build the computation graph and executing? Another useful aspect of this would be that you wouldn’t have to re-implement these if you want to use the solvers within Tensorflow.

I guess it would be possible, but we certainly don't want a dependency on tensorflow. Also, our usage of the graph is pretty basic (forward only, no state), so a homebrew implementation isn't that horrifying.
(Of course, the whole example can be done much more concisely and with code reuse in mind, but for the sake brevity I didn't include that boilerplate code.)

[kohr-h]

Another useful aspect of this would be that you wouldn’t have to re-implement these if you want to use the solvers within Tensorflow.

But that's an interesting point in itself. Maybe one can generalize this approach a bit and make the ad-hoc graph exportable to tensorflow (or other DL frameworks).

[olivierverdier]

I would love this. I basically rewrote special cases of this several time. In general, I find the class interface much more suitable for solvers than functions.

[kohr-h]

I also think that keeping the relevant information in one place while offering fine-grained control would be nice. We could get some inspiration from the SciPy solvers.

That said, today I think that trying to enforce an order by internal caching (and caching in the first place) is probably not a good idea. Now I would probably advocate for a more "pure" interface where users can hand inputs and output storage if they choose to do so. For instance, I would get rid of self.x and make it a parameter to the residual and step methods, along with out parameters for better efficiency.