From 0a6d7e16b7c1ca2b265322183b9095c949a78ae7 Mon Sep 17 00:00:00 2001
From: Hongyang Zhou <hyzhou@umich.edu>
Date: Sat, 3 Aug 2024 22:02:22 +0800
Subject: [PATCH] Separate _getvector from streamplot

---
 src/plot/pyplot.jl  | 14 +++++++++++---
 src/plot/utility.jl |  1 +
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/plot/pyplot.jl b/src/plot/pyplot.jl
index cc14177e..9b945624 100644
--- a/src/plot/pyplot.jl
+++ b/src/plot/pyplot.jl
@@ -651,10 +651,19 @@ end
     streamplot(data, var, ax=nothing; plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1,
        kwargs...)
 
-Wrapper over `streamplot` in matplotlib .
+Wrapper over `streamplot` in matplotlib.
 """
 function PyPlot.streamplot(bd::BATLData{2, T}, var::AbstractString, ax=nothing;
    plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1, kwargs...) where T
+   xi, yi, v1, v2 = _getvector(bd, var; plotrange, plotinterval)
+
+   if isnothing(ax) ax = plt.gca() end
+
+   ax.streamplot(xi, yi, v1, v2; kwargs...)
+end
+
+function _getvector(bd::BATLData{2, T}, var::AbstractString;
+   plotrange=[-Inf,Inf,-Inf,Inf], plotinterval=0.1) where T
    x, w = bd.x, bd.w
    varstream = split(var, ";")
    wnames = lowercase.(bd.head.wnames)
@@ -701,9 +710,8 @@ function PyPlot.streamplot(bd::BATLData{2, T}, var::AbstractString, ax=nothing;
          v2 = [interp2(i, j) for j in yi, i in xi]
       end
    end
-   if isnothing(ax) ax = plt.gca() end
 
-   ax.streamplot(xi, yi, v1, v2; kwargs...)
+   xi, yi, v1, v2
 end
 
 """
diff --git a/src/plot/utility.jl b/src/plot/utility.jl
index 7f3921c8..3a5d4c08 100644
--- a/src/plot/utility.jl
+++ b/src/plot/utility.jl
@@ -122,6 +122,7 @@ function interpolate2d_generalized_coords(X::T, Y::T, W::T,
    xi = range(plotrange[1], stop=plotrange[2], step=plotinterval)
    yi = range(plotrange[3], stop=plotrange[4], step=plotinterval)
    itp = NN.interpolate(X, Y, W)
+   #TODO Interpolate as a whole using predicates with multithreading
    Wi = [itp(x, y; method=NN.Triangle()) for y in yi, x in xi]::Matrix{eltype(W)}
 
    xi, yi, Wi