NaN is not supported in `curved_quiver`

[veenstrajelmer]

I finally got to comparing Ultraplots new curved_quiver to its source dfmt.velovect(). I will discontinue the function from dfm_tools, but I did notice one thing that is not supported now in curved_quiver, nan-values in U and V. This raises _CurvedQuiverTerminateTrajectory. However, nans are legitimate values when converting unstructured grids to structured grids in order to quiver them.

Simple reproducible example:

import numpy as np
import matplotlib.pyplot as plt
plt.close('all')
import ultraplot

x = np.linspace(-4,4,120)
y = np.linspace(-3,3,100)
X,Y = np.meshgrid(x,y)
U = -1 - X**2 + Y
V = 1 + X - Y**2
U[-10:,-10:] = np.nan
V[-10:,-10:] = np.nan
speed = np.sqrt(U*U + V*V)

fig,(ax1,ax2) = plt.subplots(1,2,figsize=(10,5),sharex=True,sharey=True)

# streamplot
strm = ax1.streamplot(X, Y, U, V, color=speed, cmap='winter', arrowstyle='fancy', linewidth=speed/5, integration_direction='both')
fig.colorbar(strm.lines)
ax1.set_title("matplotlib streamplot")

# curved_quiver
strm = ultraplot.axes.PlotAxes.curved_quiver(
 self=ax2, x=X, y=Y, u=U, v=V, color=speed, cmap='winter',
 linewidth=speed/5,
 arrow_at_end=True,
 density=5,
 scale=5,
 arrowsize=1,
 grains=15,
 )
fig.colorbar(strm.lines)
ax2.set_title("UltraPlot curved_quiver")

Traceback:

_CurvedQuiverTerminateTrajectory Traceback (most recent call last)
File c:\data\checkouts\dfm_tools\tests\examples\untitled2.py:34
 30 fig.colorbar(strm.lines)
 31 ax1.set_title("matplotlib streamplot")
 32 
 33 # curved_quiver
---> 34 strm = ultraplot.axes.PlotAxes.curved_quiver(
 35 self=ax2, x=X, y=Y, u=U, v=V, color=speed, cmap='winter',
 36 linewidth=speed/5,
 37 arrow_at_end=True,

File ~\AppData\Local\miniforge3\envs\dfm_tools_env\Lib\site-packages\ultraplot\axes\plot.py:2029, in PlotAxes.curved_quiver(self, x, y, u, v, linewidth, color, cmap, norm, arrowsize, arrowstyle, transform, zorder, start_points, scale, grains, density, arrow_at_end)
 2027 for xs, ys in sp2:
 2028 xg, yg = solver.domain_map.data2grid(xs, ys)
-> 2029 t = integrate(xg, yg)
 2030 if t is not None:
 2031 trajectories.append(t[0])

File ~\AppData\Local\miniforge3\envs\dfm_tools_env\Lib\site-packages\ultraplot\axes\plot_types\curved_quiver.py:232, in CurvedQuiverSolver.get_integrator.<locals>.integrate(x0, y0)
 230 self.domain_map.start_trajectory(x0, y0)
 231 self.domain_map.reset_start_point(x0, y0)
--> 232 stotal, x_traj, y_traj, m_total, hit_edge = self.integrate_rk12(
 233 x0, y0, forward_time, resolution, magnitude
 234 )
 236 if len(x_traj) > 1:
 237 return (x_traj, y_traj), hit_edge

File ~\AppData\Local\miniforge3\envs\dfm_tools_env\Lib\site-packages\ultraplot\axes\plot_types\curved_quiver.py:301, in CurvedQuiverSolver.integrate_rk12(self, x0, y0, f, resolution, magnitude)
 299 xf_traj.append(xi)
 300 yf_traj.append(yi)
--> 301 m_total.append(self.interpgrid(magnitude, xi, yi))
 303 try:
 304 k1x, k1y = f(xi, yi)

File ~\AppData\Local\miniforge3\envs\dfm_tools_env\Lib\site-packages\ultraplot\axes\plot_types\curved_quiver.py:403, in CurvedQuiverSolver.interpgrid(self, a, xi, yi)
 401 if not isinstance(xi, np.ndarray):
 402 if np.ma.is_masked(ai):
--> 403 raise _CurvedQuiverTerminateTrajectory
 404 return ai

_CurvedQuiverTerminateTrajectory: 

Please consider supporting nan-values.

[cvanelteren]

Done!