Solving coupled ODE problem with explicit fft terms of the independent variables

Hi all,

I have been trying to translate my MATLAB code that solves a set of coupled ODEs into Julia. I have only just started using Julia so am still quite unfamiliar with its types and the intricacies of getting some functions to behave nicely with said types. Below is the section of code I have written in Julia and below that, the code from MATLAB.

## Constants
L=10.0    #  propagation distance
steps=1500.0  #  number of steps to take along z 
dz=L/steps #  Discrete step between consecutive points along the propagation direction.
beta_2F = -1.0
beta_2S = -1.0
gamma_c = 2.0
gamma_s = 1.0
gamma_f = 1.0
s1 = 0.0
s2 = 2.0
s3 = 0.0
s4 = beta_2F/abs(beta_2S)
s5 = gamma_c/gamma_s
s6 = gamma_f/gamma_s
alpha = (test .* domega).^2.0
Omega = 0.5.*(s4.*omega .^ 2.0) .- s1.*omega
q = 250 ##

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

psi_0 = sqrt(2.0*q)*sech.(sqrt(2.0*q).*(t.-30))

F_0 = f.(t,10,10,0.01).*exp.(1im*12*t)

phi_0 = fft(psi_0)
phiF_0 = fft(F_0)

ini = [phi_0 phiF_0]
z = range(0, L, length = Int64(steps + 1.0))

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

    A = X[1:N,1] .* exp_term_psi # Do I need to specify using X[1:N,1], or can I just use X[1]?
    B = X[1:N,2] .* exp_term_phi
    psi = ifft(A)
    phi = ifft(B)
    rhs1 = 0.5 * exp(-1im * s3 * t) * s2 .* phi .* phi + psi .* abs.(psi) .* abs.(psi) + psi .* s5 .* abs.(phi) .* abs.(phi)
    rhs2 = exp(1im*s3*t) * s2 .* psi .* conj(phi) + phi * s5 .* abs.(psi) .* abs.(psi) + phi .* s6 .* abs.(phi) .* abs.(phi)
    rhs1 = fft(rhs1)
    rhs2 = fft(rhs2)
    dX[1] = 1im .* rhs1./exp_term_psi
    dX[2] = 1im .* rhs2./exp_term_phi
    #= dX[1] = 1im.*fft.((0.5.*s2.*(phi.^2.0).*exp.(-1im*s3.*t) .+ (abs.(psi).^2.0).*psi .+ s5.*(abs.(phi).^2.0).*psi))./exp_term_psi
    dX[2] = 1im.*fft.((s2.*psi.*conj(phi).*exp.(1im*s3.*t) .+ s5.*(abs.(psi).^2.0).*phi .+ s6.*(abs.(phi).^2.0).*phi))./exp_term_phi
p = [s2, s3, s5, s6, alpha, Omega, N];
prob = ODEProblem(ODE!,ini,(0.0,L),p);
sol = solve(prob, Tsit5());

The current error is:
LoadError: MethodError: Cannot convert an object of type Vector{ComplexF64} to an object of type ComplexF64
I’m not sure where this is an error and my debugging efforts in VSCode haven’t been very helpful at locating the problem. Am I specifying the vector variables correctly using X[1:N,1]? In MATLAB the solution is a vector of solved ODEs that I can use deval on to evaluate for any given z value I like. I’m not sure if Julia works in a similar way.
Thanks for any help in advance!

Below is the related MATLAB code, the constants are identical between the Julia and MATLAB versions:

function D = f(t,X, s2, s3, s5, s6, alpha, N, Omega)
    exp_term_psi = exp(1i*alpha.*t); % For S equation
    exp_term_phi = exp(1i*Omega.*t); % For F equation
    psi = ifft(X(1:N).*exp_term_psi,N);
    phi = ifft(X(N+1:2*N).*exp_term_phi,N);
    sol = [(1i*fft(0.5*s2*(phi.^2).*exp(-1i*s3*t) + (abs(psi).^2).*psi + s5*(abs(phi).^2).*psi)./exp_term_psi),... 
            (1i*fft(s2*psi.*conj(phi).*exp(1i*s3.*t) + s5*(abs(psi).^2).*phi + s6*(abs(phi).^2).*phi)./exp_term_phi)];
    D =[sol(:,1);sol(:,2)];

Don’t broadcast the plan. P_rhs1 = plan_fft(rhs1) You want to apply it to the array, not element-wise.

Note also that the whole point of plan_fft is to do it once and re-use it, not to recompute the plan on every timestep.

Ah yes of course thank you, a silly mistake on my half. I have now updated my code and will edit my submission with yours and @stevengj’s suggestions.

Note also that instead of exp.(1im .* x) it is more efficient to use cis.(.- x), and simpler to do @. cis(-x)

Great thank you, another useful function to know!

cis = \cos + i \sin for your memory :wink:

1 Like

I was wondering where cis came from… Thanks for clarifying!