├── tests ├── __init__.py ├── test_numpy │ ├── __init__.py │ ├── test_casadi_numpy │ │ ├── test_logicals.py │ │ ├── test_calculus.py │ │ ├── test_conditionals.py │ │ ├── test_linalg.py │ │ ├── test_arithmetic.py │ │ ├── test_finite_difference_operators.py │ │ ├── test_linalg_top_level.py │ │ ├── test_determine_type.py │ │ ├── test_array.py │ │ └── test_all_operations_run.py │ ├── test_interp.py │ └── test_lumos_numpy.py ├── test_models │ ├── __init__.py │ ├── test_tires │ │ ├── __init__.py │ │ ├── test_utils.py │ │ └── test_perantoni_tire.py │ ├── test_drone_model.py │ ├── test_aero │ │ └── test_aero.py │ ├── test_simlpe_vehicle_on_track.py │ ├── test_tracks.py │ ├── test_kinematics.py │ └── test_simple_vehicle.py ├── test_optimal_control │ ├── __init__.py │ ├── test_ltc_sweep.py │ ├── test_config.py │ ├── test_transcription.py │ ├── test_collocation.py │ └── test_ocp_logging.py ├── test_mex │ ├── compile_mex_and_run.m │ └── export_model.py └── test_simulations │ └── test_laptime_simulation.py ├── examples ├── __init__.py ├── test_examples │ ├── __init__.py │ ├── test_drone_simulation.py │ ├── test_laptime_simulation.py │ └── test_brachistochrone.py ├── drone_example.py └── laptime_simulation_example.py ├── lumos ├── models │ ├── aero │ │ ├── __init__.py │ │ └── aero.py │ ├── tires │ │ ├── __init__.py │ │ ├── utils.py │ │ ├── base.py │ │ └── perantoni.py │ ├── vehicles │ │ └── __init__.py │ ├── __init__.py │ ├── ml.py │ ├── drone_model.py │ ├── kinematics.py │ ├── simple_vehicle_on_track.py │ ├── tracks.py │ └── test_utils.py ├── optimal_control │ ├── __init__.py │ ├── fixed_mesh_ocp.py │ ├── collocation.py │ ├── convolution.py │ ├── config.py │ └── transcription.py ├── numpy │ ├── casadi_numpy │ │ ├── integrate.py │ │ ├── README.md │ │ ├── conditionals.py │ │ ├── __init__.py │ │ ├── interpolate.py │ │ ├── determine_type.py │ │ ├── trig.py │ │ ├── LICENSE.txt │ │ ├── calculus.py │ │ ├── rotations.py │ │ ├── linalg_top_level.py │ │ ├── arithmetic.py │ │ ├── spacing.py │ │ ├── logicals.py │ │ ├── linalg.py │ │ └── finite_difference_operators.py │ └── __init__.py ├── __init__.py └── simulations │ ├── __init__.py │ └── drone_simulation.py ├── imgs ├── logo.png └── vscode_open_in_container.png ├── pyproject.toml ├── environment.yml ├── data └── tracks │ └── README.md ├── .gitignore ├── docker-compose.gpu.yml ├── .env ├── docker ├── scripts │ └── default_entrypoint.sh ├── base │ ├── Dockerfile-ci │ ├── Dockerfile-base │ └── Dockerfile-dev └── docker-compose-build.yml ├── .github └── workflows │ ├── publish_to_pypi.yml │ ├── test_with_conda.yml │ ├── ltc_regression_test.yml │ └── test_mex_with_octave.yml ├── setup.py ├── docker-compose.yml ├── LICENSE ├── .devcontainer └── devcontainer.json └── regression_tests └── run_benchmark.py /tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_numpy/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lumos/models/aero/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lumos/models/tires/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lumos/models/vehicles/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /lumos/optimal_control/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /examples/test_examples/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_optimal_control/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tests/test_models/test_tires/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /imgs/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/numagic/lumos/HEAD/imgs/logo.png -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/integrate.py: -------------------------------------------------------------------------------- 1 | # A mirror scipy.interpolate 2 | 3 | # TODO -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- 1 | [build-system] 2 | requires = ["setuptools>=42"] 3 | build-backend = "setuptools.build_meta" -------------------------------------------------------------------------------- /imgs/vscode_open_in_container.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/numagic/lumos/HEAD/imgs/vscode_open_in_container.png -------------------------------------------------------------------------------- /environment.yml: -------------------------------------------------------------------------------- 1 | channels: 2 | - conda-forge 3 | - defaults 4 | dependencies: 5 | - cyipopt 6 | - pip 7 | - pip: 8 | - jax[cpu] 9 | - casadi 10 | - pyarrow 11 | - pandas 12 | 13 | 14 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/README.md: -------------------------------------------------------------------------------- 1 | # Introduction 2 | This pakckage is adapted from Peter Sharpe's [AeroSandbox](https://github.com/peterdsharpe/AeroSandbox), with its original license 3 | included in the current folder. -------------------------------------------------------------------------------- /data/tracks/README.md: -------------------------------------------------------------------------------- 1 | # Introduction 2 | The racetrack data are taken from [TUMFTM/racetrack-database](https://github.com/TUMFTM/racetrack-database) without modification, with its license include in the current folder as well. 3 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # VSCode artifacts 2 | .vscode/ 3 | 4 | # Python3 bytecode 5 | **/__pycache__ 6 | 7 | # Mac folder artifact 8 | .DS_store 9 | 10 | # temporary files 11 | *.tmp 12 | 13 | # experiment results 14 | /results 15 | 16 | # binaries 17 | *.so -------------------------------------------------------------------------------- /lumos/__init__.py: -------------------------------------------------------------------------------- 1 | # TODO: understand why we need this. Without this, the sub-packages can't be found! 2 | # It's most likely a import 'trap' introduced after 3.3, but what's strange is that it 3 | # causes a problem with 3.7 but not with 3.9 4 | import lumos.numpy as lnp 5 | -------------------------------------------------------------------------------- /docker-compose.gpu.yml: -------------------------------------------------------------------------------- 1 | version: '3' 2 | services: 3 | lumos-dev: 4 | deploy: 5 | resources: 6 | # limits: 7 | # cpus: 1 8 | reservations: 9 | devices: 10 | - capabilities: [gpu] -------------------------------------------------------------------------------- /examples/test_examples/test_drone_simulation.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | from examples.drone_example import main as drone_main 4 | 5 | 6 | class TestDroneExample(TestCase): 7 | def test_drone_simulation(self): 8 | """Simple smoke test for drone simulation example""" 9 | drone_main() 10 | -------------------------------------------------------------------------------- /examples/test_examples/test_laptime_simulation.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | from examples.laptime_simulation_example import main as ltc_main 4 | 5 | 6 | class TestLaptimeSimulationExample(TestCase): 7 | def test_laptime_simulation(self): 8 | """Simple smoke test for laptime simulation example""" 9 | ltc_main() 10 | -------------------------------------------------------------------------------- /lumos/simulations/__init__.py: -------------------------------------------------------------------------------- 1 | from jax.config import config 2 | 3 | # By default we use 64bit as over/underflow are quite likely to happen with 32bit and 4 | # 2nd derivative autograd, without a lot of careful management... 5 | config.update("jax_enable_x64", True) 6 | 7 | # For initial phase, we also report verbosely if jax is recompiling. 8 | config.update("jax_log_compiles", 1) 9 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_logicals.py: -------------------------------------------------------------------------------- 1 | import lumos.numpy.casadi_numpy as np 2 | import pytest 3 | 4 | 5 | def test_basic_logicals_numpy(): 6 | a = np.array([True, True, False, False]) 7 | b = np.array([True, False, True, False]) 8 | 9 | assert np.all(a & b == np.array([True, False, False, False])) 10 | 11 | 12 | if __name__ == "__main__": 13 | pytest.main() 14 | -------------------------------------------------------------------------------- /.env: -------------------------------------------------------------------------------- 1 | CONDA_DIR='/var/miniconda3' 2 | MINICONDA_VERSION='latest' 3 | 4 | CONTAINER_REGISTRY="ghcr.io/numagic" 5 | 6 | LUMOS_BASE_IMAGE='debian' 7 | LUMOS_BASE_IMAGE_VERSION='10.11' 8 | 9 | LUMOS_BASE_GPU_IMAGE='nvidia/cuda' 10 | LUMOS_BASE_GPU_IMAGE_VERSION='11.3.0-devel-ubuntu18.04' 11 | 12 | LUMOS_DOCKER_BASE_VERSION='0.2' 13 | LUMOS_DOCKER_PYTHON_VERSION='3.9' 14 | 15 | LUMOS_DOCKER_DATA_DIR='/var/numagic' 16 | LUMOS_DOCKER_SRC_DIR='/src' 17 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/conditionals.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 4 | 5 | 6 | def where( 7 | condition, value_if_true, value_if_false, 8 | ): 9 | if not is_casadi_type([condition, value_if_true, value_if_false], recursive=True): 10 | return _onp.where(condition, value_if_true, value_if_false) 11 | else: 12 | return _cas.if_else(condition, value_if_true, value_if_false) 13 | -------------------------------------------------------------------------------- /tests/test_models/test_drone_model.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import unittest 4 | from parameterized import parameterized, parameterized_class 5 | 6 | from lumos.models.drone_model import DroneModel 7 | from lumos.models.test_utils import BaseStateSpaceModelTest 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | 12 | class TestDroneModel(BaseStateSpaceModelTest, unittest.TestCase): 13 | ModelClass: type = DroneModel 14 | 15 | 16 | if __name__ == "__name__": 17 | unittest.main() 18 | -------------------------------------------------------------------------------- /docker/scripts/default_entrypoint.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | set -euo pipefail 3 | 4 | # Activate conda 5 | # Needed to make conda available in non-interactive shells first 6 | # Reference: https://github.com/ContinuumIO/docker-images/issues/89#issuecomment-467287039 7 | # FIXME: hard-coded conda path 8 | source $CONDA_DIR/etc/profile.d/conda.sh 9 | 10 | # Note this is also duplicated in the docker-base file, where we add them to ~/.bashrc 11 | # for non-login shells 12 | conda deactivate 13 | conda activate lumos 14 | 15 | exec "$@" -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_calculus.py: -------------------------------------------------------------------------------- 1 | import lumos.numpy.casadi_numpy as np 2 | import pytest 3 | 4 | 5 | def test_diff(): 6 | a = np.arange(100) 7 | 8 | assert np.all(np.diff(a) == pytest.approx(1)) 9 | 10 | 11 | def test_trapz(): 12 | a = np.arange(100) 13 | 14 | assert np.diff(np.trapz(a)) == pytest.approx(1) 15 | 16 | 17 | def test_invertability_of_diff_trapz(): 18 | a = np.sin(np.arange(10)) 19 | 20 | assert np.all(np.trapz(np.diff(a)) == pytest.approx(np.diff(np.trapz(a)))) 21 | 22 | 23 | if __name__ == "__main__": 24 | pytest.main() 25 | -------------------------------------------------------------------------------- /docker/base/Dockerfile-ci: -------------------------------------------------------------------------------- 1 | ARG LUMOS_DOCKER_BASE_VERSION 2 | ARG CONTAINER_REGISTRY 3 | 4 | FROM ${CONTAINER_REGISTRY}/lumos-base:${LUMOS_DOCKER_BASE_VERSION} 5 | ENV DEBIAN_FRONTEND="noninteractive" 6 | 7 | ARG CONDA_DIR 8 | # Make conda command available 9 | ENV CONDA_DIR $CONDA_DIR 10 | ENV PATH=$CONDA_DIR/bin:$PATH 11 | RUN $CONDA_DIR/bin/conda init bash 12 | 13 | 14 | # Install additional tools for ci 15 | # Rely on -i flag to activate conda 16 | # see: https://stackoverflow.com/a/60604010 17 | SHELL ["/bin/bash", "-c", "-i"] 18 | RUN conda activate lumos \ 19 | && conda install -y pytest parameterized black flake8 -------------------------------------------------------------------------------- /.github/workflows/publish_to_pypi.yml: -------------------------------------------------------------------------------- 1 | name: publish-to-pypi 2 | on: 3 | release: 4 | types: [created] 5 | 6 | jobs: 7 | build-and-publish-to-pypi: 8 | runs-on: "ubuntu-latest" 9 | steps: 10 | - uses: actions/checkout@v2 11 | - uses: actions/setup-python@v2 12 | - name: Install dependencies 13 | run: | 14 | python -m pip install --upgrade pip 15 | pip install setuptools wheel twine 16 | - name: Build and publish 17 | env: 18 | TWINE_USERNAME: __token__ 19 | TWINE_PASSWORD: ${{ secrets.PYPI_SECRET }} 20 | run: | 21 | python setup.py sdist bdist_wheel 22 | twine upload dist/* -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_conditionals.py: -------------------------------------------------------------------------------- 1 | import lumos.numpy.casadi_numpy as np 2 | import casadi as cas 3 | import pytest 4 | 5 | 6 | def test_where_numpy(): 7 | a = np.ones(4) 8 | b = 2 * np.ones(4) 9 | 10 | c = np.where(np.array([True, False, True, False]), a, b) 11 | 12 | assert np.all(c == np.array([1, 2, 1, 2])) 13 | 14 | 15 | def test_where_casadi(): 16 | a = cas.GenDM_ones(4) 17 | b = 2 * cas.GenDM_ones(4) 18 | 19 | c = np.where(cas.DM([1, 0, 1, 0]), a, b) 20 | 21 | assert np.all(c == cas.DM([1, 2, 1, 2])) 22 | 23 | 24 | # def test_if_else_mixed(): # TODO write this 25 | 26 | 27 | if __name__ == "__main__": 28 | pytest.main() 29 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_linalg.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import lumos.numpy.casadi_numpy as np 3 | import casadi as cas 4 | 5 | 6 | def test_norm_vector(): 7 | a = np.array([1, 2, 3]) 8 | cas_a = cas.DM(a) 9 | 10 | assert np.linalg.norm(a) == np.linalg.norm(cas_a) 11 | 12 | 13 | def test_norm_2D(): 14 | a = np.arange(9).reshape(3, 3) 15 | cas_a = cas.DM(a) 16 | 17 | assert np.linalg.norm(cas_a) == np.linalg.norm(a) 18 | 19 | assert np.all(np.linalg.norm(cas_a, axis=0) == np.linalg.norm(a, axis=0)) 20 | 21 | assert np.all(np.linalg.norm(cas_a, axis=1) == np.linalg.norm(a, axis=1)) 22 | 23 | 24 | if __name__ == "__main__": 25 | pytest.main() 26 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/__init__.py: -------------------------------------------------------------------------------- 1 | ### Import everything from NumPy 2 | 3 | from numpy import * 4 | 5 | ### Overwrite some functions 6 | from .array import * 7 | from .arithmetic import * 8 | from .calculus import * 9 | from .conditionals import * 10 | 11 | from .finite_difference_operators import * 12 | from .integrate import * 13 | from .interpolate import * 14 | from .linalg_top_level import * 15 | import lumos.numpy.casadi_numpy.linalg as linalg 16 | from .logicals import * 17 | from .rotations import * 18 | from .spacing import * 19 | from .trig import * 20 | 21 | ### Force-overwrite built-in Python functions. 22 | 23 | from numpy import round # TODO check that min, max are properly imported 24 | -------------------------------------------------------------------------------- /lumos/models/__init__.py: -------------------------------------------------------------------------------- 1 | from lumos.models.composition import ModelMaker 2 | from lumos.models.simple_vehicle_on_track import SimpleVehicleOnTrack 3 | from lumos.models.tires.pacejka import MF52 4 | from lumos.models.tires.perantoni import PerantoniTire 5 | from lumos.models.vehicles.simple_vehicle import SimpleVehicle 6 | from lumos.models.kinematics import TrackPosition2D 7 | from lumos.models.aero.aero import ConstAero, GPAero, MLPAero 8 | 9 | ModelMaker.add_to_registry(MF52) 10 | ModelMaker.add_to_registry(PerantoniTire) 11 | 12 | ModelMaker.add_to_registry(SimpleVehicle) 13 | ModelMaker.add_to_registry(TrackPosition2D) 14 | ModelMaker.add_to_registry(SimpleVehicleOnTrack) 15 | 16 | ModelMaker.add_to_registry(ConstAero) 17 | ModelMaker.add_to_registry(GPAero) 18 | ModelMaker.add_to_registry(MLPAero) 19 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import setuptools 2 | 3 | with open("README.md", "r", encoding="utf-8") as fh: 4 | long_description = fh.read() 5 | 6 | setuptools.setup( 7 | name="numagic-lumos", 8 | version="0.0.2rc7", 9 | author="Yunlong Xu", 10 | author_email="yunlong@numagic.io", 11 | description="lumos - scalable accelerated optimal control", 12 | long_description=long_description, 13 | long_description_content_type="text/markdown", 14 | url="https://github.com/numagic/lumos", 15 | project_urls={"Bug Tracker": "https://github.com/numagic/lumos/issues"}, 16 | classifiers=[ 17 | "Programming Language :: Python :: 3", 18 | "Operating System :: OS Independent", 19 | ], 20 | packages=setuptools.find_packages(include=["lumos", "lumos.*"]), 21 | python_requires=">=3.7", 22 | ) 23 | -------------------------------------------------------------------------------- /tests/test_models/test_tires/test_utils.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | from lumos.models.tires.utils import read_tir_file 3 | 4 | 5 | class TestReadTirFile(TestCase): 6 | def test_read(self): 7 | """Read the default file and assert the contents""" 8 | 9 | # TODO: we need better management of test data 10 | params = read_tir_file("data/tires/default.tir") 11 | 12 | # Check if all read items are floats 13 | self.assertTrue(all([isinstance(v, float) for v in params.values()])) 14 | 15 | # Check if some selected float items are correctly read 16 | self.assertAlmostEqual(params["PCX1"], 1.3605) 17 | self.assertAlmostEqual(params["PEX2"], -0.475) 18 | 19 | # Check if non-float items are not read 20 | self.assertRaises(KeyError, lambda: params["FILE_TYPE"]) 21 | self.assertRaises(KeyError, lambda: params["PROPERTY_FILE_FORMAT"]) 22 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/interpolate.py: -------------------------------------------------------------------------------- 1 | import casadi as cas 2 | from lumos.numpy.casadi_numpy.conditionals import where 3 | 4 | 5 | def interp(x, xp, fp, left=None, right=None, period=None): 6 | """One-dimensional linear interpolation, analogous to numpy.interp(). 7 | 8 | does not handle period yet. 9 | see: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.interp.html 10 | """ 11 | if period is not None: 12 | raise NotImplemented("Handling of period are not implemented yet") 13 | # Casadi supports 'bspline' and 'linear' 14 | lut = cas.interpolant("LUT", "linear", [xp], fp) 15 | f = lut(x) 16 | 17 | # Left and right default to end values, like jax.numpy 18 | if left is None: 19 | left = fp[0] 20 | 21 | if right is None: 22 | right = fp[-1] 23 | 24 | f = where(x < xp[0], left, f) 25 | f = where(x > xp[-1], right, f) 26 | 27 | return f 28 | -------------------------------------------------------------------------------- /lumos/models/ml.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | import lumos.numpy as lnp 4 | 5 | 6 | def sqeuclidean_distance(x, y): 7 | return lnp.sum((x - y) ** 2, axis=1) 8 | 9 | 10 | def rbf_kernel(x, y): 11 | return lnp.exp(-sqeuclidean_distance(x, y)) 12 | 13 | 14 | def tanh_kernel(x, y): 15 | return lnp.tanh(-sqeuclidean_distance(x, y)) 16 | 17 | 18 | def poly_kernel(x, y): 19 | return lnp.sum((x - y) ** 2, axis=1) 20 | 21 | 22 | def gp(x, gp_points, alpha): 23 | """Gaussian processes aero model""" 24 | 25 | # NOTE: need to do the vector_tile here because casadi does not use broadcast by 26 | # default. 27 | num_points = gp_points.shape[0] 28 | kxy = rbf_kernel(gp_points, lnp.vector_tile(x, num_points)) 29 | return alpha @ kxy 30 | 31 | 32 | def mlp(x, weights): 33 | """Dense fully connected neuralnet""" 34 | 35 | y = x 36 | for w in weights: 37 | y = lnp.tanh(w @ y) 38 | 39 | return y 40 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/determine_type.py: -------------------------------------------------------------------------------- 1 | import casadi as cas 2 | 3 | 4 | def is_casadi_type(array_like, recursive=True) -> bool: 5 | """ 6 | Returns a boolean of whether an object is a CasADi data type or not. If the recursive flag is True, iterates recursively. 7 | 8 | Args: 9 | 10 | object: The object to evaluate. 11 | 12 | recursive: If the object is a list or tuple, recursively iterate through every subelement. If any of the 13 | subelements are a CasADi type, return True. Otherwise, return False 14 | 15 | Returns: A boolean if the object is a CasADi data type. 16 | 17 | """ 18 | if recursive and (isinstance(array_like, list) or isinstance(array_like, tuple)): 19 | for element in array_like: 20 | if is_casadi_type(element, recursive=True): 21 | return True 22 | 23 | return ( 24 | isinstance(array_like, cas.MX) or 25 | isinstance(array_like, cas.DM) or 26 | isinstance(array_like, cas.SX) 27 | ) 28 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_arithmetic.py: -------------------------------------------------------------------------------- 1 | import lumos.numpy.casadi_numpy as np 2 | import casadi as cas 3 | import pytest 4 | 5 | 6 | def test_sum(): 7 | a = np.arange(101) 8 | 9 | assert np.sum(a) == 5050 # Gauss would be proud. 10 | 11 | 12 | def test_sum2(): 13 | # Check it returns the same results with casadi and numpy 14 | a = np.array([[1, 2, 3], [1, 2, 3]]) 15 | b = cas.SX(a) 16 | 17 | assert np.all(np.sum(a) == cas.DM(np.sum(b))) 18 | assert np.all(np.sum(a, axis=1) == cas.DM(np.sum(b, axis=1))) 19 | 20 | 21 | def test_mean(): 22 | a = np.linspace(0, 10, 50) 23 | 24 | assert np.mean(a) == pytest.approx(5) 25 | 26 | 27 | def test_cumsum(): 28 | n = np.arange(6).reshape((3, 2)) 29 | c = cas.DM(n) 30 | 31 | assert np.all(np.cumsum(n) == np.array([0, 1, 3, 6, 10, 15])) 32 | # assert np.all( # TODO add casadi testing here 33 | # np.cumsum(c) == np.array([0, 1, 3, 6, 10, 15]) 34 | # ) 35 | 36 | 37 | if __name__ == "__main__": 38 | pytest.main() 39 | -------------------------------------------------------------------------------- /tests/test_models/test_aero/test_aero.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import unittest 3 | 4 | import numpy as np 5 | 6 | from lumos.models.aero.aero import ConstAero, GPAero, MLPAero 7 | from lumos.models.test_utils import BaseModelTest 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | 12 | class TestConstAero(BaseModelTest, unittest.TestCase): 13 | ModelClass: type = ConstAero 14 | 15 | def test_outputs_are_correct(self): 16 | inputs = dict( 17 | zip(self.model.names.inputs, self.model.make_random_vector("inputs")) 18 | ) 19 | model_return = self.model.forward(inputs) 20 | 21 | # Check constant values are the same as in the params 22 | for k, v in model_return.outputs.items(): 23 | self.assertAlmostEqual(v, self.model._params[k]) 24 | 25 | 26 | class TestGPAero(BaseModelTest, unittest.TestCase): 27 | ModelClass: type = GPAero 28 | 29 | 30 | class TestMLPAero(BaseModelTest, unittest.TestCase): 31 | ModelClass: type = MLPAero 32 | 33 | 34 | if __name__ == "__name__": 35 | unittest.main() 36 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/trig.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | from numpy import pi as _pi 3 | 4 | 5 | def degrees(x): 6 | """Converts an input x from radians to degrees""" 7 | return x * 180 / _pi 8 | 9 | 10 | def radians(x): 11 | """Converts an input x from degrees to radians""" 12 | return x * _pi / 180 13 | 14 | 15 | def sind(x): 16 | """Returns the sin of an angle x, given in degrees""" 17 | return _onp.sin(radians(x)) 18 | 19 | 20 | def cosd(x): 21 | """Returns the cos of an angle x, given in degrees""" 22 | return _onp.cos(radians(x)) 23 | 24 | 25 | def tand(x): 26 | """Returns the tangent of an angle x, given in degrees""" 27 | return _onp.tan(radians(x)) 28 | 29 | 30 | def arcsind(x): 31 | """Returns the arcsin of an x, in degrees""" 32 | return degrees(_onp.arcsin(x)) 33 | 34 | 35 | def arccosd(x): 36 | """Returns the arccos of an x, in degrees""" 37 | return degrees(_onp.arccos(x)) 38 | 39 | 40 | def arctan2d(y, x): 41 | """Returns the angle associated with arctan(y, x), in degrees""" 42 | return degrees(_onp.arctan2(y, x)) 43 | -------------------------------------------------------------------------------- /docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: '3' 2 | services: 3 | lumos-dev: 4 | image: '${CONTAINER_REGISTRY}/lumos-dev-${USER}:${LUMOS_DOCKER_BASE_VERSION}' 5 | container_name: lumos_dev 6 | hostname: lumos_dev 7 | build: 8 | context: '.' 9 | dockerfile: 'docker/base/Dockerfile-dev' 10 | args: 11 | - 'CONTAINER_REGISTRY' 12 | - 'LUMOS_DOCKER_BASE_VERSION' 13 | - 'LUMOS_DOCKER_USER_NAME=${USER}' 14 | volumes: 15 | - '.:/home/${USER}/numagic/lumos' 16 | # share git credentials and config 17 | - '/Users/${USER}/.ssh:/home/${USER}/.ssh' 18 | - '/Users/${USER}/.gitconfig:/home/${USER}/.gitconfig' 19 | environment: 20 | - 'PYTHONPATH=/home/${USER}/numagic/lumos' 21 | # hard-coded ssh key name is the same 22 | - 'GIT_SSH_COMMAND=ssh -i /home/${USER}/.ssh/id_ed25519' 23 | restart: unless-stopped 24 | 25 | # Temporary hack to aid development. 26 | ports: 27 | - '8000:8000' 28 | - '81:8050' -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 numagic GmbH 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /examples/test_examples/test_brachistochrone.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | from parameterized import parameterized 4 | 5 | from examples.brachistochrone import solve_brachistochrone 6 | 7 | 8 | class TestBrachistochrone(TestCase): 9 | """Test Brachistochrone example runs and returns correct answer. 10 | 11 | Correct answer taken from: https://scipython.com/blog/the-brachistochrone-problem/ 12 | """ 13 | 14 | def test_distance_domain(self): 15 | """Test distance based formulation with fixed grid. 16 | 17 | This is harder to converge, and then to be less robust. 18 | """ 19 | t, info = solve_brachistochrone(1.0, -0.65, time_domain=False) 20 | self.assertEqual(info["status"], 0) 21 | self.assertAlmostEqual(t, 0.566, places=3) 22 | 23 | @parameterized.expand(["jax", "casadi", "custom"]) 24 | def test_time_domain(self, backend: str): 25 | """Test time based formulation, with scaled grid.""" 26 | t, info = solve_brachistochrone(1.0, -0.65, backend=backend) 27 | self.assertEqual(info["status"], 0) 28 | self.assertAlmostEqual(t, 0.566, places=3) 29 | -------------------------------------------------------------------------------- /tests/test_numpy/test_interp.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | from typing import Callable, List 3 | 4 | 5 | import casadi as cas 6 | import numpy as onp 7 | from parameterized import parameterized 8 | 9 | import lumos.numpy as lnp 10 | 11 | 12 | class TestInterp(TestCase): 13 | @parameterized.expand([["casadi"], ["jax"]]) 14 | def test_interp(self, backend: str): 15 | num_points = 11 16 | xp = onp.linspace(0, 1, num_points) 17 | fp = onp.sin(xp) 18 | 19 | # NOTE: we test also out of range behaviour 20 | for x in [-1.0, 0.0, 0.314, 1.0, 2.0]: 21 | expected = onp.interp(x, xp, fp) 22 | with lnp.use_backend(backend): 23 | if backend == "jax": 24 | actual = lnp.interp(x, xp, fp) 25 | elif backend == "casadi": 26 | # For casadi, we test the one with symbolic evaluation 27 | x_sym = cas.MX.sym("x", 1) 28 | fun = cas.Function("fun", [x_sym], [lnp.interp(x_sym, xp, fp)]) 29 | actual = float(fun(x)) 30 | self.assertAlmostEqual(actual, expected) 31 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2019-2021 Peter Sharpe 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. -------------------------------------------------------------------------------- /tests/test_mex/compile_mex_and_run.m: -------------------------------------------------------------------------------- 1 | %% Compile exported c-code into mex, execute with saved inputs and compare to saved outputs 2 | 3 | % Takes one commandline input for the file_name 4 | args = argv(); 5 | 6 | c_file = sprintf("%s.c", args{1}); 7 | data_file = sprintf("%s.mat", args{1}); 8 | 9 | fprintf("compiling %s and testing with %s\n", c_file, data_file) 10 | disp("compiling mex function ...") 11 | cmd_str = sprintf("mex %s -DMATLAB_MEX_FILE", c_file); 12 | eval(cmd_str); 13 | 14 | disp("loading test data ...") 15 | data = load(data_file); 16 | 17 | disp("executing model with test data ...") 18 | cmd = sprintf("[states_dot, outputs, con_outputs, residuals] = %s(data.states, data.inputs, data.mesh, data.params);", args{1}) 19 | eval(cmd); 20 | 21 | disp("check results against test data ...") 22 | % Check results, pay attention to the transpose as python is row major 23 | 24 | disp("checking states_dot ...") 25 | assert(states_dot, data.states_dot') 26 | 27 | disp("checking outputs ...") 28 | assert(outputs, data.outputs') 29 | % assert(con_outputs, data.con_outputs') % don't compare because it's empty 30 | 31 | disp("checking residuals ...") 32 | assert(residuals, data.residuals') -------------------------------------------------------------------------------- /docker/docker-compose-build.yml: -------------------------------------------------------------------------------- 1 | version: '3' 2 | services: 3 | lumos-base: 4 | image: '${CONTAINER_REGISTRY}/lumos-base:${LUMOS_DOCKER_BASE_VERSION}' 5 | build: 6 | context: '..' 7 | cache_from: 8 | # If cache is not found, won't fail the build. 9 | - '${CONTAINER_REGISTRY}/lumos-base:cache' 10 | dockerfile: 'docker/base/Dockerfile-base' 11 | args: 12 | - 'LUMOS_BASE_IMAGE' 13 | - 'LUMOS_BASE_IMAGE_VERSION' 14 | - 'CONDA_DIR' 15 | - 'MINICONDA_VERSION' 16 | - 'LUMOS_DOCKER_DATA_DIR' 17 | - 'LUMOS_DOCKER_PYTHON_VERSION' 18 | - 'CONTAINER_REGISTRY' 19 | lumos-ci: 20 | image: '${CONTAINER_REGISTRY}/lumos-ci:${LUMOS_DOCKER_BASE_VERSION}' 21 | build: 22 | context: '..' 23 | cache_from: 24 | # If cache is not found, won't fail the build. 25 | - '${CONTAINER_REGISTRY}/lumos-ci:cache' 26 | dockerfile: 'docker/base/Dockerfile-ci' 27 | args: 28 | - 'CONTAINER_REGISTRY' 29 | - 'LUMOS_DOCKER_BASE_VERSION' 30 | - 'CONDA_DIR' 31 | -------------------------------------------------------------------------------- /lumos/numpy/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import numpy as np 3 | from typing import Any, Union 4 | 5 | import casadi 6 | import jax.numpy as jnp 7 | import lumos.numpy.casadi_numpy as cnp 8 | from jax.config import config 9 | 10 | from lumos.numpy.utils import ( 11 | is_jax_array, 12 | vector_concat, 13 | vector_split, 14 | assert_allclose, 15 | use_backend, 16 | vector_tile, 17 | cmap, 18 | lmap, 19 | ) 20 | 21 | logger = logging.getLogger(__name__) 22 | 23 | # Enable double precision by default 24 | config.update("jax_enable_x64", True) 25 | 26 | __BACKEND_MAP = {"jax": jnp, "numpy": np, "casadi": cnp} 27 | # Use numpy as default backend 28 | __BACKEND = "numpy" 29 | 30 | ndarray = Union[np.ndarray, jnp.ndarray, casadi.SX, casadi.MX] 31 | 32 | 33 | def get_backend() -> str: 34 | return __BACKEND 35 | 36 | 37 | def set_backend(backend: str) -> None: 38 | """Set a numerical backend for model computation.""" 39 | 40 | global __BACKEND 41 | if backend in __BACKEND_MAP: 42 | __BACKEND = backend 43 | else: 44 | logger.warn( 45 | f"{backend} is not a valid backend. Staying with {__BACKEND} backend" 46 | ) 47 | 48 | 49 | def __getattr__(name: str): 50 | """Switcheable backend by overriding the namespace. 51 | 52 | FIXME: The price paid is no more auto-completion. 53 | """ 54 | 55 | return getattr(__BACKEND_MAP[__BACKEND], name) 56 | -------------------------------------------------------------------------------- /examples/drone_example.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import numpy as np 3 | from lumos.optimal_control.config import LoggingConfig 4 | from lumos.simulations.drone_simulation import DroneSimulation 5 | 6 | logger = logging.getLogger(__name__) 7 | 8 | 9 | def main(): 10 | # h_approx = "limited-memory" 11 | h_approx = "exact" 12 | 13 | ocp = DroneSimulation( 14 | sim_config=DroneSimulation.get_sim_config( 15 | num_intervals=99, 16 | hessian_approximation=h_approx, 17 | backend="casadi", 18 | transcription="LGR", 19 | is_condensed=False, 20 | logging_config=LoggingConfig( 21 | sim_name="drone", results_dir="results", log_every_nth_iter=0 22 | ), 23 | ) 24 | ) 25 | 26 | x0 = ocp.get_init_guess() 27 | 28 | solution, info = ocp.solve( 29 | x0, 30 | max_iter=200, 31 | print_level=5, 32 | print_timing_statistics="no", 33 | derivative_test="none", 34 | ) 35 | 36 | total_time = ocp.get_total_time(solution) 37 | final_theta = ocp.dec_var_operator.get_var( 38 | solution, group="states", name="theta", stage=-1 39 | ) 40 | 41 | logger.info(f"Maneuver time {total_time:.3f} sec") 42 | logger.info(f"Final theta {final_theta:.2f} rad") 43 | logger.info(f"Final sin(theta) {np.sin(final_theta):.2f}") 44 | 45 | 46 | if __name__ == "__main__": 47 | logging.basicConfig(level=logging.INFO) 48 | main() 49 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/calculus.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 4 | 5 | 6 | def diff(a, n=1, axis=-1): 7 | """ 8 | Calculate the n-th discrete difference along the given axis. 9 | 10 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.diff.html 11 | """ 12 | if not is_casadi_type(a): 13 | return _onp.diff(a, n=n, axis=axis) 14 | 15 | else: 16 | if axis != -1: 17 | raise NotImplementedError( 18 | "This could be implemented, but haven't had the need yet." 19 | ) 20 | 21 | result = a 22 | for i in range(n): 23 | result = _cas.diff(a) 24 | return result 25 | 26 | 27 | def trapz(x, modify_endpoints=False): # TODO unify with NumPy trapz, this is different 28 | """ 29 | Computes each piece of the approximate integral of `x` via the trapezoidal method with unit spacing. 30 | Can be viewed as the opposite of diff(). 31 | 32 | Args: 33 | x: The vector-like object (1D np.ndarray, cas.MX) to be integrated. 34 | 35 | Returns: A vector of length N-1 with each piece corresponding to the mean value of the function on the interval 36 | starting at index i. 37 | 38 | """ 39 | integral = (x[1:] + x[:-1]) / 2 40 | if modify_endpoints: 41 | integral[0] = integral[0] + x[0] * 0.5 42 | integral[-1] = integral[-1] + x[-1] * 0.5 43 | 44 | return integral 45 | -------------------------------------------------------------------------------- /.devcontainer/devcontainer.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "numagic_dev", 3 | "dockerComposeFile": "../docker-compose.yml", 4 | "service": "lumos-dev", 5 | // Using local or container env vars in devcontainer.json 6 | // see: https://code.visualstudio.com/docs/remote/devcontainerjson-reference 7 | "workspaceFolder": "/home/${localEnv:USER}/numagic/lumos", 8 | "extensions": [ 9 | "ms-python.python", 10 | "ms-python.vscode-pylance", 11 | "njpwerner.autodocstring", 12 | "visualstudioexptteam.vscodeintellicode" 13 | ], 14 | "settings": { 15 | // Somehow this defaultInterpeterpath is no longer picked-up automatically. The 16 | // user needs to select python interpreter manually once, and then it's set. 17 | // Failing to set the python interpreter would only affect vscode python 18 | // extension functionalities. 19 | "python.defaultInterpreterPath": "${containerEnv:CONDA_DIR}/envs/lumos/bin/python", 20 | "python.languageServer": "Pylance", 21 | "python.formatting.provider": "black", 22 | "editor.formatOnSave": true, 23 | "editor.formatOnSaveMode": "file", 24 | "python.linting.enabled": true, 25 | "python.linting.flake8Enabled": true, 26 | "python.linting.flake8Args": [ 27 | "--ignore=E203", 28 | "--ignore=E266", 29 | "--ignore=E501", 30 | "--ignore=W503", 31 | "--max-line-length=88", 32 | "--select=B,C,E,F,W,T4,B9", 33 | "--max-complexity=18" 34 | ], 35 | }, 36 | } -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_finite_difference_operators.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import lumos.numpy.casadi_numpy as np 3 | 4 | 5 | def test_uniform_forward_difference_first_degree(): 6 | assert np.finite_difference_coefficients( 7 | x=np.arange(2), x0=0, derivative_degree=1 8 | ) == pytest.approx(np.array([-1, 1])) 9 | assert np.finite_difference_coefficients( 10 | x=np.arange(9), x0=0, derivative_degree=1 11 | ) == pytest.approx( 12 | np.array([-761 / 280, 8, -14, 56 / 3, -35 / 2, 56 / 5, -14 / 3, 8 / 7, -1 / 8]) 13 | ) 14 | 15 | 16 | def test_uniform_forward_difference_higher_order(): 17 | assert np.finite_difference_coefficients( 18 | x=np.arange(5), x0=0, derivative_degree=3 19 | ) == pytest.approx(np.array([-5 / 2, 9, -12, 7, -3 / 2])) 20 | 21 | 22 | def test_uniform_central_difference(): 23 | assert np.finite_difference_coefficients( 24 | x=[-1, 0, 1], x0=0, derivative_degree=1 25 | ) == pytest.approx(np.array([-0.5, 0, 0.5])) 26 | assert np.finite_difference_coefficients( 27 | x=[-1, 0, 1], x0=0, derivative_degree=2 28 | ) == pytest.approx(np.array([1, -2, 1])) 29 | assert np.finite_difference_coefficients( 30 | x=[-2, -1, 0, 1, 2], x0=0, derivative_degree=2 31 | ) == pytest.approx(np.array([-1 / 12, 4 / 3, -5 / 2, 4 / 3, -1 / 12])) 32 | 33 | 34 | def test_nonuniform_difference(): 35 | assert np.finite_difference_coefficients( 36 | x=[-1, 2], x0=0, derivative_degree=1 37 | ) == pytest.approx(np.array([-1 / 3, 1 / 3])) 38 | 39 | 40 | if __name__ == "__main__": 41 | pytest.main() 42 | -------------------------------------------------------------------------------- /.github/workflows/test_with_conda.yml: -------------------------------------------------------------------------------- 1 | name: test-with-conda 2 | 3 | on: 4 | push: 5 | branches: 6 | - 'main' 7 | pull_request: 8 | # by default only opened, synchronize and reopened activity would trigger this event 9 | # see: https://docs.github.com/en/actions/learn-github-actions/events-that-trigger-workflows#pull_request 10 | # and def on synchronize: https://docs.github.com/en/developers/webhooks-and-events/webhooks/webhook-events-and-payloads#pull_request 11 | branches: 12 | - 'main' 13 | 14 | jobs: 15 | run-unit-tests-with-miniconda: 16 | runs-on: ${{ matrix.os }} 17 | strategy: 18 | fail-fast: false 19 | matrix: 20 | os: ["ubuntu-latest", "macos-latest"] # jax needs custom installation on windows. 21 | python-version: ["3.7", "3.8", "3.9"] 22 | steps: 23 | - uses: actions/checkout@v2 24 | - uses: conda-incubator/setup-miniconda@v2 25 | with: 26 | python-version: ${{ matrix.python-version }} 27 | auto-update-conda: true 28 | environment-file: environment.yml 29 | - name: Run unittest 30 | shell: bash -l {0} 31 | run: | 32 | python3 -m pip install . # use python3 here as python seems to point to python2, which is different to linux 33 | conda install -y pytest parameterized 34 | pytest tests/ examples/ 35 | finished-unit-test-matrix-with-miniconda: 36 | # A workaround to allow using one status check to summarise a matrix job 37 | runs-on: ubuntu-latest 38 | needs: run-unit-tests-with-miniconda 39 | steps: 40 | - run: echo Done! 41 | -------------------------------------------------------------------------------- /lumos/models/drone_model.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from typing import Any, Dict 4 | 5 | import lumos.numpy as lnp 6 | from lumos.models.base import state_space_io, StateSpaceModel, StateSpaceModelReturn 7 | 8 | logger = logging.getLogger(__name__) 9 | 10 | 11 | @state_space_io( 12 | states=("x", "x_dot", "z", "z_dot", "theta"), 13 | inputs=("f", "omega"), 14 | outputs=("sin_theta", "f_omega"), 15 | ) 16 | class DroneModel(StateSpaceModel): 17 | def __init__( 18 | self, model_config: Dict[str, Any] = {}, params: Dict[str, Any] = {}, 19 | ): 20 | super().__init__(model_config=model_config, params=params) 21 | 22 | def forward( 23 | self, 24 | states: Dict[str, float], 25 | inputs: Dict[str, float], 26 | mesh: float = 0.0, # time invariant model 27 | ) -> StateSpaceModelReturn: 28 | params = self._params 29 | theta = states["theta"] 30 | x_dot_dot = inputs["f"] * lnp.sin(theta) 31 | z_dot_dot = inputs["f"] * lnp.cos(theta) - params["gravity"] 32 | 33 | # Assemble result 34 | states_dot = self.make_dict( 35 | "states_dot", 36 | x=states["x_dot"], 37 | x_dot=x_dot_dot, 38 | z=states["z_dot"], 39 | z_dot=z_dot_dot, 40 | theta=inputs["omega"], 41 | ) 42 | 43 | outputs = self.make_dict( 44 | "outputs", sin_theta=lnp.sin(theta), f_omega=inputs["omega"] * inputs["f"], 45 | ) 46 | 47 | return StateSpaceModelReturn(states_dot=states_dot, outputs=outputs,) 48 | 49 | @classmethod 50 | def get_default_params(self) -> Dict[str, Any]: 51 | return {"gravity": 9.81} 52 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_linalg_top_level.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | import lumos.numpy.casadi_numpy as np 3 | import casadi as cas 4 | 5 | 6 | def test_cross_1D_input(): 7 | a = np.array([1, 1, 1]) 8 | b = np.array([1, 2, 3]) 9 | 10 | cas_a = cas.DM(a) 11 | cas_b = cas.DM(b) 12 | 13 | correct_result = np.cross(a, b) 14 | cas_correct_result = cas.DM(correct_result) 15 | 16 | assert np.all(np.cross(a, cas_b) == cas_correct_result) 17 | assert np.all(np.cross(cas_a, b) == cas_correct_result) 18 | assert np.all(np.cross(cas_a, cas_b) == cas_correct_result) 19 | 20 | 21 | def test_cross_2D_input_last_axis(): 22 | a = np.tile(np.array([1, 1, 1]), (3, 1)) 23 | b = np.tile(np.array([1, 2, 3]), (3, 1)) 24 | 25 | cas_a = cas.DM(a) 26 | cas_b = cas.DM(b) 27 | 28 | correct_result = np.cross(a, b) 29 | cas_correct_result = cas.DM(correct_result) 30 | 31 | assert np.all(np.cross(a, cas_b) == cas_correct_result) 32 | assert np.all(np.cross(cas_a, b) == cas_correct_result) 33 | assert np.all(np.cross(cas_a, cas_b) == cas_correct_result) 34 | 35 | 36 | def test_cross_2D_input_first_axis(): 37 | a = np.tile(np.array([1, 1, 1]), (3, 1)).T 38 | b = np.tile(np.array([1, 2, 3]), (3, 1)).T 39 | 40 | cas_a = cas.DM(a) 41 | cas_b = cas.DM(b) 42 | 43 | correct_result = np.cross(a, b, axis=0) 44 | cas_correct_result = cas.DM(correct_result) 45 | 46 | assert np.all(np.cross(a, cas_b, axis=0) == cas_correct_result) 47 | assert np.all(np.cross(cas_a, b, axis=0) == cas_correct_result) 48 | assert np.all(np.cross(cas_a, cas_b, axis=0) == cas_correct_result) 49 | 50 | 51 | if __name__ == "__main__": 52 | pytest.main() 53 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/rotations.py: -------------------------------------------------------------------------------- 1 | from lumos.numpy.casadi_numpy import linalg 2 | from lumos.numpy.casadi_numpy.array import array 3 | import numpy as _onp 4 | 5 | 6 | def rotation_matrix_2D(angle,): 7 | """ 8 | Gives the 2D rotation matrix associated with a counterclockwise rotation about an angle. 9 | Args: 10 | angle: Angle by which to rotate. Given in radians. 11 | 12 | Returns: The 2D rotation matrix 13 | 14 | """ 15 | sintheta = _onp.sin(angle) 16 | costheta = _onp.cos(angle) 17 | rotation_matrix = array([[costheta, -sintheta], [sintheta, costheta]]) 18 | return rotation_matrix 19 | 20 | 21 | def rotation_matrix_3D(angle, axis, _axis_already_normalized=False): 22 | """ 23 | Gives the 3D rotation matrix from an angle and an axis. 24 | An implmentation of https://en.wikipedia.org/wiki/Rotation_matrix#Rotation_matrix_from_axis_and_angle 25 | :param angle: can be one angle or a vector (1d ndarray) of angles. Given in radians. # TODO note deprecated functionality; must be scalar 26 | Direction corresponds to the right-hand rule. 27 | :param axis: a 1d numpy array of length 3 (x,y,z). Represents the angle. 28 | :param _axis_already_normalized: boolean, skips normalization for speed if you flag this true. 29 | :return: 30 | * If angle is a scalar, returns a 3x3 rotation matrix. 31 | * If angle is a vector, returns a 3x3xN rotation matrix. 32 | """ 33 | if not _axis_already_normalized: 34 | axis = axis / linalg.norm(axis) 35 | 36 | sintheta = _onp.sin(angle) 37 | costheta = _onp.cos(angle) 38 | cpm = array( 39 | [[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0],] 40 | ) # The cross product matrix of the rotation axis vector 41 | outer_axis = linalg.outer(axis, axis) 42 | 43 | rot_matrix = costheta * _onp.eye(3) + sintheta * cpm + (1 - costheta) * outer_axis 44 | return rot_matrix 45 | -------------------------------------------------------------------------------- /lumos/models/tires/utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import re 3 | from typing import Any, Dict 4 | 5 | logger = logging.getLogger(__name__) 6 | 7 | 8 | def read_tir_file(file_path: str) -> Dict[str, Any]: 9 | """Create tire data dictionary from Pacejka .tir file. 10 | 11 | Args: 12 | file_path (str): tir file path. 13 | Returns: 14 | Dict[str, Any]: flat parameter dicitonary contianing only numeric data. 15 | 16 | Typical data file looks like: 17 | 18 | . 19 | . 20 | . 21 | [OVERTURNING_COEFFICIENTS] 22 | QSX2 = 0.6038 $Camber induced overturning couple 23 | QSX3 = 0.025405 $Fy induced overturning couple 24 | . 25 | . 26 | . 27 | """ 28 | 29 | params = {} 30 | with open(file_path, "r") as f: 31 | tir_data = f.readlines() 32 | 33 | for line in tir_data: 34 | # TODO: this is pretty flaky and would break if somebody puts a "=" into 35 | # comments 36 | if "=" in line: 37 | name, val, *_ = re.split("[=$]", line.replace(" ", "")) 38 | try: 39 | params[name] = float(val) 40 | except ValueError: 41 | logger.debug(f"{name} is not a numeric value and is discarded.") 42 | 43 | return params 44 | 45 | 46 | # TODO: at the moment, the params can actually have different fields depending on the 47 | # form of the tir file! This makes jitting difficult (because the input pytree changes) 48 | def create_params_from_tir_file(tir_file): 49 | """Augment tir file params with some compute params.""" 50 | params = read_tir_file(tir_file) 51 | params.update( 52 | { 53 | # Used to avoid low speed singularity, [Eqn (4.E6a) Page 178 - Book] 54 | "epsilonv": 1e-6, 55 | "epsilonx": 1e-3, 56 | "epsilonk": 1e-6, 57 | "epsilony": 1e-3, 58 | } 59 | ) 60 | return params 61 | -------------------------------------------------------------------------------- /docker/base/Dockerfile-base: -------------------------------------------------------------------------------- 1 | ARG LUMOS_BASE_IMAGE 2 | ARG LUMOS_BASE_IMAGE_VERSION 3 | 4 | FROM ${LUMOS_BASE_IMAGE}:${LUMOS_BASE_IMAGE_VERSION} as base 5 | ENV DEBIAN_FRONTEND="noninteractive" 6 | 7 | # These must be specified after 'FROM' 8 | # ref: https://docs.docker.com/compose/compose-file/compose-file-v3/ 9 | ARG CONDA_DIR 10 | ARG MINICONDA_VERSION 11 | ARG LUMOS_DOCKER_DATA_DIR 12 | ARG LUMOS_DOCKER_PYTHON_VERSION 13 | 14 | # We need CONDA_DIR for scripts inside our docker image 15 | ENV CONDA_DIR $CONDA_DIR 16 | 17 | # install miniconda 18 | # ref: https://towardsdatascience.com/conda-pip-and-docker-ftw-d64fe638dc45 19 | RUN apt-get update \ 20 | && apt-get install -y wget vim \ 21 | && mkdir ${LUMOS_DOCKER_DATA_DIR} \ 22 | && wget --quiet https://repo.anaconda.com/miniconda/Miniconda3-${MINICONDA_VERSION}-Linux-x86_64.sh -O ${LUMOS_DOCKER_DATA_DIR}/miniconda.sh \ 23 | && chmod +x ${LUMOS_DOCKER_DATA_DIR}/miniconda.sh \ 24 | && ${LUMOS_DOCKER_DATA_DIR}/miniconda.sh -b -p ${CONDA_DIR} \ 25 | && rm ${LUMOS_DOCKER_DATA_DIR}/miniconda.sh 26 | 27 | # Make conda command available 28 | ENV PATH=$CONDA_DIR/bin:$PATH 29 | RUN $CONDA_DIR/bin/conda init bash 30 | 31 | # Create conda environment 32 | # See: https://towardsdatascience.com/conda-pip-and-docker-ftw-d64fe638dc45 33 | COPY environment.yml /tmp/ 34 | RUN conda create python=${LUMOS_DOCKER_PYTHON_VERSION} --name lumos \ 35 | # Make sure conda is activated in interactive shell via .bashrc 36 | && echo "conda deactivate" >> ~/.bashrc \ 37 | && echo "conda activate lumos" >> ~/.bashrc 38 | 39 | # Activate env and install gcc --> for real production, we probably don't need this 40 | SHELL ["/bin/bash", "-c", "-i"] 41 | RUN conda activate lumos \ 42 | && conda env update --file /tmp/environment.yml \ 43 | && apt upgrade \ 44 | && apt install -y build-essential 45 | 46 | # TODO: too much hard-code? 47 | COPY docker/scripts/default_entrypoint.sh /usr/local/bin/entrypoint.sh 48 | RUN chmod +x /usr/local/bin/entrypoint.sh 49 | 50 | ENTRYPOINT ["/usr/local/bin/entrypoint.sh"] 51 | 52 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/linalg_top_level.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.arithmetic import sum, abs 4 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 5 | 6 | 7 | def dot(a, b): 8 | """ 9 | Dot product of two arrays. 10 | 11 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.dot.html 12 | """ 13 | if not is_casadi_type([a, b], recursive=True): 14 | return _onp.dot(a, b) 15 | 16 | else: 17 | return _cas.dot(a, b) 18 | 19 | 20 | def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None): 21 | """ 22 | Return the cross product of two (arrays of) vectors. 23 | 24 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.cross.html 25 | """ 26 | if not is_casadi_type([a, b], recursive=True): 27 | return _onp.cross(a, b, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis) 28 | 29 | else: 30 | if axis is not None: 31 | if not (axis == -1 or axis == 0 or axis == 1): 32 | raise ValueError("`axis` must be -1, 0, or 1.") 33 | axisa = axis 34 | axisb = axis 35 | axisc = axis 36 | 37 | if axisa == -1 or axisa == 1: 38 | if not is_casadi_type(a): 39 | a = _cas.DM(a) 40 | a = a.T 41 | elif axisa == 0: 42 | pass 43 | else: 44 | raise ValueError("`axisa` must be -1, 0, or 1.") 45 | 46 | if axisb == -1 or axisb == 1: 47 | if not is_casadi_type(b): 48 | b = _cas.DM(b) 49 | b = b.T 50 | elif axisb == 0: 51 | pass 52 | else: 53 | raise ValueError("`axisb` must be -1, 0, or 1.") 54 | 55 | # Compute the cross product, horizontally (along axis 1 of a 2D array) 56 | c = _cas.cross(a, b) 57 | 58 | if axisc == -1 or axisc == 1: 59 | c = c.T 60 | elif axisc == 0: 61 | pass 62 | else: 63 | raise ValueError("`axisc` must be -1, 0, or 1.") 64 | 65 | return c 66 | -------------------------------------------------------------------------------- /lumos/models/tires/base.py: -------------------------------------------------------------------------------- 1 | from abc import abstractmethod 2 | from typing import Any, Dict, Optional, Tuple 3 | 4 | import lumos.numpy as lnp 5 | from lumos.models.base import Model, ModelReturn 6 | 7 | 8 | class BaseTire(Model): 9 | """Abstract description of a tire model. 10 | 11 | Args: 12 | params (Dict[str, Any], optional): Tire model parameters. Defaults to None. 13 | model_config (Dict[str, Any], optional): Model configuration data. Defaults to None. 14 | """ 15 | 16 | def __init__(self, params: Dict[str, Any] = {}, model_config: Dict[str, Any] = {}): 17 | super().__init__(model_config=model_config, params=params) 18 | 19 | @abstractmethod 20 | def forward( 21 | self, 22 | inputs: Dict[str, float], 23 | params: Optional[Dict[str, Any]] = None, 24 | ) -> ModelReturn: 25 | """Common tire model interface. 26 | 27 | The ISO sign convention is used, see page 29 of 28 | https://functionbay.com/documentation/onlinehelp/Documents/Tire/MFTyre-MFSwift_Help.pdf 29 | 30 | Args: 31 | inputs (lnp.ndarray): The inputs to the model. 32 | gamma: camber angle of the tire [rad] 33 | vx: long. wheel centre velocity in the tire coordinate frame [m/s] 34 | alpha: slip angle [rad] Opposite sign as force. 35 | kappa: slip ratio [-] Same sign as force 36 | Fz: normal load on the tire [N] 37 | params (Optional[Dict[str, Any]], optional): Defaults to None. 38 | Overwrite model parameters with user specified ones. If not 39 | provided, the parameters in the model attribute will be used. 40 | Returns: 41 | ModelReturn: with the 'outputs' field a vector containing: 42 | Fx: long. tire force [N] 43 | Fy: lat. tire force [N] 44 | Mx: overturning moment [Nm] 45 | My: rolling resistance moment [Nm] 46 | Mz: self-aligning moment [Nm] 47 | """ 48 | pass 49 | -------------------------------------------------------------------------------- /tests/test_mex/export_model.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | import scipy.io as sio 4 | from lumos.models.composition import ModelMaker 5 | 6 | 7 | def export_c_code_and_data(file_name: str, options, includes): 8 | """Export a state space model as c-code and data from recorded I/O""" 9 | # Create the model 10 | model = ModelMaker.make_model_from_name("SimpleVehicle") 11 | 12 | # Export c-code 13 | model.export_c_code(f"{file_name}.c", options=options, includes=includes) 14 | 15 | # Export some data that can be used for testing (to check if execution in mex gives 16 | # the same results as in python) 17 | params = model.get_recursive_params() 18 | flat_params, _ = params.tree_ravel() 19 | 20 | vx, vy, yaw_rate = 40.0, 0.0, 0.0 21 | no_slip_omega = vx / model._params["rolling_radius"] 22 | states = model.make_vector( 23 | "states", 24 | vx=vx, 25 | vy=vy, 26 | yaw_rate=yaw_rate, 27 | wheel_speed_fl=no_slip_omega, 28 | wheel_speed_fr=no_slip_omega, 29 | wheel_speed_rl=no_slip_omega, 30 | wheel_speed_rr=no_slip_omega, 31 | ) 32 | 33 | inputs = model.make_vector( 34 | "inputs", throttle=0.2, brake=0.0, steer=0.01, ax=0, ay=0 35 | ) 36 | mesh = 0.0 37 | 38 | model_return = model.forward_with_arrays(states, inputs, mesh) 39 | 40 | export_dict = model_return._asdict() 41 | export_dict.update( 42 | {"states": states, "inputs": inputs, "mesh": mesh, "params": flat_params} 43 | ) 44 | 45 | sio.savemat(f"{file_name}.mat", export_dict) 46 | 47 | 48 | if __name__ == "__main__": 49 | # Takes 1 commandline input for the file name 50 | # sys.argv is what follows 'python3' 51 | 52 | # For matlab mex function 53 | options = {"mex": True} 54 | includes = [] 55 | 56 | # For s-function as described in: https://web.casadi.org/blog/s-function/ 57 | # options = { 58 | # "casadi_real": "real_T", 59 | # "casadi_int": "int_T", 60 | # "with_header": True, 61 | # } 62 | # includes = ["simstruc.h"] 63 | export_c_code_and_data(sys.argv[1], options, includes) 64 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_determine_type.py: -------------------------------------------------------------------------------- 1 | from lumos.numpy.casadi_numpy.determine_type import * 2 | import pytest 3 | import numpy as np 4 | import casadi as cas 5 | 6 | 7 | def test_int(): 8 | assert is_casadi_type(5, recursive=True) == False 9 | assert is_casadi_type(5, recursive=False) == False 10 | 11 | 12 | def test_float(): 13 | assert is_casadi_type(5.0, recursive=True) == False 14 | assert is_casadi_type(5.0, recursive=False) == False 15 | 16 | 17 | def test_numpy(): 18 | assert is_casadi_type(np.array([1, 2, 3]), recursive=True) == False 19 | assert is_casadi_type(np.array([1, 2, 3]), recursive=False) == False 20 | 21 | 22 | def test_casadi(): 23 | assert is_casadi_type(cas.GenMX_ones(5), recursive=False) == True 24 | assert is_casadi_type(cas.GenMX_ones(5), recursive=True) == True 25 | 26 | 27 | def test_numpy_list(): 28 | assert is_casadi_type([np.array(5), np.array(7)], recursive=False) == False 29 | assert is_casadi_type([np.array(5), np.array(7)], recursive=True) == False 30 | 31 | 32 | def test_casadi_list(): 33 | assert ( 34 | is_casadi_type([cas.GenMX_ones(5), cas.GenMX_ones(5)], recursive=False) == False 35 | ) 36 | assert ( 37 | is_casadi_type([cas.GenMX_ones(5), cas.GenMX_ones(5)], recursive=True) == True 38 | ) 39 | 40 | 41 | def test_mixed_list(): 42 | assert is_casadi_type([np.array(5), cas.GenMX_ones(5)], recursive=False) == False 43 | assert is_casadi_type([np.array(5), cas.GenMX_ones(5)], recursive=True) == True 44 | 45 | 46 | def test_multi_level_contaminated_list(): 47 | a = [[1 for _ in range(10)] for _ in range(10)] 48 | 49 | assert is_casadi_type(a, recursive=False) == False 50 | assert is_casadi_type(a, recursive=True) == False 51 | 52 | a[5][5] = cas.MX(1) 53 | 54 | assert is_casadi_type(a, recursive=False) == False 55 | assert is_casadi_type(a, recursive=True) == True 56 | 57 | a[5][5] = np.array(cas.DM(1), dtype="O") 58 | 59 | assert is_casadi_type(a, recursive=False) == False 60 | assert is_casadi_type(a, recursive=True) == False 61 | 62 | 63 | if __name__ == "__main__": 64 | pytest.main([__file__]) 65 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/arithmetic.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 4 | 5 | 6 | def sum(x, axis: int = None): 7 | """ 8 | Sum of array elements over a given axis. 9 | 10 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.sum.html 11 | """ 12 | if not is_casadi_type(x): 13 | return _onp.sum(x, axis=axis) 14 | 15 | else: 16 | if axis == 0: 17 | return _cas.sum1(x).T 18 | 19 | elif axis == 1: 20 | return _cas.sum2(x) 21 | elif axis is None: 22 | return sum(sum(x, axis=0), axis=0) 23 | else: 24 | raise ValueError( 25 | "CasADi types can only be up to 2D, so `axis` must be None, 0, or 1." 26 | ) 27 | 28 | 29 | def mean(x, axis: int = None): 30 | """ 31 | Compute the arithmetic mean along the specified axis. 32 | 33 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.mean.html 34 | """ 35 | if not is_casadi_type(x): 36 | return _onp.mean(x, axis=axis) 37 | 38 | else: 39 | if axis == 0: 40 | return sum(x, axis=0) / x.shape[0] 41 | elif axis == 1: 42 | return sum(x, axis=1) / x.shape[1] 43 | elif axis is None: 44 | return mean(mean(x, axis=0), axis=1) 45 | else: 46 | raise ValueError( 47 | "CasADi types can only be up to 2D, so `axis` must be None, 0, or 1." 48 | ) 49 | 50 | 51 | def abs(x): 52 | if not is_casadi_type(x): 53 | return _onp.abs(x) 54 | 55 | else: 56 | return _cas.fabs(x) 57 | 58 | 59 | # TODO trace() 60 | 61 | # def cumsum(x, axis: int = None): 62 | # """ 63 | # Return the cumulative sum of the elements along a given axis. 64 | # 65 | # See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.cumsum.html 66 | # """ 67 | # 68 | # if not is_casadi_type(x): 69 | # return _onp.cumsum(x, axis=axis) 70 | # 71 | # else: 72 | # raise NotImplementedError 73 | # if axis is None: 74 | # return _cas.cumsum(_onp.flatten(x)) 75 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/spacing.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 4 | 5 | 6 | def linspace(start: float = 0.0, stop: float = 1.0, num: int = 50): 7 | """ 8 | Returns evenly spaced numbers over a specified interval. 9 | 10 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.linspace.html 11 | """ 12 | if not is_casadi_type([start, stop, num], recursive=True): 13 | return _onp.linspace(start, stop, num) 14 | else: 15 | return _cas.linspace(start, stop, num) 16 | 17 | 18 | def cosspace(start: float = 0.0, stop: float = 1.0, num: int = 50): 19 | """ 20 | Makes a cosine-spaced vector. 21 | 22 | Cosine spacing is useful because these correspond to Chebyshev nodes: https://en.wikipedia.org/wiki/Chebyshev_nodes 23 | 24 | To learn more about cosine spacing, see this: https://youtu.be/VSvsVgGbN7I 25 | 26 | Args: 27 | start: Value to start at. 28 | end: Value to end at. 29 | num: Number of points in the vector. 30 | """ 31 | mean = (stop + start) / 2 32 | amp = (stop - start) / 2 33 | return mean + amp * _onp.cos(linspace(_onp.pi, 0, num)) 34 | 35 | 36 | def logspace(start: float = 0.0, stop: float = 1.0, num: int = 50): 37 | """ 38 | Return numbers spaced evenly on a log scale. 39 | 40 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.logspace.html 41 | """ 42 | if not is_casadi_type([start, stop, num], recursive=True): 43 | return _onp.logspace(start, stop, num) 44 | else: 45 | return 10 ** linspace(start, stop, num) 46 | 47 | 48 | def geomspace(start: float = 1.0, stop: float = 10.0, num: int = 50): 49 | """ 50 | Return numbers spaced evenly on a log scale (a geometric progression). 51 | 52 | This is similar to logspace, but with endpoints specified directly. 53 | 54 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.geomspace.html 55 | """ 56 | if not is_casadi_type([start, stop, num], recursive=True): 57 | return _onp.geomspace(start, stop, num) 58 | else: 59 | if start <= 0 or stop <= 0: 60 | raise ValueError("Both start and stop must be positive!") 61 | return _onp.log10(10 ** linspace(start, stop, num)) 62 | -------------------------------------------------------------------------------- /lumos/models/kinematics.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import Any, Dict 3 | 4 | import lumos.numpy as lnp 5 | from lumos.models.base import state_space_io, StateSpaceModel, StateSpaceModelReturn 6 | 7 | logger = logging.getLogger(__name__) 8 | 9 | 10 | @state_space_io( 11 | states=("time", "n", "eta"), 12 | inputs=("vx", "vy", "yaw_rate", "track_curvature", "track_heading"), 13 | outputs=("yaw_angle", "curvature"), 14 | ) 15 | class TrackPosition2D(StateSpaceModel): 16 | """2D kinematics Model with velocities and yaw rate as yaw angle as inputs 17 | 18 | Sign convention: 19 | - road: curvature: +ve for LHS turn, rotation +ve anti-clockwise. Distance across is 20 | +ve if the vehicle is to the left of the centerline. 21 | - vehicle body coordinate: x forward, y to the left, z up 22 | """ 23 | 24 | def __init__( 25 | self, params: Dict[str, Any] = {}, model_config: Dict[str, Any] = None, 26 | ): 27 | super().__init__(model_config=model_config, params=params) 28 | 29 | def forward( 30 | self, states: Dict[str, float], inputs: Dict[str, float], mesh: float = 0.0, 31 | ) -> StateSpaceModelReturn: 32 | """ 33 | Sign convention: 34 | x-y in vehicle coordiante: 35 | x forward 36 | y to the left of the vehicle as seen by the driver 37 | """ 38 | 39 | n = states["n"] 40 | eta = states["eta"] 41 | vx = inputs["vx"] 42 | vy = inputs["vy"] 43 | yaw_rate = inputs["yaw_rate"] 44 | 45 | curvature = inputs["track_curvature"] 46 | track_heading = inputs["track_heading"] 47 | yaw_angle = states["eta"] + track_heading 48 | 49 | ds_dt = (vx * lnp.cos(eta) - vy * lnp.sin(eta)) / (1 - n * curvature) 50 | dn_dt = vx * lnp.sin(eta) + vy * lnp.cos(eta) 51 | deta_dt = yaw_rate - curvature * ds_dt 52 | 53 | # Pure kinematics related calculation. This is where vehicle models will be 54 | # executed in the future 55 | states_dot = self.make_dict( 56 | "states_dot", time=1 / ds_dt, n=dn_dt / ds_dt, eta=deta_dt / ds_dt, 57 | ) 58 | outputs = self.make_dict("outputs", yaw_angle=yaw_angle, curvature=curvature) 59 | return StateSpaceModelReturn(states_dot=states_dot, outputs=outputs) 60 | 61 | @classmethod 62 | def get_default_params(self) -> Dict[str, Any]: 63 | return {} 64 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/logicals.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 4 | 5 | 6 | def clip(x, min, max): 7 | """ 8 | Clip a value to a range. 9 | Args: 10 | x: Value to clip. 11 | min: Minimum value to clip to. 12 | max: Maximum value to clip to. 13 | 14 | Returns: 15 | 16 | """ 17 | return _onp.fmin(_onp.fmax(x, min), max) 18 | 19 | 20 | def minimum(x1, x2): 21 | """Element wise minimum.""" 22 | return _cas.fmin(x1, x2) 23 | 24 | 25 | def maximum(x1, x2): 26 | """Element wise maximum.""" 27 | return _cas.fmax(x1, x2) 28 | 29 | 30 | def logical_and(x1, x2): 31 | """ 32 | Compute the truth value of x1 AND x2 element-wise. 33 | 34 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.logical_and.html 35 | """ 36 | if not is_casadi_type([x1, x2], recursive=True): 37 | return _onp.logical_and(x1, x2) 38 | 39 | else: 40 | return _cas.logic_and(x1, x2) 41 | 42 | 43 | def logical_or(x1, x2): 44 | """ 45 | Compute the truth value of x1 OR x2 element-wise. 46 | 47 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.logical_or.html 48 | """ 49 | if not is_casadi_type([x1, x2], recursive=True): 50 | return _onp.logical_or(x1, x2) 51 | 52 | else: 53 | return _cas.logic_or(x1, x2) 54 | 55 | 56 | def logical_not(x): 57 | """ 58 | Compute the truth value of NOT x element-wise. 59 | 60 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.logical_not.html 61 | """ 62 | if not is_casadi_type(x, recursive=False): 63 | return _onp.logical_not(x) 64 | 65 | else: 66 | return _cas.logic_not(x) 67 | 68 | 69 | def all(a): # TODO add axis functionality 70 | """ 71 | Test whether all array elements along a given axis evaluate to True. 72 | 73 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.all.html 74 | """ 75 | if not is_casadi_type(a, recursive=False): 76 | return _onp.all(a) 77 | 78 | else: 79 | return _cas.logic_all(a) 80 | 81 | 82 | def any(a): # TODO add axis functionality 83 | """ 84 | Test whether any array element along a given axis evaluates to True. 85 | 86 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.any.html 87 | """ 88 | if not is_casadi_type(a, recursive=False): 89 | return _onp.any(a) 90 | 91 | else: 92 | return _cas.logic_any(a) 93 | -------------------------------------------------------------------------------- /docker/base/Dockerfile-dev: -------------------------------------------------------------------------------- 1 | ARG LUMOS_DOCKER_BASE_VERSION 2 | ARG CONTAINER_REGISTRY 3 | 4 | FROM ${CONTAINER_REGISTRY}/lumos-ci:${LUMOS_DOCKER_BASE_VERSION} 5 | ENV DEBIAN_FRONTEND="noninteractive" 6 | 7 | # Add non-root user, make the UID and GID the same as host to avoid permission issue on 8 | # linux: https://vsupalov.com/docker-shared-permissions/#build-the-right-image 9 | ARG LUMOS_DOCKER_USER_NAME 10 | ARG UID 11 | ARG GID 12 | 13 | # Install git and gnupg2 to make git integration work (ssh and gpg) 14 | RUN apt-get update \ 15 | && apt-get install -y vim git gnupg2 \ 16 | # HACK: the gitconfig points to /usr/local/bin (mac installation location) but 17 | # linux installs gpg in /usr/bin, so we work around this with a soft link 18 | && ln -s /usr/bin/gpg /usr/local/bin/gpg \ 19 | # install ssh server 20 | && apt-get -y --no-install-suggests --no-install-recommends install openssh-server \ 21 | && mkdir /run/sshd \ 22 | && chmod 0755 /run/sshd 23 | 24 | # Create the user: https://code.visualstudio.com/remote/advancedcontainers/add-nonroot-user 25 | # We make one change: to use /bin/bash as the login shell 26 | # group 20 on mac already exists in the base image, so use the -f option (which cancels 27 | # the gid add if exists) 28 | # 29 | # Also installs libgfortran4 to make testing with HSL libraries (if available) possible. 30 | RUN groupadd -f --gid $GID $LUMOS_DOCKER_USER_NAME \ 31 | && useradd --uid $UID --gid $GID -m $LUMOS_DOCKER_USER_NAME -s /bin/bash \ 32 | && apt-get update \ 33 | && apt-get install -y sudo \ 34 | && apt-get install -y libgfortran4 \ 35 | && echo $LUMOS_DOCKER_USER_NAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$LUMOS_DOCKER_USER_NAME \ 36 | && chmod 0440 /etc/sudoers.d/$LUMOS_DOCKER_USER_NAME 37 | 38 | # Let the user own the conda env 39 | RUN sudo chown -R $UID:$GID $CONDA_DIR/envs 40 | 41 | USER $LUMOS_DOCKER_USER_NAME 42 | 43 | # Now for the user, conda needs to be made available, so we do the following which was done 44 | # for the base as well (but that was not for the same user) 45 | # This also updates the bashrc 46 | RUN $CONDA_DIR/bin/conda init bash \ 47 | # Make sure the right conda env is activated in interactive shell via .bashrc 48 | && echo "conda deactivate" >> ~/.bashrc \ 49 | && echo "conda activate lumos" >> ~/.bashrc 50 | 51 | # Install additional tools to help dev and debug 52 | RUN conda activate lumos \ 53 | && pip install dash-bootstrap-components flask-caching 54 | 55 | # Start SSH server at entry point 56 | EXPOSE 22 57 | ENTRYPOINT ["/usr/local/bin/entrypoint.sh", "/usr/sbin/sshd", "-D"] -------------------------------------------------------------------------------- /tests/test_simulations/test_laptime_simulation.py: -------------------------------------------------------------------------------- 1 | from unittest import TestCase 2 | 3 | from lumos.models.composition import ModelMaker 4 | from lumos.models.tracks import RaceTrack 5 | from lumos.simulations.laptime_simulation import LTCConfig 6 | from lumos.simulations.laptime_simulation import LaptimeSimulation 7 | 8 | 9 | class TestLTCConfig(TestCase): 10 | def test_to_dict_and_back(self): 11 | """Test values remain the same when converting to dictionary and back""" 12 | config = LTCConfig(track="Barcelona") 13 | 14 | config_dict = config.to_dict() 15 | self.assertEqual(config, LTCConfig(**config_dict)) 16 | self.assertDictEqual(config_dict, LTCConfig(**config_dict).to_dict()) 17 | 18 | def test_set_track(self): 19 | """Should raise ValueError if no track is defined""" 20 | with self.assertRaises(ValueError): 21 | config = LTCConfig() 22 | 23 | expected_str = "/tracks/Barcelona" 24 | config = LTCConfig(track=expected_str) 25 | self.assertEqual(config.track, expected_str) 26 | 27 | def test_set_cyclic(self): 28 | """Ensure time and track heading are not made cyclic""" 29 | 30 | # If not specified, must also automtically add the reuqired non_cyclic_vars 31 | always_non_cyclic = ["states.time", "inputs.track_heading"] 32 | inputs_to_test = [["var1", "var2"], always_non_cyclic, ["states.time", "var2"]] 33 | for inputs in inputs_to_test: 34 | config = LTCConfig(track="Barcelona", non_cyclic_vars=inputs) 35 | for v in set(inputs + always_non_cyclic): 36 | self.assertIn(v, config.non_cyclic_vars) 37 | self.assertIn(v, config.non_cyclic_vars) 38 | 39 | 40 | # TODO: currently test_optimal_control/test_fixed_grid.py already tests LTC. Maybe we 41 | # should test that one with a dummy model, and test ltc solve here? 42 | class TestLaptimeSimulationWithoutSolve(TestCase): 43 | def test_set_and_change_track(self): 44 | """Test if the ocp property linked to the track is updated correctly.""" 45 | track_file = "data/tracks/Catalunya.csv" 46 | track = RaceTrack.from_tum_csv(track_file) 47 | 48 | ltc_config = LaptimeSimulation.get_sim_config(track=track_file) 49 | model_config = ModelMaker.make_config("SimpleVehicleOnTrack") 50 | ltc = LaptimeSimulation(model_config=model_config, sim_config=ltc_config) 51 | 52 | # Make sure the track length is correct 53 | self.assertAlmostEqual(ltc._mesh_scale, track.total_distance) 54 | 55 | # Change track and ensure that the distance is correctly updated. 56 | track_file = "data/tracks/Silverstone.csv" 57 | track = RaceTrack.from_tum_csv(track_file) 58 | ltc.set_track(track_file) 59 | self.assertAlmostEqual(ltc._mesh_scale, track.total_distance) 60 | -------------------------------------------------------------------------------- /tests/test_optimal_control/test_ltc_sweep.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import numpy as np 4 | from parameterized import parameterized_class 5 | 6 | from lumos.optimal_control.config import BoundaryConditionConfig 7 | from lumos.simulations.laptime_simulation import LaptimeSimulation 8 | from lumos.models.simple_vehicle_on_track import SimpleVehicleOnTrack 9 | 10 | 11 | @parameterized_class([{"backend": "casadi"}, {"backend": "jax"}]) 12 | class TestLTCSweep(unittest.TestCase): 13 | """Sweep mass parameter and solve 14 | 15 | This test mainly tests two things: 16 | 1) the laptime converges to a reasonable optimum, that it monotonically increases 17 | with mass (robustness and accuracy) 18 | 2) the parameter changing mechanism is working as expected 19 | 20 | TODO: 21 | Arguably the 2nd point can be made into a smaller unit test which only calls the nlp 22 | functions without doing the solve 23 | """ 24 | 25 | @classmethod 26 | def setUpClass(cls) -> None: 27 | """Construct an LTC problem""" 28 | model_config = SimpleVehicleOnTrack.get_recursive_default_model_config() 29 | model = SimpleVehicleOnTrack(model_config=model_config) 30 | params = model.get_recursive_default_params() 31 | 32 | sim_config = LaptimeSimulation.get_sim_config( 33 | track="data/tracks/Catalunya.csv", 34 | num_intervals=100, 35 | hessian_approximation="exact", 36 | is_cyclic=True, 37 | backend=cls.backend, 38 | transcription="LGR", 39 | boundary_conditions=(BoundaryConditionConfig(0, "states", "time", 0.0),), 40 | ) 41 | 42 | cls.ocp = LaptimeSimulation( 43 | sim_config=sim_config, model_params=params, model_config=model_config 44 | ) 45 | 46 | def _set_param_and_solve(self, path, value): 47 | self.ocp.modify_model_param(path, value) 48 | 49 | solution, info = self.ocp.solve( 50 | init_guess=self.ocp.get_init_guess(), 51 | print_level=5, 52 | max_iter=500, 53 | dual_inf_tol=1e-3, 54 | constr_viol_tol=1e-3, 55 | ) 56 | 57 | total_time = self.ocp.dec_var_operator.get_var( 58 | solution, group="states", name="time", stage=-1 59 | ) 60 | return total_time, info["status"] 61 | 62 | def test_mass_sweep(self): 63 | path = "vehicle.vehicle_mass" 64 | values = np.linspace(1600, 2400, 5) 65 | 66 | laptimes = [] 67 | for v in values: 68 | t, status = self._set_param_and_solve(path, v) 69 | self.assertEqual(status, 0, msg="solve status not optimal for mass={v:.1f}") 70 | laptimes.append(t) 71 | 72 | print(laptimes) 73 | 74 | self.assertTrue( 75 | np.all(np.diff(laptimes) > 0.01), 76 | msg="laptime not monotonically increasing with mass", 77 | ) 78 | -------------------------------------------------------------------------------- /lumos/models/simple_vehicle_on_track.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Dict, Optional 2 | 3 | import lumos.numpy as lnp 4 | from lumos.models.base import StateSpaceModel, state_space_io, StateSpaceModelReturn 5 | from lumos.models.kinematics import TrackPosition2D 6 | from lumos.models.vehicles.simple_vehicle import SimpleVehicle 7 | 8 | 9 | # Combine the signals to create the names. TODO: can we make it more automatic? 10 | @state_space_io( 11 | states=TrackPosition2D.get_direct_group_names("states") 12 | + SimpleVehicle.get_direct_group_names("states"), 13 | inputs=SimpleVehicle.get_direct_group_names("inputs") 14 | + ("track_curvature", "track_heading"), 15 | ) 16 | class SimpleVehicleOnTrack(StateSpaceModel): 17 | def __init__(self, *args, **kwargs): 18 | super().__init__(*args, **kwargs) 19 | 20 | @classmethod 21 | def get_default_submodel_config(self): 22 | 23 | return { 24 | "vehicle": "SimpleVehicle", 25 | "kinematics": "TrackPosition2D", 26 | } 27 | 28 | def forward(self, states: Dict[str, float], inputs: Dict[str, float], mesh: float): 29 | 30 | # Pick out the vehicle inputs 31 | vehicle_inputs = { 32 | k: inputs[k] for k in self.get_submodel("vehicle").get_group_names("inputs") 33 | } 34 | 35 | # Pick out vehicle states 36 | vehicle_states = { 37 | k: states[k] for k in self.get_submodel("vehicle").get_group_names("states") 38 | } 39 | 40 | # Pick out vehicle params. NOT DONE! NOT EASY! 41 | vehicle_return = self.call_submodel( 42 | "vehicle", states=vehicle_states, inputs=vehicle_inputs 43 | ) 44 | 45 | # Call Kinematics model 46 | # Pick out states 47 | kinematic_states = { 48 | k: states[k] 49 | for k in self.get_submodel("kinematics").get_group_names("states") 50 | } 51 | 52 | # pick out inputs 53 | # NOTE: this step is very custom, because the inputs come from vehicle model 54 | # outputs 55 | inputs_from_vehicle = {k: vehicle_states[k] for k in ("vx", "vy", "yaw_rate")} 56 | 57 | track_inputs = {k: inputs[k] for k in ["track_curvature", "track_heading"]} 58 | 59 | kinematic_inputs = dict(**track_inputs, **inputs_from_vehicle) 60 | 61 | # Pick out vehicle params. NOT DONE! NOT EASY! 62 | kinematics_return = self.call_submodel( 63 | "kinematics", states=kinematic_states, inputs=kinematic_inputs, mesh=mesh, 64 | ) 65 | 66 | # Convert to distance domain derivatives 67 | dt_ds = kinematics_return.states_dot["time"] 68 | states_dot = { 69 | **kinematics_return.states_dot, 70 | **{k: v * dt_ds for k, v in vehicle_return.states_dot.items()}, 71 | } 72 | 73 | # Assemble final outputs - there are no direct outputs from the current one 74 | outputs = self.make_outputs_dict() 75 | 76 | return StateSpaceModelReturn(states_dot=states_dot, outputs=outputs,) 77 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/linalg.py: -------------------------------------------------------------------------------- 1 | import numpy as _onp 2 | import casadi as _cas 3 | from lumos.numpy.casadi_numpy.arithmetic import sum, abs 4 | from lumos.numpy.casadi_numpy.determine_type import is_casadi_type 5 | 6 | 7 | def inner(x, y): 8 | """ 9 | Inner product of two arrays. 10 | 11 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.inner.html 12 | """ 13 | if not is_casadi_type([x, y], recursive=True): 14 | return _onp.inner(x, y) 15 | 16 | else: 17 | return _cas.dot(x, y) 18 | 19 | 20 | def outer(x, y): 21 | """ 22 | Compute the outer product of two vectors. 23 | 24 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.outer.html 25 | """ 26 | if not is_casadi_type([x, y], recursive=True): 27 | return _onp.outer(x, y) 28 | 29 | else: 30 | if len(y.shape) == 1: # Force y to be transposable if it's not. 31 | y = _onp.expand_dims(y, 1) 32 | return x @ y.T 33 | 34 | 35 | def solve(A, b): # TODO get this working 36 | """ 37 | Solve the linear system Ax=b for x. 38 | Args: 39 | A: A square matrix. 40 | b: A vector representing the RHS of the linear system. 41 | 42 | Returns: The solution vector x. 43 | 44 | """ 45 | if not is_casadi_type([A, b]): 46 | return _onp.linalg.solve(A, b) 47 | 48 | else: 49 | return _cas.solve(A, b) 50 | 51 | 52 | def norm(x, ord=None, axis=None): 53 | """ 54 | Matrix or vector norm. 55 | 56 | See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.linalg.norm.html 57 | """ 58 | if not is_casadi_type(x): 59 | return _onp.linalg.norm(x, ord=ord, axis=axis) 60 | 61 | else: 62 | 63 | # Figure out which axis, if any, to take a vector norm about. 64 | if axis is not None: 65 | if not (axis == 0 or axis == 1 or axis == -1): 66 | raise ValueError("`axis` must be -1, 0, or 1 for CasADi types.") 67 | elif x.shape[0] == 1: 68 | axis = 1 69 | elif x.shape[1] == 1: 70 | axis = 0 71 | 72 | if ord is None: 73 | if axis is not None: 74 | ord = 2 75 | else: 76 | ord = "fro" 77 | 78 | if ord == 1: 79 | # return _cas.norm_1(x) 80 | return sum(abs(x), axis=axis) 81 | elif ord == 2: 82 | # return _cas.norm_2(x) 83 | return sum(x ** 2, axis=axis) ** 0.5 84 | elif ord == "fro": 85 | return _cas.norm_fro(x) 86 | elif np.isinf(ord): 87 | return _cas.norm_inf() 88 | else: 89 | try: 90 | return sum(abs(x) ** ord, axis=axis) ** (1 / ord) 91 | except Exception as e: 92 | print(e) 93 | raise ValueError( 94 | "Couldn't interpret `ord` sensibly! Tried to interpret it as a floating-point order " 95 | "as a last-ditch effort, but that didn't work." 96 | ) 97 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_array.py: -------------------------------------------------------------------------------- 1 | from lumos.numpy.casadi_numpy.array import * 2 | import pytest 3 | import lumos.numpy.casadi_numpy as np 4 | import casadi as cas 5 | 6 | 7 | def test_array_numpy_equivalency_1D(): 8 | inputs = [1, 2] 9 | 10 | a = array(inputs) 11 | a_np = np.array(inputs) 12 | 13 | assert np.all(a == a_np) 14 | 15 | 16 | def test_array_numpy_equivalency_2D(): 17 | inputs = [[1, 2], [3, 4]] 18 | 19 | a = array(inputs) 20 | a_np = np.array(inputs) 21 | 22 | assert np.all(a == a_np) 23 | 24 | 25 | def test_array_casadi_1D_shape(): 26 | a = array([cas.DM(1), cas.DM(2)]) 27 | assert length(a) == 2 28 | 29 | 30 | def test_can_convert_DM_to_ndarray(): 31 | c = cas.DM([1, 2, 3]) 32 | n = np.array(c) 33 | 34 | assert np.all(n == np.array([1, 2, 3])) 35 | 36 | 37 | def test_length(): 38 | assert length(5) == 1 39 | assert length(5.0) == 1 40 | assert length([1, 2, 3]) == 3 41 | 42 | assert length(np.array(5)) == 1 43 | assert length(np.array([5])) == 1 44 | assert length(np.array([1, 2, 3])) == 3 45 | assert length(np.ones((3, 2))) == 3 46 | 47 | assert length(cas.GenMX_ones(5)) == 5 48 | 49 | 50 | def test_concatenate(): 51 | n = np.arange(10) 52 | c = cas.DM(n) 53 | 54 | assert concatenate((n, n)).shape == (20,) 55 | assert concatenate((n, c)).shape == (20, 1) 56 | assert concatenate((c, n)).shape == (20, 1) 57 | assert concatenate((c, c)).shape == (20, 1) 58 | 59 | assert concatenate((n, n, n)).shape == (30,) 60 | assert concatenate((c, c, c)).shape == (30, 1) 61 | 62 | 63 | def test_stack(): 64 | n = np.arange(10) 65 | c = cas.DM(n) 66 | 67 | assert stack((n, n)).shape == (2, 10) 68 | assert stack((n, n), axis=-1).shape == (10, 2) 69 | assert stack((n, c)).shape == (2, 10) 70 | assert stack((n, c), axis=-1).shape == (10, 2) 71 | assert stack((c, c)).shape == (2, 10) 72 | assert stack((c, c), axis=-1).shape == (10, 2) 73 | 74 | with pytest.raises(Exception): 75 | assert stack((n, n), axis=2) 76 | 77 | with pytest.raises(Exception): 78 | stack((c, c), axis=2) 79 | 80 | 81 | def test_roll_onp(): 82 | a = [1, 2, 3] 83 | b = [3, 1, 2] 84 | 85 | assert np.all(np.roll(a, 1) == b) 86 | 87 | 88 | def test_roll_casadi(): 89 | b = np.array([[3, 1, 2]]) 90 | a = cas.SX(b) 91 | 92 | assert np.all(cas.DM(np.roll(a, 1)) == b) 93 | 94 | 95 | def test_roll_casadi_2d(): 96 | a = np.array([[1, 2, 3], [4, 5, 6]]) 97 | b = cas.SX(a) 98 | 99 | assert np.all(cas.DM(np.roll(b, 1, axis=1)) == np.roll(a, 1, axis=1)) 100 | 101 | 102 | def test_max(): 103 | a = cas.SX([1, 2, 3]) 104 | b = [1, 2, 3] 105 | 106 | assert int(np.max(a)) == int(np.max(b)) 107 | 108 | 109 | def test_min(): 110 | a = cas.SX([1, 2, 3]) 111 | b = [1, 2, 3] 112 | 113 | assert int(np.min(a)) == int(np.min(b)) 114 | 115 | 116 | if __name__ == "__main__": 117 | test_can_convert_DM_to_ndarray() 118 | pytest.main() 119 | -------------------------------------------------------------------------------- /tests/test_models/test_simlpe_vehicle_on_track.py: -------------------------------------------------------------------------------- 1 | from mimetypes import init 2 | import unittest 3 | 4 | import numpy as np 5 | 6 | from lumos.models.test_utils import BaseStateSpaceModelTest 7 | from lumos.models.simple_vehicle_on_track import SimpleVehicleOnTrack 8 | 9 | 10 | class TestSimpleVehicleOnTrack(BaseStateSpaceModelTest, unittest.TestCase): 11 | ModelClass = SimpleVehicleOnTrack 12 | 13 | def _get_initial_states(self, vx=40.0): 14 | vy = 0.0 15 | yaw_rate = 0.0 16 | # FIXME: this is really bad. How do we get the rolling radius estimate? 17 | no_slip_omega = vx / 0.33 18 | 19 | return self.model.make_dict( 20 | group="states", 21 | time=0, 22 | n=0, 23 | eta=0, 24 | vx=vx, 25 | vy=vy, 26 | yaw_rate=yaw_rate, 27 | wheel_speed_fl=no_slip_omega, 28 | wheel_speed_fr=no_slip_omega, 29 | wheel_speed_rl=no_slip_omega, 30 | wheel_speed_rr=no_slip_omega, 31 | ) 32 | 33 | def test_turn_left(self): 34 | init_states = self._get_initial_states() 35 | inputs = self.model.make_dict( 36 | group="inputs", 37 | throttle=0.0, 38 | brake=0, 39 | ax=0.0, 40 | ay=0.0, 41 | steer=np.deg2rad(3.0), 42 | track_curvature=0.0, 43 | track_heading=0.0, 44 | ) 45 | 46 | model_return = self.model.forward(init_states, inputs, mesh=0.0) 47 | 48 | # should directly get +ve lateral acceleration, +ve yaw acceleration and derivative of vy 49 | self.assertGreater(model_return.states_dot["vy"], 0) 50 | self.assertGreater(model_return.states_dot["yaw_rate"], 0) 51 | self.assertGreater(model_return.outputs["vehicle.ay"], 0) 52 | 53 | # After a few timesteps: 54 | # The speed should decelerate due to cornering resistance 55 | # the vehicle should have a +ve yaw rate 56 | # The vehicle should have -ve vy (excpet for very low speed, the vehicle slips 57 | # to the inside of the turn because tire slip is heavily influenced by yaw rate) 58 | num_steps = 21 59 | dist_step = 1.0 60 | 61 | states, outputs = self._forward_euler(init_states, inputs, dist_step, num_steps) 62 | 63 | self.assertLess(states["vx"], init_states["vx"]) 64 | self.assertLess(states["vy"], 0) 65 | self.assertGreater(states["yaw_rate"], 0) 66 | 67 | # Assert position states 68 | self.assertGreater( 69 | outputs["kinematics.yaw_angle"], 70 | 0, 71 | msg="Yaw angle should be positive for turning left", 72 | ) 73 | 74 | self.assertGreater( 75 | states["n"], 76 | 0, 77 | msg="Distance across should be positive for turning left on straightline track", 78 | ) 79 | 80 | self.assertGreater( 81 | states["time"], 82 | 0, 83 | msg="Time used should be positive for turning left on straightline track", 84 | ) 85 | 86 | self.assertLess( 87 | dist_step * num_steps, 88 | init_states["vx"] * dist_step * num_steps, 89 | msg="Distance travelled should be less than going straight from the start", 90 | ) 91 | -------------------------------------------------------------------------------- /lumos/simulations/drone_simulation.py: -------------------------------------------------------------------------------- 1 | from dataclasses import dataclass, field 2 | from typing import Any, Dict, List, Tuple 3 | 4 | import numpy as np 5 | 6 | from lumos.optimal_control.config import ( 7 | BoundConfig, 8 | BoundaryConditionConfig, 9 | BoundConfig, 10 | SimConfig, 11 | ) 12 | from lumos.models.drone_model import DroneModel 13 | from lumos.optimal_control.scaled_mesh_ocp import ScaledMeshOCP 14 | 15 | 16 | def get_default_boundary_conditions(): 17 | return ( 18 | BoundaryConditionConfig(0, "states", "x", 0.0), 19 | BoundaryConditionConfig(0, "states", "x_dot", 0.0), 20 | BoundaryConditionConfig(0, "states", "z", 0.0), 21 | BoundaryConditionConfig(0, "states", "z_dot", 0.0), 22 | BoundaryConditionConfig(0, "states", "theta", 0.0), 23 | BoundaryConditionConfig(-1, "states", "x", 0.0), 24 | BoundaryConditionConfig(-1, "states", "x_dot", 0.0), 25 | BoundaryConditionConfig(-1, "states", "z", 5.0), 26 | BoundaryConditionConfig(-1, "states", "z_dot", 0.0), 27 | BoundaryConditionConfig(-1, "states", "theta", 2 * np.pi), 28 | ) 29 | 30 | 31 | def get_default_bounds(): 32 | return ( 33 | BoundConfig(group="states", name="x", values=(-50, 50)), 34 | BoundConfig(group="states", name="x_dot", values=(-50, 50)), 35 | BoundConfig(group="states", name="z", values=(-50, 50)), 36 | BoundConfig(group="states", name="z_dot", values=(-50, 50)), 37 | BoundConfig(group="states", name="theta", values=(-10 * np.pi, 10 * np.pi)), 38 | BoundConfig(group="inputs", name="f", values=(1, 20)), 39 | BoundConfig(group="inputs", name="omega", values=(-10, 10)), 40 | BoundConfig(group="global", name="mesh_scale", values=(0.1, 50)), 41 | ) 42 | 43 | 44 | @dataclass 45 | class DroneSimulationConfig(SimConfig): 46 | boundary_conditions: Tuple[BoundaryConditionConfig] = field( 47 | default_factory=get_default_boundary_conditions 48 | ) 49 | bounds: Tuple[BoundConfig] = field(default_factory=get_default_bounds) 50 | con_output_names: Tuple[str] = ("sin_theta",) 51 | 52 | 53 | class DroneSimulation(ScaledMeshOCP): 54 | ConfigClass: type = DroneSimulationConfig 55 | 56 | def __init__( 57 | self, 58 | model_params: Dict[str, Any] = {}, 59 | model_config: Dict[str, Any] = {}, 60 | sim_config: Dict[str, Any] = None, 61 | ): 62 | 63 | model = DroneModel(model_config=model_config, params=model_params,) 64 | super().__init__( 65 | model=model, sim_config=sim_config, 66 | ) 67 | 68 | def get_init_guess(self) -> np.ndarray: 69 | t_guess = 1.0 70 | 71 | inputs = np.zeros((self.num_stages, self.model.num_inputs)) + np.array( 72 | [10.0, 0] 73 | ) 74 | 75 | states = ( 76 | np.tile(self.model.make_const_vector(group="states"), (self.num_stages, 1),) 77 | + 0.1 78 | ) 79 | model_return = self.model.batched_forward( 80 | states, inputs, self.get_mesh_from_scale(t_guess), self._params 81 | ) 82 | 83 | return self.dec_var_operator.flatten_var( 84 | states=states, 85 | inputs=inputs, 86 | states_dot=model_return.states_dot, 87 | con_outputs=model_return.con_outputs, 88 | mesh_scale=t_guess, 89 | ) 90 | 91 | def get_total_time(self, x: np.array) -> float: 92 | return self._time_objective(x) 93 | -------------------------------------------------------------------------------- /.github/workflows/ltc_regression_test.yml: -------------------------------------------------------------------------------- 1 | name: ltc-regression-test 2 | on: push 3 | 4 | jobs: 5 | run-ltc-regression-test: 6 | runs-on: "ubuntu-latest" 7 | steps: 8 | - uses: actions/checkout@v2 9 | - uses: conda-incubator/setup-miniconda@v2 10 | with: 11 | python-version: 3.9 12 | auto-update-conda: true 13 | environment-file: environment.yml 14 | - name: Run LTC sweep and NLP profiling 15 | shell: bash -l {0} 16 | # This would fail if pytest fails, and will write out result files if successful 17 | # Somehow we need to use python3 -m pytest to ensure that we're running with 18 | # the correct python, and that's only for the regression tests folder... 19 | run: | 20 | conda env list 21 | python3 -m pip install . 22 | python3 regression_tests/run_benchmark.py 23 | - uses: actions/upload-artifact@v3 24 | with: 25 | path: | 26 | track_sweep_results.csv 27 | track_sweep_summary.csv 28 | summary.json 29 | if-no-files-found: error 30 | # Download previous benchmark result from cache (if exists) 31 | - name: Download previous benchmark data 32 | # Only store reults and display if push to main 33 | if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} 34 | uses: actions/cache@v1 35 | with: 36 | path: ./cache 37 | key: ${{ runner.os }}-benchmark 38 | # Run `github-action-benchmark` action for features 39 | - name: Store benchmark result for features 40 | if: ${{ github.event_name == 'push' }} 41 | uses: benchmark-action/github-action-benchmark@v1 42 | with: 43 | tool: customSmallerIsBetter 44 | # Where the output from the benchmark tool is stored 45 | output-file-path: summary.json 46 | # Workflow will fail when an alert happens 47 | fail-on-alert: False 48 | # Set auto-push to false since GitHub API token is not given 49 | auto-push: False 50 | # Need to ensure external-data-json-path is not set, otherwise won't create 51 | # gh-pages branch 52 | benchmark-data-dir-path: dev/bench-features 53 | - name: Push benchmark result for features 54 | if: ${{ github.event_name == 'push' }} 55 | run: git push 'https://${{ github.actor }}:${{ secrets.GITHUB_TOKEN }}@github.com/numagic/lumos.git' gh-pages:gh-pages 56 | # Run `github-action-benchmark` action for main only 57 | - name: Store benchmark result 58 | if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} 59 | uses: benchmark-action/github-action-benchmark@v1 60 | with: 61 | tool: customSmallerIsBetter 62 | # Where the output from the benchmark tool is stored 63 | output-file-path: summary.json 64 | # Workflow will fail when an alert happens 65 | fail-on-alert: False 66 | # Set auto-push to false since GitHub API token is not given 67 | auto-push: False 68 | # Need to ensure external-data-json-path is not set, otherwise won't create 69 | # gh-pages branch 70 | - name: Push benchmark result 71 | if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }} 72 | run: git push 'https://${{ github.actor }}:${{ secrets.GITHUB_TOKEN }}@github.com/numagic/lumos.git' gh-pages:gh-pages 73 | -------------------------------------------------------------------------------- /.github/workflows/test_mex_with_octave.yml: -------------------------------------------------------------------------------- 1 | name: test-mex-with-octave 2 | # Tests in linux, export a state-space model to c-code, and export saved I/O data. 3 | # Then in both linux and windows, compile the c-code into mex and executes it and tests 4 | # if outputs matches expected values. 5 | 6 | on: 7 | push: 8 | branches: 9 | - 'main' 10 | pull_request: 11 | branches: 12 | - 'main' 13 | 14 | # NOTE: we use 'artifact' instead of 'cache' because cache is OS specific, meaning we 15 | # can't easily generate a file in linux and share it with linxu and windows via cache 16 | jobs: 17 | generate-c-code: 18 | runs-on: ubuntu-latest 19 | outputs: 20 | artifact-name: ${{ steps.generate-artifact-name.outputs.artifact-name}} 21 | file-name: ${{ steps.generate-file-name.outputs.file-name}} 22 | steps: 23 | - name: Generate artifact name 24 | # We want to share the generated code between jobs of the same commit, but not 25 | # between commits, therefore we need to generate unique keys. 26 | id: generate-artifact-name 27 | run: | 28 | echo "::set-output name=artifact-name::codegen-${{ github.run_id }}-${{ github.run_attempt }}" 29 | shell: bash 30 | - name: Generate file name 31 | # We want to share the generated code between jobs of the same commit, but not 32 | # between commits, therefore we need to generate unique keys. 33 | id: generate-file-name 34 | run: | 35 | echo "::set-output name=file-name::forward" 36 | shell: bash 37 | - uses: actions/checkout@v2 38 | - uses: conda-incubator/setup-miniconda@v2 39 | with: 40 | python-version: 3.9 41 | auto-update-conda: true 42 | environment-file: environment.yml 43 | - name: Generate code 44 | shell: bash -l {0} 45 | run: | 46 | python3 -m pip install . # use python3 here as python seems to point to python2, which is different to linux 47 | python3 tests/test_mex/export_model.py ${{ steps.generate-file-name.outputs.file-name}} 48 | - name: Store generated code and test data as artifact 49 | uses: actions/upload-artifact@v3 50 | with: 51 | name: ${{ steps.generate-artifact-name.outputs.artifact-name }} 52 | path: | 53 | ${{ steps.generate-file-name.outputs.file-name}}.c 54 | ${{ steps.generate-file-name.outputs.file-name}}.mat 55 | 56 | compile-mex-and-run-on-linux: 57 | runs-on: ubuntu-latest 58 | needs: generate-c-code 59 | steps: 60 | - uses: actions/checkout@v2 61 | - name: Install octave 62 | run: | 63 | sudo apt -yq update 64 | sudo apt install -yq --no-install-recommends octave liboctave-dev 65 | - uses: actions/download-artifact@v3 66 | with: 67 | name: ${{ needs.generate-c-code.outputs.artifact-name }} 68 | - run: octave-cli tests/test_mex/compile_mex_and_run.m ${{ needs.generate-c-code.outputs.file-name }} 69 | 70 | compile-mex-and-run-on-windows: 71 | runs-on: windows-latest 72 | needs: generate-c-code 73 | steps: 74 | - uses: actions/checkout@v2 75 | - name: Install octave 76 | run: choco install octave.portable 77 | - uses: actions/download-artifact@v3 78 | with: 79 | name: ${{ needs.generate-c-code.outputs.artifact-name }} 80 | # On windows, 'octave' seems to dispatch into a different process, while the 81 | # the current process would exist with status 0 immediately... 82 | # Also needs to use "" for path, otherwise it won't be parsed correctly. 83 | - run: | 84 | octave-cli "tests\test_mex\compile_mex_and_run.m" ${{ needs.generate-c-code.outputs.file-name }} 85 | shell: bash 86 | 87 | successful-mex-compile-and-run-on-windows-and-linux: 88 | # A workaround to allow using one status check for multipe 89 | runs-on: ubuntu-latest 90 | needs: 91 | - compile-mex-and-run-on-windows 92 | - compile-mex-and-run-on-linux 93 | steps: 94 | - run: echo Done! 95 | -------------------------------------------------------------------------------- /tests/test_numpy/test_lumos_numpy.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import jax.numpy as jnp 4 | import numpy as np 5 | from casadi import SX 6 | 7 | 8 | import lumos.numpy as lnp 9 | 10 | 11 | class TestBasics(unittest.TestCase): 12 | def setUp(self): 13 | self._backend = lnp.get_backend() 14 | 15 | def test_set_and_get_backend(self): 16 | """ 17 | IF switched to a correct backend 18 | THEN create the backend is set and compute return type is correct. 19 | """ 20 | np_input = np.random.randn(10, 3) 21 | 22 | # numpy backend 23 | lnp.set_backend("numpy") 24 | self.assertEqual(lnp.get_backend(), "numpy") 25 | np_out = lnp.sin(np_input) 26 | self.assertFalse(lnp.is_jax_array(np_out)) 27 | 28 | # Test jax backend 29 | lnp.set_backend("jax") 30 | self.assertEqual(lnp.get_backend(), "jax") 31 | jax_out = lnp.sin(np_input) 32 | self.assertTrue(lnp.is_jax_array(jax_out)) 33 | 34 | # casadi backend 35 | lnp.set_backend("casadi") 36 | self.assertEqual(lnp.get_backend(), "casadi") 37 | casadi_out = lnp.sin(SX(np_input)) 38 | self.assertTrue(isinstance(casadi_out, SX)) 39 | 40 | def tearDown(self): 41 | lnp.set_backend(self._backend) 42 | 43 | 44 | class TestUtils(unittest.TestCase): 45 | def setUp(self): 46 | self._backend = lnp.get_backend() 47 | 48 | def test_is_jax_array(self): 49 | np_array = np.zeros((10, 3)) 50 | self.assertFalse(lnp.is_jax_array(np_array)) 51 | self.assertFalse(lnp.is_jax_array(SX(np_array))) 52 | self.assertTrue(lnp.is_jax_array(jnp.array(np_array))) 53 | 54 | def test_vector_concat(self): 55 | num_arrays, array_size = 5, 10 56 | expected_vector_size = num_arrays * array_size 57 | np_arrays = [np.random.randn(array_size) for _ in range(num_arrays)] 58 | casadi_arrays = [SX(a) for a in np_arrays] 59 | jax_arrays = [jnp.array(a) for a in np_arrays] 60 | 61 | lnp.set_backend("jax") 62 | self.assertEqual(lnp.vector_concat(jax_arrays).shape, (expected_vector_size,)) 63 | 64 | lnp.set_backend("numpy") 65 | self.assertEqual(lnp.vector_concat(np_arrays).shape, (expected_vector_size,)) 66 | 67 | # Casadi uses column vectors! 68 | lnp.set_backend("casadi") 69 | self.assertEqual( 70 | lnp.vector_concat(casadi_arrays).shape, (expected_vector_size, 1) 71 | ) 72 | 73 | def test_lmap_with_arrays(self): 74 | """lmap should work with array outputs, with different in_axes mapping""" 75 | 76 | def fn(a, b): 77 | return a + b 78 | 79 | # Mapping all inputs 80 | a_array, b_array = np.ones(10), np.ones(10) * 2 81 | c_array = lnp.lmap(fn)(a_array, b_array) 82 | np.testing.assert_allclose(c_array, a_array + b_array) 83 | 84 | # Without mapping the last input 85 | a_array, b = np.ones(10), 2 86 | c_array = lnp.lmap(fn, in_axes=[0, None])(a_array, b) 87 | np.testing.assert_allclose(c_array, a_array + b) 88 | 89 | def test_lmap_with_tuples(self): 90 | """lmap should work with tuple or namedtuple outputs, with different in_axes mapping""" 91 | 92 | def fn(a, b): 93 | return a + b, a - b 94 | 95 | # Mapping all inputs 96 | a_array, b_array = np.ones(10), np.ones(10) * 2 97 | output = lnp.lmap(fn)(a_array, b_array) 98 | np.testing.assert_allclose(output, (a_array + b_array, a_array - b_array)) 99 | 100 | # Without mapping the last input 101 | a_array, b = np.ones(10), 2 102 | actual = lnp.lmap(fn, in_axes=[0, None])(a_array, b) 103 | expected = (a_array + b, a_array - b) 104 | 105 | for act, exp in zip(actual, expected): 106 | np.testing.assert_allclose(act, exp) 107 | 108 | def tearDown(self): 109 | lnp.set_backend(self._backend) 110 | 111 | 112 | if __name__ == "__name__": 113 | unittest.main() 114 | -------------------------------------------------------------------------------- /lumos/numpy/casadi_numpy/finite_difference_operators.py: -------------------------------------------------------------------------------- 1 | from lumos.numpy.casadi_numpy.array import array, length 2 | import numpy as _onp 3 | 4 | 5 | def finite_difference_coefficients( 6 | x: _onp.ndarray, x0: float = 0, derivative_degree: int = 1, 7 | ) -> _onp.ndarray: 8 | """ 9 | Computes the weights (coefficients) in compact finite differece formulas for any order of derivative 10 | and to any order of accuracy on one-dimensional grids with arbitrary spacing. 11 | 12 | (Wording above is taken from the paper below, as are docstrings for parameters.) 13 | 14 | Modified from an implementation of: 15 | 16 | Fornberg, Bengt, "Generation of Finite Difference Formulas on Arbitrarily Spaced Grids". Oct. 1988. 17 | Mathematics of Computation, Volume 51, Number 184, pages 699-706. 18 | 19 | PDF: https://www.ams.org/journals/mcom/1988-51-184/S0025-5718-1988-0935077-0/S0025-5718-1988-0935077-0.pdf 20 | 21 | More detail: https://en.wikipedia.org/wiki/Finite_difference_coefficient 22 | 23 | Args: 24 | 25 | derivative_degree: The degree of the derivative that you are interested in obtaining. (denoted "M" in the 26 | paper) 27 | 28 | x: The grid points (not necessarily uniform or in order) that you want to obtain weights for. You must 29 | provide at least as many grid points as the degree of the derivative that you're interested in, plus 1. 30 | 31 | The order of accuracy of your derivative depends in part on the number of grid points that you provide. 32 | Specifically: 33 | 34 | order_of_accuracy = n_grid_points - derivative_degree 35 | 36 | (This is in general; can be higher in special cases.) 37 | 38 | For example, if you're evaluating a second derivative and you provide three grid points, you'll have a 39 | first-order-accurate answer. 40 | 41 | (x is denoted "alpha" in the paper) 42 | 43 | x0: The location that you are interested in obtaining a derivative at. This need not be on a grid point. 44 | 45 | Complexity is O(derivative_degree * len(x) ^ 2) 46 | 47 | Returns: A 1D ndarray corresponding to the coefficients that should be placed on each grid point. In other words, 48 | the approximate derivative at `x0` is the dot product of `coefficients` and the function values at each of the 49 | grid points `x`. 50 | 51 | """ 52 | ### Check inputs 53 | if derivative_degree < 1: 54 | return ValueError("The parameter derivative_degree must be an integer >= 1.") 55 | expected_order_of_accuracy = length(x) - derivative_degree 56 | if expected_order_of_accuracy < 1: 57 | return ValueError( 58 | "You need to provide at least (derivative_degree+1) grid points in the x vector." 59 | ) 60 | 61 | ### Implement algorithm; notation from paper in docstring. 62 | N = length(x) - 1 63 | 64 | delta = _onp.zeros(shape=(derivative_degree + 1, N + 1, N + 1), dtype="O") 65 | 66 | delta[0, 0, 0] = 1 67 | c1 = 1 68 | for n in range( 69 | 1, N + 1 70 | ): # TODO make this algorithm more efficient; we only need to store a fraction of this data. 71 | c2 = 1 72 | for v in range(n): 73 | c3 = x[n] - x[v] 74 | c2 = c2 * c3 75 | # if n <= M: # Omitted because d is initialized to zero. 76 | # d[n, n - 1, v] = 0 77 | for m in range(min(n, derivative_degree) + 1): 78 | delta[m, n, v] = ( 79 | (x[n] - x0) * delta[m, n - 1, v] - m * delta[m - 1, n - 1, v] 80 | ) / c3 81 | for m in range(min(n, derivative_degree) + 1): 82 | delta[m, n, n] = ( 83 | c1 84 | / c2 85 | * ( 86 | m * delta[m - 1, n - 1, n - 1] 87 | - (x[n - 1] - x0) * delta[m, n - 1, n - 1] 88 | ) 89 | ) 90 | c1 = c2 91 | 92 | coefficients_object_array = delta[derivative_degree, -1, :] 93 | 94 | coefficients = array( 95 | [*coefficients_object_array] 96 | ) # Reconstructs using aerosandbox.numpy to intelligently type 97 | 98 | return coefficients 99 | -------------------------------------------------------------------------------- /lumos/optimal_control/fixed_mesh_ocp.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import Any, Dict, Tuple 3 | 4 | import numpy as np 5 | 6 | from lumos.models.base import StateSpaceModel 7 | from lumos.optimal_control.nlp import ( 8 | LinearConstraints, 9 | BaseObjective, 10 | ) 11 | from lumos.optimal_control.scaled_mesh_ocp import ScaledMeshOCP 12 | 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | 17 | class FixedMeshOCP(ScaledMeshOCP): 18 | """Optimal Control problem where the mesh is fixed""" 19 | 20 | # No Global variables 21 | global_var_names: Tuple[str] = () 22 | 23 | def __init__( 24 | self, 25 | model: StateSpaceModel, 26 | sim_config: Dict[str, Any] = None, 27 | mesh_scale: float = 1.0, 28 | ): 29 | 30 | # NOTE: we choose to make mesh_scale a variable instead of a config to make the 31 | # distinction of fixed_grid clearer! 32 | assert mesh_scale > 0.0, "mesh_scale must be a positive scalar!" 33 | self.set_mesh_scale(mesh_scale) 34 | 35 | super().__init__(model=model, sim_config=sim_config) 36 | 37 | def set_mesh_scale(self, mesh_scale: float): 38 | """When we set the mesh_scale, some other properties must change""" 39 | 40 | # Currently we don't need to update anything else, but in the future we might. 41 | self._mesh_scale = mesh_scale 42 | 43 | def _get_mesh_scale(self, x): 44 | """Overwrite parent method since the mesh scale is now fixed.""" 45 | return self._mesh_scale 46 | 47 | def get_mesh(self): 48 | """Helper method for FixedMesh as we don't need any input to get the mesh.""" 49 | return self.get_mesh_from_scale(self._mesh_scale) 50 | 51 | def _time_objective(self, x): 52 | idx_end = self.dec_var_operator.get_var_index_in_dec( 53 | group="states", name="time", stage=-1 54 | ) 55 | 56 | idx_start = self.dec_var_operator.get_var_index_in_dec( 57 | group="states", name="time", stage=0 58 | ) 59 | return x[idx_end] - x[idx_start] 60 | 61 | def _time_gradient(self, x): 62 | idx_end = self.dec_var_operator.get_var_index_in_dec( 63 | group="states", name="time", stage=-1 64 | ) 65 | 66 | idx_start = self.dec_var_operator.get_var_index_in_dec( 67 | group="states", name="time", stage=0 68 | ) 69 | grad = np.zeros_like(x) 70 | grad[idx_end] = 1 71 | grad[idx_start] = -1 72 | return grad 73 | 74 | def _build_continuity_cons(self): 75 | """Overwrite base class to create linear constraints for continiuty. 76 | 77 | Since for fixed grid, the continuity is just a linear constraint. 78 | 79 | However here we do not cache the jacobian, because when the jacobian value is 80 | cached, then if the problem changes (for example, when the mesh_scale changes), this continuity jacobian is no 81 | longer valid! 82 | 83 | The performance cost is < 1ms per call for 250 intervals with 3 stages per 84 | interval, so pretty much negligible for any problem with a non-trivial model. 85 | """ 86 | 87 | continuity_cons = LinearConstraints( 88 | constraints=self._continuity_constraints, 89 | num_in=self.num_dec, 90 | num_con=self.num_continuity_cons, 91 | jacobian=self._continuity_jacobian, 92 | jacobian_structure=self._continuity_jacobianstructure(), 93 | cache_jacobian=False, 94 | ) 95 | 96 | self.add_constraints("continuity", continuity_cons) 97 | 98 | def _build_objective(self): 99 | # Common objective regardless of the problem 100 | time_objective = BaseObjective( 101 | num_in=self.num_dec, 102 | objective=lambda x: self._time_objective(x), 103 | gradient=lambda x: self._time_gradient(x), 104 | hessian=lambda x: np.array([]), 105 | hessian_structure=( 106 | np.array([], dtype=np.int32), 107 | np.array([], dtype=np.int32), 108 | ), 109 | ) 110 | self.add_objective("time", time_objective) 111 | -------------------------------------------------------------------------------- /lumos/models/tires/perantoni.py: -------------------------------------------------------------------------------- 1 | import logging 2 | from typing import Any, Dict 3 | 4 | import numpy as np 5 | 6 | import lumos.numpy as lnp 7 | from lumos.models.base import model_io, ModelReturn 8 | from lumos.models.tires.base import BaseTire 9 | 10 | logger = logging.getLogger(__name__) 11 | 12 | 13 | @model_io( 14 | inputs=( 15 | "Fz", # vertical load 16 | "kappa", # slip ratio 17 | "alpha", # slip angle 18 | "vx", # x-velocity in tire coordinate 19 | "gamma", # inclination angle 20 | ), 21 | outputs=( 22 | "Fx", 23 | "Fy", 24 | "Mx", 25 | "My", 26 | "Mz", 27 | ), 28 | ) 29 | class PerantoniTire(BaseTire): 30 | """Tire model used in Perantoni Paper 31 | Optimal Control of F1 with variable parameters. 32 | 33 | Perantoni papers sticks to some werid sign convention (Fz -ve, omega -ve). 34 | We make slight modifications to adopt iso convention, which is the same as used in 35 | MF5.2. See page 29 of: 36 | https://functionbay.com/documentation/onlinehelp/Documents/Tire/MFTire-MFSwift_Help.pdf 37 | 38 | As a result: 39 | - longitudinal slip and force are the same sign 40 | - lateral slip and force are the opposite sign. 41 | """ 42 | 43 | def __init__(self, *args, **kwargs): 44 | super().__init__(*args, **kwargs) 45 | 46 | def forward(self, inputs: Dict[str, float]): 47 | """Perantoni Apendix A, page 32""" 48 | params = self._params 49 | 50 | # unpack inputs 51 | kappa = inputs["kappa"] 52 | alpha = inputs["alpha"] 53 | Fz = inputs["Fz"] 54 | 55 | # unpack parameters 56 | Fz1 = params["Fz1"] 57 | Fz2 = params["Fz2"] 58 | mux1 = params["mux1"] 59 | mux2 = params["mux2"] 60 | muy1 = params["muy1"] 61 | muy2 = params["muy2"] 62 | kappa1 = params["kappa1"] 63 | kappa2 = params["kappa2"] 64 | alpha1 = params["alpha1"] 65 | alpha2 = params["alpha2"] 66 | Qx = params["Qx"] 67 | Qy = params["Qy"] 68 | 69 | # NOTE: here there are some inconsistency in the paper regarding parameter names 70 | # in the parameters, we have mux1, mux2, but in equations we have mux_max1, mux_max2 71 | # TODO: these are just linear interpolation, could abstract out. 72 | interp_val = (Fz - Fz1) / (Fz2 - Fz1) 73 | mux_max = interp_val * (mux2 - mux1) + mux1 74 | muy_max = interp_val * (muy2 - muy1) + muy1 75 | kappa_max = interp_val * (kappa2 - kappa1) + kappa1 76 | alpha_max = interp_val * (alpha2 - alpha1) + alpha1 77 | 78 | kappa_n = kappa / kappa_max 79 | alpha_n = alpha / alpha_max 80 | 81 | # NOTE: we need to add a jitter here to avoid divide by zero error at 0 slip 82 | jitter = 1e-6 83 | rho = lnp.sqrt(alpha_n**2 + kappa_n**2 + jitter) 84 | 85 | Sx = np.pi / 2 / lnp.arctan(Qx) 86 | Sy = np.pi / 2 / lnp.arctan(Qy) 87 | 88 | mux = mux_max * lnp.sin(Qx * lnp.arctan(Sx * rho)) 89 | muy = muy_max * lnp.sin(Qy * lnp.arctan(Sy * rho)) 90 | 91 | Fx = mux * Fz * kappa_n / rho 92 | Fy = -muy * Fz * alpha_n / rho 93 | 94 | # Fill moments with 0.0 95 | outputs = dict(Fx=Fx, Fy=Fy, Mx=0.0, My=0.0, Mz=0.0) 96 | return ModelReturn(outputs=outputs) 97 | 98 | @classmethod 99 | def get_default_params(cls) -> Dict[str, Any]: 100 | 101 | # parameters from Perantoni page 34 102 | params = { 103 | "Fz1": 2000, # reference load 1 104 | "Fz2": 6000, # reference load 2 105 | "mux1": 1.75, # peak longitudinal friction coefficient at load 1 106 | "mux2": 1.4, # peak longitudinal friction coefficient at load 2 107 | "kappa1": 0.11, # slip coefficient for the friction peak at load 1 108 | "kappa2": 0.10, # slip coefficient for the friction peak at load 2 109 | "muy1": 1.80, # peak lateral friction coefficient at load 1 110 | "muy2": 1.45, # peak lateral friction coefficient at load 2 111 | "alpha1": np.deg2rad(9.0), # slip angle for the friction peak at load 1 112 | "alpha2": np.deg2rad(8.0), # slip angle for the friction peak at load 2 113 | "Qx": 1.9, # longitudinal shape factor 114 | "Qy": 1.9, # lateral shape factor 115 | } 116 | 117 | return params 118 | -------------------------------------------------------------------------------- /tests/test_optimal_control/test_config.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from unittest import TestCase 3 | 4 | from lumos.optimal_control.config import ( 5 | BoundaryConditionConfig, 6 | BoundConfig, 7 | ScaleConfig, 8 | SimConfig, 9 | ) 10 | 11 | 12 | class TestBoundConfig(TestCase): 13 | def test_stage_var_bound_construction(self): 14 | """Test both correct and incorrect constructions of BoundConfig for stage vars""" 15 | # both scalars, should work 16 | kwargs = {"group": "states", "name": "x", "values": (1.0, 2.0)} 17 | bounds = BoundConfig(**kwargs) 18 | for k, v in kwargs.items(): 19 | self.assertEqual(getattr(bounds, k), v) 20 | 21 | # both arrays, should work 22 | kwargs = { 23 | "group": "states", 24 | "name": "x", 25 | "values": (np.ones(3), np.ones(3) * 2), 26 | } 27 | bounds = BoundConfig(**kwargs) 28 | for k, v in kwargs.items(): 29 | self.assertEqual(getattr(bounds, k), v) 30 | 31 | # one scalar, one array, should fail 32 | with self.assertRaises(TypeError): 33 | bounds = BoundConfig("states", "x", (1.0, np.ones(10))) 34 | with self.assertRaises(TypeError): 35 | bounds = BoundConfig("states", "x", (np.ones(10), 1.0)) 36 | 37 | # 2d array, should fail 38 | with self.assertRaises(AssertionError): 39 | bounds = BoundConfig("states", "x", (np.ones((10, 1)), np.ones(10))) 40 | 41 | # arrays of different sizes, should fail 42 | with self.assertRaises(AssertionError): 43 | bounds = BoundConfig("states", "x", (np.ones(10), np.ones(11))) 44 | 45 | # scalar incorrect range, should fail 46 | with self.assertRaises(AssertionError): 47 | bounds = BoundConfig("states", "x", (1.0, 0.0)) 48 | 49 | # scalar incorrect range, should fail 50 | with self.assertRaises(AssertionError): 51 | bounds = BoundConfig( 52 | "states", "x", (np.array([1.0, 2.0, 3.0]), np.array([2.0, 1.0, 2.0])) 53 | ) 54 | 55 | def test_global_var_bound_construction(self): 56 | """Test both correct and incorrect constructions of BoundConfig for global vars""" 57 | # both scalars, should work 58 | kwargs = {"group": "global", "name": "mesh_scale", "values": (1.0, 2.0)} 59 | bounds = BoundConfig(**kwargs) 60 | for k, v in kwargs.items(): 61 | self.assertEqual(getattr(bounds, k), v) 62 | 63 | # both arrays, should fail 64 | with self.assertRaises(TypeError): 65 | bounds = BoundConfig( 66 | "mesh_scale", (np.array([1.0, 2.0, 3.0]), np.array([2.0, 1.0, 2.0])) 67 | ) 68 | 69 | # scalar incorrect range, should fail 70 | with self.assertRaises(AssertionError): 71 | bounds = BoundConfig("global", "mesh_scale", (1.0, 0.0)) 72 | 73 | 74 | class TestSimConfig(TestCase): 75 | def test_mutable_fields_are_not_shared(self): 76 | """Ensure that mutable fields don't use shared memory across objects""" 77 | config1 = SimConfig() 78 | config2 = SimConfig() 79 | 80 | config1.non_cyclic_vars += ["a", "b"] 81 | 82 | # Check config2 is not affected 83 | self.assertFalse(config2.non_cyclic_vars) 84 | 85 | def test_to_dict_and_back(self): 86 | """Test values remain the same when converting to dictionary and back""" 87 | sim_config = SimConfig( 88 | num_intervals=39, 89 | transcription=("LGR", {"num_stages": 3}), 90 | is_cyclic=True, 91 | non_cyclic_vars=["a", "b"], 92 | bounds=( 93 | BoundConfig("states", "x", (0.0, 2.0)), 94 | BoundConfig("global", "mesh_scale", (0.0, 2.0)), 95 | ), 96 | scales=( 97 | ScaleConfig("states", "x", 2.0), 98 | ScaleConfig("global", "mesh_scale", 10.0), 99 | ), 100 | boundary_conditions=( 101 | BoundaryConditionConfig(0, "states", "x", 1.0), 102 | BoundaryConditionConfig(0, "states", "y", 1.0), 103 | ), 104 | ) 105 | 106 | # config -> dict -> config -> dict reproduces the same dict 107 | sim_dict = sim_config.to_dict() 108 | self.assertEqual(sim_config, SimConfig(**sim_dict)) 109 | self.assertDictEqual(sim_dict, SimConfig(**sim_dict).to_dict()) 110 | -------------------------------------------------------------------------------- /lumos/optimal_control/collocation.py: -------------------------------------------------------------------------------- 1 | from enum import auto, Enum 2 | from typing import List 3 | 4 | import numpy as np 5 | from numpy.polynomial.legendre import Legendre 6 | from numpy.polynomial.polynomial import Polynomial 7 | from scipy.interpolate import lagrange 8 | 9 | 10 | class CollocationEnum(Enum): 11 | """CollocationEnum Scheme enumeration. 12 | 13 | A collcation scheme consists of two main ingredients: 14 | - the position of the collocation point 15 | - the interpolating polynomial used 16 | 17 | See: https://hal.archives-ouvertes.fr/hal-01615132/document 18 | """ 19 | 20 | LG = auto() 21 | LGR = auto() 22 | LGL = auto() 23 | 24 | 25 | # FIXME: this is really just getting the Legendere root points, not collocation points 26 | def get_collocation_points( 27 | num_points: int, scheme: CollocationEnum = CollocationEnum.LGR 28 | ) -> np.ndarray: 29 | """Retrieve 1D collocation points for a given scheme in the interval of [-1, 1] 30 | 31 | Note: Some methods include the end points of -1, 1, while some others don't. 32 | """ 33 | 34 | if not isinstance(scheme, CollocationEnum): 35 | raise TypeError( 36 | f"Expected scheme to be of type {CollocationEnum}, but got {type(scheme)} instead" 37 | ) 38 | 39 | # TODO: would this be different for different schemes? 40 | assert num_points >= 1 41 | 42 | # TODO: perhaps we should just combine everything into a Collocation Class 43 | if scheme == CollocationEnum.LGR: 44 | # root of Legendre polynomial, whrre P is a lgrange function 45 | # P_{N-1} + P_N = 0, but flipped around zero to include t=1 46 | coefficients = [0] * (num_points - 1) + [1, 1] 47 | characteristic_polynomial = Legendre(coefficients) 48 | 49 | # solve and flip 50 | collocation_points = -characteristic_polynomial.roots()[::-1] 51 | elif scheme == CollocationEnum.LG: 52 | # roots of P_N = 0 53 | coefficients = [0] * num_points + [1] 54 | characteristic_polynomial = Legendre(coefficients) 55 | collocation_points = characteristic_polynomial.roots() 56 | elif scheme == CollocationEnum.LGL: 57 | # roots of P_dot_{N-1} + [-1, 1] 58 | coefficients = [0] * (num_points - 1) + [1] 59 | characteristic_polynomial = Legendre(coefficients).deriv() 60 | collocation_points = np.append( 61 | np.insert(characteristic_polynomial.roots(), 0, -1), 1 62 | ) 63 | else: 64 | raise NotImplemented(scheme) 65 | return collocation_points 66 | 67 | 68 | def make_lagrange_polynomial(support: np.ndarray, index: int) -> Polynomial: 69 | """Create the i-th lagrange polynomial""" 70 | 71 | weights = np.zeros_like(support) 72 | weights[index] = 1 73 | 74 | # NOTE: lagrange returns coef in decending power order, which is opposite to numpy 75 | # Polynomials. 76 | coefficients = lagrange(support, weights).coef[::-1] 77 | return Polynomial(coefficients) 78 | 79 | 80 | def make_lagrange_basis(support: np.ndarray) -> List[Polynomial]: 81 | """Create a list of lagrange basis of varying order on the same support.""" 82 | 83 | return [make_lagrange_polynomial(support, index) for index in range(len(support))] 84 | 85 | 86 | def build_lagrange_differential_matrix( 87 | support: np.ndarray, evaluation_points: np.ndarray 88 | ) -> np.ndarray: 89 | """Differential matrix for computing the derivative of a lagrange polynomial using 90 | linear matrix vector multiplication""" 91 | 92 | lagrange_basis = make_lagrange_basis(support) 93 | polynomials = [ 94 | lagrange_polynomial.deriv() for lagrange_polynomial in lagrange_basis 95 | ] 96 | 97 | return np.array([p(evaluation_points) for p in polynomials]).T 98 | 99 | 100 | def build_lagrange_integration_matrix( 101 | support: np.ndarray, evaluation_points: np.ndarray 102 | ) -> np.ndarray: 103 | """Integration matrix for computing the definitive integral of a lagrange polynomial using 104 | linear matrix vector multiplication""" 105 | 106 | lagrange_basis = make_lagrange_basis(support) 107 | polynomials = [ 108 | lagrange_polynomial.integ() for lagrange_polynomial in lagrange_basis 109 | ] 110 | 111 | # NOTE: eq39, the integral polynomial is a definite integral, so need to call the 112 | # integral twice to remove the part for tau < support[0] 113 | return np.array([p(evaluation_points) - p(support[0]) for p in polynomials]).T 114 | -------------------------------------------------------------------------------- /tests/test_models/test_tracks.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from unittest import TestCase 3 | 4 | from parameterized import parameterized, parameterized_class 5 | 6 | from lumos.models.tracks import ( 7 | RaceTrack, 8 | cartesian_to_curvilinear, 9 | curvilinear_to_cartesian, 10 | ) 11 | 12 | 13 | def generate_circle(curvature: float, num_points: int = 1000): 14 | 15 | radius = 1 / curvature 16 | curvature = np.ones(num_points) * curvature 17 | s = np.pi * 2 * np.abs(radius) * np.linspace(0, 1, num_points) 18 | 19 | # polar angle 20 | theta = curvature * s 21 | 22 | # circle starts at 3 o'lock position 23 | heading = theta + np.pi / 2 24 | x = np.cos(theta) * radius 25 | y = np.sin(theta) * radius 26 | 27 | assert heading[-1] - heading[0] == np.pi * 2 * (np.sign(radius)) 28 | 29 | return {"x": x, "y": y, "s": s, "curvature": curvature, "heading": heading} 30 | 31 | 32 | class TestTrackUtils(TestCase): 33 | @parameterized.expand([(1e-3,), (-1e-3,)]) 34 | def test_cartesian_to_curvilinear(self, expected_curvature): 35 | """Test cartesian_to_curvilinear utility against a tractable track: a circle""" 36 | 37 | expected_results = generate_circle(curvature=expected_curvature) 38 | 39 | s, curvature, heading = cartesian_to_curvilinear( 40 | x=expected_results["x"], y=expected_results["y"] 41 | ) 42 | 43 | # NOTE: the accuracy is affected by how many points we have on the circle 44 | np.testing.assert_allclose(s, expected_results["s"], rtol=1e-3, atol=1e-3) 45 | # first and last couple of elements of curvature would be off due to finite 46 | # difference, so we ignore them. 47 | np.testing.assert_allclose( 48 | curvature[2:-2], 49 | expected_results["curvature"][2:-2], 50 | rtol=1e-3, 51 | atol=1e-4, 52 | ) 53 | np.testing.assert_allclose( 54 | heading, expected_results["heading"], rtol=1e-2, atol=1e-2 55 | ) 56 | 57 | @parameterized.expand([(1e-3,), (-1e-3,)]) 58 | def test_curvilinear_to_cartesian(self, expected_curvature): 59 | """Test curvilinear_to_cartesian utility against a tractable track: a circle""" 60 | expected_results = generate_circle(curvature=expected_curvature) 61 | 62 | x, y, heading = curvilinear_to_cartesian( 63 | s=expected_results["s"], 64 | curvature=expected_results["curvature"], 65 | x0=expected_results["x"][0], 66 | y0=expected_results["y"][0], 67 | heading0=expected_results["heading"][0], 68 | ) 69 | 70 | # NOTE: the accuracy is affected by how many points we have on the circle 71 | np.testing.assert_allclose(x, expected_results["x"], rtol=1e-2, atol=1e-2) 72 | np.testing.assert_allclose(y, expected_results["y"], rtol=1e-2, atol=1e-2) 73 | np.testing.assert_allclose( 74 | heading, expected_results["heading"], rtol=1e-2, atol=1e-2 75 | ) 76 | 77 | 78 | @parameterized_class( 79 | [ 80 | {"expected_curvature": 1e-3}, 81 | {"expected_curvature": -1e-3}, 82 | ] 83 | ) 84 | class TestTrack(TestCase): 85 | expected_curvature = 1e-3 86 | 87 | def _test_against_expected(self, track): 88 | # first and last couple of elements of curvature would be off due to finite 89 | # difference, so we ignore them. 90 | np.testing.assert_allclose( 91 | track.curvature_at(self.expected_results["s"][2:-2]), 92 | self.expected_results["curvature"][2:-2], 93 | rtol=1e-3, 94 | atol=1e-4, 95 | ) 96 | np.testing.assert_allclose( 97 | track.heading_at(self.expected_results["s"]), 98 | self.expected_results["heading"], 99 | rtol=1e-2, 100 | atol=1e-2, 101 | ) 102 | 103 | def setUp(self): 104 | self.expected_results = generate_circle(curvature=self.expected_curvature) 105 | 106 | def test_from_cartesian(self): 107 | track = RaceTrack.from_cartesian( 108 | x=self.expected_results["x"], y=self.expected_results["y"] 109 | ) 110 | 111 | self._test_against_expected(track) 112 | 113 | def test_from_curvilinear(self): 114 | track = RaceTrack.from_curvilinear( 115 | s=self.expected_results["s"], 116 | curvature=self.expected_results["curvature"], 117 | heading0=self.expected_results["heading"][0], 118 | ) 119 | 120 | self._test_against_expected(track) 121 | -------------------------------------------------------------------------------- /lumos/models/aero/aero.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | 3 | from lumos.models.ml import gp, mlp 4 | from lumos.models.base import model_io, Model, ModelReturn 5 | 6 | # NOTE: we just put some dummy inputs here for now. 7 | @model_io( 8 | inputs=("front_ride_height", "rear_ride_height", "yaw"), outputs=("Cx", "Cy", "Cz") 9 | ) 10 | class AeroModel(Model): 11 | # FIXME: there are a few implicit limitations: 12 | # 2) the parameter size for some model (gp, mlp), depends on the I/O size, we should 13 | # implement checks. 14 | 15 | def __init__(self, *args, **kwargs): 16 | super().__init__(*args, **kwargs) 17 | 18 | 19 | class ConstAero(AeroModel): 20 | @classmethod 21 | def get_default_params(self): 22 | return {"Cx": 0.6, "Cy": 0.0, "Cz": 1.9} 23 | 24 | def forward(self, inputs): 25 | return ModelReturn(outputs=self._params) 26 | 27 | 28 | class GPAero(AeroModel): 29 | # FIXME: we should pass these as configuration parameters for the model instance. 30 | num_points = 1024 31 | 32 | @classmethod 33 | def get_default_params(cls): 34 | # NOTE: we don't use the property methods here because special care is needed to 35 | # use property method inside a class method 36 | # see: https://stackoverflow.com/questions/128573/using-property-on-classmethods 37 | # 38 | # While we can make these class method properties for now, but in the future 39 | # we might want to make the I/O names dynamic that are only defined after the 40 | # object is instantiated (for example, passing all submodel outputs as a part of 41 | # parent model outputs) 42 | num_inputs = len(cls.get_direct_group_names("inputs")) 43 | num_outputs = len(cls.get_direct_group_names("outputs")) 44 | return { 45 | "gp_points": np.random.randn(cls.num_points, num_inputs), 46 | "alpha": np.random.randn(num_outputs, cls.num_points), 47 | } 48 | 49 | def forward(self, inputs): 50 | array_inputs = self.make_vector("inputs", **inputs) 51 | coeff = gp(array_inputs, self._params["gp_points"], self._params["alpha"]) 52 | 53 | outputs = self.make_dict( 54 | "outputs", **{n: coeff[idx] for idx, n in enumerate(self.names.outputs)} 55 | ) 56 | 57 | return ModelReturn(outputs=outputs) 58 | 59 | 60 | class MLPAero(AeroModel): 61 | num_layers = 3 62 | layer_dim = 128 63 | 64 | weights_prefix = "weights_layer" # parameter prefix for a layer's weights 65 | 66 | @classmethod 67 | def _list_to_dict(cls, weights: list): 68 | return {f"weights_layer_{i}": w for i, w in enumerate(weights)} 69 | 70 | @classmethod 71 | def _dict_to_list(cls, weights: dict): 72 | # NOTE: this requires ordered dict, which is default in python for 3.6+ 73 | # We rely on the ordered dict instead of the number in name 74 | return [v for k, v in weights.items() if k.startswith(cls.weights_prefix)] 75 | 76 | @classmethod 77 | def get_default_params(cls): 78 | num_inputs = len(cls.get_direct_group_names("inputs")) 79 | num_outputs = len(cls.get_direct_group_names("outputs")) 80 | 81 | layers_dim = [cls.layer_dim] * cls.num_layers 82 | ins = [num_inputs] + layers_dim[:-1] 83 | outs = layers_dim 84 | # NOTE: we make the default parameters small so that it doesn't affect vehicle 85 | # behaviour too much (so we can more easilyi assess runtime and convergence of 86 | # a random aero map with MLP) 87 | weights = [1e-3 * np.random.randn(od, id) for id, od in zip(ins, outs)] 88 | 89 | # Final layer 90 | weights.append(np.random.randn(num_outputs, layers_dim[-1])) 91 | 92 | weights_dict = cls._list_to_dict(weights) 93 | 94 | params = {"air_density": 1.225} 95 | params.update(weights_dict) 96 | return params 97 | 98 | def forward(self, inputs): 99 | # NOTE: at the moment we only support parameters of a flat dict where the key is 100 | # a string, and the value a scalar or an array. And since Casadi only supports 101 | # up to 2d array (matrices), this present an issue for MLP params. 102 | # 103 | # MLP params are at least 3 dimensional, we could do list of weights (if weights) 104 | # are 2d only, but that breaks the flat dict requirement on values. 105 | array_inputs = self.make_vector("inputs", **inputs) 106 | 107 | coeff = mlp(array_inputs, weights=self._dict_to_list(self._params)) 108 | outputs = self.make_dict( 109 | "outputs", **{n: coeff[idx] for idx, n in enumerate(self.names.outputs)} 110 | ) 111 | return ModelReturn(outputs=outputs) 112 | -------------------------------------------------------------------------------- /tests/test_optimal_control/test_transcription.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | import numpy as np 4 | from scipy.interpolate import lagrange 5 | 6 | from lumos.optimal_control.transcription import ( 7 | get_transcription_options, 8 | make_transcription, 9 | ForwardEuler, 10 | Trapezoidal, 11 | LGR, 12 | LGRIntegral, 13 | ) 14 | 15 | 16 | class TestMakeTranscription(unittest.TestCase): 17 | def test_make_LGR(self): 18 | transcription = make_transcription("LGR", {"num_stages": 3}) 19 | self.assertIsInstance(transcription, LGR) 20 | self.assertEqual(transcription.num_stages_per_interval, 3) 21 | 22 | def test_make_ForwardEuler(self): 23 | transcription = make_transcription("ForwardEuler") 24 | self.assertIsInstance(transcription, ForwardEuler) 25 | self.assertEqual(transcription.num_stages_per_interval, 2) 26 | 27 | def test_make_Trapezoidal(self): 28 | transcription = make_transcription("Trapezoidal") 29 | self.assertIsInstance(transcription, Trapezoidal) 30 | self.assertEqual(transcription.num_stages_per_interval, 2) 31 | 32 | def test_make_LGRIntegral(self): 33 | transcription = make_transcription("LGRIntegral", {"num_stages": 5}) 34 | self.assertIsInstance(transcription, LGRIntegral) 35 | self.assertEqual(transcription.num_stages_per_interval, 5) 36 | 37 | def test_make_something_wrong(self): 38 | # Type not supported 39 | with self.assertRaises(RuntimeError): 40 | transcription = make_transcription("LGL") 41 | 42 | # Unwanted input arguments, the constructor should raise TypeError 43 | with self.assertRaises(TypeError): 44 | transcription = make_transcription("ForwardEuler", {"someargs": 1.0}) 45 | 46 | def test_get_and_make_transcriptions(self): 47 | options = get_transcription_options() 48 | 49 | for name in options: 50 | transcription = make_transcription(name) 51 | self.assertEqual(transcription.__class__.__name__, name) 52 | 53 | 54 | class TestTranscription(unittest.TestCase): 55 | def test_LGR_continuity(self): 56 | """Test LGR differentiation and residual computations are correct""" 57 | num_stages = 5 58 | lgr = make_transcription("LGR", {"num_stages": num_stages}) 59 | A_diff, B_diff = lgr.get_continuity_matrices() 60 | x = np.random.randn(num_stages) 61 | s = lgr.interp_points 62 | poly = lagrange(s, x) 63 | 64 | # Test derivatives are correct. 65 | x_dot_expected = poly.deriv()(s[1:]) 66 | x_dot_actual = A_diff @ x 67 | np.testing.assert_array_almost_equal(x_dot_actual, x_dot_expected) 68 | 69 | # Test with the correct x and x_dot, continuity equations are satisfied 70 | # need to add an x_dot to the first point to make size match. The first point of 71 | # x_dot is not collocated, so we can add any random number.. 72 | res = A_diff @ x - B_diff @ np.insert(x_dot_expected, 0, 0) 73 | np.testing.assert_array_almost_equal(res, np.zeros_like(res)) 74 | 75 | def test_LGRIntegral_continuity(self): 76 | """Test LGRIntegral integral and residual computations are correct""" 77 | num_stages = 5 78 | lgr_int = make_transcription("LGRIntegral", {"num_stages": num_stages}) 79 | A, B = lgr_int.get_continuity_matrices() 80 | x = np.random.randn(num_stages) 81 | s = lgr_int.interp_points 82 | poly = lagrange(s, x) 83 | x_dot_actual = poly.deriv()(s) 84 | 85 | # Test integrals are correct. 86 | x_actual = x[0] + B @ x_dot_actual 87 | x_expected = x[1:] 88 | np.testing.assert_array_almost_equal(x_actual, x_expected) 89 | 90 | # Test with the correct x and x_dot, continuity equations are satisfied 91 | res = A @ x - B @ x_dot_actual 92 | np.testing.assert_array_almost_equal(res, np.zeros_like(res)) 93 | 94 | def test_LGR_invertible(self): 95 | """LGR scheme can be inverted to reconstruct states from derivatives. 96 | 97 | See equation 72-73 on Garg, Rao, et al, 2017, https://hal.archives-ouvertes.fr/hal-01615132 98 | An overview of three pseudospectral methods for the numerical solution of 99 | optimal control problems 100 | """ 101 | 102 | num_stages = 5 103 | lgr = make_transcription("LGR", {"num_stages": num_stages}) 104 | A, B = lgr.get_continuity_matrices() 105 | 106 | x = np.sin(5 * np.linspace(0, 1, num_stages)) 107 | x_dot = A @ x 108 | 109 | x_reconstructed = np.linalg.solve(A[:, 1:], x_dot) + x[0] 110 | 111 | np.testing.assert_array_almost_equal(x_reconstructed, x[1:]) 112 | -------------------------------------------------------------------------------- /examples/laptime_simulation_example.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | 4 | from lumos.models.composition import ModelMaker 5 | from lumos.models.simple_vehicle_on_track import SimpleVehicleOnTrack 6 | from lumos.models.tires.utils import create_params_from_tir_file 7 | from lumos.optimal_control.config import LoggingConfig 8 | from lumos.simulations.laptime_simulation import LaptimeSimulation 9 | 10 | logger = logging.getLogger(__name__) 11 | 12 | 13 | # FIXME: hardcoded relative track directory 14 | TRACK_DIR = "data/tracks" 15 | 16 | 17 | def main(): 18 | is_cyclic = True 19 | track = "Catalunya" 20 | track_file = os.path.join(TRACK_DIR, track + ".csv") 21 | 22 | sim_config = LaptimeSimulation.get_sim_config( 23 | num_intervals=250, 24 | hessian_approximation="exact", 25 | is_cyclic=is_cyclic, 26 | is_condensed=False, 27 | backend="casadi", 28 | track=track_file, 29 | transcription="LGR", 30 | logging_config=LoggingConfig( 31 | sim_name=track, results_dir="results", log_every_nth_iter=0 32 | ), 33 | ) 34 | 35 | model_config = SimpleVehicleOnTrack.get_recursive_default_model_config() 36 | 37 | # EXAMPLE: change tire model 38 | # model_config.replace_subtree( 39 | # "vehicle.tire", ModelMaker.make_config("PerantoniTire") 40 | # ) 41 | 42 | # EXMAPLE: change an aero model 43 | # model_config.replace_subtree("vehicle.aero", ModelMaker.make_config("MLPAero")) 44 | 45 | model = SimpleVehicleOnTrack(model_config=model_config) 46 | params = model.get_recursive_default_params() 47 | 48 | # Example of changing model parameters 49 | # TODO: an issue here is that we need to instantiate the model first to get params 50 | # but that's unavoidable because without the model, we don't even know the tree 51 | # structure of all the submodels, let alone the default parameters. 52 | # params.set_param("vehicle.vehicle_mass", 1700) 53 | 54 | # Example: change tire parameters 55 | sharpened_params = create_params_from_tir_file("data/tires/sharpened.tir") 56 | # FIXME: here we're using private methods. We should probably add a method to change 57 | # the parameters of an entire node in the ParameterTree 58 | for c in ["fl", "fr", "rl", "rr"]: 59 | submodel_path = "vehicle.tire_" + c 60 | tire_params = params._get_subtree(submodel_path) 61 | tire_params._data = sharpened_params 62 | params.replace_subtree(submodel_path, tire_params) 63 | 64 | ocp = LaptimeSimulation( 65 | model_params=params, model_config=model_config, sim_config=sim_config 66 | ) 67 | 68 | x0 = ocp.get_init_guess() 69 | ocp.profile(x0, repeat=10, hessian=True) 70 | 71 | solution, info = ocp.solve( 72 | x0, 73 | max_iter=200, 74 | print_level=5, 75 | print_timing_statistics="yes", 76 | print_info_string="yes", 77 | derivative_test="none", 78 | dual_inf_tol=1e-3, 79 | constr_viol_tol=1e-3, 80 | ) 81 | total_time = ocp.dec_var_operator.get_var( 82 | solution, group="states", name="time", stage=-1 83 | ) 84 | logger.info(f"Maneuver time {total_time:.3f} sec") 85 | 86 | return 87 | 88 | # Change param and solve again 89 | ocp.modify_model_param("vehicle.vehicle_mass", 2100.0) 90 | solution, info = ocp.solve( 91 | solution, 92 | max_iter=200, 93 | print_level=5, 94 | print_timing_statistics="yes", 95 | print_info_string="yes", 96 | derivative_test="none", 97 | dual_inf_tol=1e-3, 98 | constr_viol_tol=1e-3, 99 | ) 100 | total_time = ocp.dec_var_operator.get_var( 101 | solution, group="states", name="time", stage=-1 102 | ) 103 | logger.info(f"Maneuver time {total_time:.3f} sec") 104 | 105 | # # Tighten slip bound, solve again 106 | # for c in ["fl", "fr", "rl", "rr"]: 107 | # for s in range(ocp.num_stages): 108 | # ocp.set_con_output_bounds("slip_ratio_" + c, stage=s, bounds=(-0.1, 0.1)) 109 | # ocp.set_con_output_bounds( 110 | # "slip_angle_" + c, stage=s, bounds=(-np.deg2rad(10), np.deg2rad(10)) 111 | # ) 112 | 113 | # solution, info = ocp.solve( 114 | # solution, 115 | # lagrange=info["mult_g"], 116 | # zl=info["mult_x_L"], 117 | # zu=info["mult_x_U"], 118 | # max_iter=200, 119 | # print_level=5, 120 | # print_timing_statistics="yes", 121 | # print_info_string="yes", 122 | # derivative_test="none", 123 | # dual_inf_tol=1e-3, 124 | # constr_viol_tol=1e-3, 125 | # ) 126 | # total_time = ocp.dec_var_operator.get_var( 127 | # solution, group="states", name="time", stage=-1 128 | # ) 129 | 130 | # logger.info(f"Maneuver time {total_time:.3f} sec") 131 | 132 | 133 | if __name__ == "__main__": 134 | logging.basicConfig(level=logging.INFO) 135 | main() 136 | -------------------------------------------------------------------------------- /tests/test_numpy/test_casadi_numpy/test_all_operations_run.py: -------------------------------------------------------------------------------- 1 | import casadi as cas 2 | import pytest 3 | 4 | import lumos.numpy.casadi_numpy as np 5 | 6 | ### Cause all NumPy warnings to raise exceptions, to make this bulletproof 7 | np.seterr(all="raise") 8 | 9 | 10 | @pytest.fixture 11 | def types(): 12 | ### NumPy data types 13 | scalar_np = np.array(1) 14 | vector_np = np.array([1, 1]) 15 | matrix_np = np.ones((2, 2)) 16 | 17 | ### CasADi data types 18 | scalar_cas = cas.SX.sym("cas_scalar", 1) 19 | vector_cas = cas.SX.sym("cas_vector", 2) 20 | 21 | ### Dynamically-typed data type creation (i.e. type depends on inputs) 22 | vector_dynamic = np.array( 23 | [scalar_cas, scalar_cas] 24 | ) # vector as a dynamic-typed array 25 | matrix_dynamic = np.array( 26 | [ # matrix as an dynamic-typed array 27 | [scalar_cas, scalar_cas], 28 | [scalar_cas, scalar_cas], 29 | ] 30 | ) 31 | 32 | ### Create lists of possible variable types for scalars, vectors, and matrices. 33 | scalar_options = [scalar_cas, scalar_np] 34 | vector_options = [vector_cas, vector_np, vector_dynamic] 35 | matrix_options = [matrix_np, matrix_dynamic] 36 | 37 | return { 38 | "scalar": scalar_options, 39 | "vector": vector_options, 40 | "matrix": matrix_options, 41 | "all": scalar_options + vector_options + matrix_options, 42 | } 43 | 44 | 45 | def test_indexing_1D(types): 46 | for x in types["vector"] + types["matrix"]: 47 | x[0] # The first element of x 48 | x[-1] # The last element of x 49 | x[1:] # All but the first element of x 50 | x[:-1] # All but the last element of x 51 | x[::2] # Every other element of x 52 | x[::-1] # The elements of x, but in reversed order 53 | 54 | 55 | def test_indexing_2D(types): 56 | for x in types["matrix"]: 57 | x[0, :] # The first row of x 58 | x[:, 0] # The first column of x 59 | x[1:, :] # All but the first row of x 60 | 61 | 62 | def test_basic_math(types): 63 | for x in types["all"]: 64 | for y in types["all"]: 65 | ### Arithmetic 66 | x + y 67 | x - y 68 | x * y 69 | x / y 70 | np.sum(x) # Sum of all entries of array-like object x 71 | 72 | ### Exponentials & Powers 73 | x ** y 74 | np.power(x, y) 75 | np.exp(x) 76 | np.log(x) 77 | np.log10(x) 78 | np.sqrt(x) # Note: do x ** 0.5 rather than np.sqrt(x). 79 | 80 | ### Trig 81 | np.sin(x) 82 | np.cos(x) 83 | np.tan(x) 84 | np.arcsin(x) 85 | np.arccos(x) 86 | np.arctan(x) 87 | np.arctan2(y, x) 88 | np.sinh(x) 89 | np.cosh(x) 90 | np.tanh(x) 91 | np.arcsinh(x) 92 | np.arccosh(x) 93 | np.arctanh(x - 0.5) # `- 0.5` to give valid argument 94 | 95 | 96 | def test_logic(types): 97 | for option_set in [ 98 | types["scalar"], 99 | types["vector"], 100 | types["matrix"], 101 | ]: 102 | for x in option_set: 103 | for y in option_set: 104 | ### Comparisons 105 | """ 106 | Note: if warnings appear here, they're from `np.array(1) == cas.MX(1)` - 107 | sensitive to order, as `cas.MX(1) == np.array(1)` is fine. 108 | 109 | However, checking the outputs, these seem to be yielding correct results despite 110 | the warning sooo... 111 | """ 112 | x == y # Warnings coming from here 113 | x != y # Warnings coming from here 114 | x > y 115 | x >= y 116 | x < y 117 | x <= y 118 | 119 | ### Conditionals 120 | np.where(x > 1, x ** 2, 0) 121 | 122 | ### Elementwise min/max 123 | np.fmax(x, y) 124 | np.fmin(x, y) 125 | 126 | for x in types["all"]: 127 | np.fabs(x) 128 | np.floor(x) 129 | np.ceil(x) 130 | np.clip(x, 0, 1) 131 | 132 | 133 | def test_vector_math(types): 134 | for x in types["vector"]: 135 | for y in types["vector"]: 136 | x.T 137 | np.linalg.inner(x, y) 138 | np.linalg.outer(x, y) 139 | np.linalg.norm(x) 140 | 141 | 142 | def test_matrix_math(types): 143 | for v in types["vector"]: 144 | for m in types["matrix"]: 145 | m = m + np.eye(2) # Regularize the matrix so it's not singular 146 | 147 | m.T 148 | m @ v 149 | for m_other in types["matrix"]: 150 | m @ m_other 151 | np.linalg.solve(m, v) 152 | 153 | 154 | def test_spacing(types): 155 | for x in types["scalar"]: 156 | np.linspace(x - 1, x + 1, 10) 157 | np.cosspace(x - 1, x + 1, 10) 158 | 159 | 160 | def test_rotation_matrices(types): 161 | for angle in types["scalar"]: 162 | for axis in types["vector"]: 163 | np.rotation_matrix_2D(angle) 164 | np.rotation_matrix_3D(angle, np.array([axis[0], axis[1], axis[0]])) 165 | 166 | 167 | if __name__ == "__main__": 168 | pytest.main() 169 | -------------------------------------------------------------------------------- /tests/test_models/test_kinematics.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import unittest 3 | 4 | import numpy as np 5 | 6 | from lumos.models.kinematics import TrackPosition2D 7 | from lumos.models.test_utils import BaseStateSpaceModelTest 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | 12 | class TestTrackPosition2D(BaseStateSpaceModelTest, unittest.TestCase): 13 | ModelClass: type = TrackPosition2D 14 | 15 | def _build_track_and_drive( 16 | self, 17 | curvature: float, 18 | start_heading: float, 19 | speed: float, 20 | yaw_rate: float = 0.0, 21 | distance_step: float = 1.0, 22 | num_steps: int = 20, 23 | ): 24 | 25 | init_states = self.model.make_dict(group="states", n=0.0, time=0.0, eta=0.0) 26 | states = init_states 27 | for step in range(num_steps): 28 | 29 | track_heading = start_heading + step * distance_step * curvature 30 | inputs = self.model.make_dict( 31 | group="inputs", 32 | vx=speed, 33 | vy=0.0, 34 | yaw_rate=yaw_rate, 35 | track_curvature=curvature, 36 | track_heading=track_heading, 37 | ) 38 | model_return = self.model.forward(states=states, inputs=inputs) 39 | 40 | for k in states: 41 | states[k] += model_return.states_dot[k] * distance_step 42 | 43 | # Call the model again on the final states to ensure the outputs and states are 44 | # sync'd 45 | inputs = self.model.make_dict( 46 | group="inputs", 47 | vx=speed, 48 | vy=0.0, 49 | yaw_rate=yaw_rate, 50 | track_curvature=curvature, 51 | track_heading=start_heading + num_steps * distance_step * curvature, 52 | ) 53 | 54 | model_return = self.model.forward(states=states, inputs=inputs) 55 | 56 | return states, model_return.outputs 57 | 58 | def test_drive_straight_on_straight_track(self): 59 | """WHEN speed is in x-direction on a straightline track 60 | THEN distance across should stay at zero 61 | """ 62 | 63 | speed = 5.0 64 | num_steps = 20 65 | distance_step = 0.1 66 | final_states, final_outputs = self._build_track_and_drive( 67 | curvature=0.0, 68 | start_heading=0.0, 69 | speed=speed, 70 | distance_step=distance_step, 71 | num_steps=num_steps, 72 | ) 73 | 74 | # 0-distance to centerline 75 | self.assertAlmostEqual(final_states["n"], 0) 76 | # distance travelled is correct 77 | self.assertAlmostEqual(final_states["time"], num_steps * distance_step / speed) 78 | 79 | def test_drive_curve_on_straight_track(self): 80 | """WHEN speed is in x-direction on a straightline track 81 | AND yaw rate is +ve 82 | THEN distance across should be positive (deviate to left) 83 | THEN yaw angle should be consistent 84 | """ 85 | 86 | # Sign convention 87 | # curvature +ve to the left (along distance) 88 | # distance_across +ve to the left of centerline 89 | # vehicle yaw +ve to the left as seen by the driver 90 | speed = 5.0 91 | yaw_rate = np.deg2rad(10) 92 | num_steps = 20 93 | distance_step = 0.1 94 | final_states, final_outputs = self._build_track_and_drive( 95 | curvature=0.0, 96 | start_heading=0.0, 97 | speed=speed, 98 | yaw_rate=yaw_rate, 99 | distance_step=distance_step, 100 | num_steps=num_steps, 101 | ) 102 | # centerline distance greater than 0 103 | # TODO: assert greater than 0 isn't that of a great test... Can we be more precise?) 104 | self.assertGreater(final_states["n"], 0) 105 | 106 | # time spent would be larger than if driving on straightl line 107 | self.assertGreater(final_states["time"], num_steps * distance_step / speed) 108 | 109 | # assert heading is correct 110 | self.assertAlmostEqual( 111 | final_outputs["yaw_angle"], 112 | num_steps * distance_step / speed * yaw_rate, 113 | places=3, 114 | ) 115 | 116 | def test_drive_curve_on_curve_track(self): 117 | """WHEN speed is in x-direction on a curved track 118 | AND yaw rate is +ve 119 | AND yaw rate and speed matches track curvature 120 | THEN should follow centerline 121 | """ 122 | 123 | # Sign convention 124 | # curvature +ve to the left (along distance) 125 | # distance_across +ve to the left of centerline 126 | # vehicle yaw +ve to the left as seen by the driver 127 | speed = 5.0 128 | yaw_rate = np.deg2rad(10) 129 | curvature = yaw_rate / speed 130 | num_steps = 20 131 | distance_step = 0.1 132 | final_states, final_outputs = self._build_track_and_drive( 133 | curvature=curvature, 134 | start_heading=0.0, 135 | speed=speed, 136 | yaw_rate=yaw_rate, 137 | distance_step=distance_step, 138 | num_steps=num_steps, 139 | ) 140 | 141 | # We should stay on centerline 142 | self.assertAlmostEqual(final_states["n"], 0) 143 | 144 | # time spent would be the same as driving on straightl line 145 | self.assertAlmostEqual(final_states["time"], num_steps * distance_step / speed) 146 | 147 | # assert heading should be the same as track heading 148 | self.assertAlmostEqual(final_states["eta"], 0) 149 | 150 | 151 | if __name__ == "__name__": 152 | unittest.main() 153 | -------------------------------------------------------------------------------- /tests/test_models/test_tires/test_perantoni_tire.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import unittest 3 | from typing import Dict 4 | 5 | import numpy as np 6 | 7 | from lumos.models.tires.perantoni import PerantoniTire 8 | from lumos.models.test_utils import BaseModelTest 9 | 10 | logger = logging.getLogger(__name__) 11 | 12 | 13 | class TestPerantoniTire(BaseModelTest, unittest.TestCase): 14 | ModelClass: type = PerantoniTire 15 | 16 | def _get_fx_from_kappa(self, kappa: float, input_dict: Dict[str, float]) -> float: 17 | """Helper to perform longitudinal sweep""" 18 | input_dict["kappa"] = kappa 19 | outputs, *_ = self.model.forward(inputs=input_dict) 20 | 21 | return outputs["Fx"] 22 | 23 | def _get_fy_from_alpha(self, alpha: float, input_dict: Dict[str, float]) -> float: 24 | input_dict["alpha"] = alpha 25 | outputs, *_ = self.model.forward(inputs=input_dict) 26 | 27 | return outputs["Fy"] 28 | 29 | def test_pure_longitudinal(self): 30 | """Test characteristics in pure longitudinal load 31 | 32 | WHEN: load is pure longitudinal 33 | THEN: zero force at zero slip 34 | and force is monotonically increasing between the peaks 35 | """ 36 | vx = 10.0 37 | input_dict = { 38 | "Fz": 2000.0, 39 | "kappa": 0.0, 40 | "alpha": 0.0, 41 | "vx": vx, 42 | "gamma": 0.0, 43 | } 44 | 45 | # No force at zero slip 46 | fx0 = self._get_fx_from_kappa(0.0, input_dict=input_dict) 47 | self.assertAlmostEqual(fx0, 0, msg="Longitudinal Force non-zero at 0 slip") 48 | 49 | # Longitudianl force and slip are the same sign 50 | # Ensure the force is monotonically increasing between the -ve and +ve peaks 51 | kappa_sweep = np.linspace(-0.3, 0.3, 101) 52 | fx_sweep = np.array( 53 | [self._get_fx_from_kappa(k, input_dict=input_dict) for k in kappa_sweep] 54 | ) 55 | idx_min = np.argmin(fx_sweep) 56 | idx_max = np.argmax(fx_sweep) 57 | fx_should_be_monotonic = fx_sweep[idx_min : idx_max + 1] 58 | self.assertGreater(len(fx_should_be_monotonic), 0) 59 | self.assertTrue( 60 | np.all(np.diff(fx_should_be_monotonic) > 0), 61 | msg="Longitudinal force should be monotonically increasing between the two peaks", 62 | ) 63 | 64 | def test_pure_lateral(self): 65 | """Test characteristics in pure lateral load 66 | 67 | WHEN: load is pure lateral 68 | THEN: zero force at zero slip 69 | and force is monotonically increasing between the peaks 70 | """ 71 | vx = 10.0 72 | input_dict = { 73 | "Fz": 2000.0, 74 | "kappa": 0.0, 75 | "alpha": 0.0, 76 | "vx": vx, 77 | "gamma": 0.0, 78 | } 79 | 80 | # No force at zero slip 81 | fy0 = self._get_fy_from_alpha(0.0, input_dict=input_dict) 82 | self.assertAlmostEqual(fy0, 0, msg="Lateral Force non-zero at 0 slip") 83 | 84 | # lateral slip and force are opposite signs 85 | # Ensure the force is monotonically decreasing between the -ve and +ve peaks 86 | alpha_sweep = np.linspace(-0.3, 0.3, 101) 87 | fy_sweep = np.array( 88 | [self._get_fy_from_alpha(a, input_dict=input_dict) for a in alpha_sweep] 89 | ) 90 | idx_min = np.argmin(fy_sweep) 91 | idx_max = np.argmax(fy_sweep) 92 | fy_should_be_monotonic = fy_sweep[idx_max : idx_min + 1] 93 | # Must also ensure the array is not empty! 94 | self.assertGreater(len(fy_should_be_monotonic), 0) 95 | self.assertTrue( 96 | np.all(np.diff(fy_should_be_monotonic) < 0), 97 | msg="Lateral force should be monotonically decreasing between the two peaks", 98 | ) 99 | 100 | def test_combined(self): 101 | """Test combined load characteristics 102 | WHEN combined slip is introduced 103 | THEN longitudinal force in combined is smaller than in pure 104 | and lateral force in combined is smaller than in pure 105 | """ 106 | vx = 10.0 107 | base_input_dict = { 108 | "Fz": 2000.0, 109 | "kappa": 0.0, 110 | "alpha": 0.0, 111 | "vx": vx, 112 | "gamma": 0.0, 113 | } 114 | 115 | # check pure longitudinal load is larger than that in combined. 116 | kappa = 0.02 117 | alpha_sweep = np.linspace(-0.3, 0.3, 20) 118 | input_dict = base_input_dict.copy() 119 | fx_in_pure = self._get_fx_from_kappa(kappa, input_dict) 120 | fx_in_combined = np.array( 121 | [ 122 | self._get_fx_from_kappa(kappa, dict(input_dict, alpha=alpha)) 123 | for alpha in alpha_sweep 124 | ] 125 | ) 126 | self.assertTrue( 127 | np.all((fx_in_pure - fx_in_combined) >= 0), 128 | "Longitudinal load should be greater in pure longitudinal load than combined.", 129 | ) 130 | 131 | # check pure lateral load magnitude is larger than that in combined. 132 | # For positive slip, this means the force is smaller (more negative) 133 | alpha = 0.02 134 | kappa_sweep = np.linspace(-0.3, 0.3, 20) 135 | input_dict = base_input_dict.copy() 136 | fy_in_pure = self._get_fy_from_alpha(alpha, input_dict) 137 | fy_in_combined = np.array( 138 | [ 139 | self._get_fy_from_alpha(alpha, dict(input_dict, kappa=kappa)) 140 | for kappa in kappa_sweep 141 | ] 142 | ) 143 | self.assertTrue( 144 | np.all((fy_in_pure - fy_in_combined) <= 0), 145 | "Lateral load should be greater in pure lateral load than combined.", 146 | ) 147 | 148 | 149 | if __name__ == "__name__": 150 | unittest.main() 151 | -------------------------------------------------------------------------------- /lumos/optimal_control/convolution.py: -------------------------------------------------------------------------------- 1 | from typing import Any, Callable, Dict 2 | 3 | import numpy as np 4 | 5 | from lumos.optimal_control.nlp import MappedConstraints, CasMappedConstraints 6 | from lumos.optimal_control.utils import ( 7 | DecVarOperator, 8 | batch_conv1d, 9 | create_offset_structure, 10 | ) 11 | from lumos.models.tree_utils import ParameterTree 12 | 13 | 14 | class ConvConstraints(MappedConstraints): 15 | def __init__( 16 | self, 17 | unit_problem: MappedConstraints, 18 | dec_var_op: DecVarOperator, 19 | normalized_mesh: np.ndarray, 20 | mesh_scale_fn: Callable, 21 | params: Dict[str, Any], 22 | ): 23 | """An NLP problem formulated by convoluting some unit problems over the decision variable. 24 | 25 | To parallelize the convolution is done in a batched fashion, by transforming the 26 | inputs into a batched array, which is then evaluated by all the NLP functions in 27 | batches. 28 | 29 | The width of the convolution is defined by the unit_probem's number of inputs. 30 | 31 | Args: 32 | unit_problem (UnitProblem): The unit_problem that acts as a convolution kernel 33 | normalized mesh: the normalized mesh [0, 1] that is needed for the unit 34 | problem, with lead dimension = batch. 35 | mesh_scale_fn: a function that takes the decision variable and returns the 36 | mesh scale. 37 | 38 | FIXME: Currently for non-fixed grid problems, the derivative w.r.t. mesh is not 39 | taken into account! This means that it only works for 'time-invariant' ODE/DAE 40 | at the moment. 41 | 42 | NOTE: we could create set_con_scales method here that just repeats the scales of 43 | the unit problem. However since that would require the unit problem to be set 44 | already at the construction of the ConvProblem, it is tricky to do. 45 | """ 46 | 47 | self._unit_problem = unit_problem 48 | self._normalized_mesh = normalized_mesh 49 | self._mesh_scale_fn = mesh_scale_fn 50 | self._op = dec_var_op 51 | self._batch = dec_var_op.num_stages 52 | self._stride = dec_var_op.num_var_stage 53 | 54 | self._width = unit_problem.num_in 55 | self.set_params(params) 56 | 57 | def _constraints(x): 58 | transformed_vars = self._transform_inputs(x) 59 | return np.ravel( 60 | unit_problem.mapped_constraints( 61 | transformed_vars, self._get_mesh(x), self._params 62 | ) 63 | ) 64 | 65 | def _jacobian(x): 66 | transformed_vars = self._transform_inputs(x) 67 | return np.ravel( 68 | unit_problem.mapped_jacobian( 69 | transformed_vars, self._get_mesh(x), self._params 70 | ) 71 | ) 72 | 73 | def _hessian(x, lagrange): 74 | transformed_vars = self._transform_inputs(x) 75 | 76 | # Reshape lagrange to get ready for vmap 77 | shape = (self._batch, unit_problem.num_con) 78 | if np.prod(shape) != len(lagrange): 79 | raise ValueError( 80 | f"wrong length of lagrangian. Expect {np.prod(shape)}, " 81 | f"but got {len(lagrange)}" 82 | ) 83 | lagrange = np.reshape(lagrange, shape) 84 | 85 | hess = unit_problem.mapped_hessian( 86 | transformed_vars, lagrange, self._get_mesh(x), self._params 87 | ) 88 | 89 | return np.ravel(hess) 90 | 91 | jac_struct, hess_struct = self._build_jac_and_hess_struct(unit_problem) 92 | 93 | super().__init__( 94 | num_in=dec_var_op.num_dec, 95 | num_con=self._batch * unit_problem.num_con, 96 | constraints=_constraints, 97 | jacobian=_jacobian, 98 | hessian=_hessian, 99 | jacobian_structure=jac_struct, 100 | hessian_structure=hess_struct, 101 | ) 102 | 103 | def set_params(self, params): 104 | # Casadi functions requires flat array as inputs. 105 | # TODO: perhaps we should move this into another wrapper on CasMappedConstraints and 106 | # leave ConvConstraints backend-agnostic. The advantage of doing it like now is 107 | # to 1) use minimum code in CasMappedConstraints and 2) only run this parameter 108 | # flattening once, instead of everytime a function is called. 109 | if isinstance(self._unit_problem, CasMappedConstraints): 110 | self._params, _ = params.tree_ravel() 111 | else: 112 | self._params = params 113 | 114 | def _transform_inputs(self, x): 115 | stage_vars, _ = self._op.split_stage_and_global_vars(x) 116 | 117 | return batch_conv1d(stage_vars, width=self._width, stride=self._stride,) 118 | 119 | def _get_mesh(self, x): 120 | return self._normalized_mesh * self._mesh_scale_fn(x) 121 | 122 | def _build_jac_and_hess_struct(self, unit_problem): 123 | 124 | # Jacobian 125 | unit_rows, unit_cols = unit_problem.jacobianstructure() 126 | 127 | jac_struct = create_offset_structure( 128 | base_rows=np.ravel(unit_rows), 129 | base_cols=np.ravel(unit_cols), 130 | row_offset=unit_problem.num_con, 131 | col_offset=self._stride, 132 | num_blocks=self._batch, 133 | ) 134 | 135 | # Hessian 136 | unit_rows, unit_cols = unit_problem.hessianstructure() 137 | 138 | hess_struct = create_offset_structure( 139 | base_rows=np.ravel(unit_rows), 140 | base_cols=np.ravel(unit_cols), 141 | row_offset=self._stride, 142 | col_offset=self._stride, 143 | num_blocks=self._batch, 144 | ) 145 | 146 | return jac_struct, hess_struct 147 | -------------------------------------------------------------------------------- /tests/test_optimal_control/test_collocation.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import unittest 3 | 4 | import numpy as np 5 | from numpy.testing import assert_allclose 6 | from parameterized import parameterized 7 | 8 | from lumos.optimal_control.collocation import ( 9 | build_lagrange_differential_matrix, 10 | CollocationEnum, 11 | build_lagrange_integration_matrix, 12 | get_collocation_points, 13 | ) 14 | 15 | 16 | class LG: 17 | """Hand-coded helper Lengendre polynomials. 18 | 19 | see: https://en.wikipedia.org/wiki/Legendre_polynomials 20 | """ 21 | 22 | def p8(x): 23 | return ( 24 | 1 25 | / 128 26 | * (6435 * x ** 9 - 12012 * x ** 6 + 6930 * x ** 4 - 1260 * x ** 2 + 35) 27 | ) 28 | 29 | def p9(x): 30 | return ( 31 | 1 32 | / 128 33 | * ( 34 | 12155 * x ** 9 35 | - 25740 * x ** 7 36 | + 18018 * x ** 5 37 | - 4620 * x ** 3 38 | + 315 * x 39 | ) 40 | ) 41 | 42 | def p9_dot(x): 43 | return ( 44 | 1 45 | / 128 46 | * ( 47 | 9 * 12155 * x ** 8 48 | - 7 * 25740 * x ** 6 49 | + 5 * 18018 * x ** 4 50 | - 3 * 4620 * x ** 2 51 | + 315 52 | ) 53 | ) 54 | 55 | def p10(x): 56 | return ( 57 | 1 58 | / 256 59 | * ( 60 | 46189 * x ** 10 61 | - 109395 * x ** 8 62 | + 90090 * x ** 6 63 | - 30030 * x ** 4 64 | + 3465 * x ** 2 65 | - 63 66 | ) 67 | ) 68 | 69 | def p10_dot(x): 70 | return ( 71 | 1 72 | / 256 73 | * ( 74 | 10 * 46189 * x ** 9 75 | - 8 * 109395 * x ** 7 76 | + 6 * 90090 * x ** 5 77 | - 4 * 30030 * x ** 3 78 | + 2 * 3465 * x 79 | ) 80 | ) 81 | 82 | 83 | class TestCollocation(unittest.TestCase): 84 | @parameterized.expand(tuple(itertools.product(CollocationEnum, [2, 10]))) 85 | def test_get_collocation_points(self, collocation_method, num_collocation_points): 86 | """Test collocation points by: 87 | - check against known roots of charateerization polynomials 88 | - comparing to hard-coded ground truth. 89 | """ 90 | collocation_points = get_collocation_points( 91 | num_collocation_points, collocation_method 92 | ) 93 | 94 | # Check correct size and correctly ordered 95 | assert len(set(collocation_points)) == num_collocation_points 96 | assert np.all(collocation_points == sorted(collocation_points)) 97 | 98 | if collocation_method == CollocationEnum.LGR: 99 | if num_collocation_points == 2: 100 | # This characteristic polynomial is a quadratic, we could solve by hand. 101 | # Q = 3/2*x^2 + x - 0.5 = 0 ==> x = -1 or 1/3 102 | # then flip to -1/3 and 1 103 | assert_allclose(collocation_points, [-1 / 3, 1]) 104 | elif num_collocation_points == 10: 105 | residuals = [LG.p10(-x) + LG.p9(-x) for x in collocation_points] 106 | assert_allclose(residuals, np.zeros_like(residuals), atol=1e-9) 107 | elif collocation_method == CollocationEnum.LG: 108 | # Q = 3/2*x^2 - 0.5 = 0 ==> solve by hand 109 | if num_collocation_points == 2: 110 | sol = np.sqrt(4 * 0.5 * 3 / 2) / 3 111 | assert_allclose(collocation_points, [-sol, sol]) 112 | elif num_collocation_points == 10: 113 | residuals = [LG.p10(x) for x in collocation_points] 114 | assert_allclose(residuals, np.zeros_like(residuals), atol=1e-9) 115 | elif collocation_method == CollocationEnum.LGL: 116 | if num_collocation_points == 2: 117 | assert_allclose(collocation_points, [-1, 1]) 118 | elif num_collocation_points == 10: 119 | # First and last points are included in LGL 120 | self.assertAlmostEqual(collocation_points[0], -1) 121 | self.assertAlmostEqual(collocation_points[-1], 1) 122 | 123 | residuals = [LG.p9_dot(x) for x in collocation_points[1:-1]] 124 | assert_allclose(residuals, np.zeros_like(residuals), atol=1e-9) 125 | else: 126 | raise NotImplementedError(collocation_method) 127 | 128 | @unittest.skip("To be implemented") 129 | def test_lagrange_interpolation(self): 130 | "Test if the interpolation goes through the points exactly." 131 | 132 | # TODO: to be implemented 133 | raise NotImplementedError 134 | 135 | def test_build_lagrange_differential_matrix(self): 136 | """Test differential matrix against manually derived ground truth.""" 137 | 138 | support = np.array([0.0, 2.0, 5]) 139 | evaluation_points = np.array([2.0, 5]) 140 | differential_matrix = build_lagrange_differential_matrix( 141 | support, evaluation_points 142 | ) 143 | 144 | # Manually derived derivatives 145 | derivatives = [ 146 | lambda x: -7 / 10 + 1 / 5 * x, 147 | lambda x: 5 / 6 - 1 / 3 * x, 148 | lambda x: -2 / 15 + 2 / 15 * x, 149 | ] 150 | 151 | expected = np.array([[f(x) for f in derivatives] for x in evaluation_points]) 152 | assert_allclose(differential_matrix, expected) 153 | 154 | def test_build_lagrange_build_lagrange_integration_matrix(self): 155 | """Test integration matrix against manually derived ground truth.""" 156 | 157 | support = np.array([0.0, 2.0, 5]) 158 | evaluation_points = np.array([2.0, 5]) 159 | integration_matrix = build_lagrange_integration_matrix( 160 | support, evaluation_points 161 | ) 162 | 163 | # Manually derived integrals 164 | integrals = [ 165 | lambda x: x - 7 / 20 * x ** 2 + 1 / 30 * x ** 3, 166 | lambda x: 5 / 12 * x ** 2 - 1 / 18 * x ** 3, 167 | lambda x: -1 / 15 * x ** 2 + 1 / 45 * x ** 3, 168 | ] 169 | 170 | expected = np.array([[f(x) for f in integrals] for x in evaluation_points]) 171 | assert_allclose(integration_matrix, expected) 172 | -------------------------------------------------------------------------------- /tests/test_models/test_simple_vehicle.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import unittest 3 | 4 | import numpy as np 5 | 6 | from lumos.models.test_utils import BaseStateSpaceModelTest, use_backends 7 | from lumos.models.vehicles.simple_vehicle import SimpleVehicle 8 | 9 | 10 | class TestSimpleVehicle(BaseStateSpaceModelTest, unittest.TestCase): 11 | ModelClass: type = SimpleVehicle 12 | 13 | def _get_initial_states(self, vx=40.0): 14 | vy = 0.0 15 | yaw_rate = 0.0 16 | no_slip_omega = vx / self.model._params["rolling_radius"] 17 | 18 | return dict( 19 | vx=vx, 20 | vy=vy, 21 | yaw_rate=yaw_rate, 22 | wheel_speed_fl=no_slip_omega, 23 | wheel_speed_fr=no_slip_omega, 24 | wheel_speed_rl=no_slip_omega, 25 | wheel_speed_rr=no_slip_omega, 26 | ) 27 | 28 | # TODO: there seems to be some common patterns which we might be able to abstract 29 | # (but maybe it's also not worth it) 30 | # 1) run model at a given point 31 | # 2) assert something 32 | # 3) solve IVP with a fixed input 33 | # 4) assert some other stuff 34 | 35 | # NOTE: as the model is implicit now, we set the algebraic variables ax and ay to 0, 36 | # which should just remove the load transfer effect 37 | 38 | def test_accel_straight(self): 39 | init_states = self._get_initial_states() 40 | inputs = dict(throttle=0.6, brake=0, steer=0, ax=0, ay=0) 41 | 42 | model_return = self.model.forward(init_states, inputs) 43 | 44 | # Rear wheel should accelerate 45 | self.assertGreater(model_return.states_dot["wheel_speed_rl"], 0) 46 | self.assertGreater(model_return.states_dot["wheel_speed_rr"], 0) 47 | 48 | # After a few timesteps, the vehicle should be faster and still driving straight 49 | num_steps = 21 50 | time_step = 0.05 51 | 52 | # Careful, numpy arrays are mutable! 53 | states, outputs = self._forward_euler(init_states, inputs, time_step, num_steps) 54 | 55 | # Speed should have increased 56 | # TODO: should we put a threshold on 'sufficient increase'? 57 | self.assertGreater(states["vx"], init_states["vx"]) 58 | 59 | # lateral speed and yaw rate should stay at 0 60 | self.assertAlmostEqual(states["vy"], 0) 61 | self.assertAlmostEqual(states["yaw_rate"], 0) 62 | 63 | def test_brake_straight(self): 64 | init_states = self._get_initial_states() 65 | inputs = dict(throttle=0.0, brake=0.1, steer=0, ax=0, ay=0) 66 | 67 | model_return = self.model.forward(init_states, inputs) 68 | 69 | # All wheels should decelerate 70 | for c in ("fl", "fr", "rl", "rr"): 71 | self.assertLess(model_return.states_dot[f"wheel_speed_{c}"], 0) 72 | 73 | # After a few timesteps, the vehicle should be slower and still driving straight 74 | num_steps = 21 75 | time_step = 0.05 76 | 77 | states, outputs = self._forward_euler(init_states, inputs, time_step, num_steps) 78 | 79 | # Speed should have increased 80 | # TODO: should we put a threshold on 'sufficient increase'? 81 | self.assertLess(states["vx"], init_states["vx"]) 82 | 83 | # lateral speed and yaw rate should stay at 0 84 | self.assertAlmostEqual(states["vy"], 0) 85 | self.assertAlmostEqual(states["yaw_rate"], 0) 86 | 87 | def test_turn_left(self): 88 | init_states = self._get_initial_states() 89 | inputs = dict(throttle=0.0, brake=0, steer=0.2, ax=0, ay=0) 90 | 91 | model_return = self.model.forward(init_states, inputs) 92 | 93 | # should directly get +ve lateral acceleration, +ve yaw acceleration and derivative of vy 94 | self.assertGreater(model_return.states_dot["vy"], 0) 95 | self.assertGreater(model_return.states_dot["yaw_rate"], 0) 96 | self.assertGreater(model_return.outputs["ay"], 0) 97 | 98 | # After a few timesteps: 99 | # The speed should decelerate due to cornering resistance 100 | # the vehicle should have a +ve yaw rate 101 | # The vehicle should have -ve vy (excpet for very low speed, the vehicle slips 102 | # to the inside of the turn because tire slip is heavily influenced by yaw rate) 103 | num_steps = 21 104 | time_step = 0.05 105 | 106 | states, outputs = self._forward_euler(init_states, inputs, time_step, num_steps) 107 | 108 | self.assertLess(states["vx"], init_states["vx"]) 109 | self.assertLess(states["vy"], 0) 110 | self.assertGreater(states["yaw_rate"], 0) 111 | 112 | # TODO: we could also check that the vehicle gets into kind of a steady-state turn 113 | 114 | def test_turn_right(self): 115 | init_states = self._get_initial_states() 116 | inputs = dict(throttle=0.0, brake=0, steer=-0.2, ax=0, ay=0) 117 | 118 | model_return = self.model.forward(init_states, inputs) 119 | 120 | # should directly get +ve lateral acceleration, +ve yaw acceleration and derivative of vy 121 | self.assertLess(model_return.states_dot["vy"], 0) 122 | self.assertLess(model_return.states_dot["yaw_rate"], 0) 123 | self.assertLess(model_return.outputs["ay"], 0) 124 | 125 | num_steps = 21 126 | time_step = 0.05 127 | 128 | states, outputs = self._forward_euler(init_states, inputs, time_step, num_steps) 129 | 130 | # All the other assertions change sign compared to left turn, except for this one 131 | self.assertLess(states["vx"], init_states["vx"]) 132 | self.assertGreater(states["vy"], 0) 133 | self.assertLess(states["yaw_rate"], 0) 134 | 135 | def test_delta_torque(self): 136 | """When there is a delta speed 137 | There is a delta torque of the opposite sign. 138 | """ 139 | 140 | states = self._get_initial_states() 141 | wheel_speed_rr = states["wheel_speed_rr"] 142 | 143 | for delta_speed in [-1.0, 1.0]: 144 | wheel_speed_rl = wheel_speed_rr + delta_speed 145 | 146 | states["wheel_speed_rl"] = wheel_speed_rl 147 | 148 | inputs = self.model.make_dict( 149 | "inputs", throttle=0.3, brake=0, steer=0.0, ax=0, ay=0 150 | ) 151 | 152 | model_return = self.model.forward(states, inputs) 153 | delta_torque = ( 154 | model_return.outputs["drive_torque_rl"] 155 | - model_return.outputs["drive_torque_rr"] 156 | ) 157 | 158 | # delta speed and torque should be opposite signs 159 | self.assertLess(delta_speed * delta_torque, 0.0) 160 | -------------------------------------------------------------------------------- /regression_tests/run_benchmark.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import json 3 | import logging 4 | import os 5 | 6 | import pandas as pd 7 | 8 | from lumos.models.composition import ModelMaker 9 | from lumos.simulations.laptime_simulation import LaptimeSimulation 10 | 11 | logger = logging.getLogger() 12 | 13 | 14 | def _create_metric(name, unit, value): 15 | return {"name": name, "unit": unit, "value": value} 16 | 17 | 18 | class MetricGenerator: 19 | def create_metrics(self): 20 | # Run jit profiling 21 | metrics = profile_jax_jit() 22 | 23 | # Run nlp profiling 24 | for backend, num_intervals in itertools.product( 25 | ["casadi", "jax"], [100, 1000, 10000] 26 | ): 27 | logger.info( 28 | f"profiling nlp with backend={backend}, num_intervals={num_intervals}" 29 | ) 30 | metrics += profile_nlp(backend, num_intervals) 31 | 32 | # Run track sweep 33 | metrics += sweep_tracks() 34 | 35 | with open("summary.json", "w+") as outfile: 36 | json.dump(metrics, outfile) 37 | 38 | 39 | def sweep_tracks(): 40 | TRACK_DIR = "data/tracks" 41 | 42 | track_files = [ 43 | os.path.join(TRACK_DIR, f) for f in os.listdir(TRACK_DIR) if f.endswith(".csv") 44 | ] 45 | 46 | # Quickest way is to run with jax and only compiling for the first solve 47 | sim_config = LaptimeSimulation.get_sim_config( 48 | num_intervals=250, 49 | hessian_approximation="exact", 50 | is_cyclic=True, 51 | is_condensed=False, 52 | backend="jax", 53 | track=track_files[0], 54 | transcription="LGR", 55 | ) 56 | 57 | model_config = ModelMaker.make_config("SimpleVehicleOnTrack") 58 | 59 | ocp = LaptimeSimulation(model_config=model_config, sim_config=sim_config) 60 | 61 | def _solve_with_track(track_file): 62 | logger.info(f"solving ltc for: {track_file}") 63 | ocp.set_track(track_file) 64 | 65 | x0 = ocp.get_init_guess() 66 | logger.info(f"Starting solve for {track_file}") 67 | solution, info = ocp.solve( 68 | x0, 69 | max_iter=200, 70 | print_level=4, 71 | print_timing_statistics="no", 72 | derivative_test="none", 73 | dual_inf_tol=1e-3, 74 | constr_viol_tol=1e-3, 75 | ) 76 | info["laptime"] = ocp.dec_var_operator.get_var( 77 | solution, group="states", name="time", stage=-1 78 | ) 79 | info["track"] = track_file 80 | logger.info(f"Maneuver time {info['laptime']:.3f} sec") 81 | 82 | COLUMNS_TO_EXTRACT = ( 83 | "track", 84 | "status", 85 | "num_iter", 86 | "laptime", 87 | "status_msg", 88 | ) 89 | 90 | result = {k: info[k] for k in COLUMNS_TO_EXTRACT} 91 | return result 92 | 93 | # Summary statistics 94 | results = pd.DataFrame.from_records([_solve_with_track(f) for f in track_files]) 95 | is_success = results["status"] == 0 96 | num_success = is_success.sum() 97 | summary = pd.DataFrame( 98 | [ 99 | { 100 | "num_success": num_success, 101 | "num_total": len(results), 102 | "success_pct": num_success / len(results) * 100, 103 | "avg_success_iter": results.loc[is_success, "num_iter"].mean(), 104 | } 105 | ] 106 | ) 107 | 108 | # Artifacts files for the record 109 | results.to_csv("track_sweep_results.csv") 110 | summary.to_csv("track_sweep_summary.csv") 111 | 112 | # Create benchmark result for traclomg amd github page visualization 113 | # See: https://github.com/benchmark-action/github-action-benchmark 114 | # using customSmallerIsBetter, so need to make metrics also better when smaller 115 | metrics = [ 116 | _create_metric("num_total", "-", len(results)), 117 | _create_metric("failure_pct", "%", (1 - num_success / len(results)) * 100), 118 | _create_metric( 119 | "avg_success_iter", "iter", results.loc[is_success, "num_iter"].mean() 120 | ), 121 | ] 122 | 123 | return metrics 124 | 125 | 126 | def profile_jax_jit(): 127 | sim_config = LaptimeSimulation.get_sim_config( 128 | num_intervals=250, 129 | hessian_approximation="exact", 130 | is_cyclic=True, 131 | is_condensed=False, 132 | backend="jax", 133 | track="data/tracks/Catalunya.csv", 134 | transcription="LGR", 135 | ) 136 | 137 | model_config = ModelMaker.make_config("SimpleVehicleOnTrack") 138 | ocp = LaptimeSimulation(model_config=model_config, sim_config=sim_config) 139 | 140 | x0 = ocp.get_init_guess() 141 | # Call just once, so approximately all time is jax jit time 142 | results = ocp.profile(x0, repeat=1, hessian=True, call_once_before=False) 143 | 144 | metrics = [] 145 | # NOTE: the nlp.constraints etc are called after model_algebra, which means jitting 146 | # has already been triggered and completed. So we need to catch jitting time from 147 | # the model_algebra calls. 148 | for name in [ 149 | "model_algebra.constraints", 150 | "model_algebra.jacobian", 151 | "model_algebra.hessian", 152 | ]: 153 | metrics.append( 154 | _create_metric(".".join(["jax_jit", name]), "sec", results[name]) 155 | ) 156 | 157 | return metrics 158 | 159 | 160 | def profile_nlp(backend: str, num_intervals: int): 161 | # Quickest way is to run with jax and only compiling for the first solve 162 | sim_config = LaptimeSimulation.get_sim_config( 163 | num_intervals=num_intervals, 164 | hessian_approximation="exact", 165 | is_cyclic=True, 166 | is_condensed=False, 167 | backend=backend, 168 | track="data/tracks/Catalunya.csv", 169 | transcription="LGR", 170 | ) 171 | 172 | model_config = ModelMaker.make_config("SimpleVehicleOnTrack") 173 | ocp = LaptimeSimulation(model_config=model_config, sim_config=sim_config) 174 | 175 | x0 = ocp.get_init_guess() 176 | results = ocp.profile(x0, repeat=10, hessian=True) 177 | 178 | metrics = [] 179 | for name in [ 180 | "nlp.objective", 181 | "nlp.gradient", 182 | "nlp.constraints", 183 | "nlp.jacobian", 184 | "nlp.hessian", 185 | ]: 186 | metrics.append( 187 | _create_metric( 188 | ".".join([backend, str(num_intervals), name]), "sec", results[name] 189 | ) 190 | ) 191 | 192 | return metrics 193 | 194 | 195 | if __name__ == "__main__": 196 | logging.basicConfig(level=logging.INFO) 197 | generator = MetricGenerator() 198 | generator.create_metrics() 199 | -------------------------------------------------------------------------------- /lumos/models/tracks.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | from abc import abstractmethod 4 | from typing import Any, Dict 5 | 6 | import numpy as np 7 | import pandas as pd 8 | 9 | logger = logging.getLogger(__name__) 10 | 11 | 12 | class Track: 13 | @abstractmethod 14 | def heading_at(self, s: float) -> float: 15 | """Heading function""" 16 | pass 17 | 18 | @abstractmethod 19 | def curvature_at(self, s: float) -> float: 20 | pass 21 | 22 | @abstractmethod 23 | def left_distance_at(self, s: float) -> float: 24 | """Distance to left bound.""" 25 | pass 26 | 27 | @abstractmethod 28 | def right_distance_at(self, s: float) -> float: 29 | """Distance to right bound.""" 30 | pass 31 | 32 | 33 | def _trapz(y, x): 34 | dx = np.diff(x) 35 | deltas = (y[:-1] + y[1:]) / 2 * dx 36 | return np.cumsum(np.insert(deltas, 0, 0.0)) 37 | 38 | 39 | def cartesian_to_curvilinear(x, y): 40 | ds = np.sqrt(np.diff(x) ** 2 + np.diff(y) ** 2) 41 | s = np.cumsum(np.insert(ds, 0, 0.0)) 42 | 43 | # first derivative 44 | dx = np.gradient(x, s) 45 | dy = np.gradient(y, s) 46 | 47 | # second derivatives 48 | d2x = np.gradient(dx, s) 49 | d2y = np.gradient(dy, s) 50 | 51 | # calculation of curvature from the typical formula 52 | curvature = (dx * d2y - d2x * dy) / (dx * dx + dy * dy) ** 1.5 53 | 54 | initial_heading = np.arctan2(dy[0], dx[0]) 55 | heading = _trapz(curvature, s) + initial_heading 56 | return s, curvature, heading 57 | 58 | 59 | def curvilinear_to_cartesian( 60 | s, curvature, x0: float = 0, y0: float = 0, heading0: float = 0 61 | ): 62 | heading = _trapz(curvature, s) + heading0 63 | x = _trapz(np.cos(heading), s) + x0 64 | y = _trapz(np.sin(heading), s) + y0 65 | 66 | return x, y, heading 67 | 68 | 69 | class RaceTrack(Track): 70 | def __init__( 71 | self, 72 | x: np.ndarray, 73 | y: np.ndarray, 74 | s: np.ndarray, 75 | curvature: np.ndarray, 76 | heading: np.ndarray, 77 | left_distance: np.ndarray = None, 78 | right_distance: np.ndarray = None, 79 | ): 80 | # Create constant distance to the sides if unspecified 81 | # FIXME: 5m is arbitrary, should we make these values compulsory? 82 | if left_distance is None: 83 | left_distance = np.ones_like(s) * 5 84 | if right_distance is None: 85 | right_distance = np.ones_like(s) * 5 86 | 87 | # Check if the left distance and right distance are incompatible with what the 88 | # curvature allows. 89 | # if curvature * distance from centerline > 1 -> the boundary will fold on itself. 90 | # this is valid for both +ve and -ve distance from centerline 91 | # NOTE: we just check, but don't do any modifications yet. 92 | b_violate_left = curvature * left_distance > 1 93 | b_violate_right = curvature * -right_distance > 1 94 | 95 | logger.info(f"left distance violation: {np.sum(b_violate_left)}") 96 | logger.info(f"right distance violation: {np.sum(b_violate_right)}") 97 | 98 | self._track_data = { 99 | "s": s, 100 | "curvature": curvature, 101 | "heading": heading, 102 | "x": x, 103 | "y": y, 104 | "left_distance": left_distance, 105 | "right_distance": right_distance, 106 | } 107 | 108 | # Offset distance to always start at 0 109 | self._track_data["s"] = self._track_data["s"] - self._track_data["s"][0] 110 | 111 | @staticmethod 112 | def from_cartesian(x, y, left_distance=None, right_distance=None): 113 | s, curvature, heading = cartesian_to_curvilinear(x, y) 114 | return RaceTrack( 115 | x=x, 116 | y=y, 117 | s=s, 118 | curvature=curvature, 119 | heading=heading, 120 | left_distance=left_distance, 121 | right_distance=right_distance, 122 | ) 123 | 124 | @staticmethod 125 | def from_curvilinear( 126 | s, curvature, heading0=0, left_distance=None, right_distance=None 127 | ): 128 | x, y, heading = curvilinear_to_cartesian(s, curvature, heading0=heading0) 129 | return RaceTrack( 130 | x=x, 131 | y=y, 132 | s=s, 133 | curvature=curvature, 134 | heading=heading, 135 | left_distance=left_distance, 136 | right_distance=right_distance, 137 | ) 138 | 139 | @staticmethod 140 | def from_tum_csv(track_file): 141 | """Create track from TUM track data format 142 | 143 | see:https://github.com/TUMFTM/racetrack-database 144 | """ 145 | df = pd.read_csv(track_file) 146 | return RaceTrack.from_cartesian( 147 | x=df["# x_m"], 148 | y=df["y_m"], 149 | left_distance=df["w_tr_left_m"], 150 | right_distance=df["w_tr_right_m"], 151 | ) 152 | 153 | def curvature_at(self, s: float) -> float: 154 | """Return the curvature of the track at given distance""" 155 | return np.interp(s, self._track_data["s"], self._track_data["curvature"]) 156 | 157 | def heading_at(self, s: float) -> float: 158 | """Return the heading angle of the track at given distance""" 159 | return np.interp(s, self._track_data["s"], self._track_data["heading"]) 160 | 161 | def left_distance_at(self, s: float) -> float: 162 | """Constant halfwidth for now""" 163 | return np.interp(s, self._track_data["s"], self._track_data["left_distance"]) 164 | 165 | def right_distance_at(self, s: float) -> float: 166 | """Constant halfwidth for now""" 167 | return np.interp(s, self._track_data["s"], self._track_data["right_distance"]) 168 | 169 | @property 170 | def total_distance(self): 171 | return self._track_data["s"][-1] 172 | 173 | 174 | class ConstCurvTrack(Track): 175 | _curvature: float 176 | _start_heading: float 177 | 178 | def __init__( 179 | self, curvature: float = 0.01, start_heading: float = 0.0, half_width: float = 2 180 | ): 181 | self._curvature = curvature 182 | self._start_heading = start_heading 183 | self._half_wdith = half_width 184 | 185 | def heading_at(self, s: float) -> float: 186 | """Heading function 187 | 188 | dtheta/ds = k 189 | 190 | For constant k 191 | theta = start_heading + k*s 192 | 193 | """ 194 | return self._start_heading + self._curvature * s 195 | 196 | def curvature_at(self, s: float) -> float: 197 | return self._curvature 198 | 199 | def half_width_at(self, s: float) -> float: 200 | """Constant halfwidth for now""" 201 | return self._half_wdith 202 | -------------------------------------------------------------------------------- /tests/test_optimal_control/test_ocp_logging.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import glob 3 | import os 4 | import unittest 5 | from tempfile import TemporaryDirectory 6 | 7 | import numpy as np 8 | import pandas as pd 9 | from pyarrow import parquet as pq 10 | from parameterized import parameterized 11 | 12 | from lumos.optimal_control.config import LoggingConfig 13 | from lumos.simulations.drone_simulation import DroneSimulation 14 | from lumos.simulations.laptime_simulation import LaptimeSimulation 15 | 16 | 17 | class TestOCPLogging(unittest.TestCase): 18 | num_iter: int = 5 19 | SimClass: type = DroneSimulation 20 | 21 | def _run_and_log( 22 | self, 23 | results_dir: str, 24 | log_every_nth_iter: int, 25 | drone_or_ltc: str = "drone", 26 | is_condensed: bool = False, 27 | ): 28 | logging_config = LoggingConfig( 29 | sim_name="temp", 30 | results_dir=results_dir, 31 | log_every_nth_iter=log_every_nth_iter, 32 | ) 33 | if drone_or_ltc == "drone": 34 | sim_config = DroneSimulation.get_sim_config( 35 | backend="casadi", 36 | num_intervals=100, 37 | is_condensed=is_condensed, 38 | logging_config=logging_config, 39 | ) 40 | ocp = DroneSimulation(sim_config=sim_config) 41 | else: 42 | sim_config = LaptimeSimulation.get_sim_config( 43 | track="data/tracks/Catalunya.csv", 44 | backend="casadi", 45 | num_intervals=100, 46 | is_condensed=is_condensed, 47 | logging_config=logging_config, 48 | ) 49 | ocp = LaptimeSimulation(sim_config=sim_config) 50 | 51 | x0 = ocp.get_init_guess() 52 | solution, info = ocp.solve(x0, max_iter=self.num_iter, print_level=0,) 53 | 54 | return ocp.logging_dir 55 | 56 | @parameterized.expand( 57 | [[0, 2, 0], [1, 2, 1], [10, 2, 1],] 58 | ) 59 | def test_correct_files_are_created( 60 | self, log_every_nth_iter, expected_num_csv, expected_num_parquet 61 | ): 62 | with TemporaryDirectory() as temp_dir: 63 | logging_dir = self._run_and_log( 64 | results_dir=temp_dir, log_every_nth_iter=log_every_nth_iter 65 | ) 66 | 67 | # Check the number of log files are correct 68 | # one csv for final results and 1 csv for metrics history should always be 69 | # written 70 | # log_every_nth_iter = 0 -> no parquet files for all iterations 71 | csv_files = glob.glob(os.path.join(logging_dir, "*.csv")) 72 | self.assertEqual(len(csv_files), expected_num_csv) 73 | 74 | parquet_files = glob.glob(os.path.join(logging_dir, "*.parquet")) 75 | self.assertEqual(len(parquet_files), expected_num_parquet) 76 | 77 | def test_no_logging(self): 78 | """Just a smoke test. Kind of difficult to verify it doesn't produce anything""" 79 | sim_config = DroneSimulation.get_sim_config( 80 | backend="casadi", num_intervals=100, 81 | ) 82 | ocp = DroneSimulation(sim_config=sim_config) 83 | x0 = ocp.get_init_guess() 84 | solution, info = ocp.solve(x0, max_iter=self.num_iter, print_level=0) 85 | 86 | @parameterized.expand([[2], [3], [10]]) 87 | def test_last_iter_is_logged(self, log_every_nth_iter: int): 88 | """Ensure that when we log iterations, the last one is alwasy logged. 89 | 90 | We also test that the last iteration logged results are correct. 91 | """ 92 | with TemporaryDirectory() as temp_dir: 93 | sim_config = DroneSimulation.get_sim_config( 94 | backend="casadi", 95 | num_intervals=100, 96 | logging_config=LoggingConfig( 97 | sim_name="temp", 98 | results_dir=temp_dir, 99 | log_every_nth_iter=log_every_nth_iter, 100 | ), 101 | ) 102 | ocp = DroneSimulation(sim_config=sim_config) 103 | x0 = ocp.get_init_guess() 104 | solution, info = ocp.solve(x0, max_iter=self.num_iter, print_level=0) 105 | 106 | combined_df = pq.read_pandas( 107 | os.path.join(ocp.logging_dir, "all_iters.parquet") 108 | ).to_pandas() 109 | 110 | # Check last iteration is logged 111 | self.assertEqual(combined_df["iter"].max(), self.num_iter) 112 | 113 | # Check the values are correct 114 | last_iter_df = combined_df.loc[combined_df["iter"] == self.num_iter] 115 | op = ocp.dec_var_operator 116 | for group in ["states", "inputs", "con_outputs"]: 117 | for name in ocp.model.get_group_names(group): 118 | np.testing.assert_allclose( 119 | last_iter_df[group + "." + name], 120 | op.get_var(solution, group, name), 121 | ) 122 | 123 | @parameterized.expand(itertools.product(("drone", "ltc"), (True, False))) 124 | def test_log_every_iter_logs_correctly(self, drone_or_ltc, is_condensed): 125 | with TemporaryDirectory() as temp_dir: 126 | logging_dir = self._run_and_log( 127 | is_condensed=is_condensed, 128 | results_dir=temp_dir, 129 | log_every_nth_iter=1, 130 | drone_or_ltc=drone_or_ltc, 131 | ) 132 | 133 | # Check the number of log files are correct 134 | # There should be: 135 | # - 1 all_iters parquet 136 | # - 1 final_iter csv 137 | # - 1 metrics history csv 138 | expected_num_csv = 2 139 | expected_num_parquet = 1 140 | csv_files = glob.glob(os.path.join(logging_dir, "*.csv")) 141 | self.assertEqual(len(csv_files), expected_num_csv) 142 | 143 | parquet_files = glob.glob(os.path.join(logging_dir, "*.parquet")) 144 | self.assertEqual(len(parquet_files), expected_num_parquet) 145 | 146 | # Ensure the max constraiont violation agrees with inf_pr 147 | combined_df = pq.read_pandas( 148 | os.path.join(logging_dir, "all_iters.parquet") 149 | ).to_pandas() 150 | metrics_history_df = pd.read_csv( 151 | os.path.join(logging_dir, "metrics_history.csv") 152 | ) 153 | 154 | # FIXME: it's not nice to get constraint names from here... 155 | con_names = [c for c in combined_df.columns if "_con." in c] 156 | for it in range(self.num_iter): 157 | iter_df = combined_df.loc[combined_df["iter"] == it] 158 | idx, column = iter_df[con_names].abs().stack().idxmax() 159 | self.assertAlmostEqual( 160 | metrics_history_df["inf_pr"][it], np.abs(iter_df[column][idx]) 161 | ) 162 | self.assertEqual( 163 | metrics_history_df["max_violation_con_name"][it], column 164 | ) 165 | self.assertAlmostEqual( 166 | metrics_history_df["max_violation_mesh"][it], iter_df["mesh"][idx], 167 | ) 168 | -------------------------------------------------------------------------------- /lumos/optimal_control/config.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from dataclasses import asdict, dataclass, field 3 | from typing import Any, Dict, List, Optional, Tuple, Union 4 | 5 | 6 | @dataclass 7 | class BoundConfig: 8 | """Config for the bounds on a stage variable 9 | 10 | The bound values must be a tuple of floats or a tuple of 1d numpy arrays. The upper 11 | bounds must be greater than or equal to the lower bounds. 12 | """ 13 | 14 | group: str 15 | name: str 16 | values: Union[Tuple[float, float], Tuple[np.ndarray, np.ndarray]] 17 | 18 | def __post_init__(self): 19 | """Perform checks on bounds. 20 | 21 | 1) bounds should be both scalars or both 1d np.ndarray 22 | 2) upper bounds should be larger than or equal to lower bounds. 23 | """ 24 | lb, ub = self.values 25 | 26 | pretext = f"Failed to set bounds for {self.group}:{self.name}. " 27 | if np.isscalar(lb) and np.isscalar(ub): 28 | assert ub >= lb, pretext + "lower bounds larger than upper bounds" 29 | elif isinstance(lb, np.ndarray) and isinstance(ub, np.ndarray): 30 | assert lb.ndim == 1 and ub.ndim == 1, pretext + "Bound arrays must be 1d." 31 | assert len(lb) == len(ub), pretext + "Bound arrays must have the same size." 32 | assert np.all(ub >= lb), ( 33 | pretext 34 | + f"lower bounds larger than upper bounds at stage {np.where(lb > ub)}" 35 | ) 36 | 37 | else: 38 | raise TypeError( 39 | pretext + "lower and upper bounds must be both scalars or " 40 | "both 1d numpy arrays of the same size." 41 | ) 42 | 43 | 44 | @dataclass 45 | class ScaleConfig: 46 | """Config for the scale on variables 47 | 48 | TODO: check all values are postiive! 49 | """ 50 | 51 | group: str 52 | name: str 53 | value: float 54 | 55 | 56 | @dataclass 57 | class BoundaryConditionConfig: 58 | """Boundary condition config to constraint a stage variable to a given value.""" 59 | 60 | stage: int 61 | group: str 62 | name: str 63 | value: float 64 | 65 | 66 | @dataclass 67 | class TranscriptionConfig: 68 | pass 69 | 70 | 71 | @dataclass 72 | class LoggingConfig: 73 | """Controls if and where the results as well as debug info are written to""" 74 | 75 | sim_name: str = "simulation" 76 | results_dir: str = "results" 77 | log_every_nth_iter: int = 0 78 | 79 | 80 | @dataclass 81 | class SimConfig: 82 | """Simulation Configuraration for Optimal Control 83 | 84 | This class is inheritted from dataclass, as such any customization to it need to be 85 | done in __post_init__ 86 | 87 | required fields are: 88 | num_intervals (int): number of intervals for the OCP, default to 99. 89 | interval_points (Optional[np.ndarray]): when given, uses custon interval points 90 | to set up the normalized_mesh. The points must be in [0, 1] and covers both 91 | end points. They must also be monotonically increasing, and have a size of 92 | num_intervals + 1. Default to None, which uses uniform intervals. See also: 93 | ScaledMeshOCP.set_normalized_mesh 94 | transcription (str, Tuple[str, Dict[str, Any]]): the transcription method to use 95 | can be either a single string, and must be one of the following: 96 | ["ForwardEuler", "Trapezoidal", "LGR", "LGRIntegral"]. Optionally, one can pass 97 | in additional parameter, for example, 'num_stages' for LGR, which defines number 98 | of stages used in collocation. So all of the following are valid: 99 | - transcription="ForwardEuler" (use ForwardEuler, no configurable parameters) 100 | - transcirption="LGR" (use LGR, with default 3 stages per interval) 101 | - transcription = ("LGR", {"num_stages": 5}) (use LGR, with custom parameter of 102 | 5 stages per interval) 103 | is_cyclic [bool]: true will set the problem as cyclic, default to False. 104 | non_cyclic_vars [List[str]]: when the problem is cyclic, the variables that can 105 | still be non-cyclic, default to empty. 106 | is_condensed [bool]: if True, will condense the problem to remove "states_dot" 107 | from decision variables, which would reduce IPOPT solve overhead, with 108 | potential cost on function call overhead and convergence. Default to False. 109 | backend [str]: numerical backend to use, must be "jax", "casadi" or "cuistom". 110 | if custom, the user must provide the derivatives calls 111 | (see examples/brachistochrone.py). Default to "jax" 112 | hessian_approximation [str]: "limited-memory" or "exact", default to "exact" 113 | boundary_conditions [Tuple[BoundaryConditionConfig]]: tuple of boundary cond 114 | configurations. Default to empty 115 | bounds [Tuple[BoundConfig]]: bounds all on variables, default to empty. 116 | scales [Tuple[ScaleConfig]]: custom scaling for all variables, defaul to empty, 117 | which uses a scale of 1 118 | con_output_names [Tuple[str]]: names of the outputs to be added to the dec vars, 119 | so that we can add bounds on them to constrain them. 120 | logging_config [Optional[LoggingConfig]]: logging configuration, default to None. 121 | 122 | """ 123 | 124 | num_intervals: int = 99 125 | interval_points: Optional[np.ndarray] = None 126 | transcription: Union[str, Tuple[str, Optional[Dict[str, Any]]]] = "Trapezoidal" 127 | is_cyclic: bool = False 128 | non_cyclic_vars: List[str] = field(default_factory=list) 129 | is_condensed: bool = False 130 | backend: str = "jax" 131 | hessian_approximation: str = "exact" 132 | boundary_conditions: Tuple[BoundaryConditionConfig] = () 133 | bounds: Tuple[BoundConfig] = () 134 | scales: Tuple[ScaleConfig] = () 135 | con_output_names: Tuple[str] = () 136 | logging_config: Optional[LoggingConfig] = None 137 | 138 | def __post_init__(self): 139 | # We allow a single string argument (with no additional arugments or relying on 140 | # default arguments) for transcription definition, eg "LGR", ("LGR", ) and 141 | # ("LGR", {"num_stages": 3}}) are all valid 142 | # It almost feels like that the SimConfig needs to be a tree as well, where 143 | # transcription can have its own config. 144 | if isinstance(self.transcription, str): 145 | self.transcription = (self.transcription,) 146 | 147 | # Additional operations for desrialization for config fields 148 | if self.bounds and isinstance(self.bounds[0], dict): 149 | self.bounds = tuple(BoundConfig(**d) for d in self.bounds) 150 | 151 | if self.scales and isinstance(self.scales[0], dict): 152 | self.scales = tuple(ScaleConfig(**d) for d in self.scales) 153 | 154 | if self.boundary_conditions and isinstance(self.boundary_conditions[0], dict): 155 | self.boundary_conditions = tuple( 156 | BoundaryConditionConfig(**d) for d in self.boundary_conditions 157 | ) 158 | 159 | def to_dict(self) -> Dict[str, Any]: 160 | """Convert a nested config to a dictionary.""" 161 | return asdict(self) 162 | -------------------------------------------------------------------------------- /lumos/optimal_control/transcription.py: -------------------------------------------------------------------------------- 1 | from abc import ABC, abstractmethod 2 | from typing import Any, Dict, List, Tuple 3 | 4 | import numpy as np 5 | from scipy.sparse import diags 6 | 7 | import lumos.numpy as lnp 8 | 9 | from lumos.optimal_control.collocation import ( 10 | build_lagrange_differential_matrix, 11 | build_lagrange_integration_matrix, 12 | CollocationEnum, 13 | get_collocation_points, 14 | ) 15 | 16 | 17 | class Transcription(ABC): 18 | """Transcription method turning a continuous time problem into a discrete one. 19 | 20 | 21 | # TODO: should we eturn this into A*x + B*x_dot to make it more conventional? 22 | It constructs the linear continuity constraints: A*x - B*x_dot*T = 0 23 | 24 | assuming an m-stage interval, and d states 25 | 26 | 27 | A and B are both [m-1, m], while x and x_dot are both [m, d]. (m-stage interval, d 28 | states) T is a constant of the interval time. This is natural for integration scheme 29 | , but for differential schemes usually the interval time is combined with A in the 30 | form of 1/T. 31 | 32 | We unify it here to: 33 | - make the interface consistent for differential and integration schcheme 34 | - ensure the continuity constraint is in the order of the states instead of swtching 35 | between states and state derivatives. 36 | """ 37 | 38 | num_stages_per_interval: int 39 | num_constraints_per_interval: int 40 | 41 | @property 42 | def _cont_matrix_shape(self) -> Tuple[int, int]: 43 | return self.num_constraints_per_interval, self.num_stages_per_interval 44 | 45 | def get_continuity_matrices(self) -> Tuple[np.ndarray, np.ndarray]: 46 | return self._get_A_matrix(), self._get_B_matrix() 47 | 48 | def continuity_con( 49 | self, x: lnp.ndarray, x_dot: lnp.ndarray, interval_length: float 50 | ) -> lnp.ndarray: 51 | A, B = self.get_continuity_matrices() 52 | continuity_con = A @ x - B @ x_dot * interval_length 53 | 54 | return continuity_con 55 | 56 | @abstractmethod 57 | def _get_A_matrix(self) -> np.ndarray: 58 | pass 59 | 60 | @abstractmethod 61 | def _get_B_matrix(self) -> np.ndarray: 62 | pass 63 | 64 | 65 | class ForwardEuler(Transcription): 66 | """x_{i+1} - x_{i} - x_dot_{i}*dt = 0""" 67 | 68 | num_stages_per_interval: int = 2 69 | num_constraints_per_interval: int = 1 70 | 71 | def _get_A_matrix(self): 72 | return diags([-1, 1], [0, 1], shape=self._cont_matrix_shape).toarray() 73 | 74 | def _get_B_matrix(self): 75 | return diags([1], [0], shape=self._cont_matrix_shape).toarray() 76 | 77 | 78 | class Trapezoidal(Transcription): 79 | """x_{i+1} - x_{i} - (x_dot_{i+1} + x_dot_{i}) * dt/2 = 0""" 80 | 81 | num_stages_per_interval: int = 2 82 | num_constraints_per_interval: int = 1 83 | 84 | def _get_A_matrix(self): 85 | return diags([-1, 1], [0, 1], shape=self._cont_matrix_shape).toarray() 86 | 87 | def _get_B_matrix(self): 88 | return diags([0.5, 0.5], [0, 1], shape=self._cont_matrix_shape).toarray() 89 | 90 | 91 | class Collocation(Transcription): 92 | """Transcription with Legendre collocation 93 | 94 | Interval of collocation is converted to [0, 1] from the standard of [-1, 1] to make 95 | downstream computations easier. 96 | 97 | More details, refer to: AN OVERVIEW OF THREE PSEUDOSPECTRAL METHODS FOR THE 98 | NUMERICAL SOLUTION OF OPTIMAL CONTROL PROBLEMS 99 | https://hal.archives-ouvertes.fr/hal-01615132/document 100 | 101 | """ 102 | 103 | interp_points: np.ndarray 104 | collocation_points: np.ndarray 105 | 106 | def __init__(self, num_stages: int): 107 | self.num_stages_per_interval: int = num_stages 108 | 109 | self._set_collocation_points(num_stages) 110 | self._set_interp_points() 111 | 112 | self.d_matrix: np.ndarray = build_lagrange_differential_matrix( 113 | support=self.interp_points, evaluation_points=self.collocation_points 114 | ) 115 | 116 | @property 117 | def num_constraints_per_interval(self): 118 | return len(self.collocation_points) 119 | 120 | @abstractmethod 121 | def _set_collocation_points(self, num_stages: int) -> None: 122 | pass 123 | 124 | @abstractmethod 125 | def _set_interp_points(self) -> None: 126 | pass 127 | 128 | def _get_A_matrix(self) -> np.ndarray: 129 | # multiply by two because collocation is in the domain of [-1, 1], 130 | # so to map to [0, 1] (easy to scale to interval time), we first need to multiply 131 | # by two. 132 | return self.d_matrix 133 | 134 | def _get_B_matrix(self) -> np.ndarray: 135 | return diags([1], [1], shape=self._cont_matrix_shape).toarray() 136 | 137 | 138 | class LGR(Collocation): 139 | """Transcription with LGR collocation""" 140 | 141 | def __init__(self, num_stages: int = 3): 142 | super().__init__(num_stages=num_stages) 143 | 144 | def _set_collocation_points(self, num_stages: int) -> None: 145 | # map collocation points from [-1, 1] to [0, 1] 146 | self.collocation_points = ( 147 | get_collocation_points( 148 | num_points=num_stages - 1, scheme=CollocationEnum.LGR 149 | ) 150 | + 1 151 | ) / 2 152 | 153 | def _set_interp_points(self) -> None: 154 | # Add the 0 to the interp_points 155 | self.interp_points = np.insert(self.collocation_points, 0, 0) 156 | 157 | 158 | class LGRIntegral(LGR): 159 | """Integral variant of the LGR scheme""" 160 | 161 | def __init__(self, num_stages: int = 3): 162 | super().__init__(num_stages=num_stages) 163 | 164 | # Now we fit the polynomial on the derivatives (so on collocaiton points) 165 | # And then evaluate the ingral at the interpretation points (except for the 1st 166 | # point) 167 | self.i_matrix = build_lagrange_integration_matrix( 168 | support=self.interp_points, evaluation_points=self.collocation_points, 169 | ) 170 | 171 | def _get_A_matrix(self) -> np.ndarray: 172 | return np.hstack( 173 | [ 174 | -np.ones((self.num_stages_per_interval - 1, 1)), 175 | np.eye(self.num_stages_per_interval - 1), 176 | ] 177 | ) 178 | 179 | def _get_B_matrix(self) -> np.ndarray: 180 | return self.i_matrix 181 | 182 | 183 | TRANSCRIPTION_OPTIONS = { 184 | t.__name__: t for t in (ForwardEuler, Trapezoidal, LGR, LGRIntegral) 185 | } 186 | 187 | 188 | def get_transcription_options() -> List[str]: 189 | """Return names of available transcription classes. 190 | 191 | Returns: 192 | List[str]: a list of names of the available Transcription classes. 193 | """ 194 | return [n for n in TRANSCRIPTION_OPTIONS.keys()] 195 | 196 | 197 | def make_transcription(name: str, kwargs: Dict[str, Any] = None) -> Transcription: 198 | """Create a Transcription object from name and keyword arguments. 199 | 200 | Args: 201 | name (str): name of the transcription class. 202 | kwargs (Dict[str, Any], optional): additional kwargs to be passed to the 203 | transcription constructtor. Defaults to None, which will be set to empty. 204 | 205 | Raises: 206 | RuntimeError: if the transcription required is not a valid option. 207 | 208 | Returns: 209 | Transcription: Transcription object that defines a descritization scheme. 210 | """ 211 | if not kwargs: 212 | kwargs = {} 213 | 214 | if name in TRANSCRIPTION_OPTIONS: 215 | return TRANSCRIPTION_OPTIONS[name](**kwargs) 216 | else: 217 | raise RuntimeError( 218 | "name is not a valid transcription type. " 219 | f"Valid options are {get_transcription_options()}" 220 | ) 221 | -------------------------------------------------------------------------------- /lumos/models/test_utils.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | import unittest 4 | import uuid 5 | from functools import wraps 6 | from typing import List 7 | 8 | import numpy as np 9 | from casadi import MX, Function 10 | 11 | import lumos.numpy as lnp 12 | 13 | 14 | logger = logging.getLogger(__name__) 15 | 16 | 17 | def use_backends(backends: List[str]): 18 | """Decortator to run tests with different backends. 19 | 20 | This is a workaround of an issue with parameterized_class + inheritance, 21 | 22 | Parameterized_class works by doing the following: 23 | 1) add new classes with subfix to current module 24 | 2) remove 'test_' method from the original class, so the original class doesn't get to run. 25 | see: https://github.com/wolever/parameterized/blob/0403e891d9a6ec5fa77c0e200f31e1298fcacbc9/parameterized/parameterized.py#L606 26 | Problem is: if the original class is inherited, the deletion doesn't happen because 27 | those those 'test_' methods are defined in its parent class! 28 | 29 | TODO: currently there are still issues compared to using parameterized 30 | 1) using 1 test instead of n tests, so as soon as one backend fails, we stop running 31 | the rest of the tests 32 | 2) more difficult to tell which backend fails and why 33 | 34 | """ 35 | 36 | def decorator(fn): 37 | @wraps(fn) 38 | def wrapped(self, *args, **kwargs): 39 | original_backend = lnp.get_backend() 40 | for i, backend in enumerate(backends): 41 | with self.subTest(i=i): 42 | lnp.set_backend(backend) 43 | logger.info(f"Running with {backend} backend") 44 | try: 45 | fn(self, *args, **kwargs) 46 | except AssertionError as ae: 47 | # Catch assertion failure seperately so it counts as 'failed' 48 | # instead of 'errors' 49 | raise ae 50 | except Exception as e: 51 | # Throw a mesage to highlight with which backend the test fails 52 | # (This is something we would get from parameterlized as it creates 53 | # separate tests) 54 | # using raise from, see: https://stackoverflow.com/a/29442282 55 | raise Exception(f"Error with {backend} backend") from e 56 | finally: 57 | # Return to original backend 58 | # TODO: maybe this should be a context manager? 59 | lnp.set_backend(original_backend) 60 | 61 | return wrapped 62 | 63 | return decorator 64 | 65 | 66 | class BaseModelTest: 67 | ModelClass: type 68 | # this defines the variable groups that need to be passed to model forward call 69 | forward_arguments: List[str] = ["inputs"] 70 | need_mesh_input: bool = False 71 | 72 | @classmethod 73 | def setUpClass(cls): 74 | cls.model = cls.ModelClass() 75 | cls.params = cls.model.get_recursive_params() 76 | 77 | cls.args_dict = { 78 | g: {n: 0.1 for n in cls.model.get_group_names(g)} 79 | for g in cls.forward_arguments 80 | } 81 | 82 | if cls.need_mesh_input: 83 | cls.args_dict["mesh"] = 0.0 84 | 85 | @use_backends(backends=["jax", "numpy", "casadi"]) 86 | def test_forward(self): 87 | """ 88 | smoke test to see if forward call throws an error. 89 | """ 90 | 91 | # TODO: Casadi behaves more like 'typebroadcast', it can take in float and 92 | # return float, but that's not what we need. We want to make sure t works with 93 | # Casadi symbolic variables. (But maybe if the test works with float, then it 94 | # will work with symbolic variables as well?) 95 | def _nested_dict_to_mx(d): 96 | if isinstance(d, dict): 97 | return {k: _nested_dict_to_mx(v) for k, v in d.items()} 98 | else: 99 | return MX(d) 100 | 101 | if lnp.get_backend() == "casadi": 102 | args_dict = _nested_dict_to_mx(self.args_dict) 103 | else: 104 | args_dict = self.args_dict 105 | 106 | # Calling without external parameters 107 | out = self.model.forward(**args_dict) 108 | 109 | logger.debug(f"outputs is of type {type(out.outputs)}") 110 | 111 | # Calling with external parameters 112 | _ = self.model.apply_and_forward(**args_dict, params=self.params) 113 | 114 | # TODO: we should impement a model auto-diff test to check if auto-diff works 115 | # (should we also check the autodiff accuracy?) 116 | @unittest.skip("To be implemented") 117 | def test_backward(self): 118 | # TODO: should backward be abstracted to model level? But numpy model wouldn't 119 | # have an autodiff tool. In addition, the original target is to only use jax 120 | # for autodiff (and casadi could be used for benchmark only) 121 | pass 122 | 123 | # TODO: code generation should only be possible for casadi backend. Where should we 124 | # put the test? 125 | # TODO: maybe for everymodel, we just want to test if it can generate code with 126 | # casadi, but don't touch anything else to do with Casadi? 127 | @use_backends(backends=["casadi"]) 128 | def test_code_generation(self): 129 | # Code generation requires using the array inputs interface 130 | args_dict = { 131 | k: MX.sym("mesh") if k == "mesh" else MX.sym(k, self.model.get_num_vars(k)) 132 | for k, v in self.args_dict.items() 133 | } 134 | 135 | # TODO: need to add parameter to the function 136 | f = Function( 137 | "f", 138 | list(args_dict.values()), 139 | list(self.model.forward_with_arrays(**args_dict)), 140 | ) 141 | 142 | # call the function with concrete values 143 | args_dict = { 144 | k: v if k == "mesh" else self.model.make_vector(k, **v) 145 | for k, v in self.args_dict.items() 146 | } 147 | _ = f.call(list(args_dict.values())) 148 | 149 | # Note: name must start with a letter, have no underscore, doesn't 150 | # overlap with some special names. So can't pass in full path name. 151 | # see: https://github.com/casadi/casadi/blob/8d0f80a4d0fe2054384bfb9748f7a0f6bae540ff/casadi/core/function.cpp#L1167 152 | # Due to such restrictions, we must generate and fix the name by hand, and 153 | # operate in current directory. 154 | 155 | file_name = f.fix_name(uuid.uuid4().hex) + ".c" 156 | f.generate(file_name) 157 | # At the moment only check if the generated code has more than 10 lines. 158 | # TODO:compile and run generated code. Need to configure compiler 159 | with open(file_name, "r") as f: 160 | file_length = len(f.read()) 161 | self.assertTrue(file_length > 10) 162 | logger.debug(f"file length = {file_length} lines") 163 | 164 | # Manually clean up 165 | os.remove(file_name) 166 | 167 | 168 | class BaseStateSpaceModelTest(BaseModelTest): 169 | forward_arguments: List[str] = ["inputs", "states"] 170 | need_mesh_input = True 171 | 172 | def _forward_euler(self, init_states, inputs, time_step, num_steps): 173 | """Helper to run forward euler with a fixed inputs for a few steps""" 174 | 175 | # Ensure we don't modify the initial states 176 | states = dict(init_states) 177 | for step in range(num_steps): 178 | model_return = self.model.forward(states, inputs, step * time_step) 179 | for k in states: 180 | states[k] += model_return.states_dot[k] * time_step 181 | 182 | return states, model_return.outputs 183 | 184 | 185 | if __name__ == "__name__": 186 | unittest.main() 187 | --------------------------------------------------------------------------------