Is there a way to improve the line quality when exporting streamplots from matplotlib?

Question

I am drawing streamplots using matplotlib, and exporting them to a vector format. However, I find the streamlines are exported as a series of separate lines - not joined objects. This has the effect of reducing the quality of the image, and making for an unwieldy file for further manipulation. An example; the following images are of a pdf generated by exportfig and viewed in Acrobat Reader:

This is the entire plot

and this is a zoom of the center.

Interestingly, the length of these short line segments is affected by 'density' - increasing the density decreases the length of the lines. I get the same behavior whether exporting to svg, pdf or eps.

Is there a way to get a streamplot to export streamlines as a single object, preferably as a curved line?

MWE

import matplotlib.pyplot as plt
import numpy as np

square_size = 101

x = np.linspace(-1,1,square_size)
y = np.linspace(-1,1,square_size)

u, v = np.meshgrid(-x,y)

fig, axis = plt.subplots(1, figsize = (4,3))
axis.streamplot(x,y,u,v)
fig.savefig('YourDirHere\\test.pdf')

Answer 1

A quick solution to this issue is to change the default cap styles of those tiny segments drawn by the streamplot function. In order to do this, follow the below steps.

1. Extract all the segments from the stream plot.
2. Bundle these segments through LineCollection function.
3. Set the collection's cap style to round.
4. Set the collection's zorder value smaller than the stream plot's default 2. If it is higher than the default value, the arrows of the stream plot will be overdrawn by the lines of the new collection.
5. Add the collection to the figure.

The solution of the example code is presented below.

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.collections import LineCollection # Import LineCollection function.

square_size = 101

x = np.linspace(-1,1,square_size)
y = np.linspace(-1,1,square_size)

u, v = np.meshgrid(-x,y)

fig, axis = plt.subplots(1, figsize = (4,3))
strm = axis.streamplot(x,y,u,v)

# Extract all the segments from streamplot.
strm_seg = strm.lines.get_segments()

# Bundle segments with round capstyle. The `zorder` value should be less than 2 to not
# overlap streamplot's arrows.
lc = LineCollection(strm_seg, zorder=1.9, capstyle='round')

# Add the bundled segment to the subplot.
axis.add_collection(lc)

fig.savefig('streamline.pdf')

Additionally, if you want to have streamlines their line widths changing throughout the graph, you have to extract them and append this information to LineCollection.

strm_lw = strm.lines.get_linewidths()
lc = LineCollection(strm_seg, zorder=1.9, capstyle='round', linewidths=strm_lw)

Sadly, the implementation of a color map is not as straight as the above solution. Therefore, using a color map with above approach will not be very pleasing. You can still automate the coloring process, as shown below.

strm_col = strm.lines.get_color()
lc = LineCollection(strm_seg, zorder=1.9, capstyle='round', color=strm_col)

Lastly, I opened a pull request to change the default capstyle option in the matplotlib repository, it can be seen here. You can apply this commit using below code too. If you prefer to do so, you do not need any tricks explained above.

diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
index 95ce56a512..0229ae107c 100644
--- a/lib/matplotlib/streamplot.py
+++ b/lib/matplotlib/streamplot.py
@@ -222,7 +222,7 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
         arrows.append(p)
 
     lc = mcollections.LineCollection(
-        streamlines, transform=transform, **line_kw)
+        streamlines, transform=transform, **line_kw, capstyle='round')
     lc.sticky_edges.x[:] = [grid.x_origin, grid.x_origin + grid.width]
     lc.sticky_edges.y[:] = [grid.y_origin, grid.y_origin + grid.height]
     if use_multicolor_lines:

Answer 2

In the end, it seemed like the best solution was to extract the lines from the streamplot object, and plot them using axis.plot. The lines are stored as individual segments with no clue as to which line they belong, so it is necessary to stitch them together into continuous lines.

Code follows:

import matplotlib.pyplot as plt
import numpy as np

def extract_streamlines(sl):
    
    # empty list for extracted lines, flag
    new_lines = []
    
    for line in sl:

        #ignore zero length lines
        if np.array_equiv(line[0],line[1]):
            continue

        ap_flag = 1

        for new_line in new_lines:
            
            #append the line segment to either start or end of exiting lines, if either the star or end of the segment is close.
            if np.allclose(line[0],new_line[-1]):
                new_line.append(list(line[1]))
                ap_flag = 0
                break

            elif np.allclose(line[1],new_line[-1]):
                new_line.append(list(line[0]))
                ap_flag = 0
                break
            
            elif np.allclose(line[0],new_line[0]):
                new_line.insert(0,list(line[1]))
                ap_flag = 0
                break

            elif np.allclose(line[1],new_line[0]):
                new_line.insert(0,list(line[0]))
                ap_flag = 0
                break
                
        # otherwise start a new line
        if ap_flag:
            new_lines.append(line.tolist())

    return [np.array(line) for line in new_lines]

square_size = 101

x = np.linspace(-1,1,square_size)
y = np.linspace(-1,1,square_size)

u, v = np.meshgrid(-x,y)

fig_stream, axis_stream = plt.subplots(1, figsize = (4,3))
stream = axis_stream.streamplot(x,y,u,v)

np_new_lines = extract_streamlines(stream.lines.get_segments())

fig, axis = plt.subplots(1, figsize = (4,4))

for line in np_new_lines:
    axis.plot(line[:,0], line[:,1])

fig.savefig('YourDirHere\\test.pdf')