Allow me to preface this by saying I have limited programming experience beyond the relatively short time I’ve been using Julia, so if you see I’ve made some kind of coding faux pas, please point it out!
I have a non-linear system of ODEs that I’m trying to simulate. The initial conditions are pre-determined from a chosen Poincare section, so I thought this would be an ideal scenario for parallelization. However, I’m encountering problems which are probably due to my inexperience, but I seem to figure out from the documentation. There are multiple callbacks, and the means to generate the initial conditions are a little involved so I wrapped everything before the definition of the ensemble problem in an
using Distributed addprocs(4) @everywhere begin using DifferentialEquations function eom(du, u, p, t) μ, k = p x, y, xd, yd, Ev = u r1(u,p) = √((x+μ)^2+y^2) r2(u,p) = √((x-1+μ)^2+y^2) du = dx = xd du = dy = yd du = dxd = (2*yd + x - k/r1(u,p)^2 * (xd-y) - ((1-μ)/r1(u,p)^3 * (x+1-μ) + μ/r2(u,p)^3 * (x-1+μ))) du = dyd = (-2*xd + y - k/r1(u,p)^2 * (yd+x) - ((1-μ)/r1(u,p)^3 + μ/r2(u,p)^3)*y) du = -k/r1(u,p)*(xd^2-xd*y+x*yd+yd^2) return nothing end function eom_jac(J,u,p,t) μ, k = p x, y, xd, yd = u r1(u,p) = √((x+μ)^2+y^2) r2(u,p) = √((x-1+μ)^2+y^2) J[1,:] = [0 0 1 0 0] J[2,:] = [0 0 0 1 0] J[3,1] = (1 - (1-μ)/r1(u,p)^3 + 3(x+μ)*(x-1+μ)*(1-μ)/r1(u,p)^5 - μ/r2(u,p)^3 + 3μ*(x+1-μ)^2 / r2(u,p)^5 + 2k*(x+μ)*(xd-y)/r1(u,p)^4) J[3,2] = 3(x-1+μ)*(1-μ)*y/r1(u,p)^5 + 3(x+1-μ)*μ*y/r2(u,p)^5 + 2k*y*(xd-y)/r1(u,p)^4 J[3,3] = -k/r1(u,p)^4 J[3,4] = 2 J[3,5] = 0 J[4,1] = 3(1-μ)*(x+μ)*y/r1(u,p)^5 + 3*(x-1+μ)*μ*y/r2(u,p)^5 + 2k*(x+μ)*(xd-y)/r1(u,p)^4 J[4,2] = (1 + 3(1-μ)*y^2/r1(u,p)^5 + 3μ*y^2/r2(u,p)^5 - (1-μ)/r1(u,p)^3 - μ/r2(u,p)^3 + 2k*y*(xd-y)/r1(u,p)^4) J[4,3] = -2 -k/r1(u,p)^4 J[4,4] = 0 J[4,5] = 0 J[5,1] = 4k*(x+μ)/r1(u,p)^3 * (xd^2+yd^2-xd*y+x*yd)-k/r1(u,p)^2 * yd J[5,2] = 4k*y/r1(u,p)^3 * (xd^2+yd^2-xd*y+x*yd) + k/r1(u,p)^2 * xd J[5,3] = -k/r1(u,p)^2 * (2xd-y) J[5,4] = -k/r1(u,p)^2 * (2yd+x) J[5,5] = 0 nothing end # ff = ODEFunction(eom; jac = eom_jac) p = [0.5,0.0] #Primary locations P1 = [-p;0] P2 = [1-p;0] #System Radii R1 = 1e-4 R2 = 1e-4 Rsys = 10.0 tspan = (0.0,1000.0) ## Event Handling import LinearAlgebra.norm function condition(out,u,t,integrator) out = sqrt(u^2+u^2)-Rsys #particle escapes out = R1-norm(u[1:2]-P1) #particle collides with the sun out = R2-norm(u[1:2]-P2) #particle collides with the planet end function affect!(integrator, event_index) terminate!(integrator) end function manifold(resid,u,p,t) #when dissipation is off, system is Hamiltonian resid = 0.5*(u^2+u^2) + V(u,u,p) - E resid = 0 resid = 0 resid = 0 resid = 0 end cbv = VectorContinuousCallback(condition,affect!,3) cbm = ManifoldProjection(manifold) cbs = CallbackSet(cbv,cbm) ## initial conditions #generate an array of initial conditions from a Poincare section with #E = -1.375, rdot = 0, phidot < 0 x = range(-2.0,2.0,length = 1000) y = range(-2.0,2.0,length = 1000) E = -1.375 r1(x,y,μ) = √((x+μ)^2+y^2) r2(x,y,μ) = √((x-1+μ)^2+y^2) R(x,y) = √(x^2 + y^2) V(x,y,μ) = -μ/r1(x,y,μ) - (1-μ)/r2(x,y,μ) - 1/2 * (x^2+y^2) #Jacobi Potential g(x,y,μ) = √(-2*V(x,y,μ)+2*E) xd0(x,y,μ) = y/R(x,y)*g(x,y,μ) #x initial velocity yd0(x,y,μ) = -x/R(x,y)*g(x,y,μ) #y initial velocity xy = [[x[i] y[j] xd0(x[i],y[j],p) yd0(x[i],y[j],p) E] for i=1:length(x), j=1:length(y)] ## Parallelize! prob = ODEProblem(ff,xy,tspan,p) #base problem # solver_args = (abstol=1e-8,reltol=1e-8) # @time sol = solve(prob,Vern7(),callback=cbs;solver_args...) function prob_func(prob,i,repeat) ODEProblem(prob.f,xy[i],prob.tspan,prob.p) end end #end @everywhere EnsProb = EnsembleProblem(prob,prob_func=prob_func) sim = solve(EnsProb,Vern7(),EnsembleDistributed(),trajectories=length(xy), callback = cbs, abstol = 1e-8, reltol = 1e-8)
When I format the code to solve serially (i.e, constructed a new ODEProblem for each initial condition and solving it), the code runs fine, but it’s obviously a very unsatisfying way to go about it. When I format as an
EnsembleProblem, the code either causes the computer to hang (99% CPU & MEM usage) or takes an extremely long time (longer than it would take single-threaded.) I’ve tried using both
EnsembleDistributed() and both have similar issues.
What I’ve come to ask is: is there any obvious mistake I’m making here?
Additionally, it would be nice if it were possible to write this in a way that’s able to use
@DiffEqGPU, but I’m encountering many errors when exploring this, and so any insights direction would also be appreciated.