Example #1
0
    def __init__(self, **kwargs):
        SystemModel.__init__(self, **kwargs)

        x_1 = self.create_state("x_0")
        x_2 = self.create_state("x_1")
        u = self.create_control("u")

        ode = [(1 - x_2**2) * x_1 - x_2 + u, x_1]

        self.include_equations(ode=ode)
Example #2
0
    def __init__(self, **kwargs):
        SystemModel.__init__(self, **kwargs)

        a = DM([[-1, -2], [5, -1]])
        b = DM([[1, 0], [0, 1]])

        x = self.create_state("x", 2)
        u = self.create_control("u", 2)

        self.include_equations(ode=mtimes(a, x) + mtimes(b, u))
Example #3
0
    def __init__(self, **kwargs):
        SystemModel.__init__(self, **kwargs)

        x_1 = self.create_state("x_0")
        x_2 = self.create_state("x_1")
        y = self.create_algebraic_variable("y")
        u = self.create_control("u")

        ode = [y + u, x_1]
        alg = [(1 - x_2**2) * x_1 - x_2 - y]

        self.include_equations(ode=ode, alg=alg)
    def setUp(self):
        # create ode system
        self.ode_model = SystemModel(name="ode_sys")
        x = self.ode_model.create_state("x", 3)
        u = self.ode_model.create_input("u", 3)
        self.ode_model.include_equations(ode=-x + u)

        self.dae_model = SystemModel(name="dae_sys")
        x = self.dae_model.create_state("x", 3)
        y = self.dae_model.create_algebraic_variable("y", 3)
        u = self.dae_model.create_input("u", 3)
        self.dae_model.include_equations(ode=-x + u, alg=y - x + u**2)
Example #5
0
    def __init__(self, name="dae_system", **kwargs):
        SystemModel.__init__(self,
                             name=name,
                             model_name_as_prefix=True,
                             **kwargs)

        x = self.create_state("x", 2)
        y = self.create_algebraic_variable("y", 2)
        u = self.create_control("u", 2)
        a = self.create_parameter("a")

        ode = [-a * x[0] + y[0], -x[1] + y[1] + u[0]]
        alg = [-y[0] - x[1] + u[1], -y[1] - x[0]]

        self.include_equations(ode=ode, alg=alg)
