<commit_before>"""
Time stepping solvers.
"""
from __future__ import absolute_import
import numpy as nm

from sfepy.base.base import (get_default, output, assert_,
                             Struct, IndexedStruct, basestr)
from sfepy.linalg.utils import output_array_stats
from sfepy.solvers.solvers import SolverMeta, TimeSteppingSolver
from sfepy.solvers.ts import TimeStepper, VariableTimeStepper
import six

class NewmarkState(Struct):

    def __init__(self, uname, vname, aname):
        Struct.__init__(self, uname=uname, vname=vname, aname=aname)

    def unpack(self, state):
        parts = state.get_parts()
        vvs = state.variables
        u = vvs[self.uname].get_reduced(parts[self.uname])
        v = vvs[self.vname].get_reduced(parts[self.vname])
        a = vvs[self.aname].get_reduced(parts[self.aname])

        return u, v, a

    def pack(self, state, u, v, a):
        vec = nm.r_[u, v, a]

        state.set_reduced(vec)

def gen_multi_vec_packing(size, num):
    assert_((size % num) == 0)
    ii = size // num

    def unpack(vec):
        return [vec[ir:ir+ii] for ir in range(0, size, ii)]

    def pack(*args):
        return nm.concatenate(args)

    return unpack, pack

def _cache(obj, attr, dep):
    def decorate(fun):
        def new_fun(*args, **kwargs):
            if dep:
                val = getattr(obj, attr)
                if val is None:
                    val = fun(*args, **kwargs)
                    setattr(obj, attr, val)

            else:
                val = fun(*args, **kwargs)

            return val
        return new_fun
    return decorate

