A Simple Dymos Example

We can also use OpenMDAO.jl Components within a Dymos ODE. This example will implement the Brachistochrone example from the Dymos docs using Julia.

Preamble

Let's make things easier on ourselves and import the VarData and PartialsData names into our local namespace:

using OpenMDAOCore: OpenMDAOCore, VarData, PartialsData

The `AbstractExplicitComp` `struct`

The ODE component for the brachistochrone has two options:

num_nodes, the number of nodes used to descritize the trajectory of the bead.
static_gravity, a flag to indicate whether gravity should vary along the trajectory (and thus have length num_nodes) or if it should be constant (and thus be a scalar).

We'll use a default value of false for static_gravity, just like the Python implementation in the Dymos docs:

struct BrachistochroneODE <: OpenMDAOCore.AbstractExplicitComp
    num_nodes::Int
    static_gravity::Bool
end

# `static_gravity` set to `false` by default.
BrachistochroneODE(; num_nodes, static_gravity=false) = BrachistochroneODE(num_nodes, static_gravity)

Main.BrachistochroneODE

`OpenMDAOCore.setup`

Next we'll define the OpenMDAOCore.setup function:

function OpenMDAOCore.setup(self::BrachistochroneODE)
	nn = self.num_nodes

	# Inputs
    input_data = Vector{VarData}()
    push!(input_data, VarData("v"; val=zeros(nn), units="m/s"))
	if self.static_gravity
        push!(input_data, VarData("g", val=9.80665, units="m/s/s", tags=["dymos.static_target"]))
	else
        push!(input_data, VarData("g", val=9.80665 * ones(nn), units="m/s/s"))
    end
    push!(input_data, VarData("theta", val=ones(nn), units="rad"))

    # Outputs
    output_data = Vector{VarData}()
    push!(output_data, VarData("xdot", val=zeros(nn), units="m/s", tags=["dymos.state_rate_source:x", "dymos.state_units:m"]))
    push!(output_data, VarData("ydot", val=zeros(nn), units="m/s", tags=["dymos.state_rate_source:y", "dymos.state_units:m"]))
    push!(output_data, VarData("vdot", val=zeros(nn), units="m/s**2", tags=["dymos.state_rate_source:v", "dymos.state_units:m/s"]))
    push!(output_data, VarData("check", val=zeros(nn), units="m/s"))

    # Setup partials
    arange = 0:nn-1
    partials_data = Vector{PartialsData}()
    push!(partials_data, PartialsData("vdot", "theta"; rows=arange, cols=arange))

    push!(partials_data, PartialsData("xdot", "v"; rows=arange, cols=arange))
    push!(partials_data, PartialsData("xdot", "theta"; rows=arange, cols=arange))

    push!(partials_data, PartialsData("ydot", "v"; rows=arange, cols=arange))
    push!(partials_data, PartialsData("ydot", "theta"; rows=arange, cols=arange))

    push!(partials_data, PartialsData("check", "v"; rows=arange, cols=arange))
    push!(partials_data, PartialsData("check", "theta"; rows=arange, cols=arange))

	if self.static_gravity
		c = zeros(Int, self.num_nodes)
        push!(partials_data, PartialsData("vdot", "g"; rows=arange, cols=c))
	else
        push!(partials_data, PartialsData("vdot", "g"; rows=arange, cols=arange))
    end

    return input_data, output_data, partials_data
end

This is probably the most complicated setup we've seen yet. A few things to note:

We can change the size of the g (gravity) input and its sub-Jacobian using the static_gravity option in the BrachistochroneODE struct.
The VarData calls use tags, which are passed to the ExplicitComponent.add_input method using the tags keyword argument in OpenMDAO.
As we'll see, the job of a Dymos ODE is to compute the state rates from the states and controls. It turns out that in many (all) ODEs, these state rates for a given trajectory node only depend on the state and controls at that particular node. This implies that the Jacobian of the ODE calculation will be sparse. This example, like the original Python implementation, passes the sparsity pattern of the various sub-Jacobians to the PartialsData structs using the rows and cols keywords.

`OpenMDAOCore.compute!` and `OpenMDAOCore.compute_partials!`

The OpenMDAOCore.compute! method for our ODE is fairly straightforward:

function OpenMDAOCore.compute!(self::BrachistochroneODE, inputs, outputs)
	theta = inputs["theta"]
	cos_theta = cos.(theta)
	sin_theta = sin.(theta)
	g = inputs["g"]
	v = inputs["v"]

	@. outputs["vdot"] = g * cos_theta
	@. outputs["xdot"] = v * sin_theta
	@. outputs["ydot"] = -v * cos_theta
	@. outputs["check"] = v / sin_theta

    return nothing
end

The @. macro tells Julia to use broadcasting for the array calculations (similar to NumPy broadcasting).

The compute_partials! method is also quite similar to the original Python implementation:

function OpenMDAOCore.compute_partials!(self::BrachistochroneODE, inputs, partials)
	theta = inputs["theta"]
	cos_theta = cos.(theta)
	sin_theta = sin.(theta)
	g = inputs["g"]
	v = inputs["v"]

	@. partials["vdot", "g"] = cos_theta
	@. partials["vdot", "theta"] = -g * sin_theta

	@. partials["xdot", "v"] = sin_theta
	@. partials["xdot", "theta"] = v * cos_theta

	@. partials["ydot", "v"] = -cos_theta
	@. partials["ydot", "theta"] = v * sin_theta

	@. partials["check", "v"] = 1 / sin_theta
	@. partials["check", "theta"] = -v * cos_theta / sin_theta ^ 2

	return nothing
end

The run script

