pr-adding-tensors

.gitignore

@@ -28,3 +28,6 @@ Manifest.toml
 
 # Local reports (analysis, status reports, previews) should not be tracked
 reports/
+
+# claude
+CLAUDE.local.md

Project.toml

@@ -1,6 +1,6 @@
 name = "CTParser"
 uuid = "32681960-a1b1-40db-9bff-a1ca817385d1"
-version = "0.8.0"
+version = "0.7.9"
 authors = ["Jean-Baptiste Caillau <[email protected]>"]
 
 [deps]

src/onepass.jl

@@ -96,19 +96,6 @@ function e_prefix!(p)
     return nothing
 end
 
-# Utils
-
-"""
-$(TYPEDSIGNATURES)
-
-Generate a fresh symbol by concatenating the given components and a
-`gensym()` suffix.
-
-This is used throughout the parser to create unique internal names that
-do not collide with user-defined identifiers.
-"""
-__symgen(s...) = Symbol(s..., gensym())
-
 """
 $(TYPEDEF)
 
@@ -200,28 +187,6 @@ __wrap(e, n, line) = quote
     end
 end
 
-"""
-$(TYPEDSIGNATURES)
-
-Return `true` if `x` represents a range.
-
-This predicate is specialised for `AbstractRange` values and for
-expressions of the form `i:j` or `i:p:j`.
-"""
-is_range(x) = false
-is_range(x::T) where {T<:AbstractRange} = true
-is_range(x::Expr) = (x.head == :call) && (x.args[1] == :(:))
-
-"""
-$(TYPEDSIGNATURES)
-
-Return `x` itself if it is a range, or a one-element array `[x]`.
-
-This is a normalisation helper used when interpreting constraint
-indices.
-"""
-as_range(x) = is_range(x) ? x : [x]
-
 # Main code
 
 """
@@ -580,7 +545,10 @@ function p_state_exa!(p, p_ocp, x, n, xx; components_names=nothing)
     ))
     code = __wrap(code, p.lnum, p.line)
     dyn_con = Symbol(:dyn_con, x) # name for the constraints associated with the dynamics
-    code = :($x = $code; $dyn_con = Vector{$pref.Constraint}(undef, $n)) # affectation must be done outside try ... catch (otherwise declaration known only to try local scope)
+    code = quote
+        $x = $code
+        $dyn_con = Vector{$pref.Constraint}(undef, $n) # affectation must be done outside try ... catch (otherwise declaration known only to try local scope)
+    end
     return code
 end
 
@@ -727,18 +695,21 @@ function p_constraint_exa!(p, p_ocp, e1, e2, e3, c_type, label)
             code = :(length($e1) == length($e3) == 1 || throw("this constraint must be scalar")) # (vs. __throw) since raised at runtime
             x0 = __symgen(:x0)
             xf = __symgen(:xf)
+            k = __symgen(:k)
             e2 = replace_call(e2, p.x, p.t0, x0)
             e2 = replace_call(e2, p.x, p.tf, xf)
             e2 = subs2(e2, x0, p.x, 0)
+            e2 = subs(e2, x0, :([$(p.x)[$k, 0] for $k ∈ 1:$(p.dim_x)]))
             e2 = subs2(e2, xf, p.x, :grid_size)
-            concat(code, :($pref.constraint($p_ocp, $e2; lcon=($e1), ucon=($e3))))
+            e2 = subs(e2, xf, :([$(p.x)[$k, grid_size] for $k ∈ 1:$(p.dim_x)]))
+            concat(code, :($pref.constraint($p_ocp, $e2; lcon=($e1), ucon=($e3)))) # debug: to vectorise
         end
         (:initial, rg) => begin
             if isnothing(rg)
                 rg = :(1:($(p.dim_x))) # x(t0) implies rg == nothing but means x[1:p.dim_x](t0)
                 e2 = subs(e2, p.x, :($(p.x)[$rg]))
             elseif !is_range(rg)
