integrator.py

import numpy as np
from typing import Optional
import casadi as cas

from name_and_value_fun import assign_values_to_model
from base import Model


class Initializer:
    def __init__(self, name: str, model: Model, solver='kinsol', options={}):
        """ Initialize ODE by yourself instead of relying on DAE/ODE solver

            Given known values, then decide the unknown values by solving algebraic equations;

            init_prob: dict with keys 'sym', 'val', 'alg'

            solver: 'newton', 'kinsol' or 'nlpsol';
               'nlpsol' is the most robust one, but options should be given to specify nlp solver,
                example: z0 = ode_fun.initialize(z0, np.concatenate([x0, p0]), 'nlpsol', {'nlpsol': 'ipopt'})
        """

        self.name = name
        self.model = model
        self.solver = solver
        self.options = options

        model.add_init_eq()
        model.classify_and_vectorize_fixed()
        self.init_eq = cas.vcat(self.model.inits)
        self.fun = self.get_initialize_fun()

    def initialize(self, update_fixed: bool = True):
        """
        update_fixed is to determine whether we update the values of fixed and free variables or not
        """
        if self.fun is None:
            return

        if update_fixed:
            self.model.collect_var_val(['fixed', 'free'])

        free_val = self.fun(self.model.vars_val['free'],  self.model.vars_val['fixed'])
        self.model.vars_val['free'] = free_val

        assign_values_to_model(self.model, self.model.vars_code['free'], free_val)

    def get_initialize_fun(self):
        """ Get the CasADi function to generate free variable values from fixed variable values"""
        init_fixed, init_free = self.model.vars_symbol['fixed'], self.model.vars_symbol['free']

        alg_fun = cas.Function('alg_fun', [init_free, init_fixed], [self.init_eq])
        sol_fun = cas.rootfinder('initializer', self.solver, alg_fun, self.options)
        return sol_fun


class Integrator:
    def __init__(self, name, solver, dae_prob, t0=0., tf=1., options=None,
                 initializer: Optional[Initializer] = None):
        self.name = name
        self.solver = solver
        self.dae_prob = dae_prob
        self.options = options
        self.initializer = initializer

        self.dae_fun = cas.integrator(name, solver, dae_prob, t0, tf, options)

    def update_dae_fun(self, t0:Optional[float]=None, tf:Optional[float]=None):
        """Even time the solver, the dae_prob or the options has been updated, we should run the function"""
        if t0 is None:
            if 't0' in self.options:
                t0 = self.options['t0']
            else:
                t0 = 0.

        if tf is None:
            if 'tf' in self.options:
                tf = self.options['tf']
            else:
                tf = 1.
        self.dae_fun = cas.integrator(self.name, self.solver, self.dae_prob, t0, tf, self.options)

    def solve_steady_state(self, z_guess, x0p, rootfinder='newton'):
        pass