Visualization of 2D rod-shaped particles in python

Question

I have rod-liked particles in 2D which consist of a middle rectangular part of length(l) and width(d), there are two semicircular caps on both ends of diameter(d). Diameter(d) for all particles remains constant and length(l) changes. I want to visualize it in python. Can matplotlib do it with custom markers? Since there are around 5K particles in simulation I want it to be fast.

I have been using gnuplot, plotting the rods as vectors which gives me an idea of the particles, but not the actual size.

Answer 1

You may want to have a look at this post, which uses an idea from this answer.

A rod would be a line with spherical caps (solid_capstyle="round").

The linewidth, which is normally given in points is then calculated from the data units and the current data transform.

In order to add many of such lines, one may change the code from the aforementionned post to create a lot of lines. Unfortunately, a LineCollection, which would speed things up, does not allow for changing the capstyle.

Here is an example:

import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
import numpy as np

class Rods():
    def __init__(self, x,y,l,w=1, angle=0, **kwargs):
        self.ax = kwargs.pop("ax", plt.gca())
        self.color=kwargs.pop("color", "lightblue")
        self.lw = 1
        self.ax.figure.canvas.draw()
        self.ppd=72./self.ax.figure.dpi
        self.trans = self.ax.transData.transform
        self.lines = []
        self.set_data(x,y,l,w, angle)

        self.cid=self.ax.figure.canvas.mpl_connect('resize_event',self._resize)
        self.cid1=self.ax.figure.canvas.mpl_connect("motion_notify_event",self._resize)
        self.cid2=self.ax.figure.canvas.mpl_connect('button_release_event',self._resize)

    def set_data(self, x,y,l,w=1, angle=0, color=None):
        if color: self.color=color
        self.lw_data = w
        m = np.array([[np.cos(angle), -np.sin(angle)],[np.sin(angle), np.cos(angle)]])
        xy1 = np.dot(m, np.c_[np.zeros_like(x),-l/2.*np.ones_like(x)].T).T
        xy2 = np.dot(m, np.c_[np.zeros_like(x), l/2.*np.ones_like(x)].T).T

        x = np.c_[xy1[:,0] + x,xy2[:,0] + x]
        y = np.c_[xy1[:,1] + y,xy2[:,1] + y]

        if self.lines:
            for line in self.lines: line.remove()
        for i in range(len(x)):
            line = Line2D(x[i,:], y[i,:], 
                           solid_capstyle="round", color=self.color)
            self.lines.append(line)

        for line in self.lines:
            self.ax.add_artist(line)
        self._resize()

    def _resize(self, event=None):
        lw =  ((self.trans((1,self.lw_data))-self.trans((0,0)))*self.ppd)[1]
        if lw != self.lw:
            for line in self.lines:
                line.set_linewidth(lw)
            self.ax.figure.canvas.draw_idle()
            self.lw = lw


fig, ax = plt.subplots()
ax.set_aspect('equal')
ax.axis([0,6,0,6])

x = np.array([1,2,3,4])
y = np.array([2,3,5,5])
l = np.array([2,3,5,4])*0.2

r = Rods(x,y,l,w=0.1, angle=45, color="crimson")
r = Rods(y[::-1],x[::-1],l,w=0.4, angle=90, color="purple")

plt.show()