Use Enzyme in ipopt rather than FiniteDiff (#20)

* working Enzyme

# Conflicts:
#	Project.toml
#	examples/budgerigar.jl
#	src/endotherm/thermoregulation/ipopt.jl
#	src/endotherm/thermoregulation/rulebased.jl

* Track HeatExchange API renames and add smoothing kwarg to IPOPTControl

Caller-side updates for upstream renames: generated_heat_flow → metabolic_heat_flow,
per-side insulation_depth/conductivity moved under .dorsal/.ventral, shape_b → axis_ratio_b,
respiration_mass → respiration_mass_flow. Adds `smoothing` field on IPOPTControl
(default SmoothBound(1e-5)) threaded into nlp_pack so AD sees differentiable kinks.

* cleanup and deps

* more cleanup

* Apply suggestion from @rafaqz

Author: rafaqz

Project.toml

@@ -6,7 +6,7 @@ authors = ["Michael Kearney and Rafael Schouten"]
 [deps]
 BiophysicalGeometry = "06c61620-4f66-4020-a018-07befe374221"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
-FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Flatten = "4c728ea3-d9ee-5c9a-9642-b6f7d7dc04fa"
 FluidProperties = "d16f1b92-71d3-42f3-aafa-43e5bb0efda8"
 HeatExchange = "720bd838-cdd6-4aad-926b-452b962dce21"
@@ -21,7 +21,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 UnitfulMoles = "999f2bd7-36bf-5ba7-9bc1-c9473aa75374"
 
 [sources]
-BiophysicalGeometry = {rev = "characteristic-dimension-options", url = "https://github.com/BiophysicalEcology/BiophysicalGeometry.jl"}
+BiophysicalGeometry = {rev = "main", url = "https://github.com/BiophysicalEcology/BiophysicalGeometry.jl"}
 FluidProperties = {rev = "main", url = "https://github.com/BiophysicalEcology/FluidProperties.jl"}
 HeatExchange = {rev = "main", url = "https://github.com/BiophysicalEcology/HeatExchange.jl"}
 
@@ -31,7 +31,7 @@ BiophysicalGeometry = "0.1"
 CSV = "0.10"
 ConstructionBase = "1.6"
 DataFrames = "1"
-FiniteDiff = "2.30.0"
+Enzyme = "0.13.144"
 Flatten = "0.4.3"
 FluidProperties = "0.1"
 HeatExchange = "0.1"

docs/ipopt_endotherm_thermoregulation.md

@@ -208,11 +208,9 @@ performance. The exact Hessian requires O(n²) = 81 finite-difference evaluation
 iteration; L-BFGS constructs a low-rank approximation from gradient differences instead,
 reducing the cost substantially.
 
-**Why finite differences?** `Unitful.jl` units propagate through `HeatExchange.heat_balance`,
-making the residuals and objective incompatible with Julia automatic differentiation packages
-(Zygote, ForwardDiff). All gradients, Jacobians, and Hessians are computed via `FiniteDiff.jl`.
-The `hess` and `cons_h` callbacks are registered (required by `IpoptOptimizer`) but never
-called at runtime when L-BFGS is active.
+**Differentiation strategy.** Gradients and constraint Jacobians are computed via
+[`Enzyme.jl`](https://github.com/EnzymeAD/Enzyme.jl): reverse-mode for the scalar objective
+gradient and forward-mode for the constraint Jacobian.
 
 ---
 
@@ -234,7 +232,7 @@ for (air_temperature, ...) in zip(air_temperatures, ...)
     out = thermoregulate(Endotherm(), IPOPTControl(), organism, environment,
                          generated_heat_flow_ipopt, skin_temperature_ipopt, insulation_temperature_ipopt)
 
-    generated_heat_flow_ipopt    = out.energy_flows.generated_heat_flow
+    generated_heat_flow_ipopt    = out.energy_flows.metabolic_heat_flow
     skin_temperature_ipopt       = out.thermoregulation.skin_temperature
     insulation_temperature_ipopt = out.thermoregulation.insulation_temperature
 end
@@ -311,7 +309,7 @@ insulation_temperature_ipopt = air_temperatures[1]
 for (air_temperature, rh, q10) in zip(air_temperatures, ...)
     out = thermoregulate(Endotherm(), IPOPTControl(), organism, environment,
                          generated_heat_flow_ipopt, skin_temperature_ipopt, insulation_temperature_ipopt)
-    generated_heat_flow_ipopt    = out.energy_flows.generated_heat_flow
+    generated_heat_flow_ipopt    = out.energy_flows.metabolic_heat_flow
     skin_temperature_ipopt       = out.thermoregulation.skin_temperature
     insulation_temperature_ipopt = out.thermoregulation.insulation_temperature
     push!(ipopt_results, ...)
@@ -331,6 +329,6 @@ Typical tuning guidance for penalty weights:
 
 ## Limitations and Future Work
 
-- **Automatic differentiation:** `Unitful.jl` units propagate through `HeatExchange.heat_balance`, making it incompatible with `ForwardDiff` or `Zygote`. All derivatives are computed by finite differences via `FiniteDiff.jl`. Stripping units from the inner-loop heat balance computation would allow exact derivatives and potentially faster convergence.
+- **Automatic differentiation:** derivatives are computed via `Enzyme.jl` (reverse-mode for the objective gradient, forward-mode for the constraint Jacobian). The Hessian and constraint Hessian callbacks are stubs because IPOPT runs with the L-BFGS Hessian approximation; supplying exact second-order information may improve convergence further.
 - **Dorsal/ventral symmetry:** the mean-weighted body approximation merges dorsal and ventral sides. The full `solve_metabolic_rate` computes them separately. In strongly asymmetric conditions (e.g. high solar loading on dorsal surface) this may introduce small errors.
 - **Global optimality:** IPOPT finds a local optimum of the NLP. For well-posed problems the objective is approximately convex and the local optimum is unique, but unusual initial conditions or extreme parameter combinations may yield suboptimal solutions.

examples/budgerigar.jl

@@ -516,6 +516,7 @@ skin_temperature_ipopt          = metabolism_pars_init.core_temperature - 3.0u"K
 insulation_temperature_ipopt    = air_temperatures[1]
 
 @time for (air_temperature, relative_humidity, q10) in zip(air_temperatures, experimental_relative_humdities, q10s)
+    global generated_heat_flow_ipopt, skin_temperature_ipopt, insulation_temperature_ipopt
 
     environment_vars_loop = example_environment_vars(;
         air_temperature,
@@ -563,7 +564,7 @@ insulation_temperature_ipopt    = air_temperatures[1]
         insulation_temperature_dorsal           = tr.dorsal.insulation_temperature,
         insulation_temperature_ventral          = tr.ventral.insulation_temperature,
         insulation_depth_dorsal       = tr.dorsal.insulation_depth,
-        axis_ratio_b      = tr.aspect_ratio,
+        axis_ratio_b      = tr.axis_ratio_b,
         flesh_conductivity            = tr.flesh_conductivity,
         pant                          = tr.pant,
         skin_wetness                  = tr.skin_wetness,
@@ -701,6 +702,9 @@ plot(pc5, pc6, pc7, layout = (1, 3), size = (1200, 350))
 #      layout = (7, 1),
 #      size = (700, 1600))
 
-plot(pc1, pc2, pc3, pc4, pc5, pc6, pc7,
+final_plot = plot(pc1, pc2, pc3, pc4, pc5, pc6, pc7,
     layout = (4, 2),
     size = (700, 1000))
+display(final_plot)
+println("\nPlot window open. Press Enter to exit.")
+readline()

src/BiophysicalBehaviour.jl

@@ -11,7 +11,7 @@ using UnitfulMoles
 
 using BiophysicalGeometry: AbstractBody, shape
 
-using FiniteDiff
+using Enzyme
 using Optimization
 using OptimizationIpopt
 import SciMLBase

src/endotherm/thermoregulation/ipopt.jl

@@ -145,22 +145,48 @@ function _run_ipopt(
         setpoint_temperature = setpoint_temperature_K,
     )
 
-    obj_fn(x, _)    = _objective_value_weighted(x, opt_pars)
+    obj_fn(x, _)     = _objective_value_weighted(x, opt_pars)
     res_fn!(r, x, _) = _heat_balance_residuals_weighted!(r, x, nlp_pars)
-    grad_fn!(g, x, _) = FiniteDiff.finite_difference_gradient!(g, e -> _objective_value_weighted(e, opt_pars), x)
-    hess_fn!(H, x, _) = (H .= FiniteDiff.finite_difference_hessian(e -> _objective_value_weighted(e, opt_pars), x))
-    cons_j_fn!(J, x, _) = (J .= FiniteDiff.finite_difference_jacobian(
-        e -> (r = zeros(eltype(e), 4); _heat_balance_residuals_weighted!(r, e, nlp_pars); r),
-        x,
-    ))
-    function cons_h_fn!(res, x, _)
-        for i in eachindex(res)
-            res[i] .= FiniteDiff.finite_difference_hessian(
-                e -> (r = zeros(4); _heat_balance_residuals_weighted!(r, e, nlp_pars); r[i]),
-                x,
-            )
+
+    # Enzyme reverse-mode AD. Bypass DifferentiationInterface — the params
+    # NamedTuples are already captured in the closures, and DI's vector packing
+    # of `p` doesn't compose with Unitful values.
+    # hess! and cons_h! are registered (required by IpoptOptimizer) but not
+    # called at runtime when hessian_approximation="limited-memory" is set.
+    function grad_fn!(g, x, _)
+        fill!(g, 0)
+        Enzyme.autodiff(Enzyme.Reverse, _objective_value_weighted,
+                        Enzyme.Active,
+                        Enzyme.Duplicated(x, g),
+                        Enzyme.Const(opt_pars))
+        return nothing
+    end
+    function cons_j_fn!(J, x, _)
+        # Forward mode: 9 passes (one per input) instead of reverse's 4.
+        # Reverse mode currently produces silent NaN for cols 2 (skin_T)
+        # and 3 (ins_T) in this call chain — even though every called
+        # function is type-stable per JET. Forward mode gives correct
+        # finite values matching finite differences. The reverse-mode
+        # NaN is an Enzyme limitation unrelated to type stability; revisit
+        # when the Enzyme issue is resolved upstream.
+        m = size(J, 1)
+        n = length(x)
+        r  = zeros(m)
+        dr = zeros(m)
+        dx = zeros(n)
+        for j in 1:n
+            fill!(r, 0); fill!(dr, 0); fill!(dx, 0); dx[j] = 1.0
+            Enzyme.autodiff(Enzyme.Forward, _heat_balance_residuals_weighted!,
+                            Enzyme.Const,
+                            Enzyme.Duplicated(r, dr),
+                            Enzyme.Duplicated(x, dx),
+                            Enzyme.Const(nlp_pars))
+            @views J[:, j] .= dr
         end
+        return nothing
     end
+    hess_fn!(H, x, _)        = (fill!(H, 0); nothing)
+    cons_h_fn!(res, x, _)    = (foreach(r -> fill!(r, 0), res); nothing)
 
     optimization_func = OptimizationFunction(obj_fn, SciMLBase.NoAD();
         cons     = res_fn!,
@@ -343,20 +369,34 @@ function _run_ipopt(
 
     obj_fn(x, _)     = _objective_value_multisided(x, opt_pars)
     res_fn!(r, x, _) = _heat_balance_residuals_multisided!(r, x, nlp_pars)
-    grad_fn!(g, x, _) = FiniteDiff.finite_difference_gradient!(g, e -> _objective_value_multisided(e, opt_pars), x)
-    hess_fn!(H, x, _) = (H .= FiniteDiff.finite_difference_hessian(e -> _objective_value_multisided(e, opt_pars), x))
-    cons_j_fn!(J, x, _) = (J .= FiniteDiff.finite_difference_jacobian(
-        e -> (r = zeros(eltype(e), 6); _heat_balance_residuals_multisided!(r, e, nlp_pars); r),
-        x,
-    ))
-    function cons_h_fn!(res, x, _)
-        for i in eachindex(res)
-            res[i] .= FiniteDiff.finite_difference_hessian(
-                e -> (r = zeros(6); _heat_balance_residuals_multisided!(r, e, nlp_pars); r[i]),
-                x,
-            )
+
+    function grad_fn!(g, x, _)
+        fill!(g, 0)
+        Enzyme.autodiff(Enzyme.Reverse, _objective_value_multisided,
+                        Enzyme.Active,
+                        Enzyme.Duplicated(x, g),
+                        Enzyme.Const(opt_pars))
+        return nothing
+    end
+    function cons_j_fn!(J, x, _)
+        m = size(J, 1)
+        n = length(x)
+        r  = zeros(m)
+        dr = zeros(m)
+        dx = zeros(n)
+        for j in 1:n
+            fill!(r, 0); fill!(dr, 0); fill!(dx, 0); dx[j] = 1.0
+            Enzyme.autodiff(Enzyme.Forward, _heat_balance_residuals_multisided!,
+                            Enzyme.Const,
+                            Enzyme.Duplicated(r, dr),
+                            Enzyme.Duplicated(x, dx),
+                            Enzyme.Const(nlp_pars))
+            @views J[:, j] .= dr
         end
+        return nothing
     end
+    hess_fn!(H, x, _)     = (fill!(H, 0); nothing)
+    cons_h_fn!(res, x, _) = (foreach(r -> fill!(r, 0), res); nothing)
 
     optimization_func = OptimizationFunction(obj_fn, SciMLBase.NoAD();
         cons     = res_fn!,
@@ -455,7 +495,8 @@ function thermoregulate(
     evap_pars  = evaporation_pars(organism)
 
     nlp_packed = HeatExchange.nlp_pack(control.nlp_strategy, organism, environment,
-                                       skin_temperature_init, insulation_temperature_init)
+                                       skin_temperature_init, insulation_temperature_init;
+                                       smoothing = control.smoothing)
 
     _run_ipopt(nlp_packed, organism, environment, limits, metab_pars, int_cond, evap_pars,
                metabolic_heat_flow_init, skin_temperature_init, insulation_temperature_init;

src/organism.jl

@@ -89,11 +89,15 @@ pant, skin_wetness) as continuous decision variables.
 - `nlp_strategy`: NLP formulation — `WeightedMeanNLP()` (default, dorsal/ventral
   weighted-mean single body, 9 variables, 4 constraints) or `MultiSidedNLP()`
   (explicit per-side heat balance, 11 variables, 7 constraints).
+- `smoothing`: smoothing policy passed to the heat-balance physics so AD sees
+  differentiable kinks. Defaults to `SmoothBound(1.0e-5)`; pass `HardBound()` to
+  match the rule-based path's exact `abs`/`max`/`step` behaviour.
 
 Requires `Optimization.jl` and `OptimizationIpopt.jl`.
 """
-Base.@kwdef struct IPOPTControl <: AbstractControlStrategy
+Base.@kwdef struct IPOPTControl{S<:HeatExchange.SmoothingStrategy} <: AbstractControlStrategy
     nlp_strategy::HeatExchange.NLPStrategy = HeatExchange.WeightedMeanNLP()
+    smoothing::S = HeatExchange.SmoothBound(1.0e-5)
 end
 
 # =============================================================================