Simple Example

To get an intuition on how to solve an ODE problem with ozone, lets walk through a simple example.

The Problem

Let's start by defining a simple ODE that we want to integrate,

\frac{\partial{y}}{\partial{t}} = -y, \quad y_0 = 0.5.

We want to integrate it from $t = 0$ to $t = 3$ using 4th order Runge-Kutta (RK4). With the problem defined, let's see how to implement this in ozone.

Defining the ODE function.

We represent our ODE function $-y$ as a csdl Model(). For those unfamiliar, csdl is an Embedded Domain Specific Language framework for modeling systems. For details, visit the csdl website.

Our csdl Model takes in the state $y$ as an input and outputs $dydt$ :

import csdl

class ODESystemModel(csdl.Model):
    def initialize(self):
        # Required every time for ODE systems or Profile Output systems
        self.parameters.declare('num_nodes')

    def define(self):
        # Required every time for ODE systems or Profile Output systems
        n = self.parameters['num_nodes']

        # State y. All ODE states must have shape = (n, .. shape of state ...)
        y = self.create_input('y', shape=n)

        # Register output dy/dt = -y
        self.register_output('dy_dt', -y)

Creating the Solver Model

We now have to create a separate csdl model that creates our integrator.

First, let's define the inputs to the ODE: initial condition $y_0 = 0.5$ and the timestep vector. The timestep vector is a vector of stepsizes to integrate over. Let's use a timestep of 0.01s. This means the timestep vector is [0.01, 0.01, ... 0.01, 0.01] with 30 elements to go from $t = 0$ to $t = 3$ . Both of these must be csdl variables.

Next, we will create the ODE solver model. Import the ODEProblem class from ozone.api. Instantiate the class with three arguments:

Numerical Method (RK4)
Solution Approach (Timemarching)
Number of timesteps (30)

We then set the state, time vector and the ODE function. Finally, we create the solver model.

import numpy as np
from ozone.api import ODEProblem

# The CSDL Model containing the ODE integrator
class RunModel(csdl.Model):

    def define(self):
        num_times = 31
        dt = 0.1

        # Create inputs to the ODE
        # Initial condition for state
        self.create_input('y_0', 0.5)

        # Timestep vector
        h_vec = np.ones(num_times-1)*dt
        self.create_input('h', h_vec)

        # Create ODEProblem class
        ode_problem = ODEProblem('RK4', 'time-marching', num_times)
        ode_problem.add_state('y', 'dy_dt', initial_condition_name='y_0', output='y_integrated')
        ode_problem.add_times(step_vector='h')
        ode_problem.set_ode_system(ODESystemModel)

        # Create CSDL Model of solver
        self.add(ode_problem.create_solver_model())

Running the Model

We can then run our csdl model. Let's also compute the derivative of the outputs with respect to the inputs. The output y_integrated is a vector of the solved state at each timestep.

import python_csdl_backend

# Simulator object:
sim = python_csdl_backend.Simulator(RunModel(), mode='rev')

# Run and check derivatives
sim.prob.run_model()
print('y integrated:', sim.prob['y_integrated'])
sim.prob.check_totals(of=['y_integrated'], wrt=['y_0'], compact_print=True)

Simple Example

The Problem​

Defining the ODE function.​

Creating the Solver Model​

Running the Model​

The Problem

Defining the ODE function.

Creating the Solver Model

Running the Model