brocksam / pycollo

General-purpose optimal control, trajectory optimisation and parameter optimisation using direct collocation
https://brocksam.github.io/pycollo/
MIT License
8 stars 3 forks source link

Allow `sympy.physics.mechanics.dynamicsymbols` to be used as state and control variables #67

Closed brocksam closed 1 year ago

brocksam commented 1 year ago

This PR adds support for using sympy.physics.mechanics.dynamicsymbols as state and control variables in OCP definitions.

Using the Brachistochrone problem as an example, the single-phase OCP's state and control variables can be set using sympy.physics.mechanics.dynamicsymbols like:

import sympy.physics.mechanics as me

import pycollo

x, y, v, u = me.dynamicsymbols("x y v u")

problem = pycollo.OptimalControlProblem(name="Brachistochrone")
phase = problem.new_phase(name="A")

# `phase.state_variables` will return a `collections.namedtuple` of length 3 where the fields are the symbol names
phase.state_variables = [x, y, v]
assert len(phase.state_variables) == 3
assert hasattr(phase.state_variables, "x")
assert hasattr(phase.state_variables, "y")
assert hasattr(phase.state_variables, "v")
assert phase.state_variables.x == x
assert phase.state_variables.y == y
assert phase.state_variables.v == v

# `phase.control_variables` will return a `collections.namedtuple` of length 1 where the field is the symbol name
phase.control_variables = u
assert len(phase.control_variables) == 1
assert hasattr(phase.control_variables, "u")
assert phase.control_variables.u == u

Note that the (t) portion of a me.dynamicsymbols's name is trimmed off in the collections.namedtuple's field name because this is an invalid identifier.

It's possible to mix sm.Symbols with me.dynamicsymbols unless their names clash. For example, the following is invalid:

import sympy as sm
import sympy.physics.mechanics as me

import pycollo

sm_x = sm.Symbol("x")
me_x = me.dynamicsymbols("x")

problem = pycollo.OptimalControlProblem(name="Symbol name clash")
phase = problem.new_phase(name="A")

# Will raise a `ValueError` because both `sm_x` and `me_x` are identifier as `x` (even though `me_x`'s name is `x(t)`)
phase.state_variables = [sm_x, me_x]