I can confirm that @danielwe is right about the problem.
Now for some comments about style and performance. I’d recommend having a look in the manual’s performance tips section. Most of what I’ll tell you can be found there.
- Don’t use (untyped) globals inside functions
Consider yourpsi!
function psi!(du, u, p, t)
Op = - im * p.par[1] * H - p.par[2]/2 * G - p.par[3] /2.0* H
mul!(du, Op, u)
end
All of the operators (H and G) here are untyped globals. This prevents the compiler from optimizing this function. A better strategy is to put everything you need into the parameters p. Here since your Hamiltonian is not time dependent you should just compute it once and stick the full thing into the parameters.
2. It is rather important for performance to reduce allocations to a minimum. That’s why I used mul! as this is a non-allocating form of matrix multiplication. The most important thing (because it is called the most) is psi! which will be fine once we move out the computation of Op. The jump condition is also fine, but affect! could use some more work. We can avoid it in this case, but if you ever find the need for some scratch space, you can also just fit into the parameters p.
3. You defined a weird type for the parameters to work around not being able to mutate threshold but unfortunately mutable types also come with a (smallish) performance penalty. More importantly, in this case I’d prefer to use a standard NamedTuple instead, as that allows me to be much more flexible and experiment more easily. Also it will simplify the code.
4. Concerning naming conventions: Usually functions that mutate one or more of their inputs get a !. Since your jump_condition! does not mutate anything, it is better named jump_condition without the !.
Here is an improved version of your code:
numf = 500 # number of points
tf = 100
tspan=(0.0,tf)
u0 = [1.0+ im * 0.0, 0.0 + im * 0.0] # initial datum
par = [0.4,0.5,0.1] # parameters
# Jump operators
L1 = [0.0 + im * 0.0 1.0+ im * 0.0; 0.0 + im * 0.0 0.0+ im * 0.0]
L2 = [0.0 + im * 0.0 0.0 + im * 0.0; 1.0+ im * 0.0 0.0+ im * 0.0]
H = L1 * L2
G = L2 * L1
function psi!(du, u, p, t)
mul!(du, p.hamiltonian, u)
end
# Rename this
function jump_condition(u,t,integrator)
# dot(x,x) is more efficient than sum(abs2, x)
real(dot(u,u)) - integrator.p.threshold[] # note the [] because I made threshold a Ref
end
function affect!(integrator)
ψ = copyto!(integrator.p.tempψ, integrator.u) # copy this to avoid aliasing later
process_ops = integrator.p.process_operators
jump_ops = integrator.p.jump_operators
# reuse process operators to compute probabilities
jump_prob1 = real(dot(ψ, process_ops[1], ψ))
jump_prob2 = real(dot(ψ, process_ops[2], ψ))
tot_jump_prob = jump_prob1 + jump_prob2
if jump_prob1 / tot_jump_prob >= rand(Float64)
jump_ampl = sqrt(jump_prob1)
# note mul!(C,A,B,α,β) computes
# C .= α*A*B + β*C
# false acts as a "hard zero"
# note that C,A,B may not be overlap in memory -> that's why we used the temporary memory in the first line of this function
mul!(integrator.u, jump_ops[1], ψ, 1/jump_ampl, false)
else
jump_ampl = sqrt(jump_prob2)
mul!(integrator.u, jump_ops[2], ψ, 1/jump_ampl, false)
end
integrator.p.threshold[] = rand(Float64)
end
jump_callback = ContinuousCallback(jump_condition,affect!);
function prob_func(prob,i,repeat)
# it is safe to overwrite this because EnsembleProblem makes a copy before each trajectory
prob.p.threshold[] = rand()
prob
end
num_traj=1000 # trajectories in the ensemble
p = (;
hamiltonian = - im * par[1] * H - par[2]/2 * G - par[3] /2.0* H,
jump_operators = (L1, L2),
process_operators = (L1'L1, L2'L2), # is there a canonical name for these?
threshold = Ref(0.0), # use a Ref here as a single element container
tempψ = similar(u0), # some scratch space
)
sse_prob0 = ODEProblem(psi!, u0, (0.0, tf), p; callback = jump_callback);
ensemble_sse0 = EnsembleProblem(sse_prob0; prob_func = prob_func)
sol_sse0 = solve(ensemble_sse0, Vern7(), EnsembleSerial(), trajectories = num_traj)