Commit a5fe594d authored by Matthieu Boileau's avatar Matthieu Boileau
Browse files

srw: color plot

parent 1956c3c8
......@@ -2,6 +2,7 @@
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.collections import LineCollection
from math import sqrt, pi
from numba import jit
......@@ -82,23 +83,49 @@ class Walk2D:
return anim
def plot(self):
def plot(self, colorize=True):
"""Plot walk in a figure"""
x = self.x
y = self.y
fig, ax = self._init_figure()
path, = ax.plot(x, y) # a line for path
if colorize:
t = np.arange(1, x.shape[0] + 1) # time variable
# symbols for initial (squared) and final (round) position
ax.plot(x[0], y[0], 's', color=path.get_color())
ax.plot(x[-1], y[-1], 'o')
# set up an array of (x,y) points
points = np.array([x, y]).transpose().reshape(-1, 1, 2)
# set up a list of segments
segs = np.concatenate([points[:-1], points[1:]], axis=1)
def plot_walk(nstep: int):
"""plot a nstep-2D walk"""
walk = Walk2D(nstep)
# make the collection of segments
cmap = plt.get_cmap('jet')
lc = LineCollection(segs, cmap=cmap, alpha=0.5)
lc.set_array(t) # color the segments by the time parameter
# plot the collection
# line collections donnot auto-scale the plot
ax.set_xlim(x.min(), x.max())
ax.set_ylim(y.min(), y.max())
start_color = cmap(t[0])
end_color = cmap(t[-2] * 256 / t.shape[0])
path, = ax.plot(x, y) # a line for path
start_color = 'C1' # blue
end_color = 'C2' # orange
# symbols for initial (squared) and final (round) position
ax.plot(x[0], y[0], 's', ms=10,
ax.plot(x[-1], y[-1], 'o', ms=10,
class NWalk:
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment