From 799640d27a08187ac7a71c92958b59ab7a1f6298 Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Fri, 12 Jun 2026 14:43:03 +0100
Subject: [PATCH 1/7] init with new doc page

---
 docs/pages.jl                     |   1 +
 docs/src/inverse_problems/udes.md | 100 ++++++++++++++++++++++++++++++
 2 files changed, 101 insertions(+)
 create mode 100644 docs/src/inverse_problems/udes.md

diff --git a/docs/pages.jl b/docs/pages.jl
index 9e80f95acf..7cdc4dd979 100644
--- a/docs/pages.jl
+++ b/docs/pages.jl
@@ -67,6 +67,7 @@ pages = Any[
     "inverse_problems/likelihood_profiler.md",
     "inverse_problems/global_sensitivity_analysis.md",
     "inverse_problems/turing_ode_param_fitting.md",
+    "inverse_problems/udes.md",
     "Examples" => Any[
         "inverse_problems/examples/ode_fitting_oscillation.md"
     ]
diff --git a/docs/src/inverse_problems/udes.md b/docs/src/inverse_problems/udes.md
new file mode 100644
index 0000000000..a085ee9299
--- /dev/null
+++ b/docs/src/inverse_problems/udes.md
@@ -0,0 +1,100 @@
+# [Fitting universal differential equations](@id udes)
+```@raw html
+<details><summary><strong>Environment setup and package installation</strong></summary>
+```
+The following code sets up an environment for running the code on this page.
+```julia
+using Pkg
+Pkg.activate(; temp = true) # Creates a temporary environment, which is deleted when the Julia session ends.
+Pkg.add("Catalyst")
+Pkg.add("DataFrames")
+Pkg.add("Distributions")
+Pkg.add("Lux")
+Pkg.add("ModelingToolkitNeuralNets")
+Pkg.add("OrdinaryDiffEq")
+Pkg.add("Optimisers")
+Pkg.add("PEtab")
+Pkg.add("Plots")
+```
+```@raw html
+</details>
+```
+```@raw html
+<details><summary><strong>Quick-start example</strong></summary>
+```
+The following code provides a brief example of how to perform parameter fitting using the [Optimization.jl](https://github.com/SciML/Optimization.jl) package.
+```julia
+# Declare the model (a mutual activation loop).
+using Catalyst
+
+# Generate synthetic data to which we will apply the fitting procedure.
+using Distributions, OrdinaryDiffEq, Plots
+
+# Create a UDE using Lu and ModelingToolkitNeuralNets.
+using ModelingToolkitNeuralNets, Lux
+
+# Fit the model using PEtab.
+using PEtab, Optimisers
+
+# Plot the fitted simulation and recovered activation functions.
+plot()
+plot()
+```
+```@raw html
+</details>
+```
+  \
+Default stuff open meues with env and copy example.
+
+## Fitting a rate-based UDE
+
+Short introduction to rate-based UDEs goes here.
+
+Generate synthetic data.
+
+Declare UDE
+
+Fit using PEtab
+
+Plotting the results
+
+
+
+---
+
+## [Citations](@id petab_citations)
+
+If you use this functionality in your research, [in addition to Catalyst](@ref doc_index_citation), please cite the following papers to support the authors of the Lux.jl and PEtab.jl packages:
+
+```bibtex
+@software{pal2025lux,
+  author    = {Pal, Avik},
+  title     = {{Lux: Explicit Parameterization of Deep Neural Networks in Julia}},
+  publisher = {Zenodo},
+  year      = {2025},
+  doi       = {10.5281/zenodo.7808903},
+  url       = {https://doi.org/10.5281/zenodo.7808903}
+}
+```
+```bibtex
+@article{PEtabBioinformatics2025,
+  title={PEtab.jl: advancing the efficiency and utility of dynamic modelling},
+  author={Persson, Sebastian and Fr{\"o}hlich, Fabian and Grein, Stephan and Loman, Torkel and Ognissanti, Damiano and Hasselgren, Viktor and Hasenauer, Jan and Cvijovic, Marija},
+  journal={Bioinformatics},
+  volume={41},
+  number={9},
+  pages={btaf497},
+  year={2025},
+  publisher={Oxford University Press}
+}
+```
+
+---
+
+## References
+
+[^1]:
+    [Rackauckas, R et al. _Universal differential equations for scientific machine learning_, arXiv (2020).](https://arxiv.org/abs/2001.04385)
+
+[^2]:
+    [Loman, T and R, Baker. _Functional and parametric identifiability for universal differential equations applied to chemical reaction networks_, arXiv (2025).](https://arxiv.org/abs/2510.14140)

From 90d8a7c2da51571bd21638b5d787191190485b6a Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Tue, 16 Jun 2026 14:55:30 +0100
Subject: [PATCH 2/7] finish writing, udpated doc env

---
 docs/Project.toml                 |   6 +
 docs/src/inverse_problems/udes.md | 357 ++++++++++++++++++++++++++++--
 2 files changed, 343 insertions(+), 20 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index 0547d99a4b..7ed9d83d06 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -21,12 +21,15 @@ Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LikelihoodProfiler = "93acb638-a083-5915-8dce-d129bc6a3f59"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
+Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839"
+ModelingToolkitNeuralNets = "f162e290-f571-43a6-83d9-22ecc16da15f"
 NetworkLayout = "46757867-2c16-5918-afeb-47bfcb05e46a"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d"
 Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 OptimizationBBO = "3e6eede4-6085-4f62-9a71-46d9bc1eb92b"
 OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
 OptimizationLBFGSB = "22f7324a-a79d-40f2-bebe-3af60c77bd15"
@@ -79,12 +82,15 @@ JumpProcesses = "9.23"
 Latexify = "0.16.5"
 LikelihoodProfiler = "1.5.1"
 LinearSolve = "2.30, 3"
+Lux = "1.31.4"
 MCMCChains = "6, 7"
 ModelingToolkitBase = "1.17"
+ModelingToolkitNeuralNets = "2.5.1"
 NetworkLayout = "0.4"
 NonlinearSolve = "4"
 NonlinearSolveFirstOrder = "1, 2.1"
 Optim = "2"
+Optimisers = "0.4.7"
 OptimizationBBO = "0.4"
 OptimizationBase = "4, 5.0"
 OptimizationLBFGSB = "1"
diff --git a/docs/src/inverse_problems/udes.md b/docs/src/inverse_problems/udes.md
index a085ee9299..f647e05b7c 100644
--- a/docs/src/inverse_problems/udes.md
+++ b/docs/src/inverse_problems/udes.md
@@ -22,49 +22,368 @@ Pkg.add("Plots")
 ```@raw html
 <details><summary><strong>Quick-start example</strong></summary>
 ```
-The following code provides a brief example of how to perform parameter fitting using the [Optimization.jl](https://github.com/SciML/Optimization.jl) package.
+The following code provides a brief example of how to declare a UDE using Catalyst,
+[ModelingToolkitNeuralNets.jl](https://github.com/SciML/ModelingToolkitNeuralNets.jl), and
+[Lux.jl](https://github.com/LuxDL/Lux.jl), and then fit it using
+[PEtab.jl](https://github.com/sebapersson/PEtab.jl).
 ```julia
 # Declare the model (a mutual activation loop).
 using Catalyst
+rn = @reaction_network begin
+    hill(Y, v, K, n), 0 --> X
+    X, 0 --> Y
+    d, (X, Y) --> 0
+end
 
 # Generate synthetic data to which we will apply the fitting procedure.
 using Distributions, OrdinaryDiffEq, Plots
+u0 = [:X => 2.0, :Y => 0.1]
+ps_true = [:v => 1.1, :K => 2.0, :n => 3.0, :d => 0.5]
+tend = 45.0
+oprob_true = ODEProblem(rn, u0, tend, ps_true)
+sol_true = solve(oprob_true)
+sample_t = range(0.0, tend; length = 20)
+σ = 0.1
+sample_X = [rand(Normal(X, σ)) for X in sol_true(sample_t; idxs = :X)]
+sample_Y = [rand(Normal(Y, σ)) for Y in sol_true(sample_t; idxs = :Y)]
+plot(sol_true; lw = 6, label = ["X (true)" "Y (true)"])
+plot!(sample_t, sample_X, seriestype = :scatter, label = "X (data)", color = 1, ms = 6)
+plot!(sample_t, sample_Y, seriestype = :scatter, label = "Y (data)", color = 2, ms = 6)
 
-# Create a UDE using Lu and ModelingToolkitNeuralNets.
+# Create a UDE using Lux and ModelingToolkitNeuralNets.
 using ModelingToolkitNeuralNets, Lux
+nn_arch = Lux.Chain(
+    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
+)
+@SymbolicNeuralNetwork U, θ = nn_arch
+A(z) = U(z, θ)[1]
+rn_ude = @reaction_network begin
+    $A(Y), 0 --> X
+    X, 0 --> Y
+    d, (X, Y) --> 0
+end
 
-# Fit the model using PEtab.
-using PEtab, Optimisers
+# Bundle the UDE and data into a PEtab problem.
+using DataFrames, PEtab
+ps_est = [PEtabMLParameter(:θ), PEtabParameter(:d)]
+observables = [PEtabObservable(:obs_X, :X, σ), PEtabObservable(:obs_Y, :Y, σ)]
+dfX = DataFrame(obs_id = "obs_X", time = sample_t, measurement = sample_X)
+dfY = DataFrame(obs_id = "obs_Y", time = sample_t, measurement = sample_Y)
+measurements = vcat(dfX, dfY)
+petab_model = PEtabModel(rn_ude, observables, measurements, ps_est; speciemap = u0)
+petab_prob = PEtabODEProblem(petab_model)
 
-# Plot the fitted simulation and recovered activation functions.
-plot()
-plot()
+# Fit the unknown function and parameters using 5 independent starts of 10,000 Adam iterations.
+using Optimisers
+options = OptimisersOptions(iterations = 10000)
+@time res = calibrate_multistart(petab_prob, Optimisers.Adam(1e-3), 5; options)
+
+# Plot the fitted simulation and the recovered activation functions.
+plot(res, petab_prob)
+plot(res, petab_prob; plot_type = :best_function)
 ```
 ```@raw html
 </details>
 ```
   \
-Default stuff open meues with env and copy example.
 
-## Fitting a rate-based UDE
+In this tutorial, we describe how to create and fit *universal differential equation* (UDE)
+models declared through Catalyst. UDEs are a *hybrid modelling* framework: they combine
+mechanistic modelling (where models are built from first principles) with data-driven
+modelling (where models are learnt from data)[^1]. The core idea is to declare known model
+components as standard differential equations, while allowing unknown system
+dynamics to be learnt from data. In practice, this is done by inserting *universal function
+approximators* (typically neural networks) into differential equations to represent unknown dynamics.
+
+In the context of Catalyst and chemical reaction networks, UDEs can naturally be used to
+learn *reaction rates*. Non-constant reaction rates are typically modelled by
+heuristically chosen functions, and in some cases a more natural choice is to learn these
+rate functions directly from data[^2]. This can be used to, for example, learn a protein's
+production rate as a function of a transcription factor's concentration. 
+
+In this tutorial, we use [Lux.jl](https://github.com/LuxDL/Lux.jl) to declare neural
+networks, [ModelingToolkitNeuralNets.jl](https://github.com/SciML/ModelingToolkitNeuralNets.jl)
+to incorporate them into Catalyst models, and
+[PEtab.jl](https://github.com/sebapersson/PEtab.jl) to fit these models to data. More
+details on these packages can be found in their respective documentation, while we here will describe
+how they compose with each other and Catalyst.
+
+## [Fitting a rate-based UDE](@id udes_basic_example)
+
+For this example, we will use a mutual activation loop in which two species ($X$ and $Y$)
+activate each other's production.
+```@example ude_rate_based
+using Catalyst
+rn_true = @reaction_network begin
+    hill(Y, v, K, n), 0 --> X
+    X, 0 --> Y
+    d, (X, Y) --> 0
+end
+```
+Here, the production of $X$ follows a Hill function, which we will later try to recover
+through a UDE.
+
+First, we designate the true simulation conditions, from which we will generate a synthetic
+dataset for the fitting procedure.
+```@example ude_rate_based
+using Distributions, OrdinaryDiffEq, Plots
+
+# Simulate the model at the true conditions.
+u0 = [:X => 2.0, :Y => 0.1]
+ps_true = [:v => 1.1, :K => 2.0, :n => 3.0, :d => 0.5]
+tend = 45.0
+oprob_true = ODEProblem(rn_true, u0, (0.0, tend), ps_true)
+sol_true = solve(oprob_true)
 
-Short introduction to rate-based UDEs goes here.
+# Generate normally distributed noisy measurements with a fixed standard deviation.
+t_measurement = range(0.0, tend; length = 20)
+σ = 0.1
+X_observed = [rand(Normal(X, σ)) for X in sol_true(t_measurement; idxs = :X)]
+Y_observed = [rand(Normal(Y, σ)) for Y in sol_true(t_measurement; idxs = :Y)]
 
-Generate synthetic data.
+# Plots the true simulation and the synthetic data.
+plot(sol_true; lw = 4, label = ["X (true)" "Y (true)"])
+plot!(t_measurement, X_observed; seriestype = :scatter, label = "X (data)", color = 1, ms = 5)
+plot!(t_measurement, Y_observed; seriestype = :scatter, label = "Y (data)", color = 2, ms = 5)
+```
+
+Next, we declare a UDE version of the model by replacing the Hill function with a neural
+network. The neural network will be created using Lux. Here, a relatively small
+feed-forward architecture is sufficient. By using a non-negative activation function in the
+final layer (here, [softplus](https://en.wikipedia.org/wiki/Softplus)), we constrain the
+fitted function to be non-negative, which is a known physical constraint for reaction rate
+functions. Softplus is also smooth, a desirable property for functions embedded in ODE
+simulations. A full list of available activation functions can be found
+[here](https://lux.csail.mit.edu/stable/api/NN_Primitives/ActivationFunctions).
+```@example ude_rate_based
+using Lux
+nn_arch = Lux.Chain(
+    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 1, Lux.softplus, use_bias = false)
+)
+```
+We now generate a `SymbolicNeuralNetwork` using ModelingToolkitNeuralNets. This embeds the
+Lux model in a symbolic format that is compatible with Catalyst (and ModelingToolkit).
+```@example ude_rate_based
+using ModelingToolkitNeuralNets
+@SymbolicNeuralNetwork U, θ = nn_arch
+```
+Here, `U` represents the neural network structure, and `θ` represents its parameters. Additional options for using the `@SymbolicNeuralNetwork` macro are described [here](https://sciml.github.io/ModelingToolkitNeuralNets.jl/stable/api/#ModelingToolkitNeuralNets.@SymbolicNeuralNetwork) and [here](https://sciml.github.io/ModelingToolkitNeuralNets.jl/stable/symbolic_ude_tutorial/). Next, using  the following syntax, we can [interpolate](@ref dsl_advanced_options_symbolics_and_DSL_interpolation) this function approximator as  a Catalyst model rate.
+```@example ude_rate_based
+A(z) = U(z, θ)[1]
+rn_ude = @reaction_network begin
+    $A(Y), 0 --> X
+    X, 0 --> Y
+    d, (X, Y) --> 0
+end
+```
+We can inspect the parameters contained within this model:
+```@example ude_rate_based
+parameters(rn_ude)
+```
+In addition to the mechanistic parameter $d$, the neural network parameters $θ$ have also
+been incorporated into the model. The neural network architecture, $U$, is also represented
+as a parameter, but in practice this paraemter is fixed and can be ignored.
 
-Declare UDE
+Now, we use PEtab to simultaneously fit $d$ and $θ$. The workflow is almost identical to
+the standard [PEtab parameter fitting workflow](@ref petab_parameter_fitting). The main
+differences are that we declare $θ$ as a `PEtabMLParameter` (rather than a
+`PEtabParameter`), and that we choose an optimisation algorithm better suited to neural
+network training. We begin by declaring the parameters we wish to fit, the observables (and
+their noise formulas), and the measurements, and then use these to generate a
+`PEtabODEProblem`.
+```@example ude_rate_based
+using DataFrames, PEtab
+ps_est = [PEtabMLParameter(:θ), PEtabParameter(:d)]
+observables = [PEtabObservable(:obs_X, :X, σ), PEtabObservable(:obs_Y, :Y, σ)]
+dfX = DataFrame(obs_id = "obs_X", time = t_measurement, measurement = X_observed)
+dfY = DataFrame(obs_id = "obs_Y", time = t_measurement, measurement = Y_observed)
+measurements = vcat(dfX, dfY)
+petab_model = PEtabModel(rn_ude, observables, measurements, ps_est; speciemap = u0)
+petab_prob = PEtabODEProblem(petab_model)
+nothing # hide
+```
+PEtab supports many options for customising the optimisation problem. Here, we choose the
+simplest settings. These options are described further in the [Catalyst-specific PEtab
+tutorial](@ref petab_parameter_fitting), and in
+[PEtab's documentation](https://sebapersson.github.io/PEtab.jl/stable/).
 
-Fit using PEtab
+The UDE model can now be fitted to the data. Here, we use the Adam implementation from the
+[Optimisers.jl](https://github.com/FluxML/Optimisers.jl) package, running 5 independent starts 
+with 10,000 iterations each.
+```@example ude_rate_based
+using Optimisers
+options = OptimisersOptions(iterations = 10000)
+@time res = calibrate_multistart(petab_prob, Optimisers.Adam(1e-3), 5; options)
+nothing # hide
+```
+
+PEtab's plot recipes can be applied to fitted UDE models without adaptation. Here, we plot
+the fitted UDE together with the data.
+```@example ude_rate_based
+plot(res, petab_prob; lw = 4, la = 0.6, ms = 5, title = "Final fit")
+```
+In the next section, we will show different UDE-specific plotting options.
+
+### [Investigating fitted neural networks](@id udes_nn_investigation)
+
+In many cases, UDEs are used to fit unknown model *functions*. When a fitted function has a
+single input argument, the final fit can often be interpreted by directly plotting that
+function. With PEtab, this can be done using
+```@example ude_rate_based
+plot(res, petab_prob; plot_type = :best_function, lw = 6)
+```
+Since we use synthetic data, we can also compare the fitted function to the ground-truth
+function:
+```@example ude_rate_based
+true_f(Y) = Catalyst.hill(Y, 1.1, 2.0, 3.0)
+plot!(true_f, 0.1, 3.8, linestyle = :dash, lw = 3, color = :blue, label = "True activation function")
+```
+This lets us check whether the correct function has been correctly recovered.
+
+When carrying out multistart calibration, PEtab also provides the
+`plot_type = :function_ensemble` option, which plots the functional forms found by each
+optimisation run. In practice, some runs might not have converged. Here, we use the
+`loss_thres` option to only plot functions associated with optimisation runs with a similar
+loss value to the best run:
+```@example ude_rate_based
+plot(res, petab_prob; plot_type = :function_ensemble, lw = 4, xlimit = (0.0, 3.8), loss_thres = res.fmin + 2.0)
+```
+Here, we can see that all optimisation runs achieving a good model-to-data fit converge to
+the same functional form. In practice, this can be used as a simple form of
+[*identifiability analysis*](@ref likelihood_profiler_practical_identifiability), assessing
+the extent to which a single functional form can reliably be recovered from the data[^2].
+
+A more thorough description of these plots and their options can be found
+[here](https://sebapersson.github.io/PEtab.jl/dev/tutorials/sciml/sciml_plotting). To
+retrieve the fitted function, for example to evaluate it manually at selected values, use:
+```@example ude_rate_based
+fitted_oprob = get_odeproblem(res, petab_prob)[1]
+fitted_f(z) = fitted_oprob.ps[U](z, fitted_oprob.ps[θ])[1]
+fitted_f(1.0)
+```
+
+## [Alternative approaches for incorporating neural networks into models](@id udes_alternative_forms)
+Finally, we consider some additional contexts in which neural networks can be incorporated
+into Catalyst models.
+
+### [Declaring UDEs programmatically](@id udes_programmatic)
+While this tutorial has primarily created UDEs using the Catalyst `@reaction_network` DSL,
+it is also possible to create them [programmatically](@ref programmatic_CRN_construction).
+Below, we create a UDE representation of the same mutual activation loop as previously, but
+using the programmatic interface.
+```@example ude_programmatic
+using Catalyst, ModelingToolkitNeuralNets, Lux
+t = default_t()
+@species X(t) Y(t)
+@parameters d
+nn_arch = Lux.Chain(
+    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
+)
+@SymbolicNeuralNetwork U, θ = nn_arch
+rxs = [
+    Reaction(U(Y, θ)[1], [], [X]),
+    Reaction(X, [], [Y]),
+    Reaction(d, [X], []),
+    Reaction(d, [Y], []),
+]
+@named rs = ReactionSystem(rxs, t)
+rs = complete(rs)
+```
+ModelingToolkitNeuralNets also provides a function version of the `@SymbolicNeuralNetwork`
+macro, which provides some additional functionality. It can be found
+[here](https://sciml.github.io/ModelingToolkitNeuralNets.jl/stable/api/#ModelingToolkitNeuralNets.SymbolicNeuralNetwork).
+
+### [Non-rate based UDEs](@id udes_non_rate_based)
+Above, we described how a neural network can be used to learn an unknown rate function.
+However, symbolic neural networks created using `@SymbolicNeuralNetwork` are highly
+flexible: they act as normal functions and can be used wherever such functions can be used.
+For example, neural networks can be inserted into [coupled equations](@ref coupled_models_diffeqs).
+Here, we use this to model a cell growing in a nutrient ($N$). The rate at which the nutrient is
+consumed is modelled as an unknown function of both a nutrient-activated protein's active form ($Xₐ$)
+and the nutrient's concentration. Note that the unknown function has two input variables, which is
+straightforward to handle.
+```@example ude_equations
+# Create a symbolic neural network depending on two input variables.
+using ModelingToolkitNeuralNets, Lux
+nn_arch = Lux.Chain(
+    Lux.Dense(2 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
+)
+@SymbolicNeuralNetwork U, θ = nn_arch
+C(z1, z2) = U(z1, z2, θ)[1]
+
+# Create a model of a cell that consumes a limited nutrient supply.
+using Catalyst
+rs = @reaction_network begin
+    @equations D(N) ~ -$C(N, Xₐ)
+    (kₐ*N, kᵢ), Xᵢ <--> Xₐ
+end
+```
+
+### [Non-ODE UDEs](@id udes_non_odes)
+Up until now, we have assumed that all fitted universal differential equations are based on
+ODEs. However, in practice, UDEs declared through Catalyst act as normal differential 
+equations (although with an unusually large number of parameters). This means that `JumpProblem`s
+and `SDEProblem`s can be generated from Catalyst UDEs using the standard syntax, for example:
+```@example ude_sde
+# Create the UDE.
+using Catalyst, ModelingToolkitNeuralNets, Lux
+nn_arch = Lux.Chain(
+    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 1, Lux.softplus, use_bias = false)
+)
+@SymbolicNeuralNetwork U, θ = nn_arch
+A(z) = U(z, θ)[1]
+rn_ude = @reaction_network begin
+    $A(Y), 0 --> X
+    X, 0 --> Y
+    d, (X, Y) --> 0
+end
+
+# Create an SDEProblem.
+u0 = [:X => 2.0, :Y => 0.1]
+ps = [:d => 0.5]
+sprob = SDEProblem(rn_ude, u0, 45.0, ps)
+nothing # hide
+```
+generates a universal stochastic differential equation.
 
-Plotting the results
+The parameters of this model cannot currently be fitted using PEtab. However, one can use
+an approach similar to those used [here](@ref optimization_parameter_fitting) and
+[here](@ref behaviour_optimisation) to declare a loss function comparing the behaviour of
+this UDE with available data. Next, the parameters can be fitted by minimising that loss
+function. An example of how this is done for an ODE-declared UDE can be found
+[here](https://sciml.github.io/ModelingToolkitNeuralNets.jl/stable/symbolic_ude_tutorial/).
 
+### [Learning parameters or observables using neural networks](@id udes_parameters_n_observables)
+Throughout this tutorial, we have shown how neural networks can be incorporated into
+Catalyst models to learn unknown functions of system variables within the main model
+system. PEtab, however, also supports two additional ways to incorporate neural networks
+that, due to how they are inserted, can be handled more efficiently:
+- Models where the neural network exists within the mapping between the model's *states* and
+  its *observables*.
+- Models where parameters are neural-network mappings of other parameter values. Primarily,
+  this can be used to learn phenotypical parameters (such as binding rates) from measurable
+  parameters (such as protein structure data).
 
+To learn more about how to declare such models, please see the following two PEtab
+tutorials: [observable neural networks](https://sebapersson.github.io/PEtab.jl/dev/tutorials/sciml/observable)
+and [pre-simulation neural networks](https://sebapersson.github.io/PEtab.jl/dev/tutorials/sciml/pre_simulate).
 
 ---
 
-## [Citations](@id petab_citations)
+## [Citations](@id udes_citations)
 
-If you use this functionality in your research, [in addition to Catalyst](@ref doc_index_citation), please cite the following papers to support the authors of the Lux.jl and PEtab.jl packages:
+If you use this functionality in your research, [in addition to Catalyst](@ref doc_index_citation),
+please cite the following papers to support the authors of the Lux.jl and PEtab.jl packages:
 
 ```bibtex
 @software{pal2025lux,
@@ -93,8 +412,6 @@ If you use this functionality in your research, [in addition to Catalyst](@ref d
 
 ## References
 
-[^1]:
-    [Rackauckas, R et al. _Universal differential equations for scientific machine learning_, arXiv (2020).](https://arxiv.org/abs/2001.04385)
+[^1]: [Rackauckas, R et al. _Universal differential equations for scientific machine learning_, arXiv (2020).](https://arxiv.org/abs/2001.04385)
 
-[^2]:
-    [Loman, T and R, Baker. _Functional and parametric identifiability for universal differential equations applied to chemical reaction networks_, arXiv (2025).](https://arxiv.org/abs/2510.14140)
+[^2]: [Loman, T and R, Baker. _Functional and parametric identifiability for universal differential equations applied to chemical reaction networks_, arXiv (2025).](https://arxiv.org/abs/2510.14140)

From f27a056de62106caebcc92b29fd6c5242164522b Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Thu, 25 Jun 2026 11:15:48 +0100
Subject: [PATCH 3/7] optim tutorial fix

---
 docs/Project.toml                                         | 2 ++
 .../inverse_problems/optimization_ode_param_fitting.md    | 8 ++++----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/docs/Project.toml b/docs/Project.toml
index c44c1778f8..91e7b27bfa 100644
--- a/docs/Project.toml
+++ b/docs/Project.toml
@@ -32,6 +32,7 @@ Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 OptimizationBBO = "3e6eede4-6085-4f62-9a71-46d9bc1eb92b"
 OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb"
+OptimizationEvolutionary = "cb963754-43f6-435e-8d4b-99009ff27753"
 OptimizationLBFGSB = "22f7324a-a79d-40f2-bebe-3af60c77bd15"
 OptimizationNLopt = "4e6fcdb7-1186-4e1f-a706-475e75c168bb"
 OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
@@ -87,6 +88,7 @@ Optim = "2"
 Optimisers = "0.4.7"
 OptimizationBBO = "0.4.7"
 OptimizationBase = "4, 5.0"
+OptimizationEvolutionary = "0.4.9"
 OptimizationLBFGSB = "1"
 OptimizationNLopt = "0.3"
 OptimizationOptimJL = "0.4"
diff --git a/docs/src/inverse_problems/optimization_ode_param_fitting.md b/docs/src/inverse_problems/optimization_ode_param_fitting.md
index bc31075fe8..e9b6d62078 100644
--- a/docs/src/inverse_problems/optimization_ode_param_fitting.md
+++ b/docs/src/inverse_problems/optimization_ode_param_fitting.md
@@ -8,7 +8,7 @@ using Pkg
 Pkg.activate(; temp = true) # Creates a temporary environment, which is deleted when the Julia session ends.
 Pkg.add("Catalyst")
 Pkg.add("OptimizationBase")
-Pkg.add("OptimizationBBO")
+Pkg.add("OptimizationEvolutionary")
 Pkg.add("OptimizationNLopt")
 Pkg.add("OptimizationOptimJL")
 Pkg.add("OrdinaryDiffEq")
@@ -159,10 +159,10 @@ Catalyst.PNG(plot(plt; fmt = :png, dpi = 200)) # hide
 !!! note
     Here, a good exercise is to check the resulting parameter set and note that, while it creates a good fit to the data, it does not actually correspond to the original parameter set. Identifiability is a concept that studies how to deal with this problem.<!--NTS: re-add ref when identifiablity works again-->
 
-Say that we instead would like to use a [differential evolution](https://en.wikipedia.org/wiki/Differential_evolution) approach, as implemented by the [BlackBoxOptim.jl](https://github.com/SciML/BlackBoxOptim.jl) package. In this case we can run:
+Say that we instead would like to use a [genetic algorithm](https://en.wikipedia.org/wiki/Genetic_algorithm) approach, as implemented by the [Evolutionary.jl](https://github.com/wildart/Evolutionary.jl) package. In this case we can run:
 ```@example optimization_paramfit_1 
-using OptimizationBBO
-sol = solve(optprob, BBO_adaptive_de_rand_1_bin_radiuslimited())
+using OptimizationEvolutionary
+sol = solve(optprob, Evolutionary.GA())
 nothing # hide
 ```
 to solve `optprob` for this combination of solve and implementation.

From 782233e4386f290a10a3f9b956f2ca2e3a1b98db Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Sun, 12 Jul 2026 09:56:40 +0100
Subject: [PATCH 4/7] Update docs/src/inverse_problems/udes.md

Co-authored-by: Sebastian Persson <46872750+sebapersson@users.noreply.github.com>
---
 docs/src/inverse_problems/udes.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/inverse_problems/udes.md b/docs/src/inverse_problems/udes.md
index f647e05b7c..3b613f172d 100644
--- a/docs/src/inverse_problems/udes.md
+++ b/docs/src/inverse_problems/udes.md
@@ -213,7 +213,7 @@ tutorial](@ref petab_parameter_fitting), and in
 
 The UDE model can now be fitted to the data. Here, we use the Adam implementation from the
 [Optimisers.jl](https://github.com/FluxML/Optimisers.jl) package, running 5 independent starts 
-with 10,000 iterations each.
+with 10,000 iterations/epochs each.
 ```@example ude_rate_based
 using Optimisers
 options = OptimisersOptions(iterations = 10000)

From 40b328493a548d53da43e2fc35c89cfe8efdcf85 Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Sun, 12 Jul 2026 09:56:56 +0100
Subject: [PATCH 5/7] Update docs/src/inverse_problems/udes.md

Co-authored-by: Sebastian Persson <46872750+sebapersson@users.noreply.github.com>
---
 docs/src/inverse_problems/udes.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/src/inverse_problems/udes.md b/docs/src/inverse_problems/udes.md
index 3b613f172d..2f66aae119 100644
--- a/docs/src/inverse_problems/udes.md
+++ b/docs/src/inverse_problems/udes.md
@@ -366,7 +366,7 @@ function. An example of how this is done for an ODE-declared UDE can be found
 ### [Learning parameters or observables using neural networks](@id udes_parameters_n_observables)
 Throughout this tutorial, we have shown how neural networks can be incorporated into
 Catalyst models to learn unknown functions of system variables within the main model
-system. PEtab, however, also supports two additional ways to incorporate neural networks
+system. PEtab.jl, however, also supports two additional ways to incorporate neural networks
 that, due to how they are inserted, can be handled more efficiently:
 - Models where the neural network exists within the mapping between the model's *states* and
   its *observables*.

From 602d5539c8de8cbc3f8415329c942cfc22d045cc Mon Sep 17 00:00:00 2001
From: Torkel <torkel.loman@gmail.com>
Date: Sun, 12 Jul 2026 10:19:40 +0100
Subject: [PATCH 6/7] add section on mutlipell inputs or outputs.

---
 docs/src/inverse_problems/udes.md | 38 ++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/docs/src/inverse_problems/udes.md b/docs/src/inverse_problems/udes.md
index 2f66aae119..6bdd26f4a9 100644
--- a/docs/src/inverse_problems/udes.md
+++ b/docs/src/inverse_problems/udes.md
@@ -186,7 +186,7 @@ parameters(rn_ude)
 ```
 In addition to the mechanistic parameter $d$, the neural network parameters $θ$ have also
 been incorporated into the model. The neural network architecture, $U$, is also represented
-as a parameter, but in practice this paraemter is fixed and can be ignored.
+as a parameter, but in practice this parameter is fixed and can be ignored.
 
 Now, we use PEtab to simultaneously fit $d$ and $θ$. The workflow is almost identical to
 the standard [PEtab parameter fitting workflow](@ref petab_parameter_fitting). The main
@@ -266,6 +266,42 @@ fitted_f(z) = fitted_oprob.ps[U](z, fitted_oprob.ps[θ])[1]
 fitted_f(1.0)
 ```
 
+### [Neural networks with multiple inputs or outputs](@id udes_nn_multi_ins_or_outs)
+
+It is possible to use neural networks with multiple inputs. To do this, simply list all inputs in the symbolic neural network function's arguments, while also ensuring that the declared Lux chain have the correct number of inputs in the first layer. Here we use this approach to declare a feedforward loop where the production rate of $Z$ is an unknown function of both $X$ and $Y$.
+```@example ude_multi_ins_or_outs
+using Catalyst, ModelingToolkitNeuralNets, Lux
+nn_arch = Lux.Chain(
+    Lux.Dense(2 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 1, Lux.softplus, use_bias = false),
+)
+@SymbolicNeuralNetwork U, θ = nn_arch
+f(x, y) = U(x, y, θ)[1]
+multi_input_ude = @reaction_network begin
+    d, (X, Y, Z)
+    hill(X,v,K,n), 0 --> Y
+    $f(X, Y), 0 --> Z
+end
+```
+
+It is also possible to declare neural networks with multiple output. Generally, this is equivalent to fitting two unknown functions, however, if the two fitted functions are expected to have similar forms, declaring them as a single neural network with multiple outputs can be numerically advantageous. Below, we use this approach to declare a model where a single species $X$ activates the production of two different species ($Y$ and $Z$) according to two different unknown functions that are approximated using a single, two-output, neural network.
+```@example ude_multi_ins_or_outs
+nn_arch = Lux.Chain(
+    Lux.Dense(1 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 3, Lux.softplus, use_bias = false),
+    Lux.Dense(3 => 2, Lux.softplus, use_bias = false),
+)
+@SymbolicNeuralNetwork U, θ = nn_arch
+f1(x) = U(x, θ)[1]
+f2(x) = U(x, θ)[2]
+rn_ude = @reaction_network begin
+    d, (X, Y, Z)
+    $f1(X), 0 --> Y
+    $f2(X), 0 --> Z
+end
+```
+
 ## [Alternative approaches for incorporating neural networks into models](@id udes_alternative_forms)
 Finally, we consider some additional contexts in which neural networks can be incorporated
 into Catalyst models.

From d554a668e128b405f1b5da29e700c2cc111dea44 Mon Sep 17 00:00:00 2001
From: Torkel Loman <torkel.loman@gmail.com>
Date: Sun, 12 Jul 2026 22:07:20 +0100
Subject: [PATCH 7/7] fix bad reactions in ude doc

---
 docs/src/inverse_problems/udes.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/src/inverse_problems/udes.md b/docs/src/inverse_problems/udes.md
index 6bdd26f4a9..9cdbda67c7 100644
--- a/docs/src/inverse_problems/udes.md
+++ b/docs/src/inverse_problems/udes.md
@@ -279,7 +279,7 @@ nn_arch = Lux.Chain(
 @SymbolicNeuralNetwork U, θ = nn_arch
 f(x, y) = U(x, y, θ)[1]
 multi_input_ude = @reaction_network begin
-    d, (X, Y, Z)
+    d, (X, Y, Z) --> 0
     hill(X,v,K,n), 0 --> Y
     $f(X, Y), 0 --> Z
 end
@@ -296,7 +296,7 @@ nn_arch = Lux.Chain(
 f1(x) = U(x, θ)[1]
 f2(x) = U(x, θ)[2]
 rn_ude = @reaction_network begin
-    d, (X, Y, Z)
+    d, (X, Y, Z) --> 0
     $f1(X), 0 --> Y
     $f2(X), 0 --> Z
 end