Package: numagic-lumos
lumos: scalable optimal control for dynamical systemsBenchmark on main|Benchmark on featuresOptimal Control Example|Quickstart|Environment setupWhat is lumos?lumosis a scalable numerical framework for the modelling and simulation of dynamical systems. Currently lumos focus on optimal control problems only, but the ideas it implements can be easily extended to other problems such as parameter inferences and state estimations.One of the core philosophies oflumosis to leverage some fundamental technologies that enables the deep learning revolution: namely automatic differnetiation and large scale parallelization using accelerated hardware (GPU, TPU, etc).lumosadopts a numpy-like syntax to write models, and currently has two built-in backends:JAXandCasadi. In addition to the built-in backends,lumosis also extensible and can solve optimal control problems with user-defined custom model bakends, provided that a few necessary APIs are exposed.Optimal Control ExampleimportnumpyasnpfromtypingimportAny,Dictimportlumos.numpyaslnpfromlumos.models.baseimportStateSpaceModel,StateSpaceModelReturn,state_space_iofromlumos.optimal_control.configimport(BoundaryConditionConfig,BoundConfig,)fromlumos.optimal_control.scaled_mesh_ocpimportScaledMeshOCP# Create a model@state_space_io(states=("v","x","y"),inputs=("theta",),outputs=("theta",))classTimeModel(StateSpaceModel):"""The Brachistochrone model formulatd in the time domain"""def__init__(self,model_config:Dict[str,Any]={},params:Dict[str,Any]={},):super().__init__(model_config=model_config,params=params)defforward(self,states:Dict[str,float],inputs:Dict[str,float],mesh:float=0.0,# time invariant model)->StateSpaceModelReturn:params=self._paramstheta=inputs["theta"]v=states["v"]v_dot=-params["gravity"]*lnp.sin(theta)dx_dt=lnp.cos(theta)*vdy_dt=lnp.sin(theta)*v# Assemble resultstates_dot=self.make_dict(group="states_dot",v=v_dot,x=dx_dt,y=dy_dt,)outputs=self.make_dict(group="outputs",theta=theta)returnStateSpaceModelReturn(states_dot=states_dot,outputs=outputs)@classmethoddefget_default_params(self)->Dict[str,Any]:return{"gravity":9.81}# Set up the model and the problemmodel=TimeModel()sim_config=ScaledMeshOCP.get_sim_config(boundary_conditions=(BoundaryConditionConfig(0,"states","x",0.0),BoundaryConditionConfig(0,"states","y",0.0),BoundaryConditionConfig(0,"states","v",0.0),BoundaryConditionConfig(-1,"states","x",1.0),BoundaryConditionConfig(-1,"states","y",-0.6),),bounds=(BoundConfig("global","mesh_scale",(0.01,10.0)),BoundConfig("inputs","theta",(-np.pi/2,np.pi/2)),),num_intervals=49,hessian_approximation="exact",)ocp=ScaledMeshOCP(model=model,sim_config=sim_config)# Solve with a trivial initial guesssolution,info=ocp.solve(init_guess=np.zeros(ocp.num_dec),max_iter=200,print_level=4,)print(f"maneuveur time:{ocp.objective(solution):.3f}seconds")ContentsOptimal Control ExampleQuickstart: Colab in the CloudEnvirohnment setupQuickstart: Colab in the CloudBrachistochrone notebookof the Brachistochrone example with the end-to-end implementation with and without automatic differentiationLaptime simulation notebookfor optimal laptime around a race track, using GPUs.Environment setuplumosIt is recommened to develop with lumos via two ways: usingcondaor usingdocker environment, with the latter preferred, especially for developers as.At the moment,lumosis tested and supports Python 3.7, 3.8 and 3.9.Setting up with condasetting up the dependenciescondacreate-nlumospython=3.9
condaenvupdate-nlumos-fenvironment.yml
condaactivatelumosInstalllumosfrom test.pypi (temporarily)pipinstall-ihttps://test.pypi.org/simple/numagic-lumos==0.0.5a0Alternatively, installlumosfrom sourcepython3-mpipinstall.Test if it workspython3examples/drone_example.pySetting up with dockeryou can directy build the local dev container with:docker-composebuild--build-argUID=$(id-u)--build-argGID=$(id-g)After the docker container is build, start it withdocker-composeup-dThe container should now be running in the background.[Optional] download and install VSCode, and then install the 'remote containers' extension.once VSCode with the extensions are installed, open VSCode, and open the folder where the repo is. It should automatically recognize the container configuration file, and asks if you want to open it in an conatiner, click 'Reopen in Container'Once the VSCode enter the container, it will set up the other extensions and adopt settinsg defined in .devcontainer.json. After it installs the extensions, it might ask you to give permission to reload the conatiner to activate extensions like pylance language server, simply click yes.at this point, you should be able to run the examples, in a conatiner, try:(lumos)<username>@lumos_dev:~/numagic/lumos$python3examples/drone_example.pyand the simulations should run and in the end you should get something like:Numberofobjectivefunctionevaluations=225Numberofobjectivegradientevaluations=67Numberofequalityconstraintevaluations=225Numberofinequalityconstraintevaluations=0NumberofequalityconstraintJacobianevaluations=67NumberofinequalityconstraintJacobianevaluations=0NumberofLagrangianHessianevaluations=66TotalsecondsinIPOPT=3.247