-                rg = as_range(rg)
+                rg = as_range(rg) # case rg = i (vs i:j or i:p:j)
             end
             code = :(
                 length($e1) == length($e3) == length($rg) || throw("wrong bound dimension")
@@ -757,7 +728,7 @@ function p_constraint_exa!(p, p_ocp, e1, e2, e3, c_type, label)
                 rg = :(1:($(p.dim_x)))
                 e2 = subs(e2, p.x, :($(p.x)[$rg]))
             elseif !is_range(rg)
-                rg = as_range(rg)
+                rg = as_range(rg) # case rg = i (vs i:j or i:p:j)
             end
             code = :(
                 length($e1) == length($e3) == length($rg) || throw("wrong bound dimension")
@@ -776,7 +747,7 @@ function p_constraint_exa!(p, p_ocp, e1, e2, e3, c_type, label)
                 rg = :(1:($(p.dim_v)))
                 e2 = subs(e2, p.v, :($(p.v)[$rg]))
             elseif !is_range(rg)
-                rg = as_range(rg)
+                rg = as_range(rg) # case rg = i (vs i:j or i:p:j)
             end
             code_box = :(
                 length($e1) == length($e3) == length($rg) || throw("wrong bound dimension")
@@ -791,10 +762,9 @@ function p_constraint_exa!(p, p_ocp, e1, e2, e3, c_type, label)
         end
         (:state_range, rg) => begin
             if isnothing(rg)
-                rg = :(1:($(p.dim_x)))
-                e2 = subs(e2, p.x, :($(p.x)[$rg]))
+                rg = :(1:($(p.dim_x))) # NB. no need to update e2 (unused) here
             elseif !is_range(rg)
-                rg = as_range(rg)
+                rg = as_range(rg) # case rg = i (vs i:j or i:p:j)
             end
             code_box = :(
                 length($e1) == length($e3) == length($rg) || throw("wrong bound dimension")
@@ -809,10 +779,9 @@ function p_constraint_exa!(p, p_ocp, e1, e2, e3, c_type, label)
         end
         (:control_range, rg) => begin
             if isnothing(rg)
-                rg = :(1:($(p.dim_u)))
-                e2 = subs(e2, p.u, :($(p.u)[$rg]))
+                rg = :(1:($(p.dim_u))) # NB. no need to update e2 (unused here)
             elseif !is_range(rg)
-                rg = as_range(rg)
+                rg = as_range(rg) # case rg = i (vs i:j or i:p:j)
             end
             code_box = :(
                 length($e1) == length($e3) == length($rg) || throw("wrong bound dimension")
@@ -831,8 +800,11 @@ function p_constraint_exa!(p, p_ocp, e1, e2, e3, c_type, label)
             ut = __symgen(:ut)
             e2 = replace_call(e2, [p.x, p.u], p.t, [xt, ut])
             j = __symgen(:j)
+            k = __symgen(:k)
             e2 = subs2(e2, xt, p.x, j)
+            e2 = subs(e2, xt, :([$(p.x)[$k, $j] for $k ∈ 1:$(p.dim_x)]))
             e2 = subs2(e2, ut, p.u, j)
+            e2 = subs(e2, ut, :([$(p.u)[$k, $j] for $k ∈ 1:$(p.dim_u)]))
             e2 = subs(e2, p.t, :($(p.t0) + $j * $(p.dt)))
             concat(
                 code,
@@ -931,26 +903,33 @@ function p_dynamics_coord_exa!(p, p_ocp, x, i, t, e)
     j1 = __symgen(:j)
     j2 = :($j1 + 1)
     j12 = :($j1 + 0.5)
+    k = __symgen(:k)
     ej1 = subs2(e, xt, p.x, j1)
+    ej1 = subs(ej1, xt, :([$(p.x)[$k, $j1] for $k ∈ 1:$(p.dim_x)]))
     ej1 = subs2(ej1, ut, p.u, j1)
+    ej1 = subs(ej1, ut, :([$(p.u)[$k, $j1] for $k ∈ 1:$(p.dim_u)]))
     ej1 = subs(ej1, p.t, :($(p.t0) + $j1 * $(p.dt)))
     ej2 = subs2(e, xt, p.x, j2)
+    ej2 = subs(ej2, xt, :([$(p.x)[$k, $j2] for $k ∈ 1:$(p.dim_x)]))
     ej2 = subs2(ej2, ut, p.u, j2)
+    ej2 = subs(ej2, ut, :([$(p.u)[$k, $j2] for $k ∈ 1:$(p.dim_u)]))
     ej2 = subs(ej2, p.t, :($(p.t0) + $j2 * $(p.dt)))
-    ej12 = subs5(e, xt, p.x, j1)
+    ej12 = subs2m(e, xt, p.x, j1)
+    ej12 = subs(ej12, xt, :([(($(p.x)[$k, $j1] + $(p.x)[$k, $j1 + 1]) / 2) for $k ∈ 1:$(p.dim_x)]))
     ej12 = subs2(ej12, ut, p.u, j1)
+    ej12 = subs(ej12, ut, :([$(p.u)[$k, $j1] for $k ∈ 1:$(p.dim_u)]))
     ej12 = subs(ej12, p.t, :($(p.t0) + $j12 * $(p.dt)))
     dxij = :($(p.x)[$i, $j2] - $(p.x)[$i, $j1])
     code = quote
         if scheme == :euler
-            $pref.constraint($p_ocp, $dxij - $(p.dt) * $ej1 for $j1 in 0:(grid_size - 1))
+            $pref.constraint($p_ocp, $dxij - $(p.dt) * $ej1 for $j1 in 0:grid_size-1)
         elseif scheme ∈ (:euler_implicit, :euler_b) # euler_b is deprecated
-            $pref.constraint($p_ocp, $dxij - $(p.dt) * $ej2 for $j1 in 0:(grid_size - 1))
+            $pref.constraint($p_ocp, $dxij - $(p.dt) * $ej2 for $j1 in 0:grid_size-1)
         elseif scheme == :midpoint
-            $pref.constraint($p_ocp, $dxij - $(p.dt) * $ej12 for $j1 in 0:(grid_size - 1))
+            $pref.constraint($p_ocp, $dxij - $(p.dt) * $ej12 for $j1 in 0:grid_size-1)
         elseif scheme ∈ (:trapeze, :trapezoidal) # trapezoidal is deprecated
             $pref.constraint(
-                $p_ocp, $dxij - $(p.dt) * ($ej1 + $ej2) / 2 for $j1 in 0:(grid_size - 1)
+                $p_ocp, $dxij - $(p.dt) * ($ej1 + $ej2) / 2 for $j1 in 0:grid_size-1
             )
         else
             throw(
@@ -1001,22 +980,27 @@ function p_lagrange_exa!(p, p_ocp, e, type)
     j1 = __symgen(:j)
     j2 = :($j1 + 1)
     j12 = :($j1 + 0.5)
+    k = __symgen(:k)
     ej1 = subs2(e, xt, p.x, j1)
+    ej1 = subs(ej1, xt, :([$(p.x)[$k, $j1] for $k ∈ 1:$(p.dim_x)]))
     ej1 = subs2(ej1, ut, p.u, j1)
+    ej1 = subs(ej1, ut, :([$(p.u)[$k, $j1] for $k ∈ 1:$(p.dim_u)]))
     ej1 = subs(ej1, p.t, :($(p.t0) + $j1 * $(p.dt)))
-    ej12 = subs5(e, xt, p.x, j1)
+    ej12 = subs2m(e, xt, p.x, j1)
+    ej12 = subs(ej12, xt, :([(($(p.x)[$k, $j1] + $(p.x)[$k, $j1 + 1]) / 2) for $k ∈ 1:$(p.dim_x)]))
     ej12 = subs2(ej12, ut, p.u, j1)
+    ej12 = subs(ej12, ut, :([$(p.u)[$k, $j1] for $k ∈ 1:$(p.dim_u)]))
     ej12 = subs(ej12, p.t, :($(p.t0) + $j12 * $(p.dt)))
     code = quote
         if scheme == :euler
-            $pref.objective($p_ocp, $(p.dt) * $ej1 for $j1 in 0:(grid_size - 1))
+            $pref.objective($p_ocp, $(p.dt) * $ej1 for $j1 in 0:grid_size-1)
         elseif scheme ∈ (:euler_implicit, :euler_b) # euler_b is deprecated
             $pref.objective($p_ocp, $(p.dt) * $ej1 for $j1 in 1:grid_size)
         elseif scheme == :midpoint
-            $pref.objective($p_ocp, $(p.dt) * $ej12 for $j1 in 0:(grid_size - 1))
+            $pref.objective($p_ocp, $(p.dt) * $ej12 for $j1 in 0:grid_size-1)
         elseif scheme ∈ (:trapeze, :trapezoidal) # trapezoidal is deprecated
             $pref.objective($p_ocp, $(p.dt) * $ej1 / 2 for $j1 in (0, grid_size))
-            $pref.objective($p_ocp, $(p.dt) * $ej1 for $j1 in 1:(grid_size - 1))
+            $pref.objective($p_ocp, $(p.dt) * $ej1 for $j1 in 1:grid_size-1)
         else
             throw(
                 "unknown numerical scheme: $scheme (possible choices are :euler, :euler_implicit, :midpoint, :trapeze)",
@@ -1062,10 +1046,13 @@ function p_mayer_exa!(p, p_ocp, e, type)
     pref = prefix_exa()
     x0 = __symgen(:x0)
     xf = __symgen(:xf)
+    k = __symgen(:k)
     e = replace_call(e, p.x, p.t0, x0)
     e = replace_call(e, p.x, p.tf, xf)
     e = subs2(e, x0, p.x, 0)
+    e = subs(e, x0, :([$(p.x)[$k, 0] for $k ∈ 1:$(p.dim_x)]))
     e = subs2(e, xf, p.x, :grid_size)
+    e = subs(e, xf, :([$(p.x)[$k, grid_size] for $k ∈ 1:$(p.dim_x)]))
     # now, x[i](t0) has been replaced by x[i, 0] and x[i](tf) by x[i, grid_size]
     code = :($pref.objective($p_ocp, $e))
     return __wrap(code, p.lnum, p.line)
@@ -1295,7 +1282,7 @@ function def_fun(e; log=false)
         $p_ocp = $pref.PreModel()
         $code
         $pref.definition!($p_ocp, $ee)
-        $pref.time_dependence!($p_ocp; autonomous=$p.is_autonomous)
+        $pref.time_dependence!($p_ocp; autonomous=$p.is_autonomous) # not $(p.xxxx) as this info is known statically
     end
 
     if is_active_backend(:exa)
@@ -1383,7 +1370,7 @@ function def_exa(e; log=false)
             $(p.box_u) # lvar and uvar for control
             $(p.box_v) # lvar and uvar for variable (after x and u for compatibility with CTDirect)
             $p_ocp = $pref.ExaCore(
-                base_type; backend=backend, minimize=($p.criterion == :min)
+                base_type; backend=backend, minimize=($p.criterion == :min) # not $(p.xxxx) as this info is known statically
             )
             $code
             $dyn_check