Capture observed dependence on unknowns in simplified nonlinear systems

docs/src/tutorials/nonlinear.md

@@ -9,31 +9,27 @@ We use (unknown) variables for our nonlinear system.
 ```@example nonlinear
 using ModelingToolkit, NonlinearSolve
 
+# Define a nonlinear system
 @variables x y z
 @parameters σ ρ β
+eqs = [0 ~ σ * (y - x)
+       0 ~ x * (ρ - z) - y
+       0 ~ x * y - β * z]
+guesses = [x => 1.0, y => 0.0, z => 0.0]
+ps = [σ => 10.0, ρ => 26.0, β => 8 / 3]
+@mtkbuild ns = NonlinearSystem(eqs)
 
-# Define a nonlinear system
-eqs = [0 ~ σ * (y - x),
-    0 ~ x * (ρ - z) - y,
-    0 ~ x * y - β * z]
-@mtkbuild ns = NonlinearSystem(eqs, [x, y, z], [σ, ρ, β])
-
-guess = [x => 1.0,
-    y => 0.0,
-    z => 0.0]
-
-ps = [σ => 10.0
-      ρ => 26.0
-      β => 8 / 3]
+guesses = [x => 1.0, y => 0.0, z => 0.0]
+ps = [σ => 10.0, ρ => 26.0, β => 8 / 3]
 
-prob = NonlinearProblem(ns, guess, ps)
+prob = NonlinearProblem(ns, guesses, ps)
 sol = solve(prob, NewtonRaphson())
 ```
 
 We can similarly ask to generate the `NonlinearProblem` with the analytical
 Jacobian function:
 
 ```@example nonlinear
-prob = NonlinearProblem(ns, guess, ps, jac = true)
+prob = NonlinearProblem(ns, guesses, ps, jac = true)
 sol = solve(prob, NewtonRaphson())
 ```

src/systems/nonlinear/nonlinearsystem.jl

@@ -193,8 +193,12 @@ function calculate_jacobian(sys::NonlinearSystem; sparse = false, simplify = fal
         return cache[1]
     end
 
-    rhs = [eq.rhs for eq in equations(sys)]
+    # observed equations may depend on unknowns, so substitute them in first
+    # TODO: rather keep observed derivatives unexpanded, like "Differential(obs)(expr)"?
+    obs = Dict(eq.lhs => eq.rhs for eq in observed(sys))
+    rhs = map(eq -> fixpoint_sub(eq.rhs, obs), equations(sys))
     vals = [dv for dv in unknowns(sys)]
+
     if sparse
         jac = sparsejacobian(rhs, vals, simplify = simplify)
     else
@@ -215,7 +219,8 @@ function generate_jacobian(
 end
 
 function calculate_hessian(sys::NonlinearSystem; sparse = false, simplify = false)
-    rhs = [eq.rhs for eq in equations(sys)]
+    obs = Dict(eq.lhs => eq.rhs for eq in observed(sys))
+    rhs = map(eq -> fixpoint_sub(eq.rhs, obs), equations(sys))
     vals = [dv for dv in unknowns(sys)]
     if sparse
         hess = [sparsehessian(rhs[i], vals, simplify = simplify) for i in 1:length(rhs)]

test/nonlinearsystem.jl

@@ -283,3 +283,37 @@ sys = structural_simplify(ns)
 @test length(equations(sys)) == 0
 sys = structural_simplify(ns; conservative = true)
 @test length(equations(sys)) == 1
+
+# https://github.com/SciML/ModelingToolkit.jl/issues/2858
+@testset "Jacobian/Hessian with observed equations that depend on unknowns" begin
+    @variables x y z
+    @parameters σ ρ β
+    eqs = [0 ~ σ * (y - x)
+           0 ~ x * (ρ - z) - y
+           0 ~ x * y - β * z]
+    guesses = [x => 1.0, y => 0.0, z => 0.0]
+    ps = [σ => 10.0, ρ => 26.0, β => 8 / 3]
+    @mtkbuild ns = NonlinearSystem(eqs)
+
+    @test isequal(calculate_jacobian(ns), [(-1 - z + ρ)*σ -x*σ
+                                           2x*(-z + ρ) -β-(x^2)])
+    # solve without analytical jacobian
+    prob = NonlinearProblem(ns, guesses, ps)
+    sol = solve(prob, NewtonRaphson())
+    @test sol.retcode == ReturnCode.Success
+
+    # solve with analytical jacobian
+    prob = NonlinearProblem(ns, guesses, ps, jac = true)
+    sol = solve(prob, NewtonRaphson())
+    @test sol.retcode == ReturnCode.Success
+
+    # system that contains a chain of observed variables when simplified
+    @variables x y z
+    eqs = [0 ~ x^2 + 2z + y, z ~ y, y ~ x] # analytical solution x = y = z = 0 or -3
+    @mtkbuild ns = NonlinearSystem(eqs) # solve for y with observed chain z -> x -> y
+    @test isequal(expand.(calculate_jacobian(ns)), [3 // 2 + y;;])
+    @test isequal(calculate_hessian(ns), [[1;;]])
+    prob = NonlinearProblem(ns, unknowns(ns) .=> -4.0) # give guess < -3 to reach -3
+    sol = solve(prob, NewtonRaphson())
+    @test sol[x] ≈ sol[y] ≈ sol[z] ≈ -3
+end