Python streamline algorithm

Question

Goal:

I have 2 arrays vx and vy representing velocity components. I want to write a streamline algorithm:

1. Input the coordinates of a point (seed)
2. Evaluate which pixels are on the path of the input point based on its velocity components
3. Return the indices of the points in the path of the seed point

Issue/Question:

I initially wrote a Euler-forward algorithm that was solving very poorly my problem. I was advised to consider my problem as an Ordinary Differential Equation (ODE) where dx/dt = v_x(t) and dy/dt = v_y(t). I can interpolate my velocities but struggle with solving the ODE with Scipy. How could I do that ?

Homemade algorithm:

I have 2 arrays vx and vy representing velocity components. When one has a NaN, the other has one too. I have a point from which I start, the seed point. I want to track which cells this point went through based on the velocity components. I interpolate the velocity components vx and vy in order to input them in an ODE solver.

Example:

This code tests the algorithm for a 10x11 velocities array. I am blocked at the ODE solver.

import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import RegularGridInterpolator
from scipy.integrate import odeint


# Create coordinates
x = np.linspace(0, 10, 100)
y = np.linspace(11, 20, 90)
Y, X = np.meshgrid(x, y)

# Create velocity fields
vx = -1 - X**2 + Y
vy = 1 + X - Y**2


# Seed point
J = 5
I = 14

# Interpolate the velocity fields
interpvx = RegularGridInterpolator((y,x), vx)
interpvy = RegularGridInterpolator((y,x), vy)

# Solve the ODE to get the point's path, but I don't know what to put for the parameter t
#solx = odeint(interpvx, interpvx((I,J)), np.linspace(0,1,501))
#soly = odeint(interpvy, interpvx((I,J)), np.linspace(0,1,501))

Answer 1

Someone I work with came with a solution, I am not taking credit for the answer. But considering there is no streamline algorithm for Python, hopefully that will be useful to someone here:

import numpy as np
# Create coordinates
x = np.linspace(0, 1000, 100)
y = np.linspace(0, 1000, 90)
Y, X = np.meshgrid(x, y)

# Create velocity fields
vx = -1 - X**2 + Y
vy = 1 + X - Y**2


# Seed point
J = 5
I = 14

# Interpolate the velocities
from scipy.interpolate import RegularGridInterpolator
X, Y = np.meshgrid(x,y)
fx = RegularGridInterpolator((y, x), vx, bounds_error=False, fill_value=None)
fy = RegularGridInterpolator((y, x), vy, bounds_error=False, fill_value=None)

# define the velocity function to be integrated:
def f(t, y):
    return np.squeeze([fy(y), fx(y)])

# Solve for a seed point
from scipy.integrate import solve_ivp
sol = solve_ivp(f, [0, 100], [J,I], t_eval=np.arange(0,100,1))