Merge pull request #49 from Tim-Salzmann/v2

Update to v2

.github/workflows/ci_v2.yaml

@@ -0,0 +1,92 @@
+name: L4CasADi v2
+
+on:
+  push:
+    branches: [ v2 ]
+
+jobs:
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    timeout-minutes: 5
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        ref: 'v2'
+    - name: Run mypy
+      run: |
+        pip install mypy
+        mypy . --ignore-missing-imports --exclude examples
+    - name: Run flake8
+      run: |
+        pip install flake8
+        # stop the build if there are Python syntax errors or undefined names
+        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+
+  tests:
+    runs-on: ${{ matrix.runs-on }}
+    needs: [ lint ]
+    timeout-minutes: 60
+    strategy:
+      fail-fast: false
+      matrix:
+        runs-on: [ubuntu-latest, ubuntu-20.04, macos-latest]
+
+    name: Tests on ${{ matrix.runs-on }}
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v3
+      with:
+        ref: 'v2'
+        fetch-depth: 0
+
+    - name: Install Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: '>=3.9 <3.12'
+
+    - name: Install L4CasADi
+      run: |
+        python -m pip install --upgrade pip
+        pip install torch --index-url https://download.pytorch.org/whl/cpu  # Ensure CPU torch version
+        pip install -r requirements_build.txt
+        pip install . -v --no-build-isolation 
+
+    - name: Test with pytest
+      working-directory: ./tests
+      run: |
+        pip install pytest
+        pytest .
+
+  test-on-aarch:
+    runs-on: ubuntu-latest
+    needs: [ lint ]
+    timeout-minutes: 60
+
+    name: Tests on aarch64
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+        with:
+          ref: 'v2'
+          fetch-depth: 0
+      - uses: uraimo/run-on-arch-action@v2
+        name: Install and Test
+        with:
+          arch: aarch64
+          distro: ubuntu20.04
+          install: |
+            apt-get update
+            apt-get install -y --no-install-recommends python3.9 python3-pip python-is-python3
+            pip install -U pip
+            apt-get install -y build-essential
+
+          run: |
+            python -m pip install --upgrade pip
+            pip install torch --index-url https://download.pytorch.org/whl/cpu  # Ensure CPU torch version
+            pip install -r requirements_build.txt
+            pip install . -v --no-build-isolation 
+            # pip install pytest
+            # pytest .

README.md