Example #6
0
    def _create_model(self, sampled_parameters):
        sampled_parameters_p_unc = sampled_parameters[:self.socp.n_p_unc, :]
        model = SystemModel(name=self.socp.model.name + '_PCE')

        model.include_control(self.socp.model.u_sym)
        model.include_parameter(self.socp.get_p_without_p_unc())
        model.include_theta(self.socp.model.theta_sym)

        u_global = model.u_sym
        p_global = model.p_sym
        theta_global = model.theta_sym

        t_global = model.t_sym
        tau_global = model.tau_sym

        cost_list = []
        for s in range(self.n_samples):
            model_s = self.socp.model.get_deepcopy()

            # cost of sample
            cost_ode = self._create_cost_ode_of_sample(model_s)
            cost_s = model_s.create_state('cost_' + str(s))
            model_s.include_system_equations(ode=cost_ode)

            # replace the parameter variable with the sampled variable
            p_unc_s = model_s.p_sym.get(False, self.socp.get_p_unc_indices())

            model_s.replace_variable(p_unc_s, sampled_parameters_p_unc[:, s])
            model_s.remove_parameter(p_unc_s)

            # replace the model variables with the global variables
            model_s.replace_variable(model_s.u_sym, u_global)
            model_s.replace_variable(model_s.p_sym, p_global)
            model_s.replace_variable(model_s.theta_sym, theta_global)
            model_s.remove_control(model_s.u_sym)
            model_s.remove_parameter(model_s.p_sym)
            model_s.remove_theta(model_s.theta_sym)
            model_s.t_sym = t_global
            model_s.tau_sym = tau_global

            # merge the sample model in the unique model
            model.merge(model_s)

            # collect the sample cost
            cost_list.append(cost_s)

        cost_list = vertcat(*cost_list)
        return model, cost_list
    def __init__(self, index, **kwargs):
        # define default values
        self.gamma = 0.5
        self.k = 3.33e-6

        # super class method
        SystemModel.__init__(self,
                             name="pump_" + str(index),
                             model_name_as_prefix=True,
                             **kwargs)

        # create variables
        v = self.create_control("v")  # Create input pump voltage
        q = self.create_algebraic_variable("q", 2)  # Create pump flow

        # create equations:
        alg = [
            q[0] - self.gamma * self.k * v,
            q[1] - (1 - self.gamma) * self.k * v
        ]
        self.include_equations(alg=alg)
    def __init__(self, index, **kwargs):
        # define default values
        self.g = 9.8
        self.A = 28e-4
        self.a = 0.071e-4

        # super class method
        SystemModel.__init__(self,
                             name="tank_" + str(index),
                             model_name_as_prefix=True,
                             **kwargs)

        # create variables
        h = self.create_state("h")
        q_in = self.create_control("q_in")
        q_out = self.create_algebraic_variable("q_out")

        ode, alg = self.equations(h, q_in, q_out)

        # create equations:
        self.include_equations(ode, alg)
    def __init__(self, **kwargs):
        SystemModel.__init__(self, **kwargs)

        # create the pumps
        pumps = [Pump(index=0, gamma=0.7), Pump(index=1, gamma=0.6)]
        # create the tanks
        tanks = [
            Tank(index=0, n_inputs=2),
            Tank(index=1, n_inputs=2),
            Tank(index=2, n_inputs=1),
            Tank(index=3, n_inputs=1),
        ]

        # include all variables and equations from pumps and tanks
        self.include_models(pumps)
        self.include_models(tanks)

        # connections
        self.connect(tanks[0].u, pumps[0].y[0] + tanks[2].y[0])
        self.connect(tanks[3].u, pumps[1].y[1] + tanks[3].y[0])
        self.connect(tanks[1].u, pumps[1].y[0])
        self.connect(tanks[2].u, pumps[0].y[1])
Example #10
0
def get_model(name="dae_system"):
    model = SystemModel(name=name, model_name_as_prefix=True)

    x = model.create_state("x", 2)
    y = model.create_algebraic_variable("y", 2)
    u = model.create_control("u")
    a = model.create_parameter("a")

    # model.include_system_equations(ode=[
    #     -a * x[0] + y[0],
    #     -x[1] + y[1] + u[0]
    # ], alg=[
    #     -y[0] - x[1] ** 2,
    #     -y[1] - x[0] ** 1
    # ])
    model.include_equations(ode=[-a * x[0] + y[0], -x[1] + y[1] + u[0]],
                            alg=[-y[0] - x[1], -y[1] - x[0]])
    return model
Example #11
0
    def get_problem(self):
        """
            Create a single OCP which is the composition of all problems and connections.

        :return: problem
        :rtype: OptimalControlProblem
        """
        model = SystemModel(name=self.name + '_model')
        problem = OptimalControlProblem(model=model,
                                        name=self.name + '_problem')
        problem.t_0 = self.problems[0].t_0
        problem.t_f = self.problems[0].t_f
        problem.merge(self.problems)
        for edge in self.graph.edges:
            problem.connect(u=self.graph.edges[edge]['u'],
                            y=self.graph.edges[edge]['y'])

        return problem
Example #12
0
    def get_model(self):
        """
            Create a single model which is the composition of all models and the connections

        :return: model
        :rtype: SystemModel
        """
        model = SystemModel(name=self.name + '_model')
        model.include_models(self.models)
        for edge in self.graph.edges:
            model.connect(u=self.graph.edges[edge]['u'],
                          y=self.graph.edges[edge]['y'])

        return model
