src/compose.py

import numpy as np


class SumOfSeparableFunctions:
    def __init__(self, fnlist):
        self.fnlist = fnlist

    def apply(self, x):
        val = 0.0
        for f_i, x_i in zip(self.fnlist, x):
            val += f_i.apply(x_i)
        return val

    def diff(self, x):
        d = []
        for f_i, x_i in zip(self.fnlist, x):
            d.append(f_i.diff(x_i))
        return d

    def apply_and_diff(self, x):
        d = []
        val = 0.0
        for f_i, x_i in zip(self.fnlist, x):
            (a, v) = f_i.apply_and_diff(x_i)
            val += a
            d.append(v)
        return (val, d)

    def diff_lipschitz_factor(self):
        res = 0
        for f_i in self.fnlist:
            res = max(res, f_i.diff_lipschitz_factor())
        return res

    def diff_bound(self, xbound=None):
        res = 0
        for f_i in self.fnlist:
            res = max(res, f_i.diff_bound(xbound=xbound))
        return res


class ComposeFnWithOperator:
    def __init__(self, f, op, xbound=None, xbound_pair=None):
        self.f = f
        self.op = op
        self.xbound = xbound
        self.xbound_pair = xbound_pair

    def apply(self, *args):
        return self.f.apply(self.op.apply(*args))

    def diff(self, *args):
        # TODO: precalculations in apply should be used in diff_adjdir
        w = self.op.apply(*args)
        v = self.f.diff(w)
        return self.op.diff_adjdir(v, *args, apply_result=w)

    def apply_and_diff(self, *args):
        # TODO: precalculations in apply should be used in diff_adjdir
        w = self.op.apply(*args)
        (a, v) = self.f.apply_and_diff(w)
        return (a, self.op.diff_adjdir(v, *args, apply_result=w))

    def diff_lipschitz_factor(self):
        # ‖∇A(x)^*∇F(A(x)) - ∇A(y)^*∇F(A(y))‖
        # = ‖[∇A(x)^*-∇A(y)^*]∇F(A(x)) - ∇A(y)^*[∇F(A(y))-∇F(A(x))]‖
        # ≤ L_{∇A(x)} M_{∇F} + M_{∇A(y)^*} L_{∇F}L_A.
        if hasattr(self.op, "opnorm"):
            # Linear operator
            lda = 0.0  # This is zero,
            mdf = 0.0  # hence this not needed.
        else:
            mdf = self.f.diff_bound(xbound=self.op.codomain_bound(xbound=self.xbound))
            lda = self.op.diff_adj_lipschitz_factor()

        ldf = self.f.diff_lipschitz_factor()
        la = self.op.lipschitz_factor()
        mda = self.op.diff_bound(xbound=self.xbound)

        return lda * mdf + mda * ldf * la

    def diff_lipschitz_factor_pair(self):
        if self.op.hasattr("opnorm"):
            # Linear operator
            lda1, lda2 = 0.0, 0.0  # This is zero,
            mdf = 0.0  # hence this not needed.
        else:
            lda1, lda2 = self.op.diff_adj_lipschitz_factor_pair()
            mdf = self.f.diff_bound(
                xbound=self.op.codomain_bound_pair(xbound=self.xbound_pair)
            )

        ldf = self.f.diff_lipschitz_factor()
        la1, la2 = self.op.lipschitz_factor_pair()
        mda = self.op.diff_bound_pair(xbound=self.xbound_pair)

        return lda1 * mdf + mda * ldf * la1, lda2 * mdf + mda * ldf * la2


class InjectSecond:
    def __init__(self, y):
        self.y = y

    def apply(self, x):
        return (x, self.y)

    def diff_adjdir(self, j, _x, apply_result=None):
        return j[0]

    # This is not really a linear operator, but for our purposes affine behave essentially
    # the same
    def opnorm(self, *args):
        return 1.0

    def lipschitz_factor(self, *args):
        return 1.0

    def diff_adj_lipschitz_factor(self, *args):
        return 0.0

    def diff_bound(self, xbound=None):
        return 1.0

    def codomain_bound(self, xbound=None):
        if xbound is None:
            raise Exception("Linear operators have unbounded range")
        else:
            return xbound