Update userfunc.jl to account for DAE systems

.gitignore

@@ -1 +1,3 @@
-docs/build
+docs/build
+Manifest.toml
+.vscode

Project.toml

@@ -1,10 +1,13 @@
 name = "EOptInterface"
 uuid = "96d9f51d-8c41-4e2d-b1a3-7b48dea3ddd1"
-authors = ["Joseph Choi <jsphchoi@mit.edu>", "Dimitri Alston <dimitri.alston@uconn.edu>", "Pengfei Xu <pengfei.xu@uconn.edu>", "Matthew Stuber <matthew.stuber@uconn.edu>"]
 version = "0.2.0"
+authors = ["Joseph Choi <jsphchoi@mit.edu>", "Dimitri Alston <dimitri.alston@uconn.edu>", "Pengfei Xu <pengfei.xu@uconn.edu>", "Matthew Stuber <matthew.stuber@uconn.edu>"]
 
 [deps]
+CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -13,7 +16,10 @@ SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [compat]
+CSV = "0.10.17"
+DataFrames = "1.8.2"
 DocStringExtensions = "0.9"
+Ipopt = "1.15.0"
 JuMP = "1"
 ModelingToolkit = "11"
 Reexport = "1"

examples/dae_model.jl

@@ -0,0 +1,57 @@
+using EOptInterface
+using Ipopt
+using JuMP
+using ModelingToolkit
+using ModelingToolkit: t_nounits as t, D_nounits as D
+using SciCompDSL
+
+@mtkmodel TankValve begin
+    @parameters begin
+        A = 1.0     # tank cross-section area  [m²]
+        q_in = 2.0     # constant inlet flow      [m³/s]
+        k_v                # valve coefficient (FREE – no default)
+    end
+    @variables begin
+        h(t) = 1.0    # liquid level [m]  (ODE state, IC = 1)
+        q_out(t), [irreducible = true]           # outlet flow  [m³/s] (algebraic, no IC)
+    end
+    @equations begin
+        # ODE: tank mass balance
+        D(h) ~ (q_in - q_out) / A
+        # Algebraic constraint: valve equation  (written as 0 ~ rhs)
+        q_out ~ k_v * sqrt(h)
+    end
+end
+
+@mtkcompile sys = TankValve()
+
+tspan = (0.0, 10.0)
+tstep = 0.01
+
+N = Int(floor((tspan[2] - tspan[1]) / tstep)) + 1
+V = length(ModelingToolkit.unknowns(sys))
+n_dvars = length(decision_vars(sys))
+P = n_dvars - V                      # number of free parameters
+
+jump_model = JuMP.Model(Ipopt.Optimizer)
+
+JuMP.@variable(jump_model, 0.0 <= xs[1:V, 1:N] <= 100.0)
+JuMP.@variable(jump_model, 0.01 <= ps[1:P] <= 10.0)
+
+register_daesystem(jump_model, sys, tspan, tstep, "EE")
+
+h_target = 4.0
+JuMP.@objective(jump_model, Min, (xs[1, N] - h_target)^2)
+
+JuMP.optimize!(jump_model)
+
+println("═══ Results ═══")
+println("Termination : ", JuMP.termination_status(jump_model))
+println("Primal       : ", JuMP.primal_status(jump_model))
+println("Solve time   : ", round(JuMP.solve_time(jump_model), digits=4), " s")
+println("Objective    : ", round(JuMP.objective_value(jump_model), digits=6))
+
+# Extract optimal k_v
+kv_opt = JuMP.value(ps[1])
+println("k_v*         : ", round(kv_opt, digits=4))
+println("Expected k_v : 1.0  (analytical: q_in / sqrt(h_target) = 2/√4 = 1)")

src/EOptInterface.jl

@@ -31,6 +31,6 @@ module EOptInterface
     include("userfuncs.jl")
 
     # Exports
