Am I going about solving this PDE system the right way...?

General Usage
1

I have been transitioning my MATLAB code to Julia and managed to get working a sort of analogue to my MATLAB code.
My PDE system has 2 coupled (1+1)D nonlinear PDEs:
-\mathrm{i}\left(\frac{\partial \psi}{\partial X}\right) = -\frac{\mathrm{sgn}\left(\beta_{2S} \right)}{2} \! \left(\frac{\partial^{2} \psi}{\partial \tau^{2}} \right) + {| \psi|}^{2} \psi,\\ -\mathrm{i} \left(\frac{\partial \phi}{\partial X}\right) = -\frac{s_{4}}{2} \left(\frac{\partial^{2} \phi}{\partial \tau^{2}}\right)+{i} s_{1} \! \left(\frac{\partial \phi}{\partial \tau}\right) + s_{2} \psi \phi^{*} {\mathrm{exp}}[{\mathrm{i} \left(s_{3} + q \right) \! X}] + s_{5} {| \psi|}^{2} \phi.
In MATLAB, we used a UPPE style method to transform these equations into a set of coupled ODEs by taking the Fourier transform, solving the equations, and then inverse transforming them to find the original functions \psi and \phi.
The ODEs look like this:
-\mathrm{i} \left(\frac{\mathrm{d} \varphi}{\mathrm{d} X}\right) = \mathcal{F}\left[\frac{s_{2}}{2} \phi^{2} {\mathrm{exp}}\left({-\mathrm{i} s_{3} X}\right)+\left({| \psi |}^{2}+s_{5} {| \phi |}^{2}\right) \! \psi \right] {\mathrm{exp}}\left({-\mathrm{i} \alpha X}\right), \\ -\mathrm{i} \left(\frac{\mathrm{d} \theta}{\mathrm{d} X}\right) = \mathcal{F}\left[s_{2} \psi \phi^{*} {\mathrm{exp}}\left({\mathrm{i} s_{3} X}\right) +\left(s_{5} {| \psi|}^{2}+s_{6} {| \phi|}^{2}\right) \! \phi \right] {\mathrm{exp}}\left({-\mathrm{i} \Omega X}\right),

where \mathcal{F}\left(\psi\right)=\tilde{\psi}, \mathcal{F}\left(\phi\right)=\tilde{\phi}, \tilde{\psi}=\varphi {\mathrm e}^{i \alpha X}, \tilde{\phi}=\theta {\mathrm e}^{i \Omega X} and \Omega is related to the initial conditions and parameters of the system.
I have solved the system in Julia with some simple initial boundary conditions using a variety of different solvers in DifferentialEquations.jl (I found DP8 to be the most performant in my limited testing).
After all this, my question is am I going about this the right way in Julia? I have no experience with the PDE solvers available in Julia so I am not sure if they will be more performant than the method described above. I know this might be really hard to judge from what I’ve posted and I should probably just try both and compare, but if anyone has any insights into an approach that could prove more efficient and performant that would be great.
Thanks in advance!

2

A pseudospectral method, as you described, is a great way to solve that equation. We will have some things to automate that in the future, but there’s nothing that sounds wrong with what you’re doing.

3

Great thank you for the reassurance, I’ll keep experimenting with different solvers to try and improve solve times.
I had a quick look at trying to use analytic Jacobians to speed things up with ModellingToolkit but I’m not sure that it plays nicely with the Fourier transforms. Am I at the limit of techniques to improve speed and I should just be content with the performance I’m getting with the solvers?

4

It would be hard to judge without seeing the code. I would try the Runge-Kutta-Chebyshev methods as well, ROCK2 and ROCK4, since it sounds like it’s almost non-stiff, but PDEs are never really non-stiff and those methods are good for the “in-between” PDE cases. You might want to try the parallelized extrapolation methods. In fact, this might be a good test for some new methods we’re developing.

Other things are, if you’re having to solve this multiple times, parallelizing the solves effectively for that application is good. Parallelized multiple shooting if you’re doing fitting, or parallel ensembles for Monte-Carlo, etc. If it’s just one fit then there’s other things to work on.

2 Likes
5

Here is the ODE function I am using:

function ODE!(dX,X,p,t)
    s2, s3, s5, s6, alpha, Omega, N = p
    exp_term_psi = @.cis(alpha .* t) #exp.(1im .* alpha.*t)
    exp_term_phi = @.cis(Omega .* t) #exp.(1im .* Omega.*t)

    A = X[1:N] .* exp_term_psi
    B = X[N+1:2*N] .* exp_term_phi
    psi = ifft(A)
    phi = ifft(B)
    
    rhs1 = 0.5 * @.cis(-s3 * t) * s2 .* (phi.^2.0) .+ psi .* abs2.(psi) .+ psi .* s5 .* abs2.(phi)
    rhs2 = @.cis(s3 * t) * s2 .* psi .* conj(phi) .+ phi * s5 .* abs2.(psi) .+ phi .* s6 .* abs2.(phi)
    rhs1 = fft(rhs1)
    rhs2 = fft(rhs2)
    dX[1:N] = 1im .* rhs1./exp_term_psi
    dX[N+1:2*N] = 1im .* rhs2./exp_term_phi
end

In terms of multiple solves, I am changing the initial conditions of the ODE but also things like the total number of coupled equations through N, the range I am solving over and the parameters in p.
Are things like preallocating all my arrays and creating plan_ffts worthwhile?
I will try some of the solvers you have suggested above.
Thanks again!

6

This is a minor question, but why are you both using @. and broadcasting your functions?

7

Good point! It’s unnecessary, I’ve changed it now. Thanks for pointing out!

8

