dereks
dereks

Reputation: 564

How to plot 2d vectors with matplotlib automatically setting axes?

I almost wrote a code which plots vectors:

a = [2,-1]
b = [1,2]
d = [5,2]

def plot_vectors(**kwargs):
  M = []
  for key, value in kwargs.items():
    M.append(list([key]) + list(value))

  ax = plt.axes()
  ax.grid(b=True, which='major')

  for i in range(len(M)):
    l = 0

    for j in range(1, len(M[i])):
      l += M[i][j]**2

    l = l**0.5
    ax.text(M[i][1]/2, M[i][2]/2, f"{M[i][0]}={l:.2f}", size=14)
    ax.plot([0,M[i][1]], [0,M[i][2]])

  ax.set_aspect('equal', 'box')

plot_vectors(a=a, b=b, d=d)

enter image description here

The main idea is not to set up ax.set_xlim directly but do this automatically with ax.set_aspect('equal', 'box'). I didn't find how to do that with ax.quiver and ax.arrow. Can anybody suggest how to draw arrows here and modify y axis values to look like this: enter image description here

I modified the code to support 2D numpy arrays:

a = [2,-1]
b = [1,2]
d = [5,2]

def plot_vectors(**kwargs):
  M = []
  for key, value in kwargs.items():
    if isinstance(value, np.ndarray):
      value = value.tolist()

    for i, v in enumerate(value):
      if isinstance(v, np.ndarray):
        value[i] = value[i].tolist()

    if not isinstance(value[0], list):
      value = [[0,0], value]

    M.append([key] + value)

  ax = plt.axes()
  ax.grid(b=True, which='major')

  for i in range(len(M)):
    l = 0; pos = []
    for j in range(0, len(M[i][1])):
      pos.append(M[i][2][j] - M[i][1][j])
      l += (pos[j])**2
      pos[j] = pos[j] / 2 + M[i][1][j]

    l = l**0.5
    ax.plot([M[i][1][0], M[i][2][0]], [M[i][1][1], M[i][2][1]])
    ax.text(pos[0], pos[1], f"{M[i][0]}={l:.2f}", size=14)

  ax.set_aspect('equal', 'box')

plot_vectors(a=np.array(a), b=b, d=d, e=[d,np.array(b)])

enter image description here

My attempt with quiver:

a = [2,-1]
b = [1,2]
d = [5,2]

def plot_vectors(**kwargs):
  M = []
  for key, value in kwargs.items():
    if isinstance(value, np.ndarray):
      value = value.tolist()

    for i, v in enumerate(value):
      if isinstance(v, np.ndarray):
        value[i] = value[i].tolist()

    if not isinstance(value[0], list):
      value = [[0,0], value]
    M.append([key] + value)

  ax = plt.axes()
  ax.grid(b=True, which='major')
  print(M)
  for i in range(len(M)):
    l = 0; pos = []
    for j in range(0, len(M[i][1])):
      pos.append(M[i][2][j] - M[i][1][j])
      l += (pos[j])**2
      pos[j] = pos[j] / 2 + M[i][1][j]
    l = l**0.5
    ax.text(pos[0], pos[1], f"{M[i][0]}={l:.2f}", size=14)

  x, y, u, v = zip(*[(i[1][0], i[1][1], i[2][0], i[2][1]) for i in M])
  print(x, y, u, v)
  ax.quiver(x, y, u, v, scale=1)
  ax.set_aspect('equal', 'box')

plot_vectors(a=np.array(a), b=b, d=d, e=[d,np.array(b)])

returns:

[['a', [0, 0], [2, -1]], ['b', [0, 0], [1, 2]], ['d', [0, 0], [5, 2]],
['e', [5, 2], [1, 2]]]
(0, 0, 0, 5) (0, 0, 0, 2) (2, 1, 5, 1) (-1, 2, 2, 2)

enter image description here

Problems:

  1. Vector e was not drawn.
  2. How to assign visually distinct random color for every new vector?
  3. Part of the vector a was not drawn.
  4. How to remove decimal values from y axis?

Finally I did what I wanted with great help of @ImportanceOfBeingErnest:

import matplotlib.pyplot as plt
from matplotlib.ticker import MaxNLocator
import numpy as np

a = [2,-1]
b = [1,2]
d = [5,2]

def plot_vectors(**kwargs):
  M = []
  for key, value in kwargs.items():
    if isinstance(value, np.ndarray):
      value = value.tolist()

    for i, v in enumerate(value):
      if isinstance(v, np.ndarray):
        value[i] = value[i].tolist()

    if not isinstance(value[0], list):
      value = [[0,0], value]

    M.append([key] + value)

  plt.figure(figsize=(12,12))
  ax = plt.axes()
  ax.grid(b=True, which='major')
  ax.xaxis.set_major_locator(MaxNLocator(integer=True)); ax.yaxis.set_major_locator(MaxNLocator(integer=True))
  ax.set_aspect('equal', 'box')
  cmap = plt.get_cmap('nipy_spectral')
  lc = np.linspace(0.03, 0.99, 20)
  colors = cmap(np.insert(lc[::2], range(10), lc[::-2]))

  for i in range(len(M)):
    l = 0; pos = []

    for j in range(0, len(M[i][1])):
      pos.append(M[i][2][j] - M[i][1][j])
      l += (pos[j])**2
      pos[j] = pos[j] / 2 + M[i][1][j]

    l = l**0.5
    ax.text(pos[0], pos[1], f'{M[i][0]}={l:.2f}', size=18)

  x, y, u, v = zip(*[(i[1][0], i[1][1], i[2][0] - i[1][0], i[2][1] - i[1][1]) for i in M])
  ax.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=1., color=colors[:len(M)])
  ax.plot(np.array(x)+np.array(u), np.array(y)+np.array(v), np.array(x), np.array(y), visible=False)

plot_vectors(a=np.array(a), b=b, d=d, e=np.array([d,np.array(b)]), ab=[a,b])

enter image description here

Upvotes: 1

Views: 594

Answers (1)

ImportanceOfBeingErnest
ImportanceOfBeingErnest

Reputation: 339290

I think you will find solutions to 2. and 4. by searching stackoverflow a bit harder. The real problem (1./3.) is that the arrow ends do not take part in the autoscaling mechanism. This is in general expected when they are not in data coordinates, but in case they are, one could expect them to change the data limits of the plot.

In any case a workaround would be to plot an invisible plot in addition to the quiver with the points from the start and end of the vectors:

ax.quiver(x, y, u, v, angles='xy', scale_units='xy', scale=1.)
ax.plot(np.array(x)+np.array(u), np.array(y)+np.array(v), np.array(x), np.array(y), visible=False)

enter image description here

Upvotes: 2

Related Questions