Example #13
0
    def insert_intermediary_nodes(self):
        old_connections = list(self.connections)

        for (node1, node2) in old_connections:
            y = self.graph.edges[node1, node2]['y']
            u = self.graph.edges[node1, node2]['u']

            y_guess = vertcat(
                *node1.problem.y_guess)[find_variables_indices_in_vector(
                    y, node1.problem.model.y)]
            u_guess = vertcat(
                *node2.problem.u_guess)[find_variables_indices_in_vector(
                    u, node2.problem.model.u)]

            copy_y = vertcat(
                *
                [SX.sym('Dummy_' + y[ind].name()) for ind in range(y.numel())])
            copy_u = vertcat(
                *
                [SX.sym('Dummy_' + u[ind].name()) for ind in range(u.numel())])

            new_model = SystemModel(
                name='Dummy_Model_{}_to_{}'.format(node1.name, node2.name))
            new_model.include_variables(u=copy_y, y=copy_u)
            new_model.include_equations(alg=copy_u - copy_y)
            new_problem = OptimalControlProblem(
                name='OCP_Dummy_{}_to_{}'.format(node1.name, node2.name),
                model=new_model,
                t_f=node1.problem.t_f,
                y_guess=u_guess,
                u_guess=y_guess)

            new_node = Node(name='Dummy_node_{}_to_{}'.format(
                node1.name, node2.name),
                            model=new_model,
                            problem=new_problem,
                            color=0.75)

            self.include_nodes(new_node)

            self.remove_connection(node1, node2)
            self.connect(y, copy_y, node1, new_node)
            self.connect(copy_u, u, new_node, node2)

        for node in self.nodes:
            print(node.node_id, node.name)
Example #14
0
from yaocptool.methods import IndirectMethod
from yaocptool.modelling import SystemModel, OptimalControlProblem

model = SystemModel()
model.create_state('x')
model.create_control('x')

model.include_equations(ode=(-model.x[0] + model.u[0]))

problem = OptimalControlProblem(model, obj={"Q": 1, "R": 1}, x_0=[1], t_f=5.0)
# problem.u_min[0] = 0
# problem.u_max[0] = 1

# problem.x_min[0] = 0.6

solution_method = IndirectMethod(
    problem,
    degree_control=3,
    discretization_scheme="multiple-shooting",
    # discretization_scheme='collocation',
    degree=3,
    finite_elements=30,
    integrator_type="implicit",
)

solution = solution_method.solve()

solution.plot([{"x": [0, 1]}, {"u": [0]}])
Example #15
0
from __future__ import print_function
from yaocptool.modelling import SystemModel
from yaocptool.modelling import OptimalControlProblem
from yaocptool.methods import DirectMethod

# PART 1
model = SystemModel(name="simple_model")
x = model.create_state("x")  # vector of state variables
u = model.create_control("u")  # vector of control variables

# Include the dynamic equation
ode = [-x + u]
model.include_equations(ode=ode)

# Print model information
print(model)

# Part 2
problem = OptimalControlProblem(model, x_0=[1], t_f=10, obj={"Q": 1, "R": 1})

# Part 3
# Initialize a DirectMethod to solve the OCP using collocation
solution_method = DirectMethod(
    problem, finite_elements=20, discretization_scheme="collocation"
)

# Solve the problem and get the result
result = solution_method.solve()

# Make one plot with the element x[0] (the first state) and one plot with the control u[0]
result.plot([{"x": [0]}, {"u": [0]}])
Example #16
0
def _create_linear_system(n_x, n_u, a, b, name):
    model = SystemModel(name=name)
    x = model.create_state('x', n_x)
    u = model.create_control('u', n_u)
    model.include_equations(ode=mtimes(a, x) + mtimes(b, u))
    return model
