Add parametric variational inference

demos/meanfield_demo.py

+import nifty7 as ift
+from matplotlib import pyplot as plt
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+#
+# Copyright(C) 2013-2020 Max-Planck-Society
+#
+# NIFTy is being developed at the Max-Planck-Institut fuer Astrophysik.
+
+###############################################################################
+# Log-normal field reconstruction from Poissonian data with inhomogenous
+# exposure (in case for 2D mode)
+# 1D (set mode=0), 2D (mode=1), or on the sphere (mode=2)
+###############################################################################
+
+import sys
+
+import numpy as np
+
+
+# def exposure_2d(domain):
+#     # Structured exposure for 2D mode
+#     x_shape, y_shape = domain.shape
+#     exposure = np.ones(domain.shape)
+#     exposure[x_shape//3:x_shape//2, :] *= 2.
+#     exposure[x_shape*4//5:x_shape, :] *= .1
+#     exposure[x_shape//2:x_shape*3//2, :] *= 3.
+#     exposure[:, x_shape//3:x_shape//2] *= 2.
+#     exposure[:, x_shape*4//5:x_shape] *= .1
+#     exposure[:, x_shape//2:x_shape*3//2] *= 3.
+#     return ift.Field.from_raw(domain, exposure)
+
+
+if __name__ == '__main__':
+
+    # Two-dimensional regular grid with inhomogeneous exposure
+    position_space = ift.RGSpace([100])
+
+    # Define harmonic space and harmonic transform
+    harmonic_space = position_space.get_default_codomain()
+    HT = ift.HarmonicTransformOperator(harmonic_space, position_space)
+
+    # Domain on which the field's degrees of freedom are defined
+    domain = ift.DomainTuple.make(harmonic_space)
+
+    # Define amplitude (square root of power spectrum)
+    def sqrtpspec(k):
+        return 1./(1. + k**2)
+
+    p_space = ift.PowerSpace(harmonic_space)
+    pd = ift.PowerDistributor(harmonic_space, p_space)
+    a = ift.PS_field(p_space, sqrtpspec)
+    A = pd(a)
+
+    # Define sky operator
+    sky = 10*ift.exp(HT(ift.makeOp(A))).ducktape('xi')
+
+    # M = ift.DiagonalOperator(exposure)
+    GR = ift.GeometryRemover(position_space)
+    # Define instrumental response
+    # R = GR(M)
+    R = GR
+
+    # Generate mock data and define likelihood operator
+    d_space = R.target[0]
+    lamb = R(sky)
+    mock_position = ift.from_random(sky.domain, 'normal')
+    data = lamb(mock_position)
+    data = ift.random.current_rng().poisson(data.val.astype(np.float64))
+    data = ift.Field.from_raw(d_space, data)
+    likelihood = ift.PoissonianEnergy(data) @ lamb
+
+    # Settings for minimization
+    ic_newton = ift.DeltaEnergyController(
+        name='Newton', iteration_limit=1, tol_rel_deltaE=1e-8)
+    # minimizer = ift.L_BFGS(ic_newton)
+    minimizer_fc = ift.ADVIOptimizer(steps=10)
+    minimizer_mf = ift.ADVIOptimizer(steps=10)
+
+    # Compute MAP solution by minimizing the information Hamiltonian
+    H = ift.StandardHamiltonian(likelihood)
+    # initial_position = ift.from_random(domain, 'normal')
+
+    # meanfield_model = ift.MeanfieldModel(H.domain)
+    fullcov_model = ift.FullCovarianceModel(H.domain)
+    meanfield_model = ift.MeanfieldModel(H.domain)
+
+    position_fc = fullcov_model.get_initial_pos(initial_sig=0.01)
+    position_mf = meanfield_model.get_initial_pos(initial_sig=0.01)
+    KL_fc = ift.ParametricGaussianKL.make(position_fc,H,fullcov_model,3,True)
+    KL_mf = ift.ParametricGaussianKL.make(position_mf,H,meanfield_model,3,True)
+
+    # plt.figure('data')
+    # plt.imshow(sky(mock_position).val)
+    plt.pause(0.001)
+    for i in range(3000):
+        # KL = ParametricGaussianKL.make(position,H,meanfield_model,3,True)
+        KL_fc, _ = minimizer_fc(KL_fc)
+        KL_mf, _ = minimizer_mf(KL_mf)
+
+        plt.figure('result')
+        plt.cla()
+        plt.plot(sky(fullcov_model.generator(KL_fc.position)).val,'b-',label='fc')
+        plt.plot(sky(meanfield_model.generator(KL_mf.position)).val,'r-',label='mf')
+        for samp in KL_fc.samples:
+            plt.plot(sky(fullcov_model.generator(KL_fc.position + samp)).val,'b-',alpha=0.3)
+        for samp in KL_mf.samples:   
+            plt.plot(sky(meanfield_model.generator(KL_mf.position + samp)).val,'r-',alpha=0.3)
+        plt.plot(data.val,'kx')
+        plt.plot(sky(mock_position).val,'k-',label='true')
+        plt.legend()
+        plt.ylim(0,data.val.max()+10)
+        plt.pause(0.001)
\ No newline at end of file

