Add Python examples

README.md

@@ -104,7 +104,7 @@ pip install sleipnirgroup-jormungandr
 
 ## Examples
 
-See the [examples](https://github.com/SleipnirGroup/Sleipnir/tree/main/examples) and [optimization unit tests](https://github.com/SleipnirGroup/Sleipnir/tree/main/test/optimization).
+See the [examples](https://github.com/SleipnirGroup/Sleipnir/tree/main/examples), [C++ optimization unit tests](https://github.com/SleipnirGroup/Sleipnir/tree/main/test/optimization), and [Python optimization unit tests](https://github.com/SleipnirGroup/Sleipnir/tree/main/jormungandr/test/optimization).
 
 ## Build
 

examples/CurrentManager/__init__.py

@@ -0,0 +1 @@
+from .current_manager import *

examples/CurrentManager/current_manager.py

@@ -0,0 +1,80 @@
+from jormungandr.autodiff import Variable, VariableMatrix
+from jormungandr.optimization import OptimizationProblem
+
+
+class CurrentManager:
+    """
+    This class computes the optimal current allocation for a list of subsystems
+    given a list of their desired currents and current tolerances that determine
+    which subsystem gets less current if the current budget is exceeded.
+    Subsystems with a smaller tolerance are given higher priority.
+    """
+
+    def __init__(self, current_tolerances: list[float], max_current: float):
+        """
+        Constructs a CurrentManager.
+
+        Parameter ``currentTolerances``:
+            The relative current tolerance of each subsystem.
+
+        Parameter ``max_current``:
+            The current budget to allocate between subsystems.
+        """
+        self.__desired_currents = VariableMatrix(len(current_tolerances), 1)
+        self.__problem = OptimizationProblem()
+        self.__allocated_currents = self.__problem.decision_variable(
+            len(current_tolerances)
+        )
+
+        # Ensure desired_currents contains initialized Variables
+        for row in range(self.__desired_currents.rows()):
+            # Don't initialize to 0 or 1, because those will get folded by
+            # Sleipnir
+            self.__desired_currents[row] = Variable(float("inf"))
+
+        J = 0.0
+        current_sum = 0.0
+        for i in range(len(current_tolerances)):
+            # The weight is 1/tolᵢ² where tolᵢ is the tolerance between the
+            # desired and allocated current for subsystem i
+            error = self.__desired_currents[i] - self.__allocated_currents[i]
+            J += error * error / (current_tolerances[i] * current_tolerances[i])
+
+            current_sum += self.__allocated_currents[i]
+
+            # Currents must be nonnegative
+            self.__problem.subject_to(self.__allocated_currents[i] >= 0.0)
+        self.__problem.minimize(J)
+
+        # Keep total current below maximum
+        self.__problem.subject_to(current_sum <= max_current)
+
+    def calculate(self, desired_currents: list[float]) -> list[float]:
+        """
+        Returns the optimal current allocation for a list of subsystems given a
+        list of their desired currents and current tolerances that determine
+        which subsystem gets less current if the current budget is exceeded.
+        Subsystems with a smaller tolerance are given higher priority.
+
+        Parameter ``desiredCurrents``:
+            The desired current for each subsystem.
+
+        Raises ``ValueError``:
+            if the number of desired currents doesn't equal the number of
+            tolerances passed in the constructor.
+        """
+        if self.__desired_currents.rows() != len(desired_currents):
+            raise ValueError(
+                "Number of desired currents must equal the number of tolerances passed in the constructor."
+            )
+
+        for i in range(len(desired_currents)):
+            self.__desired_currents[i].set_value(desired_currents[i])
+
+        self.__problem.solve()
+
+        result = []
+        for i in range(len(desired_currents)):
+            result.append(max(self.__allocated_currents.value(i), 0.0))
+
+        return result

examples/CurrentManager/main.py

@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+
+from current_manager import CurrentManager
+
+
+def main():
+    manager = CurrentManager([1.0, 5.0, 10.0, 5.0], 40.0)
+    currents = manager.calculate([25.0, 10.0, 5.0, 0.0])
+    print(f"currents = {currents}")
+
+
+if __name__ == "__main__":
+    main()

examples/CurrentManager/test/__init__.py

@@ -0,0 +1 @@
+

examples/CurrentManager/test/current_manager_test.py

@@ -0,0 +1,55 @@
+import pytest
+
+from CurrentManager import CurrentManager
+
+
+def test_current_manager_enough_current():
+    manager = CurrentManager([1.0, 5.0, 10.0, 5.0], 40.0)
+
+    currents = manager.calculate([25.0, 10.0, 5.0, 0.0])
+
+    assert currents[0] == pytest.approx(25.0, abs=1e-3)
+    assert currents[1] == pytest.approx(10.0, abs=1e-3)
+    assert currents[2] == pytest.approx(5.0, abs=1e-3)
+    assert currents[3] == pytest.approx(0.0, abs=1e-3)
+
+
+def test_current_manager_not_enough_current():
+    manager = CurrentManager([1.0, 5.0, 10.0, 5.0], 40.0)
+
+    currents = manager.calculate([30.0, 10.0, 5.0, 0.0])
+
+    # Expected values are from the following CasADi program:
+    #
+    # #!/usr/bin/env python3
+    #
+    # import casadi as ca
+    # import numpy as np
+    #
+    # opti = ca.Opti()
+    # allocated_currents = opti.variable(4, 1)
+    #
+    # current_tolerances = np.array([[1.0], [5.0], [10.0], [5.0]])
+    # desired_currents = np.array([[30.0], [10.0], [5.0], [0.0]])
+    #
+    # J = 0.0
+    # current_sum = 0.0
+    # for i in range(4):
+    #     error = desired_currents[i, 0] - allocated_currents[i, 0]
+    #     J += error**2 / current_tolerances[i] ** 2
+    #
+    #     current_sum += allocated_currents[i, 0]
+    #
+    #     # Currents must be nonnegative
+    #     opti.subject_to(allocated_currents[i, 0] >= 0.0)
+    # opti.minimize(J)
+    #
+    # # Keep total current below maximum
+    # opti.subject_to(current_sum <= 40.0)
+    #
+    # opti.solver("ipopt")
+    # print(opti.solve().value(allocated_currents))
+    assert currents[0] == pytest.approx(29.960, abs=1e-3)
+    assert currents[1] == pytest.approx(9.007, abs=1e-3)
+    assert currents[2] == pytest.approx(1.032, abs=1e-3)
+    assert currents[3] == pytest.approx(0.0, abs=1e-3)

examples/FlywheelDirectTranscription/main.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python3
+
+import math
+
+from jormungandr.optimization import OptimizationProblem
+import numpy as np
+
+
+def main():
+    T = 5.0  # s
+    dt = 0.005  # s
+    N = int(T / dt)
+
+    # Flywheel model:
+    # States: [velocity]
+    # Inputs: [voltage]
+    A = math.exp(-dt)
+    B = 1.0 - math.exp(-dt)
+
+    problem = OptimizationProblem()
+    X = problem.decision_variable(1, N + 1)
+    U = problem.decision_variable(1, N)
+
+    # Dynamics constraint
+    for k in range(N):
+        problem.subject_to(
+            X[:, k + 1 : k + 2] == A * X[:, k : k + 1] + B * U[:, k : k + 1]
+        )
+
+    # State and input constraints
+    problem.subject_to(X[0, 0] == 0.0)
+    problem.subject_to(-12 <= U)
+    problem.subject_to(U <= 12)
+
+    # Cost function - minimize error
+    r = np.array([[10.0]])
+    J = 0.0
+    for k in range(N + 1):
+        J += (r - X[:, k : k + 1]).T * (r - X[:, k : k + 1])
+    problem.minimize(J)
+
+    problem.solve()
+
+    # The first state
+    print(f"x₀ = {X.value(0, 0)}")
+
+    # The first input
+    print(f"u₀ = {U.value(0, 0)}")
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -45,4 +45,4 @@ install_args = [ "--strip" ]
 
 [tool.pytest.ini_options]
 minversion = "6.0"
-testpaths = [ "jormungandr/test" ]
+testpaths = [ "jormungandr/test", "examples/CurrentManager/test" ]