def test_include_equations_ode_with_x(empty_model: SystemModel):
    x = empty_model.create_state('x')
    ode = -x
    empty_model.include_equations(ode=ode, x=x)
    assert empty_model.ode.numel() == x.numel()
    assert is_equal(empty_model.ode, ode, 20)
Example #18
0
# initial state and control

x_0 = DM([1, 1])
initial_control = [0.01]

# Prediction window, finite elements, and sampling time
prediction_window = 20
finite_elements = 20
t_s = prediction_window / finite_elements

######################
#      Model         #
######################
# Create a new model with 2 Tanks, the output of the first tank is connected on the second tank

model = SystemModel(name="2-Tanks")

# Get the symbolic variables
h_1, h_2 = model.create_state("h_1"), model.create_state("h_2")
u = model.create_control("u")

# Model Parameters
a = model.create_parameter("a")  # (m^2) Holes cross section
a_mean = model.create_parameter("a_mean")  # (m^2) Holes cross section
A = 28e-3  # (m^2) the tank area
g = 9.8  # gravitational acceleration

# Define the model ODEs
ode = [
    (u - a * sqrt(2 * g * h_1)) / A,
    (a * sqrt(2 * g * h_1) - a * sqrt(2 * g * h_2)) / A,
class TestSystemModel(TestCase):
    def setUp(self):
        # create ode system
        self.ode_model = SystemModel(name="ode_sys")
        x = self.ode_model.create_state("x", 3)
        u = self.ode_model.create_input("u", 3)
        self.ode_model.include_equations(ode=-x + u)

        self.dae_model = SystemModel(name="dae_sys")
        x = self.dae_model.create_state("x", 3)
        y = self.dae_model.create_algebraic_variable("y", 3)
        u = self.dae_model.create_input("u", 3)
        self.dae_model.include_equations(ode=-x + u, alg=y - x + u**2)

    def test_system_type(self):
        assert self.ode_model.system_type == "ode"
        assert self.dae_model.system_type, "dae"

    def test_x_sys_sym(self):
        self.assertTrue(is_equal(self.ode_model.x_sys_sym, self.ode_model.x))

        # with adjoints
        ode_x = self.ode_model.x[:]
        dae_x = self.dae_model.x[:]
        self.ode_model.create_state("lamb", self.ode_model.n_x)
        self.ode_model.has_adjoint_variables = True

        self.dae_model.create_state("lamb", self.dae_model.n_x)
        self.dae_model.has_adjoint_variables = True

        self.assertTrue(is_equal(self.ode_model.x_sys_sym, ode_x))
        self.assertTrue(is_equal(self.dae_model.x_sys_sym, dae_x))

    def test_lamb_sym(self):
        self.assertEqual(self.ode_model.lamb_sym.numel(), 0)
        self.assertEqual(self.dae_model.lamb_sym.numel(), 0)

        # with adjoints
        ode_lamb = self.ode_model.create_state("lamb", self.ode_model.n_x)
        self.ode_model.has_adjoint_variables = True

        dae_lamb = self.dae_model.create_state("lamb", self.dae_model.n_x)
        self.dae_model.has_adjoint_variables = True

        self.assertTrue(is_equal(self.ode_model.lamb_sym, ode_lamb))
        self.assertTrue(is_equal(self.dae_model.lamb_sym, dae_lamb))

    def test_all_sym(self):
        for model in [self.ode_model, self.dae_model]:
            answer = [
                model.t,
                model.x,
                model.y,
                model.p,
                model.theta,
                model.u_par,
            ]
            self.assertEqual(len(model.all_sym), len(answer))

            for index in range(len(model.all_sym)):
                self.assertTrue(is_equal(model.all_sym[index], answer[index]))

    def test_t(self):
        self.assertTrue(is_equal(self.ode_model.t, self.ode_model.t))
        self.assertTrue(is_equal(self.dae_model.t, self.dae_model.t))

    def test_t_setter(self):
        new_t = SX.sym("t")
        self.ode_model.t = new_t
        self.dae_model.t = new_t

        self.assertTrue(is_equal(self.ode_model.t, new_t))
        self.assertTrue(is_equal(self.dae_model.t, new_t))

    def test_tau(self):
        self.assertTrue(is_equal(self.ode_model.tau, self.ode_model.tau))
        self.assertTrue(is_equal(self.dae_model.tau, self.dae_model.tau))

    def test_tau_setter(self):
        new_tau = SX.sym("tau")
        self.ode_model.tau = new_tau
        self.dae_model.tau = new_tau

        self.assertTrue(is_equal(self.ode_model.tau, new_tau))
        self.assertTrue(is_equal(self.dae_model.tau, new_tau))

    def test_x_names(self):
        for model in [self.ode_model, self.dae_model]:
            for ind in range(model.n_x):
                self.assertEqual(model.x[ind].name(), model.x_names[ind])

    def test_y_names(self):
        for model in [self.ode_model, self.dae_model]:
            for ind in range(model.n_y):
                self.assertEqual(model.y[ind].name(), model.y_names[ind])

    def test_u_names(self):
        for model in [self.ode_model, self.dae_model]:
            for ind in range(model.n_u):
                self.assertEqual(model.u[ind].name(), model.u_names[ind])

    def test_print_variables(self):
        for model in [self.ode_model, self.dae_model]:
            model.print_variables()

    def test_create_input(self):
        n_u_initial = self.ode_model.n_u
        n_new_u = 4
        u = self.ode_model.create_input("u", n_new_u)
        self.assertEqual(self.ode_model.n_u, n_u_initial + n_new_u)
        self.assertTrue(is_equal(self.ode_model.u[-n_new_u:], u))

    def test_get_variable_by_name(self):
        model = self.dae_model.get_copy()

        # Variable does not exist:
        self.assertRaises(ValueError, model.get_variable_by_name, "other_var")

        # Multiple vars
        self.assertRaises(ValueError, model.get_variable_by_name, "x")

        # Some var that exists and is unique
        var = model.get_variable_by_name("x_1")
        self.assertTrue(is_equal(model.x[1], var))

    def test_get_variables_by_names(self):
        model = self.dae_model.get_copy()

        # Variable does not exist
        self.assertEqual(model.get_variables_by_names("other_var"), [])

        # Multiple vars
        self.assertGreater(len(model.get_variables_by_names("x")), 0)
        res = model.get_variables_by_names("x")
        for ind in range(model.n_x):
            self.assertTrue(is_equal(res[ind], model.x[ind]))

        # Some var that exists and is unique
        res = model.get_variables_by_names("x_1")
        self.assertEqual(len(res), 1)
        self.assertTrue(is_equal(model.x[1], res[0]))

        # find var with var_type
        model2 = self.dae_model.get_copy()
        p_with_name_starting_with_x = model2.create_parameter("x_ref")
        res = model2.get_variables_by_names("x", var_type="p")
        self.assertEqual(len(res), 1)
        self.assertTrue(is_equal(p_with_name_starting_with_x, res[0]))

        # find var with list of names
        res = model.get_variables_by_names(["x_1", "u_2"])
        self.assertEqual(len(res), 2)
        self.assertTrue(is_equal(model.x[1], res[0]))
        self.assertTrue(is_equal(model.u[2], res[1]))

    def test_has_variable(self):
        self.assertTrue(self.dae_model.has_variable(self.dae_model.x[0]))
        self.assertFalse(self.dae_model.has_variable(SX.sym("x_0")))

    def test_is_parametrized(self):
        model = self.dae_model.get_copy()
        self.assertFalse(model.is_parametrized())

        k = model.create_parameter("k")
        model.parametrize_control(model.u[0], -k * model.x[0], k)
        self.assertTrue(model.is_parametrized())
Example #20
0
from casadi import mtimes

from yaocptool.methods import DirectMethod
from yaocptool.modelling import SystemModel, OptimalControlProblem

# create model
model = SystemModel(name="dae_system")

x = model.create_state("x", 2)
y = model.create_algebraic_variable("y", 2)
u = model.create_control("u")
a = model.create_parameter("a")
b = model.create_theta("b")

model.include_equations(
    ode=[-a * x[0] + b * y[0], -x[1] + y[1] + u[0]],
    alg=[-y[0] - x[1]**2, -y[1] - x[0]**1],
)

# create ocp
problem = OptimalControlProblem(model)
problem.t_f = 10
# problem.L = mtimes(x.T, x) + u ** 2
problem.S = mtimes(x.T, x) + u**2 + b**2
problem.x_0 = [0, 1]
problem.set_theta_as_optimization_theta(b, -0.5, 0.5)
# problem.include_equality(problem.p_opt + 0.25)
problem.include_time_inequality(+u + x[0], when="end")

# instantiate a solution method
solution_method = DirectMethod(
Example #21
0
from yaocptool.modelling import SystemModel

# create model
model = SystemModel(name="dae_system")

x = model.create_state("x", 2)
y = model.create_algebraic_variable("y", 2)
u = model.create_control("u")
a = model.create_parameter("a")
b = model.create_theta("b")

model.include_equations(
    ode=[-a * x[0] + b * y[0], -x[1] + y[1] + u[0]],
    alg=[-y[0] - x[1]**2, -y[1] - x[0]**1],
)

x_0 = [1, 2]
sim_result = model.simulate(x_0,
                            range(1, 10),
                            0.5,
                            u=1.0,
                            p=[1],
                            theta=dict(zip(range(0, 9), [0] * 9)))
# sim_result.plot([{'x': 'all'}, {'y': 'all'}, {'u': 'all'}])

# Include data at the end of the sim_result
copy_of_sim_result = sim_result.get_copy()
sim_result2 = model.simulate(x_0,
                             range(20, 30),
                             19,
                             u=1.0,
"""
Created on Wed Nov 02 18:55:27 2016

@author: marco
"""
import sys
from os.path import dirname, abspath

from casadi import vertcat, DM, mtimes

from yaocptool.methods import DirectMethod
from yaocptool.modelling import SystemModel, OptimalControlProblem

sys.path.append(abspath(dirname(dirname(__file__))))

model = SystemModel()
x = model.create_state("x", 2)
u = model.create_control("u", 2)

a = DM([[-1, -2], [5, -1]])
b = DM([[1, 0], [0, 1]])

model.include_equations(ode=vertcat(mtimes(a, x) + mtimes(b, u)))

problem = OptimalControlProblem(
    model,
    obj={
        "Q": DM.eye(2),
        "R": DM.eye(2)
    },
    x_0=[1, 1],
Example #23
0
c_matrix = DM([[1, 0]])

# initial state and control
x_0 = DM([1, 1])
initial_control = [0.01]

# Prediction window, finite elements, and sampling time
prediction_window = 20.0
finite_elements = 20
t_s = prediction_window / finite_elements

######################
#      Model         #
######################
# Create a new model with 2 Tanks, the output of the first tank is connected on the second tank
model = SystemModel(name="2-Tanks")

# Get the symbolic variables
h = model.create_state("h", 2)
u = model.create_control("u")

# Model Parameters
A = 28e-3  # (m^2) the tank area
a = 0.071e-2  # (m^2) Holes cross section
g = 9.8  # gravitational acceleration

# Define the model ODEs
ode = [
    (u - a * sqrt(2 * g * h[0])) / A,
    (a * sqrt(2 * g * h[0]) - a * sqrt(2 * g * h[1])) / A,
]
Example #24
0
from yaocptool.modelling import SystemModel
from casadi import sqrt, DM

# Create a new model with 2 Tanks, the output of the first tank is connected on the second tank
model = SystemModel(name="2-Tanks", n_x=2, n_u=1)

# Get the symbolic variables
h_1, h_2 = model.x[0], model.x[1]
u = model.u_sym

# Model Parameters
A = 28e-4  # (m^2) the tank area
a = 0.071e-4  # (m^2) Holes cross section
g = 9.8  # gravitational acceleration

# Define the model ODEs
ode = [
    (u - a * sqrt(2 * g * h_1)) / A,
    (a * sqrt(2 * g * h_1) - a * sqrt(2 * g * h_2)) / A,
]

# Include the equations in the model
model.include_equations(ode=ode)

# Just check if the where correctly included
print(model)

# Do a simple simulation
x_0 = DM([0.101211, 0.101211])
sim_result = model.simulate(x_0=x_0,
                            t_f=[i * 0.2 for i in range(1, 201)],
Example #25
0
    def _create_augmented_model_and_p_update(self, model):
        """

        :type model: SystemModel
        """
        aug_model = SystemModel("aug_linearized_EKF_model")
        aug_model.include_state(self.model.x_sym)
        aug_model.include_state(self.model.y_sym)
        aug_model.include_control(self.model.u_sym)
        aug_model.include_parameter(self.model.p_sym)
        aug_model.include_theta(self.model.theta_sym)

        # remove u_par (self.model.u_par model.u_par)
        all_sym = list(self.model.all_sym)

        # Mean
        a_func = Function("A_matrix", all_sym,
                          [jacobian(self.model.ode, self.model.x_sym)])
        b_func = Function("B_matrix", all_sym,
                          [jacobian(self.model.ode, self.model.y_sym)])
        c_func = Function("C_matrix", all_sym,
                          [jacobian(self.model.alg, self.model.x_sym)])
        d_func = Function("D_matrix", all_sym,
                          [jacobian(self.model.alg, self.model.y_sym)])

        x_lin = aug_model.create_parameter("x_lin", self.model.n_x)
        y_lin = aug_model.create_parameter("y_lin", self.model.n_y)

        all_sym[1] = x_lin
        all_sym[2] = y_lin

        a_matrix = a_func(*all_sym)
        b_matrix = b_func(*all_sym)
        c_matrix = c_func(*all_sym)
        d_matrix = d_func(*all_sym)

        a_aug_matrix = vertcat(
            horzcat(a_matrix, b_matrix),
            horzcat(
                mtimes(-solve(d_matrix, c_matrix), a_matrix),
                mtimes(-solve(d_matrix, c_matrix), b_matrix),
            ),
        )

        x_aug = vertcat(model.x_sym, model.y_sym)
        aug_model.include_equations(ode=[mtimes(a_aug_matrix, x_aug)])

        # Covariance
        gamma = vertcat(DM.eye(self.model.n_x), -solve(d_matrix, c_matrix))

        trans_matrix_sym = SX.sym("trans_matrix", a_aug_matrix.shape)
        p_matrix_sym = SX.sym("p_matrix",
                              (model.n_x + model.n_y, model.n_x + model.n_y))
        gamma_matrix_sym = SX.sym("gamma_matrix", gamma.shape)
        q_matrix_sym = SX.sym("q_matrix", (self.model.n_x, self.model.n_x))

        p_kpp = mtimes(trans_matrix_sym,
                       mtimes(p_matrix_sym, trans_matrix_sym.T)) + mtimes(
                           gamma_matrix_sym,
                           mtimes(q_matrix_sym, gamma_matrix_sym.T))

        a_aug_matrix_func = Function("trans_matrix", all_sym, [a_aug_matrix])
        gamma_func = Function("gamma", all_sym, [gamma])
        p_update_func = Function(
            "p_update",
            [trans_matrix_sym, p_matrix_sym, gamma_matrix_sym, q_matrix_sym],
            [p_kpp],
        )

        return aug_model, a_aug_matrix_func, gamma_func, p_update_func