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):
        res = 0
        for f_i in self.fnlist:
            res = max(res, f_i.diff_bound())
        return res


class ComposeFnWithOperator:
    def __init__(self, f, op):
        self.f = f
        self.op = op
        self.override_lipschitz = None
        self.override_lipschitz_pair = None

    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_bound(self):
        mf = self.f.diff_bound()
        if hasattr(self.op, "opnorm"):
            lda = self.op.opnorm() ** 2
        else:
            lda = self.op.diff_bound()
        return lda * mf

    def diff_bound_pair(self):
        mf = self.f.diff_bound()
        lda1, lda2 = self.op.diff_bound_pair()
        return lda1 * mf, lda2 * mf

    # def lipschitz_factor(self, xbound=None):
    #     if xbound is None:
    #         xbound = self.xbound
    #     lf = self.f.lipschitz_factor(xbound=self.op.codomain_bound(xbound=xbound))
    #     la = self.op.lipschitz_factor(xbound=xbound)
    #     return lf * la

    # def lipschitz_factor_pair(self, xbound=None):
    #     if xbound is None:
    #         xbound = self.xbound
    #     lf = self.f.lipschitz_factor(xbound=self.op.codomain_bound(xbound=xbound))
    #     la1, la2 = self.op.lipschitz_factor_pair(xbound=xbound)
    #     return lf * la1, lf * la2

    def diff_lipschitz_factor(self):
        """
        Calculate the Lipschitz factor of the differential of this composed function.
        We assume that either the operator is linear and implementes `opnorm`, or it is nonlinear, and impliements `diff_chain_lipschitz_factor` to directly calculate a Lipschitz factor of $x ↦ ∇A(x)^*∇F(A(x))$, given a bound $M$
        on ∇F. The function `diff_chain_lipschitz_factor` should return the Lipschitz factor divided by $M$: we obtain $M$ through the `diff_bound` on $F$.
        """

        if self.override_lipschitz is not None:
            return self.override_lipschitz
        elif hasattr(self.op, "opnorm"):
            return self.f.diff_lipschitz_factor() * self.op.opnorm() ** 2
        else:
            mdf = self.f.diff_bound()
            lda = self.op.diff_chain_lipschitz_factor()
            # print(
            #     "LDA %s %f; MDF %f; total %f"
            #     % (type(self.op).__name__, lda, mdf, mdf * lda),
            # )
            return mdf * lda

    def diff_lipschitz_factor_pair(self):
        """
        This is similar to `diff_lipschitz_factor`, except separates the
        factor for arguments pairs.

        This requires the operator to implement `diff_chain_lipschitz_factor`;
        there is no special handling of linear opeartors.
        """

        if self.override_lipschitz_pair is not None:
            return self.override_lipschitz_pair

        mdf = self.f.diff_bound()

        lda1, lda2 = self.op.diff_chain_lipschitz_factor_pair()

        # print("LDA %s %f %f; MDF %f;  " % (type(self.op).__name__, lda1, lda2, mdf))

        return mdf * lda1, mdf * lda2


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):
    #     return 1.0

    # def diff_chain_lipschitz_factor(self, *args):
    #     return 0.0

    def diff_bound(self):
        return 1.0

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


class ScaleAuxStepHack:
    """
    This hack only scale the differential, to get different step length parameters wrt. k and c,
    without rescaling the variables themselves. Essentially it amounts to computing the gradient wrt.
    different inner products.
    """

    def __init__(self, scale):
        self.scale = scale

    def apply(self, x):
        # μ, (k, (c1, c2)) = x
        # sck, scc = self.scale
        # return μ, (k * sck, (c1 * scc, c2 * scc))
        return x

    def diff_adjdir(self, j, _x, apply_result=None):
        dμ, (dk, (dc1, dc2)) = j
        sck, scc = self.scale
        return dμ, (dk * sck, (dc1 * scc, dc2 * scc))

    def opnorm(self, *args):
        return 1

    def diff_bound(self):
        return 1