class NewmarkTS(TimeSteppingSolver):
    """
    Assumes block-diagonal matrix in `u`, `v`, `a`.

    Common settings [1]:

                                 beta gamma Omega_crit
    trapezoidal rule:    implicit 1/4  1/2  unconditional
    linear acceleration: implicit 1/6  1/2  2\sqrt{3}
    Fox-Goodwin:         implicit 1/12 1/2  \sqrt{6}
    central difference:  explicit 0    1/2  2

    All of these methods are 2-order of accuracy.

    [1] <NAME>, <NAME>: Solvers for Computational Mechanics
    """
    name = 'ts.newmark'

    __metaclass__ = SolverMeta

    _parameters = [
        ('t0', 'float', 0.0, False,
         'The initial time.'),
        ('t1', 'float', 1.0, False,
         'The final time.'),
        ('dt', 'float', None, False,
         'The time step. Used if `n_step` is not given.'),
        ('n_step', 'int', 10, False,
         'The number of time steps. Has precedence over `dt`.'),
        ('quasistatic', 'bool', False, False,
         """If True, the non-linear solver is invoked also for
         the initial time."""),
        ('is_linear', 'bool', False, False,
         'If True, the problem is considered to be linear.'),
        ('beta', 'float', 0.25, False, 'The Newmark method parameter beta.'),
        ('gamma', 'float', 0.5, False, 'The Newmark method parameter gamma.'),
    ]

    def __init__(self, conf, nls=None, context=None, **kwargs):
        TimeSteppingSolver.__init__(self, conf, nls=nls, context=context,
                                    **kwargs)

        self.ts = TimeStepper.from_conf(self.conf)

        nd = self.ts.n_digit
        format = '====== time %%e (step %%%dd of %%%dd) =====' % (nd, nd)

        self.format = format
        self.verbose = self.conf.verbose
        self.constant_matrices = None
        self.matrix = None

    def get_matrices(self, nls, vec):
        if self.conf.is_linear and self.constant_matrices is not None:
            out = self.constant_matrices

        else:
            aux = nls.fun_grad(vec)

            assert_((len(vec) % 3) == 0)
            i3 = len(vec) // 3

            K = aux[:i3, :i3]
            C = aux[i3:2*i3, i3:2*i3]
            M = aux[2*i3:, 2*i3:]

            out = (M, C, K)

            if self.conf.is_linear:
                M.eliminate_zeros()
                C.eliminate_zeros()
                K.eliminate_zeros()
                self.constant_matrices = (M, C, K)

        return out

    def get_a0(self, nls, u0, v0):
        vec = nm.r_[u0, v0, nm.zeros_like(u0)]

        aux = nls.fun(vec)
        i3 = len(u0)
        r = aux[:i3] + aux[i3:2*i3] + aux[2*i3:]

        M = self.get_matrices(nls, vec)[0]
        a0 = nls.lin_solver(-r, mtx=M)
        output_array_stats(a0, 'initial acceleration', verbose=self.verbose)
        return a0

    def create_nlst(self, nls, dt, gamma, beta, u0, v0, a0):
        nlst = nls.copy()

        dt2 = dt**2

        def v(a):
            return v0 + dt * ((1.0 - gamma) * a0 + gamma * a)

        def u(a):
            return u0 + dt * v0 + dt2 * ((0.5 - beta) * a0 + beta * a)

        def fun(at):
            vec = nm.r_[u(at), v(at), at]

            aux = nls.fun(vec)

            i3 = len(at)
            rt = aux[:i3] + aux[i3:2*i3] + aux[2*i3:]
            return rt

        @_cache(self, 'matrix', self.conf.is_linear)
        def fun_grad(at):
            vec = None if self.conf.is_linear else nm.r_[u(at), v(at), at]
            M, C, K = self.get_matrices(nls, vec)

            Kt = M + gamma * dt * C + beta * dt2 * K
            return Kt

        nlst.fun = fun
        nlst.fun_grad = fun_grad
        nlst.u = u
        nlst.v = v

        return nlst

    def __call__(self, vec0=None, conf=None, nls=None, init_fun=None,
                 prestep_fun=None, poststep_fun=None,
                 status=None, **kwargs):
        """
        Solve elastodynamics problems by the Newmark method.
        """
        conf = get_default(conf, self.conf)
        nls = get_default(nls, self.nls)

        ts = self.ts

        unpack, pack = gen_multi_vec_packing(len(vec0), 3)

        gamma = conf.gamma
        beta = conf.beta

        init_fun(ts)

        prestep_fun(ts, vec0)
        u0, v0, _ = unpack(vec0)

        ut = u0
        vt = v0
        at = self.get_a0(nls, u0, v0)

        vec = pack(ut, vt, at)
        poststep_fun(ts, vec)
        ts.advance()
        for step, time in ts.iter_from(ts.step):
            output(self.format % (time, step + 1, ts.n_step),<commit_msg>update NewmarkTS.__call__() signature: remove conf argument<commit_after>"""
Time stepping solvers.
"""
from __future__ import absolute_import
import numpy as nm

from sfepy.base.base import (get_default, output, assert_,
                             Struct, IndexedStruct, basestr)
from sfepy.linalg.utils import output_array_stats
from sfepy.solvers.solvers import SolverMeta, TimeSteppingSolver
from sfepy.solvers.ts import TimeStepper, VariableTimeStepper
import six

class NewmarkState(Struct):

    def __init__(self, uname, vname, aname):
        Struct.__init__(self, uname=uname, vname=vname, aname=aname)

    def unpack(self, state):
        parts = state.get_parts()
        vvs = state.variables
        u = vvs[self.uname].get_reduced(parts[self.uname])
        v = vvs[self.vname].get_reduced(parts[self.vname])
        a = vvs[self.aname].get_reduced(parts[self.aname])

        return u, v, a

    def pack(self, state, u, v, a):
        vec = nm.r_[u, v, a]

        state.set_reduced(vec)

def gen_multi_vec_packing(size, num):
    assert_((size % num) == 0)
    ii = size // num

    def unpack(vec):
        return [vec[ir:ir+ii] for ir in range(0, size, ii)]

    def pack(*args):
        return nm.concatenate(args)

    return unpack, pack

def _cache(obj, attr, dep):
    def decorate(fun):
        def new_fun(*args, **kwargs):
            if dep:
                val = getattr(obj, attr)
                if val is None:
                    val = fun(*args, **kwargs)
                    setattr(obj, attr, val)

            else:
                val = fun(*args, **kwargs)

            return val
        return new_fun
    return decorate

class NewmarkTS(TimeSteppingSolver):
    """
    Assumes block-diagonal matrix in `u`, `v`, `a`.

    Common settings [1]:

                                 beta gamma Omega_crit
    trapezoidal rule:    implicit 1/4  1/2  unconditional
    linear acceleration: implicit 1/6  1/2  2\sqrt{3}
    Fox-Goodwin:         implicit 1/12 1/2  \sqrt{6}
    central difference:  explicit 0    1/2  2

    All of these methods are 2-order of accuracy.

    [1] <NAME>, <NAME>: Solvers for Computational Mechanics
    """
    name = 'ts.newmark'

    __metaclass__ = SolverMeta

    _parameters = [
        ('t0', 'float', 0.0, False,
         'The initial time.'),
        ('t1', 'float', 1.0, False,
         'The final time.'),
        ('dt', 'float', None, False,
         'The time step. Used if `n_step` is not given.'),
        ('n_step', 'int', 10, False,
         'The number of time steps. Has precedence over `dt`.'),
        ('quasistatic', 'bool', False, False,
         """If True, the non-linear solver is invoked also for
         the initial time."""),
        ('is_linear', 'bool', False, False,
         'If True, the problem is considered to be linear.'),
        ('beta', 'float', 0.25, False, 'The Newmark method parameter beta.'),
        ('gamma', 'float', 0.5, False, 'The Newmark method parameter gamma.'),
    ]

    def __init__(self, conf, nls=None, context=None, **kwargs):
        TimeSteppingSolver.__init__(self, conf, nls=nls, context=context,
                                    **kwargs)

        self.ts = TimeStepper.from_conf(self.conf)

        nd = self.ts.n_digit
        format = '====== time %%e (step %%%dd of %%%dd) =====' % (nd, nd)

        self.format = format
        self.verbose = self.conf.verbose
        self.constant_matrices = None
        self.matrix = None

    def get_matrices(self, nls, vec):
        if self.conf.is_linear and self.constant_matrices is not None:
            out = self.constant_matrices

        else:
            aux = nls.fun_grad(vec)

            assert_((len(vec) % 3) == 0)
            i3 = len(vec) // 3

            K = aux[:i3, :i3]
            C = aux[i3:2*i3, i3:2*i3]
            M = aux[2*i3:, 2*i3:]

            out = (M, C, K)

            if self.conf.is_linear:
                M.eliminate_zeros()
                C.eliminate_zeros()
                K.eliminate_zeros()
                self.constant_matrices = (M, C, K)

        return out

    def get_a0(self, nls, u0, v0):
        vec = nm.r_[u0, v0, nm.zeros_like(u0)]

        aux = nls.fun(vec)
        i3 = len(u0)
        r = aux[:i3] + aux[i3:2*i3] + aux[2*i3:]

        M = self.get_matrices(nls, vec)[0]
        a0 = nls.lin_solver(-r, mtx=M)
        output_array_stats(a0, 'initial acceleration', verbose=self.verbose)
        return a0

    def create_nlst(self, nls, dt, gamma, beta, u0, v0, a0):
        nlst = nls.copy()

        dt2 = dt**2

        def v(a):
            return v0 + dt * ((1.0 - gamma) * a0 + gamma * a)

        def u(a):
            return u0 + dt * v0 + dt2 * ((0.5 - beta) * a0 + beta * a)

        def fun(at):
            vec = nm.r_[u(at), v(at), at]

            aux = nls.fun(vec)

            i3 = len(at)
            rt = aux[:i3] + aux[i3:2*i3] + aux[2*i3:]
            return rt

        @_cache(self, 'matrix', self.conf.is_linear)
        def fun_grad(at):
            vec = None if self.conf.is_linear else nm.r_[u(at), v(at), at]
            M, C, K = self.get_matrices(nls, vec)

            Kt = M + gamma * dt * C + beta * dt2 * K
            return Kt

        nlst.fun = fun
        nlst.fun_grad = fun_grad
        nlst.u = u
        nlst.v = v

        return nlst

    def __call__(self, vec0=None, nls=None, init_fun=None, prestep_fun=None,
                 poststep_fun=None, status=None, **kwargs):
        """
        Solve elastodynamics problems by the Newmark method.
        """
        conf = self.conf
        nls = get_default(nls, self.nls)

        ts = self.ts

        unpack, pack = gen_multi_vec_packing(len(vec0), 3)

        gamma = conf.gamma
        beta = conf.beta

        init_fun(ts)

        prestep_fun(ts, vec0)
        u0, v0, _ = unpack(vec0)

        ut = u0
        vt = v0
        at = self.get_a0(nls, u0, v0)

        vec = pack(ut, vt, at)
        poststep_fun(ts, vec)
        ts.advance()
        for step, time in ts.iter_from(ts.step):
            output(self.format % (time, step + 1, ts.n_step),
{'fraction_non_alphanumeric': 0.12918817072171704, 'fraction_numerical': 0.017694765298599164, 'mean_word_length': 1.8443616029822927, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4303841', 'n_tokens_mistral': 4344, 'n_tokens_neox': 4110, 'n_words': 1084}