-    export decision_vars, full_solution, register_nlsystem, register_odesystem
+    export decision_vars, full_solution, register_nlsystem, register_odesystem, register_daesystem
 
 end

src/userfuncs.jl

@@ -108,6 +108,114 @@ function register_odesystem(model::JuMP.Model, sys::ModelingToolkit.System, tspa
     return
 end
 
+"""
+    $(DocStringExtensions.TYPEDSIGNATURES)
+
+Automatically applies specified direct transcription method and registers the 
+discretized ODE `ModelingToolkit.System` as algebraic JuMP constraints besides adding constraints of algebraic 
+equations in the DAE model.
+Currently supports Explicit Euler "EE" and Implicit Euler "IE".
+"""
+function register_daesystem(model::JuMP.Model, sys::ModelingToolkit.System, tspan::Tuple{Real,Real}, tstep::Real, integrator::String)
+    if integrator != "EE" && integrator != "IE"
+        error("Available integrators: EE, IE")
+    end
+    # Number of discrete time nodes
+    N = Int(floor((tspan[2] - tspan[1]) / tstep)) + 1
+    # Number of ODE variables
+    V = length(ModelingToolkit.unknowns(sys))
+
+    init_values = copy(ModelingToolkit.initial_conditions(sys).dict)
+    param_keys = ModelingToolkit.parameters(sys)
+    param_dict = Dict()
+    dvars = decision_vars(sys)
+    sys_unknowns = ModelingToolkit.unknowns(sys)
+    for key in param_keys
+        if any(isequal(key), dvars)
+            continue
+        else
+            param_dict[key] = init_values[key]
+        end
+    end
+
+    dx_funcs = []
+    x_funcs = []
+
+
+    for j in 1:V # j is the counter for each unknown state variables (this does not include unknown parameters)
+        if string(ModelingToolkit.full_equations(sys)[j].lhs) == "0"
+            # the equation is algebraic
+            expr_j = ModelingToolkit.full_equations(sys)[j].rhs
+            expr_j = SymbolicUtils.substitute(expr_j, ModelingToolkit.bindings(sys))
+            expr_j = SymbolicUtils.substitute(expr_j, ModelingToolkit.bindings(sys))
+            # Fully substitute fixed parameters with default values
+            while ~isempty(intersect(Symbolics.get_variables(expr_j), keys(param_dict)))
+                expr_j = SymbolicUtils.substitute(expr_j, param_dict)
+            end
+            # Build runtime callable (remaining free symbols: state vars + decision vars)
+            xj_func = Symbolics.build_function(
+                expr_j,
+                dvars...,
+                expression=Val{false}
+            )
+            push!(x_funcs, xj_func)
+
+        else
+            # the equation is Differential
+            expr_j = ModelingToolkit.full_equations(sys)[j].rhs
+            expr_j = SymbolicUtils.substitute(expr_j, ModelingToolkit.bindings(sys))
+            expr_j = SymbolicUtils.substitute(expr_j, ModelingToolkit.bindings(sys))
+            # Fully substitute fixed parameters with default values
+            while ~isempty(intersect(Symbolics.get_variables(expr_j), keys(param_dict)))
+                expr_j = SymbolicUtils.substitute(expr_j, param_dict)
+            end
+
+            # Build runtime callable (remaining free symbols: state vars + decision vars)
+            dxj_func = Symbolics.build_function(
+                expr_j,
+                dvars...,
+                expression=Val{false}
+            )
+            push!(dx_funcs, dxj_func)
+        end
+
+    end
+
+    # Extract JuMP decision variable vector (pure parameters, not state vars)
+    ps = JuMP.all_variables(model)[end-length(setdiff(dvars, sys_unknowns))+1:end]
+    # Reshape remaining JuMP variables into state trajectory matrix [V × N]
+    xs = reshape(setdiff(JuMP.all_variables(model), ps), V, N)
+
+    # Extract initial conditions from the ModelingToolkit system and fix them in the JuMP model for x[1:V,1]
+
+
+
+    for (i, unk_var) in enumerate(sys_unknowns)
+        if haskey(init_values, unk_var)
+            JuMP.fix(xs[i, 1], init_values[unk_var].val, force=true)
+        end
+    end
+
+    # Formulate ODE discretization constraints
+    if integrator == "EE"
+        # Differential equations: Explicit Euler
+        JuMP.@constraint(model, [j in 1:length(dx_funcs), i in 1:(N-1)],
+            xs[j, i+1] == xs[j, i] + tstep * dx_funcs[j](xs[:, i]..., ps...))
+        # Algebraic equations: enforce at each step i
+        JuMP.@constraint(model, [j in 1:length(x_funcs), i in 1:N],
+            x_funcs[j](xs[:, i]..., ps...) == 0)
+
+    elseif integrator == "IE"
+        # Differential equations: Implicit Euler
+        JuMP.@constraint(model, [j in 1:length(dx_funcs), i in 1:(N-1)],
+            xs[j, i+1] == xs[j, i] + tstep * dx_funcs[j](xs[:, i+1]..., ps...))
+        # Algebraic equations: enforce at each step i+1
+        JuMP.@constraint(model, [j in 1:length(x_funcs), i in 1:(N-1)],
+            x_funcs[j](xs[:, i+1]..., ps...) == 0)
+    end
+
+    return
+end
 """
     $(DocStringExtensions.TYPEDSIGNATURES)
 

test/runtests.jl

@@ -232,3 +232,54 @@ end
     @test isapprox(JuMP.objective_value(model), 16796.032234817612, atol=1e-3)
 
 end
+
+@testset "DAE Model" begin
+
+    @mtkmodel TankValve begin
+        @parameters begin
+            A = 1.0     # tank cross-section area  [m²]
+            q_in = 2.0     # constant inlet flow      [m³/s]
+            k_v                # valve coefficient (FREE – no default)
+        end
+        @variables begin
+            h(t) = 1.0    # liquid level [m]  (ODE state, IC = 1)
+            q_out(t), [irreducible=true]           # outlet flow  [m³/s] (algebraic, no IC)
+        end
+        @equations begin
+            # ODE: tank mass balance
+            D(h) ~ (q_in - q_out) / A
+            # Algebraic constraint: valve equation  (written as 0 ~ rhs)
+            q_out ~ k_v * sqrt(h)
+        end
+    end
+
+    @mtkcompile system = TankValve()
+
+    tspan = (0.0, 10.0)
+    tstep = 0.01
+
+    N = Int(floor((tspan[2] - tspan[1]) / tstep)) + 1
+    V = length(ModelingToolkit.unknowns(system))
+    n_dvars = length(decision_vars(system))
+    P = n_dvars - V # number of parameters
+
+    model = JuMP.Model(Ipopt.Optimizer)
+
+    # State trajectory variables  [V × N]
+    JuMP.@variable(model, 0.0 <= xs[1:V, 1:N] <= 100.0)
+    # Free parameter variables    [P]
+    JuMP.@variable(model, 0.01 <= ps[1:P] <= 10.0)
+
+    register_daesystem(model, system, tspan, tstep, "EE")
+
+    h_target = 4.0
+    # h is the first unknown → xs[1, N]
+    JuMP.@objective(model, Min, (xs[1, N] - h_target)^2)
+
+    JuMP.optimize!(model)
+    soln_dict = EOptInterface.full_solution(model, system)
+    @test JuMP.termination_status(model) == JuMP.LOCALLY_SOLVED
+    @test JuMP.primal_status(model) == JuMP.FEASIBLE_POINT
+    @test isapprox(JuMP.objective_value(model), 0.0, atol=1e-3)
+    @test isapprox(JuMP.value(ps[1]), 0.9698, atol = 1e-3)
+end