We'll need the Dymos library to solve the Brachistochrone problem, of course:

using PythonCall: pyimport
dm = pyimport("dymos")

Python module: <module 'dymos' from '/home/runner/work/OpenMDAO.jl/OpenMDAO.jl/docs/.CondaPkg/env/lib/python3.11/site-packages/dymos/__init__.py'>

And then the rest of the script will be pretty much identical to the Python version, but written in Julia. We'll put it in a function that allows us to try out static_gravity=false and static_gravity=true.

using OpenMDAO: om, make_component, DymosifiedCompWrapper

function main(; static_gravity)
  #
  # Initialize the Problem and the optimization driver
  #
  p = om.Problem(model=om.Group())
  p.driver = om.ScipyOptimizeDriver()
  p.driver.declare_coloring()
  #
  # Create a trajectory and add a phase to it
  #
  traj = p.model.add_subsystem("traj", dm.Trajectory())

  # `Trajectory.add_phase` expects a class that it can instantiate with the number of nodes used for the phase.
  # That's easy enough to create with an anonymous function, but unfortunatly it won't work with Dymos (because of the way JuliaCall implements truthiness/falsiness of Julia callables).
  # So, as a workaround, OpenMDAO.jl has a small wrapper `struct` that will fix this for us called `DymosifiedCompWrapper`.
  # The way this works: we give DymosifiedCompWrapper the ODE type (not a type instance, the type itself), and the keyword arguments we need to instantiate the type, other than `num_nodes`.
  # Check out the docstring for an example.
  dcw = DymosifiedCompWrapper(BrachistochroneODE; static_gravity=static_gravity)
  phase = traj.add_phase("phase0",
                         dm.Phase(ode_class = dcw,
                                  transcription = dm.GaussLobatto(num_segments=10)))

  #
  # Set the variables
  #
  phase.set_time_options(fix_initial=true, duration_bounds=(.5, 10))

  phase.add_state("x", fix_initial=true, fix_final=true)

  phase.add_state("y", fix_initial=true, fix_final=true)

  phase.add_state("v", fix_initial=true, fix_final=false)

  phase.add_control("theta", continuity=true, rate_continuity=true,
                    units="deg", lower=0.01, upper=179.9)

  phase.add_parameter("g", units="m/s**2", val=9.80665)

  #
  # Minimize time at the end of the phase
  #
  phase.add_objective("time", loc="final", scaler=10)
  #
  p.model.linear_solver = om.DirectSolver()
  #
  # Setup the Problem
  #
  p.setup()

  #
  # Set the initial values
  #
  p["traj.phase0.t_initial"] = 0.0
  p["traj.phase0.t_duration"] = 2.0

  p.set_val("traj.phase0.states:x", phase.interp("x", ys=[0, 10]))
  p.set_val("traj.phase0.states:y", phase.interp("y", ys=[10, 5]))
  p.set_val("traj.phase0.states:v", phase.interp("v", ys=[0, 9.9]))
  p.set_val("traj.phase0.controls:theta", phase.interp("theta", ys=[5, 100.5]))

  #
  # Solve for the optimal trajectory
  #
  dm.run_problem(p)

  # Check the results
  println("static_gravity = $static_gravity, elapsed time = $(p.get_val("traj.phase0.timeseries.time")[-1]) (should be 1.80164719)")
end

main(; static_gravity=false)
main(; static_gravity=true)


--- Constraint Report [traj] ---
    --- phase0 ---
        None

Model viewer data has already been recorded for Driver.
Full total jacobian was computed 3 times, taking 0.845695 seconds.
Total jacobian shape: (40, 50)


Jacobian shape: (40, 50)  (19.85% nonzero)
FWD solves: 13   REV solves: 0
Total colors vs. total size: 13 vs 50  (74.0% improvement)

Sparsity computed using tolerance: 1e-25
Time to compute sparsity: 0.845695 sec.
Time to compute coloring: 0.029953 sec.
static_gravity = false, elapsed time = [1.80161673] (should be 1.80164719)
/home/runner/work/OpenMDAO.jl/OpenMDAO.jl/docs/.CondaPkg/env/lib/python3.11/site-packages/openmdao/recorders/sqlite_recorder.py:227: UserWarning:The existing case recorder file, dymos_solution.db, is being overwritten.
Optimization terminated successfully    (Exit mode 0)
            Current function value: 18.016167304638337
            Iterations: 24
            Function evaluations: 24
            Gradient evaluations: 24
Optimization Complete
-----------------------------------

--- Constraint Report [traj] ---
    --- phase0 ---
        None

Model viewer data has already been recorded for Driver.
Full total jacobian was computed 3 times, taking 0.037691 seconds.
Total jacobian shape: (40, 50)


Jacobian shape: (40, 50)  (19.85% nonzero)
FWD solves: 13   REV solves: 0
Total colors vs. total size: 13 vs 50  (74.0% improvement)

Sparsity computed using tolerance: 1e-25
Time to compute sparsity: 0.037691 sec.
Time to compute coloring: 0.029237 sec.
static_gravity = true, elapsed time = [1.80161673] (should be 1.80164719)

At the end we see we got pretty much the same answer for the elapsed time as the Python example in the Dymos docs. (But not exactly the same, which is a bummer...)

A Simple Dymos Example

Preamble

The AbstractExplicitComp struct

OpenMDAOCore.setup

OpenMDAOCore.compute! and OpenMDAOCore.compute_partials!

The run script

The `AbstractExplicitComp` `struct`

`OpenMDAOCore.setup`

`OpenMDAOCore.compute!` and `OpenMDAOCore.compute_partials!`