# How can I solve complex-valued ordinary differential equations (ODEs) using neural networks, given limitations with complex data types in libraries like Lux?

**URL:** <https://discourse.julialang.org/t/how-can-i-solve-complex-valued-ordinary-differential-equations-odes-using-neural-networks-given-limitations-with-complex-data-types-in-libraries-like-lux/110581>\
**Category:** Machine Learning\
**Tags:** ode, neural-network, lux\
**Created:** [February 22, 2024, 2:05pm UTC](https://discourse.julialang.org/t/how-can-i-solve-complex-valued-ordinary-differential-equations-odes-using-neural-networks-given-limitations-with-complex-data-types-in-libraries-like-lux/110581 "2024-02-22T14:05:15Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![sahakyan2003](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/sahakyan2003/32/207220_2.png) [@sahakyan2003](https://discourse.julialang.org/u/sahakyan2003)\
**Post date:** [February 22, 2024, 2:05pm UTC](https://discourse.julialang.org/t/how-can-i-solve-complex-valued-ordinary-differential-equations-odes-using-neural-networks-given-limitations-with-complex-data-types-in-libraries-like-lux/110581/1 "2024-02-22T14:05:15Z")

</div>

**Here I describe the code where I want to train the neural network but face the problem**

Here are the libraries I use.

```Julia
using NeuralPDE
using DifferentialEquations
using Plots
using Lux, Random

```

The function `system_of_de!` defines a system of four coupled complex-valued differential equations. The equations describe the dynamics of a physical system with four variables `u[1]` , `u[2]` , `u[3]` , and `u[4]`. This equations name is Bloch equations.

The parameters of the system are `Ω` , `Δ` , and `Γ` , which are defined outside the function.

```julia
function system_of_de!(du, u, p, t)
    Ω, Δ, Γ = p
    γ = Γ / 2

    du[1] = im * Ω * (u[3] - u[4]) + Γ * u[2]
    du[2] = -im * Ω * (u[3] - u[4]) - Γ * u[2]
    du[3] = -(γ + im * Δ) * u[3] - im * Ω * (u[2] - u[1])
    du[4] = conj(du[3])

    return nothing
end

```

**Initial Conditions, Time Span, and Parameters:**

```julia
u0 = zeros(ComplexF64, 4)
u0[1] = 1

time_span = (0.0, 7.0)

# Ω, Δ, Γ         
parameters = [100.0, 0.0, 1.0]

```

**Defining the ODE Problem**

```julia
problem = ODEProblem(system_of_de!, u0, time_span, parameters)

```

```julia
ODEProblem with uType Vector{ComplexF64} and tType Float64. In-place: true
timespan: (0.0, 7.0)
u0: 4-element Vector{ComplexF64}:

```

This part involves setting up and using the Neural Network Ordinary Differential Equation (NNODE) solver.

- `rng` is a random number generator used for initialization.
- `chain` defines the architecture of the neural network used by NNODE.
- `ps` and `st` are the initial parameters of the neural network.
- `opt` is an optimization algorithm used to train the neural network.
- `alg` is an NNODE solver object that combines the neural network and optimizer.

```julia
rng = Random.default_rng()
Random.seed!(rng, 0)
chain = Chain(Dense(1, 5, σ), Dense(5, 1))
ps, st = Lux.setup(rng, chain) |> Lux.f64

```

**Solving the ODE:**

- `sol` is the solution obtained by solving the `problem` using the NNODE solver.
- `maxiters` specifies the maximum number of iterations the solver allows.
- `saveat` defines how often the solution is saved during the simulation

```julia
using OptimizationOptimisers

opt = Adam(0.1)
alg = NNODE(chain, opt, init_params = ps)

sol = solve(problem, alg, verbose = true, maxiters = 2000, saveat = 0.01)

```

**Checking the Result:**

- `ground_truth` is the solution obtained using a traditional numerical solver (Tsit5()).
- The code plots both solutions for the first two variables (`u[1]` and `u[2]`) to compare them visually.

```julia
ground_truth = solve(problem, Tsit5(), saveat = 0.01)    

plot(ground_truth.t, real.(ground_truth[1, :]), linecolor=:blue, legend = false)
plot!(ground_truth.t, real.(ground_truth[2, :]), linecolor=:blue)

plot!(sol.t, real.(sol[1, :]), linecolor=:red)
plot!(sol.t, real.(sol[2, :]), linecolor=:red)

```

---

<div class="post-metadata">

**Author:** ![SteffenPL](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/steffenpl/32/206270_2.png) [@SteffenPL](https://discourse.julialang.org/u/SteffenPL)\
**Post date:** [February 22, 2024, 3:24pm UTC](https://discourse.julialang.org/t/how-can-i-solve-complex-valued-ordinary-differential-equations-odes-using-neural-networks-given-limitations-with-complex-data-types-in-libraries-like-lux/110581/2 "2024-02-22T15:24:38Z")

</div>

Without running the code: Is there a question or a particular error that occurs?

At the moment it looks like a tutorial on how to do it. Also, one trivial (but maybe unsatisfying) approach would be to convert the ODE to a real-valued one.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [February 22, 2024, 5:56pm UTC](https://discourse.julialang.org/t/how-can-i-solve-complex-valued-ordinary-differential-equations-odes-using-neural-networks-given-limitations-with-complex-data-types-in-libraries-like-lux/110581/3 "2024-02-22T17:56:50Z")

</div>

This is the issue [How can I solve complex-valued ordinary differential equations (ODEs) using neural networks, given limitations with complex data types in libraries like Lux? · Issue #818 · SciML/NeuralPDE.jl · GitHub](https://github.com/SciML/NeuralPDE.jl/issues/818), if interested track it there.