Solving a coupled differential equation

Hello everyone, I am new to julia and try to use it to solve my problem, which involves solving coupled differential equations. I will show my problem below:


Here is my attempted solution:

using DifferentialEquations, LinearAlgebra

# Constants
ħ = 1.0  # Reduced Planck's constant (in natural units)
m = 1.0  # Particle mass
k_max = 5.0  # Maximum momentum
N = 50  # Number of momentum points
B_span = (0.0, 10.0)  # Range of B (flow time)

# Define the momentum grid
k_grid = range(-k_max, k_max, length=N)
dk = step(k_grid)

# Define the energy function E_k
E_k = k -> ħ^2 * k^2 / (2m)

# Define the ODE system for dg_{k',k}/dB
function dg_matrix_dB(g_vec, B)
    N = 50
    g_matrix = reshape(g_vec, (N, N))  # Reshape vector to matrix
    dg_matrix = zeros(N, N)
    
    for i in 1:N  # Iterate over k'
        for j in 1:N  # Iterate over k
            E_1 = g_matrix[i,i]
            E_2 = g_matrix[j,j]

            # First term
            term1 = -((E_1 - E_2)^2) * g_matrix[i, j]

            # Summation term
            term2 = 0.0
            for p in 1:N
                E_p = g_matrix[p,p]
                term2 += (E_1 + E_2 - 2 * E_p) * g_matrix[i, p] * g_matrix[p, j]
            end

            dg_matrix[i, j] = term1 + term2 # * dk # Include dk for integration
        end
    end

    return reshape(dg_matrix, N^2)  # Return as vector
end

# Define the initial condition for g_{k',k}
c = 1.0  # Constant part of g_{k',k}
g_initial = zeros(N, N)
for i in 1:N
    for j in 1:N
        if i == j
            g_initial[i, j] = c + E_k(k_grid[i])  # Diagonal: constant + energy
        else
            g_initial[i, j] = c  # Off-diagonal: constant
        end
    end
end

# Reshape the initial condition for ODE solver
g0 = reshape(g_initial, N^2)

prob = ODEProblem(dg_matrix_dB, g0, B_span)
sol = solve(prob, Tsit5())

But it seems that the error indicates I use some sqrt function, which I do not explicitly use it, so I am very confused. My guess was it had something to do with vector to matrix conversion?

MethodError: no method matching sqrt(::Tuple{Int64}) Closest candidates are: sqrt(!Matched::Union{Float32, Float64}) at math.jl:590 sqrt(!Matched::StridedMatrix{T}) where T<:Union{Real, Complex} at /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:853 sqrt(!Matched::Union{Static.StaticBool{N}, Static.StaticFloat64{N}, Static.StaticInt{N}}) where N at ~/.julia/packages/Static/IYKUj/src/Static.jl:442

I get a totally different error from you:

julia> prob = ODEProblem(dg_matrix_dB, g0, B_span)
ERROR: All methods for the model function `f` had too few arguments. For example,
an ODEProblem `f` must define either `f(u,p,t)` or `f(du,u,p,t)`. This error
can be thrown if you define an ODE model for example as `f(u,t)`. The parameters
`p` are not optional in the definition of `f`! For more information on the required
number of arguments for the function you were defining, consult the documentation
for the `SciMLProblem` or `SciMLFunction` type that was being constructed.

For example, here is the no parameter Lorenz equation. The two valid versions
are out of place:

```julia
function lorenz(u,p,t)
  du1 = 10.0*(u[2]-u[1])
  du2 = u[1]*(28.0-u[3]) - u[2]
  du3 = u[1]*u[2] - 8/3*u[3]
  [du1,du2,du3]
 end
 u0 = [1.0;0.0;0.0]
 tspan = (0.0,100.0)
 prob = ODEProblem(lorenz,u0,tspan)
```

and in-place:

```julia
function lorenz!(du,u,p,t)
  du[1] = 10.0*(u[2]-u[1])
  du[2] = u[1]*(28.0-u[3]) - u[2]
  du[3] = u[1]*u[2] - 8/3*u[3]
 end
 u0 = [1.0;0.0;0.0]
 tspan = (0.0,100.0)
 prob = ODEProblem(lorenz!,u0,tspan)
```

