Nicholaswogan / numbalsoda

Python wrapper of LSODA (solving ODEs) which can be called from within numba functions.
MIT License
91 stars 8 forks source link

elementwise vector operation and passing size of the system #27

Open diffproject opened 1 year ago

diffproject commented 1 year ago

Hello,

I am new to numblsoda so apologies if I missed something obvious. But I am trying to solve a system of N equations with N being an user input. Is there any way to pass N to the step function(rhs in the docs) written below? Also, is there any way to do vector operations without writing for loops?

Here is my minimal working code in scipy:

N=1000
km=1+np.log(2)/3
kp=1.1*km
kbm=10**-6
T=24
t=np.arange(0,T+1,1)

d=np.array([-((kp*k/N + kbm)*(N - k) + km*k) for k in range(0,N+1)])
u=np.array([km*(k + 1) for k in range(0,N)])
l=np.array([(kp*(k - 1)/N + kbm)*(N - k + 1) for k in range(1,N+1)])

pi=np.zeros(N+1, dtype=np.float64)
pi[N]=1
dp=np.zeros(N+1)

def step(p, t):
      dp[0]=d[0]*p[0]+u[0]*p[1]
      dp[1:N]=d[1:N]*p[1:N]+u[1:]*p[2:]+l[:N-1]*p[:N-1]
      dp[N]=d[N]*p[N]+l[N-1]*p[N-1]
      return dp

sol = odeint(step, pi, t, mxstep=2000)

Thanks.

Nicholaswogan commented 1 year ago

Hi! See the section labeled "Passing data to the right-hand-side function" in the README. Regarding the vectorized operations question, for numbalsoda, it doesn't matter if you write for loops or use "vectorized" operations. Loops are just as fast as numpy vectorization (or faster) in a numba function. For example, see the benchmark below

import numpy as np
import numba as nb
import timeit

@nb.njit()
def add_numba(a, b):
    c = np.empty(a.shape[0])
    for i in range(a.shape[0]):
        c[i] = a[i] + b[i]

@nb.njit()
def test_numba():
    n = 1000
    a = np.random.rand(n)
    b = np.random.rand(n)
    c = add_numba(a, b)

def add_numpy(a, b):
    c = a + b
    return c

def test_numpy():
    n = 1000
    a = np.random.rand(n)
    b = np.random.rand(n)
    c = add_numpy(a, b)

test_numba()
test_numpy()

iters = 100000
time_numba = timeit.Timer(test_numba).timeit(number=iters)/iters
time_numpy = timeit.Timer(test_numpy).timeit(number=iters)/iters
print('numba time = %e'%time_numba)
print('numpy time = %e'%time_numpy)
diffproject commented 1 year ago

I see. Thanks. It runs, however, I get nan for solution even though scipy odeint gives me correct answer. Here is my implementation:

N=1000
km=1+np.log(2)/3
kp=1.1*km
kbm=10**-6
T=24
t=np.arange(0,T+1,1)
tsteps=np.array([0,16,19,24])

d=np.array([-((kp*k/N + kbm)*(N - k) + km*k) for k in range(0,N+1)])
u=np.array([km*(k + 1) for k in range(0,N)])
l=np.array([(kp*(k - 1)/N + kbm)*(N - k + 1) for k in range(1,N+1)])

pi=np.zeros(N+1, dtype=np.float64)
pi[N]=1
dp=np.zeros(N+1)

def make_lsoda(N,d,u,l):
    @cfunc(lsoda_sig)
    def step(t, p, dp, data):
        dp[0]=d[0]*p[0]+u[0]*p[1]

        for i in range(1,N):
            dp[i]=d[i]*p[i]+u[i]*p[i+1]+l[i-1]*p[i-1]

        dp[N]=d[N]*p[N]+l[N-1]*p[N-1]

    return step

rhs=make_lsoda(N,d,u,l)
funcptr = rhs.address

sol, success=lsoda(funcptr, pi, t)
Nicholaswogan commented 1 year ago

The problem is you are giving numbalsoda an array of integers instead of an array of floating point numbers for the times to evaluate the solution at. This is a bug in the code. You can fix you case by writing the following

t=np.arange(0,T+1,1,dtype=np.float64)
diffproject commented 1 year ago

Thanks. That works. I will be closing the issue.