ODEs and Boundary conditions generated for the "constrained brachistochrone" problem using https://github.com/thomasantony/beluga/

bc.py

import numpy as np
from math import *
def bc_func_left(_ya, _p, _aux):
    # Declare all auxiliary variables
    g = _aux['const']['g']
    _constraint1 = _aux['const']['_constraint1']
    eps_constraint1 = _aux['const']['eps_constraint1']
    [] = _p

    # Generalize to multipoint later
    # Left BCs
    [x,y,v,xi11,lamX,lamY,lamV,lamXI11,tf,] = _ya[:9]
    [theta,ue1,mu1,] = _ya[9:]

    # Declare all predefined expressions
    psi10 = _constraint1 + exp(xi11)
    psi11 = exp(xi11)

    _x0 = _aux['initial']

    _H = compute_hamiltonian(0,_ya,_p,_aux,[theta,ue1,mu1,])
    res_left = np.array([x - _x0['x'],
                    y - _x0['y'],
                    v - _x0['v'],
                    lamXI11,
                     ])

    return res_left

def bc_func_right(_yb, _p, _aux):
    # Declare all auxiliary variables
    g = _aux['const']['g']
    _constraint1 = _aux['const']['_constraint1']
    eps_constraint1 = _aux['const']['eps_constraint1']
    [] = _p

    # Right BCs
    [x,y,v,xi11,lamX,lamY,lamV,lamXI11,tf,] = _yb[:9]
    [theta,ue1,mu1,] = _yb[9:]
    # Declare all predefined expressions
    psi10 = _constraint1 + exp(xi11)
    psi11 = exp(xi11)

    _xf = _aux['terminal']

    _H = compute_hamiltonian(1,_yb,_p,_aux,[theta,ue1,mu1,])
    res_right = np.array([x - _xf['x'],
                y - _xf['y'],
                -psi10 + x + y,
                lamV,
                _H - 0,
                g*lamV*cos(theta) - lamX*v*sin(theta) + lamY*v*cos(theta) + mu1*(-v*sin(theta) + v*cos(theta)),
                2*eps_constraint1*ue1 + lamXI11 - mu1*psi11,
                -psi11*ue1 + v*sin(theta) + v*cos(theta),
                ])

    return res_right

def bc_func(_ya, _yb, _p, _aux):
    res_left = bc_func_left(_ya[:,0], _p, _aux)
    res_right = bc_func_right(_yb[:,-1], _p, _aux)

    return np.hstack((res_left, res_right))

ode.py

import numpy as np
from math import *

import scipy.linalg
import numba

def im(Z):
    return Z.imag

def compute_hamiltonian(_t,_X,_p,_aux,_u):
    [x,y,v,xi11,lamX,lamY,lamV,lamXI11,tf,] = _X[:9]
    [theta,ue1,mu1,] = _X[9:]
    # Declare all auxiliary variables
    g = _aux['const']['g']
    _constraint1 = _aux['const']['_constraint1']
    eps_constraint1 = _aux['const']['eps_constraint1']

    # Declare all quantities
    psi10 = _constraint1 + exp(xi11)
    psi11 = exp(xi11)

    return eps_constraint1*ue1**2 + g*lamV*sin(theta) + lamX*v*cos(theta) + lamXI11*ue1 + lamY*v*sin(theta) + mu1*(-psi11*ue1 + v*sin(theta) + v*cos(theta)) + 1

def compute_control(_t,_X,_p,_aux):
  return _X[9:(9+3)]

# Used to solve initial guess
def get_dhdu_func(_t,_X,_p,_aux):
    [x,y,v,xi11,lamX,lamY,lamV,lamXI11,tf,] = _X[:(9)]

    # Declare all auxiliary variables
    g = _aux['const']['g']
    _constraint1 = _aux['const']['_constraint1']
    eps_constraint1 = _aux['const']['eps_constraint1']

    def dHdu(_u):
      [theta,ue1,mu1,] = _u
      # Declare all quantities
      psi10 = _constraint1 + exp(xi11)
      psi11 = exp(xi11)

      return np.array([g*lamV*cos(theta) - lamX*v*sin(theta) + lamY*v*cos(theta) + mu1*(-v*sin(theta) + v*cos(theta)),
                2*eps_constraint1*ue1 + lamXI11 - mu1*psi11,
                -psi11*ue1 + v*sin(theta) + v*cos(theta),
                ])

    return dHdu

def deriv_func_nojit(_t,_X,_p,_const,arc_idx):
    [x,y,v,xi11,lamX,lamY,lamV,lamXI11,tf,] = _X[:9]
    [theta,ue1,mu1,] = _X[9:(9+3)]

    [] = _p

    # Declare all auxiliary variables
    g,_constraint1,eps_constraint1, = _const

    # Declare all predefined expressions
    psi10 = _constraint1 + exp(xi11)
    psi11 = exp(xi11)

    return np.array([tf*v*cos(theta),
       tf*v*sin(theta),
       g*tf*sin(theta),
       tf*ue1,
       0,
       0,
       tf*(-lamX*cos(theta) - lamY*sin(theta) - mu1*(sin(theta) + cos(theta))),
       mu1*tf*ue1*exp(xi11),
       0,
       ]+
       [(tf)*((g*(lamX*sin(theta) - lamY*cos(theta) - mu1*(-sin(theta) + cos(theta)))*sin(theta) - g*(-lamX*cos(theta) - lamY*sin(theta) - mu1*(sin(theta) + cos(theta)))*cos(theta) - (-v*sin(theta) + v*cos(theta))*(g*(-sin(theta) - cos(theta))*sin(theta) + ue1**2*exp(xi11) - (-v*sin(theta) + v*cos(theta))*(g*(lamX*sin(theta) - lamY*cos(theta) - mu1*(-sin(theta) + cos(theta)))*sin(theta) - g*(-lamX*cos(theta) - lamY*sin(theta) - mu1*(sin(theta) + cos(theta)))*cos(theta))/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta))))/(-(-v*sin(theta) + v*cos(theta))**2/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta))) - psi11**2/(2*eps_constraint1)))/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta)))),
       (tf)*(psi11*(g*(-sin(theta) - cos(theta))*sin(theta) + ue1**2*exp(xi11) - (-v*sin(theta) + v*cos(theta))*(g*(lamX*sin(theta) - lamY*cos(theta) - mu1*(-sin(theta) + cos(theta)))*sin(theta) - g*(-lamX*cos(theta) - lamY*sin(theta) - mu1*(sin(theta) + cos(theta)))*cos(theta))/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta))))/(2*eps_constraint1*(-(-v*sin(theta) + v*cos(theta))**2/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta))) - psi11**2/(2*eps_constraint1)))),
       (tf)*((g*(-sin(theta) - cos(theta))*sin(theta) + ue1**2*exp(xi11) - (-v*sin(theta) + v*cos(theta))*(g*(lamX*sin(theta) - lamY*cos(theta) - mu1*(-sin(theta) + cos(theta)))*sin(theta) - g*(-lamX*cos(theta) - lamY*sin(theta) - mu1*(sin(theta) + cos(theta)))*cos(theta))/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta))))/(-(-v*sin(theta) + v*cos(theta))**2/(-g*lamV*sin(theta) - lamX*v*cos(theta) - lamY*v*sin(theta) + mu1*(-v*sin(theta) - v*cos(theta))) - psi11**2/(2*eps_constraint1))),
       ]
    )

def deriv_func_ode45(_t,_X,_p,_aux):
    try:
        return deriv_func(_t,_X,_p,list(_aux['const'].values()),0)
    except Exception as e:
        from beluga.utils import keyboard
        keyboard()