Offending function: f
Methods:
# 1 method for generic function "dg_matrix_dB" from Main:
 [1] dg_matrix_dB(g_vec, B)
     @ REPL[11]:1

Stacktrace:
 [1] #isinplace#11
   @ ~/.julia/packages/SciMLBase/RJzsP/src/utils.jl:280
 [2] isinplace (repeats 2 times)
   @ ~/.julia/packages/SciMLBase/RJzsP/src/utils.jl:246 [inlined]
 [3] #ODEProblem#327
   @ ~/.julia/packages/SciMLBase/RJzsP/src/problems/ode_problems.jl:198
 [4] ODEProblem(f::Function, u0::Vector{Float64}, tspan::Tuple{Float64, Float64}, p::SciMLBase.NullParameters)
   @ SciMLBase ~/.julia/packages/SciMLBase/RJzsP/src/problems/ode_problems.jl:197
 [5] ODEProblem
   @ ~/.julia/packages/SciMLBase/RJzsP/src/problems/ode_problems.jl:197
 [6] top-level scope
   @ REPL[16]:1

If I fix that issue by giving dg_matrix_dB the right signature:

function dg_matrix_dB(g_vec, p, B)
    N = 50
    g_matrix = reshape(g_vec, (N, N))  # Reshape vector to matrix
    dg_matrix = zeros(N, N)
    
    for i in 1:N  # Iterate over k'
        for j in 1:N  # Iterate over k
            E_1 = g_matrix[i,i]
            E_2 = g_matrix[j,j]

            # First term
            term1 = -((E_1 - E_2)^2) * g_matrix[i, j]

            # Summation term
            term2 = 0.0
            for p in 1:N
                E_p = g_matrix[p,p]
                term2 += (E_1 + E_2 - 2 * E_p) * g_matrix[i, p] * g_matrix[p, j]
            end

            dg_matrix[i, j] = term1 + term2 # * dk # Include dk for integration
        end
    end

    return reshape(dg_matrix, N^2)  # Return as vector
end

then it runs fine for me.

Though I’ll note that you don’t need to do this reshapeing at each step, and you’ll likely get much better performance if you solve the equations in-place instead of allocating a new output each time.

Weird, I rerun it, and it is actually the same as yours. so I guess I am asking for the solution as you posted.

Thanks, I think it is the right solution. But I am not sure I can follow what the in-place functionality means, can you elaborate a little bit?

I would probably write it like this:

using DifferentialEquations, LinearAlgebra

# Constants
ħ = 1.0  # Reduced Planck's constant (in natural units)
m = 1.0  # Particle mass
k_max = 5.0  # Maximum momentum
N = 50  # Number of momentum points
B_span = (0.0, 10.0)  # Range of B (flow time)

# Define the momentum grid
k_grid = range(-k_max, k_max, length=N)
dk = step(k_grid)

# Define the energy function E_k
E_k(k) = ħ^2 * k^2 / (2m)

# Define the ODE system for dg_{k',k}/dB
function dg_dB!(dg, g, p, B)
    for j in axes(g, 2)  # Iterate over k
        E_2 = g[j,j]
        for i in axes(g, 1)  # Iterate over k′
            E_1 = g[i,i]

            # First term
            term1 = -((E_1 - E_2)^2) * g[i, j]

            # Summation term
            term2 = sum(axes(g, 1)) do p
                E_p = g[p,p]
                (E_1 + E_2 - 2 * E_p) * g[i, p] * g[p, j]
            end
            dg[i, j] = term1 + term2 # * dk # Include dk for integration
        end
    end
end

# Define the initial condition for g_{k',k}
c = 1.0  # Constant part of g_{k',k}
g0 = zeros(N, N)
for i in 1:N
    for j in 1:N
        if i == j
            g0[i, j] = c + E_k(k_grid[i])  # Diagonal: constant + energy
        else
            g0[i, j] = c  # Off-diagonal: constant
        end
    end
end

prob = ODEProblem(dg_dB!, g0, B_span)
sol = solve(prob, Tsit5())
1 Like

Thanks, it helps me a lot