Python / numpy implementation of DIIS mixing. https://en.wikipedia.org/wiki/DIIS

diis.py

import numpy as np
from numpy.linalg import multi_dot

# only necessary for FlatMixingDecorator
import w2dyn.auxiliaries.deepflatten as deepflatten

class InitialMixingDecorator(object):
    def __init__(self, init_count, init_mixer, mixer):
        self.init_count = init_count
        self.init_mixer = init_mixer
        self.mixer = mixer
    def __call__(self, *args):
        if self.init_count > 0:
            self.init_count -= 1
            return self.init_mixer(*args)
        return self.mixer(*args)

class FlatMixingDecorator(object):
    def __init__(self, mixer):
        self.mixer = mixer
    def __call__(self, *args):
        if len(args) == 1: args = args[0]
        types = deepflatten.types(args)
        shape = deepflatten.shape(args)
        x = deepflatten.flatten(args)
        x = self.mixer(x)
        x = deepflatten.restore(x, shape, types)
        return x

class RealMixingDecorator(object):
    def __init__(self, mixer):
        self.mixer = mixer
    def __call__(self, x):
        n = x.shape[0]
        x = np.concatenate([np.real(x), np.imag(x)])
        x = self.mixer(x)
        x = x[:n] + 1j*x[n:]
        return x

class NoMixingMixer(object):
    def __call__(self, *args):
        return *args

class LinearMixer(object):
    def __init__(self, old_share=0):
        self.old_share = float(old_share)
        self.old_value = None

    def __call__(self, new_value):
        if new_value is None:
            raise ValueError("new value must not be None")

        if self.oldvalue is None:
            new_trial = new_value
        else:
            new_trial = self.old_share * self.old_value + (1 - self.old_share) * new_value
        self.old_value = new_trial
        return new_trial

class DiisMixer(object):
    def __init__(self, old_share, history, period):
        self.alpha = 1 - old_share
        self.history = history
        self.period = period

        self.i = 0
        self.trials = []
        self.residuals = []

    def __call__(self, new_value):
        if self.i <= 0:
            # no history yet
            new_trial = new_value
        else:
            trial = self.trials[-1]
            residual = new_value - trial
            self.residuals.append(residual)

            # trim history
            self.trials = self.trials[-self.history:]
            self.residuals = self.residuals[-self.history:]

            if self.i <= 1 or (self.i % self.period) != 0:
                # linear mixing
                new_trial = trial + self.alpha * residual
            else:
                # pulary mixing
                R = np.array(self.trials); R = R[1:] - R[:-1]; R = R.T
                F = np.array(self.residuals); F = F[1:] - F[:-1]; F = F.T

                new_trial = trial + self.alpha * residual \
                         - np.linalg.multi_dot([R + self.alpha * F, np.linalg.pinv(F), residual])

        self.i += 1
        self.trials.append(new_trial)
        return new_trial