@@ -26,6 +26,32 @@ arXiv: [Learning for CasADi: Data-driven Models in Numerical Optimization](https
 
 Talk: [Youtube](https://youtu.be/UYdkRnGr8eM?si=KEPcFEL9b7Vk2juI&t=3348)
 
+## L4CasADi v2 Breaking Changes
+After feedback from first use-cases L4CasADi v2 is designed with efficiency and simplicity in mind.
+
+This leads to the following breaking changes:
+
+- L4CasADi v2 can leverage PyTorch's batching capabilities for increased efficiency. When passing `batched=True`,
+L4CasADi will understand the **first** input dimension as batch dimension. Thus, first and second-order derivatives
+across elements of this dimension are assumed to be **sparse-zero**. To make use of this, instead of having multiple calls to a L4CasADi function in
+your CasADi program, batch all inputs together and have a single L4CasADi call. An example of this can be seen when
+comparing the [non-batched NeRF example](examples/nerf_trajectory_optimization/nerf_trajectory_optimization.py) with the
+[batched NeRF example](examples/nerf_trajectory_optimization/nerf_trajectory_optimization_batched.py) which is faster by
+a factor of 5-10x.
+- L4CasADi v2 will not change the shape of an input anymore as this was a source of confusion. The tensor forwarded to
+the PyTorch model will resemble the **exact dimension** of the input variable by CasADi. You are responsible to make
+sure that the PyTorch model handles a **two-dimensional** input matrix! Accordingly, the parameter
+`model_expects_batch_dim` is removed.
+- By default, L4CasADi v2 will not provide the Hessian, but the Jacobian of the Adjoint. This is sufficient for most
+many optimization problems. However, you can explicitly request the generation of the Hessian by passing
+`generate_jac_jac=True`.
+
+[//]: # (L4CasADi v2 can use the new **torch compile** functionality starting from PyTorch 2.4. By passing `scripting=False`. This
+will lead to a longer compile time on first L4CasADi function call but will lead to a overall faster
+execution. However, currently this functionality is experimental and not fully stable across all models. In the long
+term there is a good chance this will become the default over scripting once the functionality is stabilized by the
+Torch developers.)
+
 ## Table of Content
 - [Projects using L4CasADi](#projects-using-l4casadi)
 - [Installation](#installation)
@@ -205,14 +231,6 @@ https://github.com/Tim-Salzmann/l4casadi/blob/421de6ef408267eed0fd2519248b2152b6
 
 ## FYIs
 
-### Batch Dimension
-
-If your PyTorch model expects a batch dimension as first dimension (which most models do) you should pass
-`model_expects_batch_dim=True` to the `L4CasADi` constructor. The `MX` input to the L4CasADi component is then expected
-to be a vector of shape `[X, 1]`. L4CasADi will add a batch dimension of `1` automatically such that the input to the
-underlying PyTorch model is of shape `[1, X]`.
-
----
 
 ### Warm Up
 

examples/acados.py

@@ -128,7 +128,7 @@ def ocp(self):
             ocp.cost.W = np.array([[1.]])
 
             # Trivial PyTorch index 0
-            l4c_y_expr = l4c.L4CasADi(lambda x: x[0], name='y_expr', model_expects_batch_dim=False)
+            l4c_y_expr = l4c.L4CasADi(lambda x: x[0], name='y_expr')
 
             ocp.model.cost_y_expr = l4c_y_expr(x)
             ocp.model.cost_y_expr_e = x[0]

examples/cpp_usage/generate.py

@@ -10,7 +10,7 @@ def forward(self, x):
 
 
 def generate():
-    l4casadi_model = l4c.L4CasADi(TorchModel(), model_expects_batch_dim=False, name='sin_l4c')
+    l4casadi_model = l4c.L4CasADi(TorchModel(), name='sin_l4c')
 
     sym_in = cs.MX.sym('x', 1, 1)
 

examples/fish_turbulent_flow/utils.py

@@ -266,7 +266,7 @@ def import_l4casadi_model(device):
     x = cs.MX.sym("x", 3)
     xn = (x - meanX) / stdX
 
-    y = l4c.L4CasADi(model, name="turbulent_model", model_expects_batch_dim=True)(xn)
+    y = l4c.L4CasADi(model, name="turbulent_model", generate_adj1=False, generate_jac_jac=True)(xn.T).T
     y = y * stdY + meanY
     fU = cs.Function("fU", [x], [y[0]])
     fV = cs.Function("fV", [x], [y[1]])

examples/matlab/export.py

@@ -10,7 +10,7 @@ def forward(self, x):
 
 
 def generate():
-    l4casadi_model = l4c.L4CasADi(TorchModel(), model_expects_batch_dim=False, name='sin_l4c')
+    l4casadi_model = l4c.L4CasADi(TorchModel(), name='sin_l4c')
     sym_in = cs.MX.sym('x', 1, 1)
     l4casadi_model.build(sym_in)
     return

examples/naive/readme.py

@@ -3,15 +3,15 @@
 
 
 naive_mlp = l4c.naive.MultiLayerPerceptron(2, 128, 1, 2, 'Tanh')
-l4c_model = l4c.L4CasADi(naive_mlp, model_expects_batch_dim=True)
+l4c_model = l4c.L4CasADi(naive_mlp)
 
-x_sym = cs.MX.sym('x', 2, 1)
+x_sym = cs.MX.sym('x', 3, 2)
 y_sym = l4c_model(x_sym)
 f = cs.Function('y', [x_sym], [y_sym])
 df = cs.Function('dy', [x_sym], [cs.jacobian(y_sym, x_sym)])
-ddf = cs.Function('ddy', [x_sym], [cs.hessian(y_sym, x_sym)[0]])
+ddf = cs.Function('ddy', [x_sym], [cs.jacobian(cs.jacobian(y_sym, x_sym), x_sym)])
 
-x = cs.DM([[0.], [2.]])
+x = cs.DM([[0., 2.], [0., 2.], [0., 2.]])
 print(l4c_model(x))
 print(f(x))
 print(df(x))

examples/nerf_trajectory_optimization/density_nerf.py

@@ -46,7 +46,7 @@ def __init__(self):
             [nn.Linear(self.input_ch, W)]
             + [
                 nn.Linear(W, W)
-                if i not in self.skips
+                if i != 4
                 else nn.Linear(W + self.input_ch, W)
                 for i in range(D - 1)
             ]
@@ -60,7 +60,7 @@ def forward(self, x):
         for i, l in enumerate(self.pts_linears):
             h = self.pts_linears[i](h)
             h = F.relu(h)
-            if i in self.skips:
+            if i == 4:
                 h = torch.cat([input_pts, h], -1)
 
         alpha = self.alpha_linear(h)

examples/nerf_trajectory_optimization/nerf_trajectory_optimization.py

@@ -10,6 +10,7 @@
 
 CASE = 1
 
+os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
 
 def polynomial(n, n_eval):
     """Generates a symbolic function for a polynomial of degree n-1"""
@@ -86,7 +87,7 @@ def trajectory_generator_solver(n, n_eval, L, warmup, threshold):
             f += cs.sum2(sk**2)
 
             # While having a maximum density (1.) of the NeRF as constraint.
-            lk = L(pk.T)
+            lk = L(pk)
             g = cs.horzcat(g, lk)
             lbg = cs.horzcat(lbg, cs.DM([-10e32]).T)
             ubg = cs.horzcat(ubg, cs.DM([threshold]).T)
@@ -175,7 +176,7 @@ def main():
         strict=False,
     )
     # -------------------------- Create L4CasADi Module -------------------------- #
-    l4c_nerf = l4c.L4CasADi(model)
+    l4c_nerf = l4c.L4CasADi(model, scripting=False)
 
     # ---------------------------------------------------------------------------- #
     #                                   NLP warmup                                   #