demos/mgvi_visualized.py

@@ -61,26 +61,27 @@ def main():
         return lh + prior
 
     z = np.exp(-1*np_ham(xx, yy))
-    plt.plot(y, np.sum(z, axis=0))
-    plt.xlabel('y')
-    plt.ylabel('unnormalized pdf')
-    plt.title('Marginal density')
-    plt.pause(2.0)
-    plt.close()
-    plt.plot(x*scale, np.sum(z, axis=1))
-    plt.xlabel('x')
-    plt.ylabel('unnormalized pdf')
-    plt.title('Marginal density')
-    plt.pause(2.0)
-    plt.close()
+    # plt.plot(y, np.sum(z, axis=0))
+    # plt.xlabel('y')
+    # plt.ylabel('unnormalized pdf')
+    # plt.title('Marginal density')
+    # plt.pause(2.0)
+    # plt.close()
+    # plt.plot(x*scale, np.sum(z, axis=1))
+    # plt.xlabel('x')
+    # plt.ylabel('unnormalized pdf')
+    # plt.title('Marginal density')
+    # plt.pause(2.0)
+    # plt.close()
 
     pos = ift.from_random(ham.domain, 'normal')
     MAP = ift.EnergyAdapter(pos, ham, want_metric=True)
     minimizer = ift.NewtonCG(
         ift.GradientNormController(iteration_limit=20, name='Mini'))
+    minimizer_mf = ift.ADVIOptimizer(10)
     MAP, _ = minimizer(MAP)
     map_xs, map_ys = [], []
-    for ii in range(10):
+    for ii in range(20):
         samp = (MAP.metric.draw_sample(from_inverse=True) + MAP.position).val
         map_xs.append(samp['a'])
         map_ys.append(samp['b'])
@@ -88,11 +89,14 @@ def main():
     minimizer = ift.NewtonCG(
         ift.GradientNormController(iteration_limit=2, name='Mini'))
     pos = ift.from_random(ham.domain, 'normal')
-    plt.figure(figsize=[12, 8])
-    for ii in range(15):
-        if ii % 3 == 0:
-            mgkl = ift.MetricGaussianKL.make(pos, ham, 40, False)
+    mf_model = ift.MeanfieldModel(ham.domain)
+    pos_mf = mf_model.get_initial_pos(initial_mean=pos)
+    mfkl = ift.ParametricGaussianKL.make(pos_mf,ham,mf_model,20,True)
 