EXIT:OptimalSolutionFound.
INFO:__main__:Maneuvertime1.300sec
INFO:__main__:Finaltheta6.28rad
INFO:__main__:Finalsin(theta)-0.00[FIXME] currently the python extension setting for vscode extensions (which does not affect running the code, but affects hwo the VSCode extensions work) does not work automatically using the setting in .devcontainer.json. Therefore one must manually set the interpreter for tools such as linting with flake8, and autoformatting with black to work.
6.1) cmd + shift + p -> type 'select interpreter', and choose 'Python: select interpreter' when the option shows up
6.2) the lumos conda env should show up as an option (something like: Python 3.9.10 ('lumos': conda)). Choose this one.
6.3) Then linting and autoformatting should work.Build the base images locally [Optional]The aforementinoed process will download the base docker images from the public github repo, alternatively one could build the images locally:<repo_root>$docker-compose--env-file.env-fdocker/docker-compose-build.ymlbuildBuilding GPU container [Optional]Note: at the moment we switched the base image OS base to a smaller size linux base, to rebuild the GPU image, one needs to:change to a GPU base image with cuda, eg: nvidia/cuda11.3.0-devel-ubuntu18.04change the jaxlib version to a corresponding one with compatible cuda version. (sometimes one might also need to change the jax version because of jax and jaxlib version compatibility)then to start the GPU containers, use multiple docker-compose seeheredocker-compose-fdocker-compose.yml-fdocker-compose.gpu.ymlup-dEnvironment setup on Mac with M1 chip (Experimental)It is possible to set up the environment on Mac with M1 chip using both the conda and the docker approach described earlier, but both will require some customization steps to tackle general problems that arise from library/OS/chip compatibilty.From the speed perspetive, using conda directly seems to have a significant advantage, probably because it can better utilize the power of the M1 chip than via docker with emulation.Setting up with Conda on M1 MacIt follows the general approach outlinedabove, but two main changes:install casadi with conda-forge instead of pipModify the environment.yml to remove cyipopt (we'll install it from source later), and move casadi from pip dependency to conda dependency. Add jax and jaxlib to conda dependencies, as shown below:channels:-conda-forge-defaultsdependencies:-casadi-jax-jaxlib-pip-pip:-casadi-pyarrow-pandasAnd then create and update the conda environment asbeforecondacreate-nlumos_m1python=3.9
condaenvupdate--fileenvironment.ymlCheck if jax and casadi are both working (in a python terminal)importjax.numpyasjnpimportnumpyasnpimportcasadiascasaa=np.random.randn(10)jnp.dot(aa,aa)cas.dot(aa,aa)install cyipopt from sourceClone the cyipopt repo, and move into repo directory:gitclonehttps://github.com/mechmotum/cyipopt.gitcdcyipoptCheck out a release tag (optional)gitcheckouttags/v1.1.0Install the dependencies required (usethisas a reference)condainstall-q-ylapack"libblas=*=*netlib"cython>=0.26"ipopt=3.14"numpy>=1.15pkg-config>=0.29.2setuptools>=39.0Install from sourcepythonsetup.pyinstall
python-mpipinstall.Check if cyipopt is working correctlypipinstalpytest
pytestTest if all components are all working correctly (using lumos examples)cd..
pythonexamples/laptime_simulation_example.pySetting up with Docker on M1 MacIt is also possible to set up the dev enviornment using the same dev conatiner, but we would need to replace the jax dependencies in the container as they were built with AVX which is not supported for docker on M1 (which uses Rosetta emulation)So inside the container, one needs to replacejaxandjaxlibwith versions that are built WITHOUT AVX, seeherecondainstall-cconda-forgejaxlib
condainstall-cconda-forgejaxAnd then all should just work as before.If one wants to build the corresponding container locally, this is equivalent to modifying theenviorenment.ymlto:channels:-conda-forge-defaultsdependencies:-cyipopt-jax-jaxlib-pip-pip:-casadi-pyarrow-pandasStargazers over time