Skip to content

Commit 34397ac

Browse files
AdwaithBatchuAdwaithBatchu
committed
Implement basic arrows3d method in Axes3D
1 parent 40a47e5 commit 34397ac

6 files changed

Lines changed: 135 additions & 1 deletion

File tree

‎doc/users/next_whats_new/arrow3d.rst‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
New ``arrow3d`` Method
2+
----------------------
3+
4+
The new ``arrow3d`` method for `~.Axes3D` allows users to plot 3D arrows,
5+
which can be used for representing vectors or directional data.

‎galleries/examples/mplot3d/arrow3d.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
"""
2+
=============
3+
3D arrow plot
4+
=============
5+
6+
Demonstrates plotting arrows in a 3D space.
7+
8+
Here we plot two arrows from the same start point to different
9+
end points. The properties of the second arrow is changed by passing
10+
additional parameters other than ``end`` and ``start`` to
11+
`.patches.FancyArrowPatch`.
12+
"""
13+
14+
import matplotlib.pyplot as plt
15+
import numpy as np
16+
17+
fig = plt.figure()
18+
ax = fig.add_subplot(111, projection='3d')
19+
20+
# Define the start and end points of the arrow
21+
start = np.array([0, 0, 0])
22+
end = np.array([1, 1, 1])
23+
24+
# Create the arrow
25+
ax.arrow3d(end, start)
26+
27+
end1 = np.array([1, 2, 3])
28+
# Passing additional keyword arguments to control properties of the arrow.
29+
# If the `start` parameter is not passed, the arrow is drawn from (0, 0, 0).
30+
ax.arrow3d(end1, mutation_scale=20, color='r', arrowstyle='->', linewidth=2)
31+
32+
plt.show()
33+
34+
# %%
35+
# .. tags::
36+
# plot-type: 3D,
37+
# level: beginner

‎lib/mpl_toolkits/mplot3d/art3d.py‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
path as mpath, rcParams)
1919
from matplotlib.collections import (
2020
Collection, LineCollection, PolyCollection, PatchCollection, PathCollection)
21-
from matplotlib.patches import Patch
21+
from matplotlib.patches import Patch, FancyArrowPatch
2222
from . import proj3d
2323

2424

@@ -1668,3 +1668,25 @@ def norm(x):
16681668
colors = np.asanyarray(color).copy()
16691669

16701670
return colors
1671+
1672+
1673+
class Arrow3D(FancyArrowPatch):
1674+
"""
1675+
3D FancyArrowPatch object.
1676+
"""
1677+
def __init__(self, xs, ys, zs, *args, **kwargs):
1678+
"""
1679+
Parameters
1680+
----------
1681+
xs, ys, zs : array-like
1682+
The x, y, and z coordinates of the arrow's start and end points.
1683+
"""
1684+
super().__init__((0,0), (0,0), *args, **kwargs)
1685+
self._verts3d = xs, ys, zs
1686+
1687+
def do_3d_projection(self, renderer=None):
1688+
"""Projects the points according to the renderer matrix."""
1689+
xs3d, ys3d, zs3d = self._verts3d
1690+
xs, ys, zs = proj3d.proj_transform(xs3d, ys3d, zs3d, self.axes.M)
1691+
self.set_positions((xs[0], ys[0]), (xs[1], ys[1]))
1692+
return np.min(zs)

‎lib/mpl_toolkits/mplot3d/axes3d.py‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4062,6 +4062,50 @@ def stem(self, x, y, z, *, linefmt='C0-', markerfmt='C0o', basefmt='C3-',
40624062

40634063
stem3D = stem
40644064

4065+
@_preprocess_data()
4066+
def arrow3d(self, end, start=None, **kwargs):
4067+
"""
4068+
3D plot of a single arrow
4069+
4070+
Parameters
4071+
----------
4072+
end : 1D array
4073+
an array of shape (3,).
4074+
4075+
start : 1D array, default: (0,0,0)
4076+
an array of shape (3,).
4077+
4078+
**kwargs
4079+
All other keyword arguments are passed on to
4080+
`~mpl_toolkits.mplot3d.art3d.Arrow3D`.
4081+
4082+
Returns
4083+
-------
4084+
arrow : `~mpl_toolkits.mplot3d.art3d.Arrow3D`
4085+
4086+
"""
4087+
had_data = self.has_data()
4088+
4089+
if start is None:
4090+
start = np.zeros_like(end)
4091+
if np.shape(end) != (3,):
4092+
raise ValueError("end must be an array of length 3")
4093+
if np.shape(start) != (3,):
4094+
raise ValueError("start must be an array of length 3")
4095+
4096+
# Set default arrow properties and update with any additional keyword args
4097+
arrow_props = dict(
4098+
mutation_scale=20, arrowstyle="-|>", shrinkA=0, shrinkB=0
4099+
)
4100+
arrow_props.update(kwargs)
4101+
4102+
xs, ys, zs = [start[0], end[0]], [start[1], end[1]], [start[2], end[2]]
4103+
arrow = art3d.Arrow3D(xs, ys, zs, **arrow_props)
4104+
self.add_artist(arrow)
4105+
self.auto_scale_xyz(xs, ys, zs, had_data)
4106+
4107+
return arrow
4108+
40654109

40664110
def get_test_data(delta=0.05):
40674111
"""Return a tuple X, Y, Z with a test data set."""

‎lib/mpl_toolkits/mplot3d/tests/baseline_images/test_axes3d/arrow3d_custom_props.png‎

52.6 KB
Loading

‎lib/mpl_toolkits/mplot3d/tests/test_axes3d.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2689,3 +2689,29 @@ def test_ndarray_color_kwargs_value_error():
26892689
ax = fig.add_subplot(111, projection='3d')
26902690
ax.scatter(1, 0, 0, color=np.array([0, 0, 0, 1]))
26912691
fig.canvas.draw()
2692+
2693+
2694+
@check_figures_equal()
2695+
def test_arrow3d_default(fig_test, fig_ref):
2696+
ax_ref = fig_ref.add_subplot(projection='3d')
2697+
start = [0, 0, 0]
2698+
end = [1, 2, 3]
2699+
ax_ref.arrow3d(end, start)
2700+
2701+
ax_test = fig_test.add_subplot(projection='3d')
2702+
ax_test.arrow3d(end)
2703+
2704+
2705+
@mpl3d_image_comparison(['arrow3d_custom_props.png'], style='mpl20')
2706+
def test_arrow3d_custom_props():
2707+
fig = plt.figure()
2708+
ax = fig.add_subplot(projection='3d')
2709+
2710+
start1 = [1, 2, 3]
2711+
end1 = [4, 5, 6]
2712+
ax.arrow3d(end1, start1,
2713+
arrowstyle="->, head_length=0.6, head_width=0.3", color='red')
2714+
2715+
start2 = [2, 5, 7]
2716+
end2 = [4, 6, -8]
2717+
ax.arrow3d(end2, start2, color='violet', ls='--', lw=2)

0 commit comments

Comments
 (0)