+    plt.figure(figsize=[12, 8])
+    for ii in range(300):
+        if ii % 1 == 0:
+            mgkl = ift.MetricGaussianKL.make(pos, ham, 20, True)
         plt.cla()
         plt.imshow(z.T, origin='lower', norm=LogNorm(), vmin=1e-3,
                    vmax=np.max(z), cmap='gist_earth_r',
@@ -105,17 +109,27 @@ def main():
             samp = (samp + pos).val
             xs.append(samp['a'])
             ys.append(samp['b'])
+        xxs, yys = [], []
+        for samp in mfkl.samples:
+            samp = mf_model.generator((samp + mfkl.position)).val
+            xxs.append(samp['a'])
+            yys.append(samp['b'])
         plt.scatter(np.array(xs)*scale, np.array(ys), label='MGVI samples')
+        plt.scatter(np.array(xxs)*scale, np.array(yys), label='MFVI samples')
+
         plt.scatter(pos.val['a']*scale, pos.val['b'], label='MGVI latent mean')
         plt.scatter(np.array(map_xs)*scale, np.array(map_ys),
                     label='Laplace samples')
         plt.scatter(MAP.position.val['a']*scale, MAP.position.val['b'],
                     label='Maximum a posterior solution')
         plt.legend()
+        plt.ylim(-4,4)
+        plt.xlim(-8,8)
         plt.draw()
-        plt.pause(1.0)
+        plt.pause(0.01)
 
         mgkl, _ = minimizer(mgkl)
+        mfkl, _ = minimizer_mf(mfkl)
         pos = mgkl.position
     ift.logger.info('Finished')
     # Uncomment the following line in order to leave the plots open

src/__init__.py

@@ -53,7 +53,7 @@ from .operators.energy_operators import (
     Squared2NormOperator, StudentTEnergy, VariableCovarianceGaussianEnergy)
 from .operators.convolution_operators import FuncConvolutionOperator
 from .operators.normal_operators import NormalTransform, LognormalTransform
-
+from .operators.multifield_flattener import MultifieldFlattener
 from .probing import probe_with_posterior_samples, probe_diagonal, \
     StatCalculator, approximation2endo
 
@@ -67,11 +67,12 @@ from .minimization.nonlinear_cg import NonlinearCG
 from .minimization.descent_minimizers import (
     DescentMinimizer, SteepestDescent, VL_BFGS, L_BFGS, RelaxedNewton,
     NewtonCG)
+from .minimization.stochastic_minimizer import ADVIOptimizer
 from .minimization.scipy_minimizer import L_BFGS_B
 from .minimization.energy import Energy
 from .minimization.quadratic_energy import QuadraticEnergy
 from .minimization.energy_adapter import EnergyAdapter
-from .minimization.metric_gaussian_kl import MetricGaussianKL
+from .minimization.gaussian_kl import MetricGaussianKL, ParametricGaussianKL
 
 from .sugar import *
 
@@ -90,6 +91,7 @@ from .library.adjust_variances import (make_adjust_variances_hamiltonian,
 from .library.nft import Gridder, FinuFFT
 from .library.correlated_fields import CorrelatedFieldMaker
 from .library.correlated_fields_simple import SimpleCorrelatedField
+from .library.variational_models import MeanfieldModel, FullCovarianceModel
 
 from . import extra
 

src/library/variational_models.py

+import numpy as np
+from ..operators.multifield_flattener import MultifieldFlattener
+from ..operators.simple_linear_operators import FieldAdapter, PartialExtractor
+from ..operators.energy_operators import EnergyOperator
+from ..operators.sandwich_operator import SandwichOperator
+from ..operators.linear_operator import LinearOperator
+from ..operators.einsum import MultiLinearEinsum
+from ..sugar import full, from_random, makeField, domain_union
+from ..linearization import Linearization
+from ..field import Field
+from ..multi_field import MultiField
+from ..domain_tuple import DomainTuple
+from ..domains.unstructured_domain import UnstructuredDomain
+
+class MeanfieldModel():
+    def __init__(self, domain):
+        self.domain = domain
+        self.Flat = MultifieldFlattener(self.domain)
+
+        self.std = FieldAdapter(self.Flat.target,'var').absolute()
+        self.latent = FieldAdapter(self.Flat.target,'latent')
+        self.mean = FieldAdapter(self.Flat.target,'mean')
+        self.generator = self.Flat.adjoint(self.mean + self.std * self.latent)
+        self.entropy = GaussianEntropy(self.std.target) @ self.std
+
+    def get_initial_pos(self, initial_mean=None,initial_sig = 1):
+        initial_pos = from_random(self.generator.domain).to_dict()
+        initial_pos['latent'] = full(self.generator.domain['latent'], 0.)
+        initial_pos['var'] = full(self.generator.domain['var'], initial_sig)
+
+        if initial_mean is None:
+            initial_mean = 0.1*from_random(self.generator.target)
+
+        initial_pos['mean'] = self.Flat(initial_mean)
+        return MultiField.from_dict(initial_pos)
+
+class FullCovarianceModel():
+    def __init__(self, domain):
+        self.domain = domain
+        self.Flat = MultifieldFlattener(self.domain)
+        one_space = UnstructuredDomain(1)
+        self.flat_domain = self.Flat.target[0]
+        N_tri = self.flat_domain.shape[0]*(self.flat_domain.shape[0]+1)//2
+        triangular_space = DomainTuple.make(UnstructuredDomain(N_tri))
+        tri = FieldAdapter(triangular_space, 'cov')
+        mat_space = DomainTuple.make((self.flat_domain,self.flat_domain))
+        lat_mat_space = DomainTuple.make((one_space,self.flat_domain))
+        lat = FieldAdapter(lat_mat_space,'latent')
+        LT = LowerTriangularProjector(triangular_space,mat_space)
+        mean = FieldAdapter(self.flat_domain,'mean')
+        cov = LT @ tri
+        co = FieldAdapter(cov.target, 'co')
+
+        matmul_setup_dom = domain_union((co.domain,lat.domain))
+        co_part = PartialExtractor(matmul_setup_dom, co.domain)
+        lat_part = PartialExtractor(matmul_setup_dom, lat.domain)
+        matmul_setup = lat_part.adjoint @ lat.adjoint @ lat + co_part.adjoint @ co.adjoint @ cov
+        MatMult = MultiLinearEinsum(matmul_setup.target,'ij,ki->jk', key_order=('co','latent'))
+
+        Resp = Respacer(MatMult.target, mean.target)
+        self.generator = self.Flat.adjoint @ (mean + Resp @ MatMult @ matmul_setup)
+        
+        Diag = DiagonalSelector(cov.target, self.Flat.target)
+        diag_cov = Diag(cov).absolute()
+        self.entropy = GaussianEntropy(diag_cov.target) @ diag_cov
+
+    def get_initial_pos(self, initial_mean = None, initial_sig = 1):
+        initial_pos = from_random(self.generator.domain).to_dict()
+        initial_pos['latent'] = full(self.generator.domain['latent'], 0.)
+        diag_tri = np.diag(np.full(self.flat_domain.shape[0],initial_sig))[np.tril_indices(self.flat_domain.shape[0])]
+        initial_pos['cov'] = makeField(self.generator.domain['cov'], diag_tri)
+        if initial_mean is None:
+            initial_mean = 0.1*from_random(self.generator.target)
+        initial_pos['mean'] = self.Flat(initial_mean)
+        return MultiField.from_dict(initial_pos)
+
+
+
+class GaussianEntropy(EnergyOperator):
+    def __init__(self, domain):
+        self._domain = domain
+
+    def apply(self, x):
+        self._check_input(x)
+        res =  -0.5* (2*np.pi*np.e*x**2).log().sum()
+        if not isinstance(x, Linearization):
+            return Field.scalar(res)
+        if not x.want_metric:
+            return res
+        return res.add_metric(SandwichOperator.make(res.jac)) #FIXME not sure about metric
+
+
+class LowerTriangularProjector(LinearOperator):
+    def __init__(self, domain, target):
+        self._domain = domain
+        self._target = target
+        self._indices=np.tril_indices(target.shape[0])
+        self._capability = self.TIMES | self.ADJOINT_TIMES
+
+    def apply(self, x, mode):
+        self._check_mode(mode)
+        if mode == self.TIMES:
+            mat = np.zeros(self._target.shape)
+            mat[self._indices] = x.val
+            return makeField(self._target,mat)
+        return makeField(self._domain, x.val[self._indices].reshape(self._domain.shape))
+
+class DiagonalSelector(LinearOperator):
+    def __init__(self, domain, target):
+        self._domain = domain
+        self._target = target
+        self._capability = self.TIMES | self.ADJOINT_TIMES
+
+    def apply(self, x, mode):
+        self._check_mode(mode)
+        if mode == self.TIMES:
+            result = np.diag(x.val)
+            return makeField(self._target,result)
+        return makeField(self._domain,np.diag(x.val).reshape(self._domain.shape))
+
+
+class Respacer(LinearOperator):
+    def __init__(self, domain, target):
+        self._domain = domain
+        self._target = target
+        self._capability = self.TIMES | self.ADJOINT_TIMES
+
+
+    def apply(self,x,mode):
+        self._check_mode(mode)
+        if mode == self.TIMES:
+            return makeField(self._target,x.val.reshape(self._target.shape))
+        return makeField(self._domain,x.val.reshape(self._domain.shape))

src/minimization/metric_gaussian_kl.py → src/minimization/gaussian_kl.py

@@ -23,12 +23,14 @@ from ..logger import logger
 from ..multi_field import MultiField
 from ..operators.endomorphic_operator import EndomorphicOperator
 from ..operators.energy_operators import StandardHamiltonian
+from ..operators.multifield_flattener import MultifieldFlattener
 from ..probing import approximation2endo
-from ..sugar import makeOp
+from ..sugar import makeOp, full, from_random
 from ..utilities import myassert
 from .energy import Energy
 
 
+
 class _KLMetric(EndomorphicOperator):
     def __init__(self, KL):
         self._KL = KL
@@ -258,3 +260,157 @@ class MetricGaussianKL(Energy):
                     yield s
                     if self._mirror_samples:
                         yield -s
+
+
+class ParametricGaussianKL(Energy):
+    """Provides the sampled Kullback-Leibler divergence between a distribution
+    and a Parametric Gaussian.
+    Notes
+    -----
+
+    See also
+
+    """
+    def __init__(self, variational_parameters, hamiltonian, variational_model, 
+                    n_samples, mirror_samples, comm,
+                        local_samples, nanisinf, _callingfrommake=False):
+        if not _callingfrommake:
+            raise NotImplementedError
+        super(ParametricGaussianKL, self).__init__(variational_parameters)
+        assert variational_model.generator.target is hamiltonian.domain
+        self._hamiltonian = hamiltonian
+        self._variational_model = variational_model
+        self._full_model = hamiltonian(variational_model.generator) + variational_model.entropy
+
+        self._n_samples = int(n_samples)
+        self._mirror_samples = bool(mirror_samples)
+        self._comm = comm
+        self._local_samples = local_samples
+        self._nanisinf = bool(nanisinf)
+
+        lin = Linearization.make_partial_var(variational_parameters, ['latent'])
+        v, g = [], []
+        for s in self._local_samples:
+            # s = _modify_sample_domain(s, variational_parameters.domain)
+            tmp = self._full_model(lin+s)
+            tv = tmp.val.val
+            tg = tmp.gradient
+            if mirror_samples:
+                tmp = self._full_model(lin-s)
+                tv = tv + tmp.val.val
+                tg = tg + tmp.gradient
+            v.append(tv)
+            g.append(tg)
+        self._val = utilities.allreduce_sum(v, self._comm)[()]/self.n_eff_samples
+        if np.isnan(self._val) and self._nanisinf:
+            self._val = np.inf
+        self._grad = utilities.allreduce_sum(g, self._comm)/self.n_eff_samples
+
+    @staticmethod
+    def make(variational_parameters, hamiltonian, variational_model, n_samples, mirror_samples,
+                    comm=None, nanisinf=False):
+        """Return instance of :class:`MetricGaussianKL`.
+
+        Parameters
+        ----------
+        mean : Field
+            Mean of the Gaussian probability distribution.
+        hamiltonian : StandardHamiltonian
+            Hamiltonian of the approximated probability distribution.
+        n_samples : integer
+            Number of samples used to stochastically estimate the KL.
+        mirror_samples : boolean
+            Whether the negative of the drawn samples are also used, as they are
+            equally legitimate samples. If true, the number of used samples
+            doubles. Mirroring samples stabilizes the KL estimate as extreme
+            sample variation is counterbalanced. Since it improves stability in
+            many cases, it is recommended to set `mirror_samples` to `True`.
+        constants : list
+            List of parameter keys that are kept constant during optimization.
+            Default is no constants.
+        napprox : int
+            Number of samples for computing preconditioner for sampling. No
+            preconditioning is done by default.
+        comm : MPI communicator or None
+            If not None, samples will be distributed as evenly as possible
+            across this communicator. If `mirror_samples` is set, then a sample and
+            its mirror image will always reside on the same task.
+        nanisinf : bool
+            If true, nan energies which can happen due to overflows in the forward
+            model are interpreted as inf. Thereby, the code does not crash on
+            these occaisions but rather the minimizer is told that the position it
+            has tried is not sensible.
+
+        Note
+        ----
+        The two lists `constants` and `point_estimates` are independent from each
+        other. It is possible to sample along domains which are kept constant
+        during minimization and vice versa.
+        """
+
+        if not isinstance(hamiltonian, StandardHamiltonian):
+            raise TypeError
+        if hamiltonian.domain is not variational_model.generator.target:
+            raise ValueError
+        if not isinstance(n_samples, int):
+            raise TypeError
+        if not isinstance(mirror_samples, bool):
+            raise TypeError
+
+        n_samples = int(n_samples)
+        mirror_samples = bool(mirror_samples)
+        local_samples = []
+        sseq = random.spawn_sseq(n_samples)
+
+        #FIXME dirty trick, many multiplications with zero
+        DirtyMaskDict = full(variational_model.generator.domain,0.).to_dict()
+        DirtyMaskDict['latent'] = full(variational_model.generator.domain['latent'], 1.)
+        DirtyMask = MultiField.from_dict(DirtyMaskDict)
+
+        for i in range(*_get_lo_hi(comm, n_samples)):
+            with random.Context(sseq[i]):
+                local_samples.append(DirtyMask * from_random(variational_model.generator.domain))
+        local_samples = tuple(local_samples)
+
+        return ParametricGaussianKL(
+            variational_parameters, hamiltonian,variational_model,n_samples, mirror_samples, comm, local_samples,
+            nanisinf, _callingfrommake=True)
+    
+    @property
+    def n_eff_samples(self):
+        if self._mirror_samples:
+            return 2*self._n_samples
+        return self._n_samples
+
+    def at(self, position):
+        return ParametricGaussianKL(
+            position, self._hamiltonian, self._variational_model, self._n_samples, self._mirror_samples,
+            self._comm, self._local_samples, self._nanisinf, True)
+            
+
+    @property
+    def value(self):
+        return self._val
+
+    @property
+    def gradient(self):
+        return self._grad
+
+    @property
+    def samples(self):
+        ntask, rank, _ = utilities.get_MPI_params_from_comm(self._comm)
+        if ntask == 1:
+            for s in self._local_samples:
+                yield s
+                if self._mirror_samples:
+                    yield -s
+        else:
+            rank_lo_hi = [utilities.shareRange(self._n_samples, ntask, i) for i in range(ntask)]
+            lo, _ = _get_lo_hi(self._comm, self._n_samples)
+            for itask, (l, h) in enumerate(rank_lo_hi):
+                for i in range(l, h):
+                    data = self._local_samples[i-lo] if rank == itask else None
+                    s = self._comm.bcast(data, root=itask)
+                    yield s
+                    if self._mirror_samples:
+                        yield -s

src/minimization/stochastic_minimizer.py

+from ..minimization.gaussian_kl import ParametricGaussianKL
+from .minimizer import Minimizer
+import numpy as np
+
+class ADVIOptimizer(Minimizer):
+
+    def __init__(self, steps,eta=1, alpha=1, tau=1, epsilon=1e-16):
+        self.alpha = alpha
+        self.eta = eta
+        self.tau=tau
+        self.epsilon = epsilon
+        self.counter = 1
+        self.steps = steps
+        self.s = None
+
+    def _step(self, position, gradient):
+        self.s = self.alpha * gradient**2 + (1-self.alpha)*self.s
+        self.rho = self.eta * self.counter**(-0.5+ self.epsilon) / (self.tau + (self.s).sqrt())
+        new_position = position - self.rho * gradient
+        self.counter += 1
+        return new_position
+
+    def __call__(self, E):
+        if self.s is None:
+            self.s = E.gradient**2
+        convergence = 0 # come up with somthing how to determine convergence
+        for i in range(self.steps):
+            x = self._step(E.position, E.gradient)
+            # maybe some KL function for resample? Should make it more generic.
+            E = ParametricGaussianKL.make(x, E._hamiltonian, E._variational_model,
+                                E._n_samples, E._mirror_samples)
+
+        return E, convergence
+
+    def reset(self):
+        self.counter = 1
+        self.s = None
\ No newline at end of file

src/operators/multifield_flattener.py

+
+
+import numpy as np
+
+from .linear_operator import LinearOperator