Can you share the full runnable thing? This will be nice for research, and maybe nice to add to the SciMLBenchmarks?

9

Of course!
For context, this is part of a masters project simulating resonant scattering and frequency conversion with temporal solitons.
The code should work as is however there are a few sections commented out that can safely be ignored.

using Logging: global_logger
using TerminalLoggers: TerminalLogger
global_logger(TerminalLogger())

using FFTW
using LinearAlgebra
using DifferentialEquations
using Plots; gr()

## define grid 

N=2048

dt=0.1
t=(-N/2.0:N/2.0-1.0)*dt 

domega=2.0*pi/(N*dt) # Conversion factor between t-space and the Fourier conjugate omega-space.
test = -N/2.0:N/2.0-1.0
omega=fftshift(test .* domega)                         
##

## Constants
L=75.0    #  propagation distance
steps=2000.0  #  number of steps to take along z 
dz=L/steps #  Discrete step between consecutive points along the propagation direction.
shift = 20.0
beta_2F = -1.0
beta_2S = -1.0
gamma_c = 2.0
gamma_s = 1.0
gamma_f = 1.0
s1 = 0.0
s2 = 0.2
s3 = -7.5
s4 = beta_2F/abs(beta_2S)
s5 = gamma_c/gamma_s
s6 = gamma_f/gamma_s
alpha = 0.5*sign(beta_2S).*omega.^2.0
Omega = 0.5.*(s4.*omega .^ 2.0) .- s1.*omega
q = 10 
##

## Initial conditions
function f(x,y,z,a)
    return a.*exp.(-((x.-y).^2.0)./(2*z^2.0))
end

function soliton(q,t,shift)
    sqrt(2.0*q)*sech.(sqrt(2.0*q).*(t.-shift))
end

psi_0 = soliton(q,t,shift)

F_0 = f.(t,-20.0,10,10^-18).*cis.(1.12157 * t)

ini = [fft(psi_0); fft(F_0)]
z = 0:L/steps:L
z = range(0, L, length = Int64(steps + 1))
#soln = Array{ComplexF64}(undef,N,2,length(z))


## Solution

function ODE!(dX,X,p,t)
    s2, s3, s5, s6, alpha, Omega, N = p
    exp_term_psi = cis.(alpha .* t)
    exp_term_phi = cis.(Omega .* t) 

    A = X[1:N] .* exp_term_psi # Critical to specify with X[1:N,1] as otherwise stores as a single complex float, not as a vector of complex floats which is what the solution should be. This is in an analogous manner to MATLAB.
    B = X[N+1:2*N] .* exp_term_phi
    psi = ifft(A)
    phi = ifft(B)
    
    rhs1 = 0.5 * cis.(-s3 * t) * s2 .* (phi.^2.0) .+ psi .* abs2.(psi) .+ psi .* s5 .* abs2.(phi)
    rhs2 = cis.(s3 * t) * s2 .* psi .* conj(phi) .+ phi * s5 .* abs2.(psi) .+ phi .* s6 .* abs2.(phi)
    rhs1 = fft(rhs1)
    rhs2 = fft(rhs2)
    dX[1:N] = 1im .* rhs1./exp_term_psi # Equally as important here to specify dX[1:N,1] here to ensure you are storing a vector of complex float solutions!
    dX[N+1:2*N] = 1im .* rhs2./exp_term_phi
end

p = [s2, s3, s5, s6, alpha, Omega, N]
prob = ODEProblem(ODE!,ini,(0.0,L),p)

sol, tim = @timed solve(prob,ROCK4(), saveat=z,progress = true, progress_steps = 1000)
sol = transpose(Array(sol))

function grids(alpha, Omega, z)
    len = length(z)
    A = transpose(alpha) .* ones(len)
    W = transpose(Omega) .* ones(len)
    Z = transpose(ones(length(alpha))) .* z
    return A, W, Z
end
A, W, Z = grids(alpha, Omega, z)

a_fourier_psi = sol[:,1:N] .*cis.(A .* Z)
a_fourier_phi = sol[:,N+1:Int64(2*N)] .* cis.(W .* Z)

function un_fourier(x)
    return abs.(ifft(x, 2)) # This computes the IFFT along the row dimension, 2 = rows, 1 = columns, maybe because Julia is column-major?
end


#sol = nothing
aout_psi = un_fourier(a_fourier_psi)
aout_phi = un_fourier(a_fourier_phi)

heaviside(x::AbstractFloat) = ifelse(x < 0, zero(x), ifelse(x > 0, one(x), 0.0))
max_func(var2,var3) = maximum(var3)/maximum(var2)

## Plots
#= zstep = 1551; 
#width = range(shift - 1.0/sqrt(2.0*q),shift + 1.0/sqrt(2.0*q); length=1000);
width = range(shift - sqrt(2.0/q)*log(2.0 + sqrt(3.0)),shift + sqrt(2.0/q)*log(2.0 + sqrt(3.0)); length=100);
incident = abs.(aout_phi[1,:] .* heaviside.(width[1] .- t));
incident_fourier = abs2.(fft(ifft(a_fourier_phi,2)[1,:] .* heaviside.(width[1] .- t)));
reflected = abs.(aout_phi[zstep,:] .* heaviside.(width[1] .- t));
reflected_fourier = abs2.(fft(ifft(a_fourier_phi,2)[zstep,:] .* heaviside.(width[1] .- t)));
transmitted = abs.(aout_phi[zstep,:] .* heaviside.(t .- width[end]));
trapped = abs.(aout_phi[zstep,:] .* (heaviside.(-t .+ width[end]) .- heaviside.(width[1] .- t))); =#

heatmap(t, z, aout_phi)
1 Like