In [2]:
%load_ext Cython

In [3]:
%%cython
import numpy as np
cimport numpy as np
cimport cython
from libc cimport math


@cython.boundscheck(False) # turn off bounds-checking for entire function
@cython.wraparound(False)  # turn off negative index wrapping for entire function
cdef simulate_path(
    double x,
    double y,
    double sd,
    double speed,
    double t_max,    
    double dt):
    
    cdef int steps
    steps = int(math.floor(t_max / dt))    
    
    # pre-generate random numbers for each step. border reaches cause bounce backs
    direction_sds_arr = np.random.normal(0.0, sd, steps)
    cdef double[:] direction_sds = direction_sds_arr
    
    positions_arr = np.zeros([steps, 2], dtype=np.float64)
    cdef double[:, :] positions = positions_arr
    positions[0, 0] = x / 2.0
    positions[0, 1] = y / 2.0
    
    cdef double prev_direction = 0.0
    cdef double cur_direction = 0.0
    cdef double new_position_x = 0.0
    cdef double new_position_y = 0.0
    cdef double cur_x = 0.0
    cdef double cur_y = 0.0
    cdef Py_ssize_t step
    # loop over timesteps
    for step in range(1, steps):
        cur_direction = direction_sds[step] + prev_direction
        cur_x = math.cos(cur_direction)
        cur_y = math.sin(cur_direction)
        new_position_x = positions[step - 1, 0] + speed * dt * cur_x
        new_position_y = positions[step - 1, 1] + speed * dt * cur_y
        
        # turn direction around if wall is hit
        if (new_position_x < 0.0) or (new_position_x > x):
            cur_x = math.cos(cur_direction + 3.1415)
            new_position_x = positions[step - 1, 0] + speed * dt * cur_x
            
        if (new_position_y < 0.0) or (new_position_y > y):
            cur_y = math.sin(cur_direction + 3.1415)
            new_position_y = positions[step - 1, 1] + speed * dt * cur_y
            
        positions[step, 0] = new_position_x
        positions[step, 1] = new_position_y   
        prev_direction = cur_direction
    
    return positions

def simulate(x, y, sd, speed, t_max, dt):
    return simulate_path(x, y, sd, speed, t_max, dt)

In [7]:
path = simulate(1.25, 1.25, 0.2, 0.4, 36000, 0.01)
np.save("rat_path_10h.npy", path)

In [8]:
path = simulate(1.25, 1.25, 0.2, 0.4, 3600, 0.01)
np.save("rat_path_1h.npy", path)

In [9]:
path = simulate(1.25, 1.25, 0.2, 0.4, 1, 0.01)
np.save("rat_path_1s.npy", path)

In [ ]: