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!