Error when solving ode that has sin and cos functions

[anbonimus]

Hi, first of all thx for implementing diffeqpy, there's not much python ode solvers on gpu now.

So i have a system of equations that involves sin and cos, and im trying to solve it on gpu, and here is part of the code i am running.

import numpy as np
import random
from diffeqpy import de
from diffeqpy import cuda
import cupy as cp

e = 1#1.602176634*10**(-19)
kb = 1#1.380649*10**(-23)
mi = 170.936323*1822.8884 #Yb+
ma =  6.0151228874*1822.8884
rmass = 1/(1/mi+1/ma)
depth = 10**(-6)*3.16*10**(-6)
c4 = 82.0563*2#10.907*10**(-57)
c8 = c4**2/depth/16

w = 2*np.pi*2.5*10**6*2.4188*10**(-17)


ax = -2.982*10**-4#-0.5*az/2
ay = ax#-0.5*az/2
az = -2*ax

qx = 0.219
qy = -qx
qz = 0

T = 10**-6*3.16*10**(-6) 
R0 = 2100#10**(-7)/(0.528177*10**-10)

def prob_func(prob,i,rep):
    de.remake(prob,u0=[random.uniform(0, 1)*u0[i] for i in range(0,3)],
            p=[random.uniform(0, 1)*p[i] for i in range(0,3)])

def atom(y,p,t):
    rx,ry,rz,drx,dry,drz,Rx,Ry,Rz,dRx,dRy,dRz,tp = y
    ax,ay,az,qx,qy,qz,w,c4,c8,mi,ma = p
    
    dis = (y[6]-y[0])**2+(y[7]-y[1])**2+(y[8]-y[2])**2
        
    ddrx = -(ax+2*qx*cp.cos(w*tp))*(w**2)*y[0]/4 + (2*c4*(y[6]-y[0])/dis**3 - 8*c8*(y[6]-y[0])/dis**5)/mi
    ddry = -(ay+2*qy*cp.cos(w*tp))*(w**2)*y[1]/4 + (2*c4*(y[7]-y[1])/dis**3 - 8*c8*(y[7]-y[1])/dis**5)/mi
    ddrz = -(az+2*qz*cp.cos(w*tp))*(w**2)*y[2]/4 + (2*c4*(y[8]-y[2])/dis**3 - 8*c8*(y[8]-y[2])/dis**5)/mi
    
    ddRx = (-2*c4*(y[6]-y[0])/dis**3 + 8*c8*(y[6]-y[0])/dis**5)/ma
    ddRy = (-2*c4*(y[7]-y[1])/dis**3 + 8*c8*(y[7]-y[1])/dis**5)/ma
    ddRz = (-2*c4*(y[8]-y[2])/dis**3 + 8*c8*(y[8]-y[2])/dis**5)/ma


p1=[ax,ay,az,qx,qy,qz,w,c4,c8,mi,ma]

r0x=0.0
r0y=0.0
r0z=0.0
dr0x=0.0
dr0y=0.0
dr0z=0.0


theta = cp.arccos(2*np.random.rand()-1)
phi = 2*np.pi*cp.random.rand()
R0x = R0*cp.sin(theta)*cp.cos(phi)+r0x
R0y = R0*cp.sin(theta)*cp.sin(phi)+r0y
R0z = R0*cp.cos(theta)+r0z


dR0r = cp.sqrt(2*T/ma)*cp.random.weibull(2)
dR0phi = cp.random.normal(loc=0.0,scale=np.sqrt(T/ma))
dR0theta = cp.random.normal(loc=0.0,scale=np.sqrt(T/ma))


dR0x = dR0r*cp.sin(theta)*cp.cos(phi)+dR0theta*cp.cos(theta)*cp.cos(phi)-dR0phi*cp.sin(phi)
dR0y = dR0r*cp.sin(theta)*cp.sin(phi)+dR0theta*cp.cos(theta)*cp.sin(phi)+dR0phi*cp.cos(phi)
dR0z = dR0r*cp.cos(theta)-dR0theta*cp.sin(theta)


ini_state = [r0x,r0y,r0z,dr0x,dr0y,dr0z,R0x,R0y,R0z,dR0x,dR0y,dR0z,0]
          
tf = 1e-9/(2.4188*10**(-17))
prob = de.ODEProblem(atom,ini_state,(0.0,tf),p1)
fast_prob = de.jit32(prob)
ensembleprob = de.EnsembleProblem(fast_prob,safetycopy=False)
sol = de.solve(ensembleprob,cuda.GPUTsit5(),cuda.EnsembleGPUKernel(cuda.CUDABackend()),trajectories=10000,saveat=0.01)

and i got the error message

Traceback (most recent call last):
  File "C:\Users\Ruiren\Desktop\temp.py", line 130, in <module>
    fast_prob = de.jit32(prob)
                ^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\any.jl", line 262, in __call__
    return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\Desktop\temp.py", line 54, in atom
    ddrx = -(ax+2*qx*cp.cos(w*tp))*(w**2)*y[0]/4 + (2*c4*(y[6]-y[0])/dis**3 - 8*c8*(y[6]-y[0])/dis**5)/mi
                     ^^^^^^^^^^^^
  File "cupy\_core\_kernel.pyx", line 1286, in cupy._core._kernel.ufunc.__call__
  File "cupy\_core\_kernel.pyx", line 159, in cupy._core._kernel._preprocess_args
  File "cupy\_core\_kernel.pyx", line 145, in cupy._core._kernel._preprocess_arg
TypeError: Unsupported type <class 'juliacall.RealValue'>

but when i change from cupy to numpy package, i got

