Solving ODE on GPU from Python with DifferentialEquations.jl

New to Julia
, ,
1

Hi everyone,

I am new in Julia. My goal is to solve stiff ODE systems (N>1000) representing chemical reaction networks up to steady-state conditions. Since the Python Scipy implementation is too slow for the purpose, I just started learning Julia as DifferentialEquations.jl allows GPU acceleration. I wrote the following piece of code where I call Julia from a Python method:

def integrate_jl(self, y0):
        from julia.api import Julia
        jl = Julia()
        from julia import Main
        Main.y0 = y0  # (Ny,1)  vector of Float64
        Main.v = self.v_dense  # (Ny,NR)  Matrix of Int8
        Main.vf = self.v_forward_dense  # (NR,Ny)  Matrix of Int8
        Main.vb = self.v_backward_dense  # (NR,Ny)  matrix of Int8
        Main.kd, Main.kr = self.kd, self.kr  # (NR,1)  vector of Float64
        Main.gas_mask = self.gas_mask  # (Ny, 1)  vector of Bool
        Main.eval("""
        using CUDA
        using DifferentialEquations
        CUDA.allowscalar(true)
        y0 = CuArray{Float64}(y0)
        v = CuArray{Int8}(v)
        kd = CuArray{Float64}(kd)
        kr = CuArray{Float64}(kr)
        vf = CuArray{Int8}(vf)
        vb = CuArray{Int8}(vb)
        gas_mask = CuArray{Bool}(gas_mask)
        p = (v = v, kd = kd, kr = kr, gas_mask = gas_mask, vf =  vf, vb = vb)
        function ode_pfr!(du, u, p, t)
            rate = p.kd .* vec(prod((u .^ p.vf')', dims=2)) .- p.kr .* vec(prod((u .^ p.vb')', dims=2))
            du = vec(p.v * net_rate)
            du[p.gas_mask] .= 0.0
        end
        prob = SteadyStateProblem(ode_pfr!, y0, p)
        sol = solve(prob, DynamicSS(QNDF()))
        CUDA.allowscalar(false)
        """)
        ...

However, I keep getting the following 'ArgumentError: cannot take the CPU address of a CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}'.

I managed to make it work on the CPU, but on the GPU I am struggling a lot.

1 Like
2

What line is throwing the error? What is the stack trace?

I can see that first du needs to be du .=, but they won’t solve your error

2 Likes
3

The line throwing the error is

sol = solve(prob, DynamicSS(QNDF()))

The stack trace I get from PyJulia is (after splitting the Julia code in multiple Main.eval calls):

File ".../my_module.py", line 275, in integrate_jl
    Main.eval("""
  File "/home/smorandi/miniconda3/envs/my_env/lib/python3.11/site-packages/julia/core.py", line 633, in eval
    ans = self._call(src)
          ^^^^^^^^^^^^^^^
  File "/home/smorandi/miniconda3/envs/my_env/lib/python3.11/site-packages/julia/core.py", line 561, in _call
    self.check_exception(src)
  File "/home/smorandi/miniconda3/envs/my_env/lib/python3.11/site-packages/julia/core.py", line 615, in check_exception
    raise JuliaError(u'Exception \'{}\' occurred while calling julia code:\n{}'
julia.core.JuliaError: Exception 'ArgumentError: cannot take the CPU address of a CuArray{Float64, 2, CUDA.Mem.DeviceBuffer}' occurred while calling julia code:

        sol = solve(prob, DynamicSS(QNDF()))

Thanks Chris for pointing out the missing dot in the du definition, as you said this does not solve the problem.

4

I just changed the stiff solver from QNDF to KenCarp4 and now works! The error source seems to be the solver algorithm.