Traceback (most recent call last):
  File "C:\Users\Ruiren\Desktop\temp.py", line 125, in <module>
    fast_prob = de.jit32(prob)
                ^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\any.jl", line 262, in __call__
    return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\Desktop\temp.py", line 49, in atom
    ddrx = -(ax+2*qx*np.cos(w*tp))*(w**2)*y[0]/4 + (2*c4*(y[6]-y[0])/dis**3 - 8*c8*(y[6]-y[0])/dis**5)/mi
                     ^^^^^^^^^^^^
ValueError: setting an array element with a sequence. The requested array would exceed the maximum number of dimension of 32.

and if i use math package i got the following, which i think its because math.sin and cos is returning 64 bit float

Traceback (most recent call last):
  File "C:\Users\Ruiren\Desktop\temp.py", line 126, in <module>
    fast_prob = de.jit32(prob)
                ^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\any.jl", line 262, in __call__
    return self._jl_callmethod($(pyjl_methodnum(pyjlany_call)), args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\Desktop\temp.py", line 50, in atom
    ddrx = -(ax+2*qx*math.cos(w*tp))*(w**2)*y[0]/4 + (2*c4*(y[6]-y[0])/dis**3 - 8*c8*(y[6]-y[0])/dis**5)/mi
                     ^^^^^^^^^^^^^^
  File "C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\number.jl", line 183, in __float__
    return self._jl_callmethod($(pyjl_methodnum(pyfloat)))
           ^^^^^^^^^^^^^^^^^^^^^^^^
juliacall.JuliaError: MethodError: no method matching Float64(::Symbolics.Num)
The type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.

Closest candidates are:
  (::Type{T})(::Real, !Matched::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:265
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:900
  Float64(!Matched::UInt8)
   @ Base float.jl:245
  ...

Stacktrace:
  [1] convert(::Type{Float64}, x::Symbolics.Num)
    @ Base .\number.jl:7
  [2] cconvert(T::Type, x::Symbolics.Num)
    @ Base .\essentials.jl:687
  [3] PyFloat_FromDouble(x1::Symbolics.Num)
    @ PythonCall.C C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\C\pointers.jl:303
  [4] pyfloat(x::Symbolics.Num)
    @ PythonCall.Core C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\Core\builtins.jl:744
  [5] _pyjl_callmethod(f::Any, self_::Ptr{PythonCall.C.PyObject}, args_::Ptr{PythonCall.C.PyObject}, nargs::Int64)
    @ PythonCall.JlWrap C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\base.jl:62
  [6] _pyjl_callmethod(o::Ptr{PythonCall.C.PyObject}, args::Ptr{PythonCall.C.PyObject})
    @ PythonCall.JlWrap.Cjl C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\C.jl:63
  [7] PyObject_CallObject
    @ C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\C\pointers.jl:303 [inlined]
  [8] macro expansion
    @ C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\Core\Py.jl:132 [inlined]
  [9] pycallargs(f::Py, args::Py)
    @ PythonCall.Core C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\Core\builtins.jl:220
 [10] pycall(::Py, ::Vector{Symbolics.Num}, ::Vararg{Any}; kwargs::@Kwargs{})
    @ PythonCall.Core C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\Core\builtins.jl:243
 [11] pycall
    @ C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\Core\builtins.jl:233 [inlined]
 [12] Py
    @ C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\Core\Py.jl:357 [inlined]
 [13] call_composed
    @ .\operators.jl:1054 [inlined]
 [14] call_composed
    @ .\operators.jl:1053 [inlined]
 [15] ComposedFunction
    @ .\operators.jl:1050 [inlined]
 [16] ODEFunction
    @ C:\Users\Ruiren\.julia\packages\SciMLBase\sYmAV\src\scimlfunctions.jl:2468 [inlined]
 [17] modelingtoolkitize(prob::SciMLBase.ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, PyList{Any}, SciMLBase.ODEFunction{false, SciMLBase.AutoSpecialize, ComposedFunction{typeof(SciMLBasePythonCallExt._pyconvert), Py}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, @Kwargs{}, SciMLBase.StandardODEProblem}; u_names::Nothing, p_names::Nothing, kwargs::@Kwargs{})
    @ ModelingToolkit C:\Users\Ruiren\.julia\packages\ModelingToolkit\YLJ0I\src\systems\diffeqs\modelingtoolkitize.jl:70
 [18] modelingtoolkitize
    @ C:\Users\Ruiren\.julia\packages\ModelingToolkit\YLJ0I\src\systems\diffeqs\modelingtoolkitize.jl:6 [inlined]
 [19] jit(x::SciMLBase.ODEProblem{Vector{Float64}, Tuple{Float64, Float64}, false, PyList{Any}, SciMLBase.ODEFunction{false, SciMLBase.AutoSpecialize, ComposedFunction{typeof(SciMLBasePythonCallExt._pyconvert), Py}, LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing, Nothing}, @Kwargs{}, SciMLBase.StandardODEProblem})
    @ Main .\none:2
 [20] pyjlany_call(self::typeof(jit), args_::Py, kwargs_::Py)
    @ PythonCall.JlWrap C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\any.jl:47
 [21] _pyjl_callmethod(f::Any, self_::Ptr{PythonCall.C.PyObject}, args_::Ptr{PythonCall.C.PyObject}, nargs::Int64)
    @ PythonCall.JlWrap C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\base.jl:73
 [22] _pyjl_callmethod(o::Ptr{PythonCall.C.PyObject}, args::Ptr{PythonCall.C.PyObject})
    @ PythonCall.JlWrap.Cjl C:\Users\Ruiren\.julia\packages\PythonCall\WMWY0\src\JlWrap\C.jl:63

Is it possible to solve this problem? im running on my pc which has a rtx 4070ti. Thank you so much in advance!