How to annotate a regression line with the proper text rotation - python

I have the following snippet of code to draw a best-fit line through a collections of points on a graph, and annotate it with the corresponding R2 value:
import matplotlib.pyplot as plt
import numpy as np
import scipy.stats
x = 50 * np.random.rand(20) + 50
y = 200 * np.random.rand(20)
plt.plot(x, y, 'o')
# k, n = np.polyfit(x, y, 1)
k, n, r, _, _ = scipy.stats.linregress(x, y)
line = plt.axline((0, n), slope=k, color='blue')
xy = line.get_xydata()
plt.annotate(
f'$R^2={r**2:.3f}$',
(xy[0] + xy[-1]) // 2,
xycoords='axes fraction',
ha='center', va='center_baseline',
rotation=k, rotation_mode='anchor',
)
plt.show()
I have tried various different (x,y) pairs, different xycoords and other keyword parameters in annotate but I haven't been able to get the annotation to properly appear where I want it. How do I get the text annotation to appear above the line with proper rotation, located either at the middle point of the line, or at either end?

1. Annotation coordinates
We cannot compute the coordinates using xydata here, as axline() just returns dummy xydata (probably due to the way matplotlib internally plots infinite lines):
print(line.get_xydata())
# array([[0., 0.],
# [1., 1.]])
Instead we can compute the text coordinates based on the xlim():
xmin, xmax = plt.xlim()
xtext = (xmin + xmax) // 2
ytext = k*xtext + n
Note that these are data coordinates, so they should be used with xycoords='data' instead of 'axes fraction'.
2. Annotation angle
We cannot compute the angle purely from the line points, as the angle will also depend on the axis limits and figure dimensions (e.g., imagine the required rotation angle in a 6x4 figure vs 2x8 figure).
Instead we should normalize the calculation to both scales to get the proper visual rotation:
rs = np.random.RandomState(0)
x = 50 * rs.rand(20) + 50
y = 200 * rs.rand(20)
plt.plot(x, y, 'o')
# save ax and fig scales
xmin, xmax = plt.xlim()
ymin, ymax = plt.ylim()
xfig, yfig = plt.gcf().get_size_inches()
k, n, r, _, _ = scipy.stats.linregress(x, y)
plt.axline((0, n), slope=k, color='blue')
# restore x and y limits after axline
plt.xlim(xmin, xmax)
plt.ylim(ymin, ymax)
# find text coordinates at midpoint of regression line
xtext = (xmin + xmax) // 2
ytext = k*xtext + n
# find run and rise of (xtext, ytext) vs (0, n)
dx = xtext
dy = ytext - n
# normalize to ax and fig scales
xnorm = dx * xfig / (xmax - xmin)
ynorm = dy * yfig / (ymax - ymin)
# find normalized annotation angle in radians
rotation = np.rad2deg(np.arctan2(ynorm, xnorm))
plt.annotate(
f'$R^2={r**2:.3f}$',
(xtext, ytext), xycoords='data',
ha='center', va='bottom',
rotation=rotation, rotation_mode='anchor',
)

Related

Python highlight user chosen area in contourf plot

What is the best solution for highlighting an area in a contourf plot?
I want the background to be opacity of 0.5 and the user chosen area to be normal. How can I achieve this?
In How to nicely plot clipped layered artists in matplotlib? Jake Vanderplas shows a way to draw a rectangle with a rectangular hole. The code can be adapted for your situation. The following example starts from a tutorial example, and highlights the third contour:
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.path import Path
from matplotlib.patches import PathPatch
def DoubleRect(xy1, width1, height1,
xy2, width2, height2, **kwargs):
base = np.array([(0, 0), (0, 1), (1, 1), (1, 0), (0, 0)])
verts = np.vstack([xy1 + (width1, height1) * base,
xy2 + (width2, height2) * base[::-1],
xy1])
codes = 2 * ([Path.MOVETO] + 4 * [Path.LINETO]) + [Path.CLOSEPOLY]
return PathPatch(Path(verts, codes), **kwargs)
origin = 'lower'
delta = 0.025
x = y = np.arange(-3.0, 3.01, delta)
X, Y = np.meshgrid(x, y)
Z1 = np.exp(-X ** 2 - Y ** 2)
Z2 = np.exp(-(X - 1) ** 2 - (Y - 1) ** 2)
Z = (Z1 - Z2) * 2
fig, ax = plt.subplots()
contours = ax.contourf(X, Y, Z, 10, cmap=plt.cm.turbo, origin=origin)
# contours.collections[2].set_color('deepskyblue') # mark one contour
# calculate (or get) the coordinates of the hole
bbox = contours.collections[2].get_paths()[0].get_extents()
hole_xy, hole_width, hole_height = bbox.p0, bbox.width, bbox.height
# find the coordinates of the surrounding rectangle
xmin, xmax = ax.get_xlim()
ymin, ymax = ax.get_ylim()
full_rect = plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, color='black', alpha=0.5)
ax.add_patch(full_rect)
# create a rectangle with a hole to clip the surrounding rectangle
mask = DoubleRect((xmin, ymin), xmax - xmin, ymax - ymin,
hole_xy, hole_width, hole_height,
facecolor='none', edgecolor='none')
ax.add_patch(mask)
full_rect.set_clip_path(mask)
plt.show()
Instead of darkening the outside region, it could also be hatched (similar to the linked post). This would set the edge color of the mask to 'black', and create the full rectangle with hatching.
full_rect = plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, facecolor='none', edgecolor='black', hatch='//')

How to plot upper and lower boundary with a LINEAR line on a scatter plot?

I have a data frame df with columns A and Q. I am using this code to draw a line of equation on it.
#Actual line of equation, which has to be plotted: Q=alpha*A^beta : ln(Q)=a+b*ln(A) : y = a+b(x)
x = np.log(df['A'])
y = np.log(df['Q'])
#deriving b,a
b,a = np.polyfit(np.log(x), y, 1)
#deriving alpha and beta. By using a = ln(alpha); b = beta -1
alpha = np.exp(a)
beta = b + 1
Q = df['Q'].values
A = df['A'].values
#equation of line
q = alpha * np.power(A,beta)
#plotting the points and line
plt.scatter(A,Q)
plt.plot(A,q, '-r')
plt.yscale('log')
plt.xscale('log')
This gives the following output, which is similar to a regression line.
But I am interested in plotting the same line of the equation as the upper and lower curve/boundary joining the farthest points(perpendicular to the green line) on both sides as shown below with the same slope as that of the continuous green line.
The idea is to first search the index of the point where the difference between the line and the plot is minimal (cf. maximal). With this point, alpha_min can be calculated such that
Q[pos_min] == alpha_min * np.power(A[pos_min], beta), thus
alpha_min = Q[pos_min] / np.power(A[pos_min], beta).
As such lines can extend quite far away from the original points, it can help to restore the x and y limits (thus clipping the plot to the original region).
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
df = pd.DataFrame()
df['A'] = 10 ** np.random.uniform(0, 1, 1000) ** 2
df['Q'] = 10 ** np.random.uniform(0, 1, 1000) ** 2
x = np.log(df['A'])
y = np.log(df['Q'])
# deriving b,a
b, a = np.polyfit(np.log(x), y, 1)
# deriving alpha and beta. By using a = ln(alpha); b = beta - 1
alpha = np.exp(a)
beta = b + 1
Q = df['Q'].values
A = df['A'].values
# plotting the points and line
plt.yscale('log')
plt.xscale('log')
plt.scatter(A, Q, color='b')
# equation of line
xmin, xmax = plt.xlim() # the limits of the x-axis for drawing the line
x = np.linspace(xmin, xmax, 50)
q = alpha * np.power(x, beta)
plt.plot(x, q, '-r')
ymin, ymax = plt.ylim() # store the limits of the scatter and line plot so they can be restored later
pos_min = np.argmin(Q / np.power(A, beta))
pos_max = np.argmax(Q / np.power(A, beta))
alpha_min = Q[pos_min] / np.power(A[pos_min], beta)
alpha_max = Q[pos_max] / np.power(A[pos_max], beta)
# plt.scatter(A[pos_min], Q[pos_min], s=100, fc='none', ec='r', lw=3)
# plt.scatter(A[pos_max], Q[pos_max], s=100, fc='none', ec='g', lw=3)
plt.plot(x, (alpha_max) * np.power(x, beta), '--r')
plt.plot(x, (alpha_min) * np.power(x, beta), '--r')
plt.xlim(xmin, xmax) # restore the limits of the scatter plot
plt.ylim(ymin, ymax)
plt.show()

how to plot a line with a slope in matplotlib using plt.loglog

Here I need to plot a frequency and a simple line with a slope of -5/3.
The problem is that the main plot is using plt.loglog() and when I want to show the line it gives me nothing or something strange.
Here are the pictures of it. I have plotted the right one and desired is the left one.
I already used np.linspace and some other things, but I was not able to solve the problem. Also, it is not clear at which points I have the first and the end of the frequency plot. That's another reason why I can not use 'np.linspace'. Can anyone help me?
Thanks a lot for your attention. I tried your code but I found out maybe there are better ways to do it with kind of my dataset. So I did this:
Change the class of dataset to np.array() and had np.log() function on it:
x = ... # type(x) = list
y = ... # type(y) = list
.
.
.
x = np.log(np.array(x))
y = np.log(np.array(y))
In this case I did not have to use plt.loglog() or np.log() and np.exp() in calculations anymore.
Locate the min, max, and mean for x and y:
ymin, ymax = ([y.min(), y.max()])
ymid = (ymin + ymax) / 2
xmin, xmax = ([x.min(), x.max()])
xmid = (xmin + xmax) / 2
Use np.linspace() for rest:
slope = - 5 / 3
x1 = np.linspace(xmin, xmax)
y1 = slope * (x1 - xmid) + ymid
ax.plot(x1, y1, 'r')
And got the result I wanted.
the result
Edited: the plot in log scale.
Because of that, it is better to use plt.loglog() kind of plots in frequency spectrums, I edited these things:
Changed back x and y to normal np.array()
x = array(...)
y = array(...)
find the middle of x and y to have the center of the line and used simple equation of a straight line and then ploted the line using np.exp():
ymin, ymax = log([y.min(), y.max()])
ymid = (ymin + ymax) / 2
xmin, xmax = log([x.min(), x.max()])
xmid = (xmin + xmax) / 2
slope = - 5 / 3
y1 = slope * (xmin - xmid) + ymid
y2 = slope * (xmax - xmid) + ymid
ax.plot(exp([xmin, xmax]), exp([y1, y2]), 'r')
plt.loglog()
the result
As you see now we have the plot in log scale.
Here is a possible approach. All calculations happen in logspace. The points are transformed back to linear space to be plotted. As matplotlib plots a line given two points always as straight, independent to the tranformation of the axes, only two points are needed.
Steps:
find the lowest and highest values of the curve (y0, y1); these define the extension of the sloped line
find a point near the center of the curve (mid_x, mid_y); this point serves as an anchor point to know where the sloped line should go
find the x values that correspond to a sloped line through (mid_x, mid_y) and go to y0 and y1
import matplotlib.pyplot as plt
import numpy as np
fig, ax = plt.subplots(figsize=(8, 8))
# first create some toy data roughly resembling the example plot
x = np.linspace(10, 2000, 1000)
y = np.random.normal(2000 / x ** np.linspace(.7, .55, x.size), 100 / x ** .7)
ax.plot(x, y)
y0, y1 = np.log([y.min(), y.max()])
# mid_x, mid_y = np.log([x[x.size // 2], y[y.size // 2]])
mid_x, mid_y = (np.log(x.min()) + np.log(x.max())) / 2, (y0 + y1) / 2
slope = -5 / 3
x0 = mid_x + slope * (y0 - mid_y)
x1 = mid_x + slope * (y1 - mid_y)
ax.plot(np.exp([x0, x1]), np.exp([y0, y1]), color='crimson')
plt.loglog()
plt.show()

Compromise between scipy.interpolate.Rbf and scipy.interpolate.griddata?

I need to generate a 2D density map given some unstructured (x, y) coordinates, and a z value which acts as the weight assigned to each point.
I can interpolate the values into a grid (see code below) using either scipy.interpolate.Rbf or scipy.interpolate.griddata. The first one is very accurate but very slow, while the second one is a lot faster but noticeably less accurate.
See image below where Rfb is to the left and griddata to the right:
Is there some way to compromise between these two functions? Something that will produce a result more similar to the Rfb function but with a speed more similar to that of the griddata function?
# Data
N = 1000
x, y, z = np.random.uniform(0.0005, 0.03, N),\
np.random.uniform(6., 10., N), np.random.uniform(1., 10., N)
xmin, xmax, ymin, ymax = min(x), max(x), min(y), max(y)
# Set up a regular grid of interpolation points
xi, yi = np.linspace(xmin, xmax, 200), np.linspace(ymin, ymax, 200)
xi, yi = np.meshgrid(xi, yi)
# Normalize data and grid.
x_new, xi_new = (x - xmin) / (xmax - xmin), (xi - xmin) / (xmax - xmin)
y_new, yi_new = (y - ymin) / (ymax - ymin), (yi - ymin) / (ymax - ymin)
# Interpolate new data with Rbf.
s = time.clock()
rbf = scipy.interpolate.Rbf(x_new, y_new, z, function='linear')
zi = rbf(xi_new, yi_new)
print(time.clock() - s)
# Plot density map.
ax1 = plt.subplot(121)
plt.pcolormesh(xi, yi, zi, cmap=plt.get_cmap('GnBu_r'), zorder=2)
plt.contour(xi, yi, zi, 2, colors='#551a8b', linewidths=0.5, zorder=3)
# Interpolate new data with griddata.
s = time.clock()
# Python 2.7
vals = zip(*[x_new, y_new])
# Python 3.6
# vals = np.array([x_new, y_new]).T
zi = scipy.interpolate.griddata(vals, z, (xi_new, yi_new), method='linear')
print(time.clock() - s)
# Plot density map.
ax2 = plt.subplot(122)
plt.imshow(
zi, extent=(xmin, xmax, ymin, ymax), origin='lower',
cmap=plt.get_cmap('GnBu_r'))
plt.contour(xi, yi, zi, 2, colors='#551a8b', linewidths=0.5, zorder=3)
plt.show()

Rotate square to be normal to a vector

Win 7, x64, Python 2.7
I'm trying to rotate a square that is initially in the xz plane so that its normal aligns with a given 3D vector. Also I am translating the square to the start of the vector but that isnt a problem.
The path I have taken is as follows,
1) Find the axis of rotation via the cross product of the given vector & the square's normal, a unit vector in the y direction in this case.
2) Find the angle of rotation via the dot product of the given vector and the square's normal.
3) Build appropriate rotation matrix.
4) Apply rotation matrix to the vertices of the square.
5) Translate to the start of the given vector.
The code..
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
na = np.array
def rotation_matrix(axis, theta):
"""
Return the rotation matrix associated with counterclockwise rotation about
the given axis by theta radians.
"""
axis = np.asarray(axis)
axis = axis/math.sqrt(np.dot(axis, axis))
a = math.cos(theta/2.0)
b, c, d = -axis*math.sin(theta/2.0)
aa, bb, cc, dd = a*a, b*b, c*c, d*d
bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
[2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
[2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])
edgeLen = 4.0 # length of square side
pos = na([2.0,2.0,2.0]) # starting point of vector
dirc = na([6.0,6.0,6.0]) # direction of vector
Ux = na([1.0,0.0,0.0]) # unit basis vectors
Uy = na([0.0,1.0,0.0])
Uz = na([0.0,0.0,1.0])
x = pos[0]
y = pos[1]
z = pos[2]
# corner vertices of square in xz plane
verts = na([[edgeLen/2.0, 0, edgeLen/2.0],
[edgeLen/2.0, 0, -edgeLen/2.0],
[-edgeLen/2.0, 0, -edgeLen/2.0],
[-edgeLen/2.0, 0, edgeLen/2.0]])
# For axis & angle of rotation
dirMag = np.linalg.norm(dirc)
axR = np.cross(dirc, Uy)
theta = np.arccos((np.dot(dirc, Uy) / dirMag))
Rax = rotation_matrix(axR, theta) # rotation matrix
# rotate vertices
rotVerts = na([0,0,0])
for v in verts:
temp = np.dot(Rax, v)
temp = na([temp[0]+x, temp[1]+y, temp[2]+z])
rotVerts = np.vstack((rotVerts, temp))
rotVerts = np.delete(rotVerts, rotVerts[0], axis=0)
# plot
# oringinal square
ax.scatter(verts[:,0], verts[:,1], verts[:,2], s=10, c='r', marker='o')
ax.plot([verts[0,0], verts[1,0]], [verts[0,1], verts[1,1]], [verts[0,2], verts[1,2]], color='g', linewidth=1.0)
ax.plot([verts[1,0], verts[2,0]], [verts[1,1], verts[2,1]], [verts[1,2], verts[2,2]], color='g', linewidth=1.0)
ax.plot([verts[2,0], verts[3,0]], [verts[2,1], verts[3,1]], [verts[2,2], verts[3,2]], color='g', linewidth=1.0)
ax.plot([verts[0,0], verts[3,0]], [verts[0,1], verts[3,1]], [verts[0,2], verts[3,2]], color='g', linewidth=1.0)
# rotated & translated square
ax.scatter(rotVerts[:,0], rotVerts[:,1], rotVerts[:,2], s=10, c='b', marker='o')
ax.plot([rotVerts[0,0], rotVerts[1,0]], [rotVerts[0,1], rotVerts[1,1]], [rotVerts[0,2], rotVerts[1,2]], color='b', linewidth=1.0)
ax.plot([rotVerts[1,0], rotVerts[2,0]], [rotVerts[1,1], rotVerts[2,1]], [rotVerts[1,2], rotVerts[2,2]], color='b', linewidth=1.0)
ax.plot([rotVerts[2,0], rotVerts[3,0]], [rotVerts[2,1], rotVerts[3,1]], [rotVerts[2,2], rotVerts[3,2]], color='b', linewidth=1.0)
ax.plot([rotVerts[0,0], rotVerts[3,0]], [rotVerts[0,1], rotVerts[3,1]], [rotVerts[0,2], rotVerts[3,2]], color='b', linewidth=1.0)
# vector
ax.plot([pos[0], pos[0]+dirc[0]], [pos[1], pos[1]+dirc[1]], [pos[1], pos[1]+dirc[1]], color='r', linewidth=1.0)
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')
This gives the following output..
The green square is the original in the xz plane, the blue square the transformed square & the given vector is in red.
As you can see its well off. After many hours pouring through similar questions & replies, I am still none the wiser as to why this does not work.
So what am I missing here?
EDIT: After pouring over the Euler Angles link, given by El Dude in the comments below, I tried the following....
Defined the square in yz plane of a static frame of reference xyz with basis vectors Ux, Uy & Uz
Used a direction vector 'dirVec' as the normal for the plane I want to rotate my square into.
I decided to use the x-convention and the ZXZ rotation matrix as discribed in Euler angles link.
Steps I have taken,
1) Create a rotated frame with Tx, Ty & Tz as basis vectors;
Tx = dirVec
Ty = Tx cross Uz (Tx not allowed to parallel to Uz)
Tz = Ty cross Tx
2) Defined a Node Line, a vector along the intersection of the planes UxUy & TxTy by taking the cross product of Uz & Tz
3) Defined the Euler angles as per the definitions in the above link
4) Defined the ZXZ rotation matrix as per the above link
5) Applied rotation matrix to coordinates of square's vertices
It doesn't work, something odd is happening, no matter what the value of 'dirVec' alpha always comes out as 0.
Is there something obvious going on that I'm just missing?
Here's the amended code...
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
na = np.array
def rotation_ZXZ(alpha=0.0, beta=0.0, gamma=0.0):
"""
Return ZXZ rotaion matrix
"""
a = alpha
b = beta
g = gamma
ca = np.cos(a)
cb = np.cos(b)
cg = np.cos(g)
sa = np.sin(a)
sb = np.sin(b)
sg = np.sin(g)
return np.array([[(ca*cg-cb*sa*sg), (-ca*sg-cb*cg*sa), sa*sb],
[(cg*sa+ca*cb*sg), (ca*cb*cg-sa*sg), -ca*sb],
[sb*sg, cg*sb, cb]])
def rotated_axes(vector=[0,1,0]):
"""
Return unit basis vectors for rotated frame
"""
vx = np.asarray(vector) / np.linalg.norm(vector)
if vx[1] != 0 or vx[2] != 0:
U = na([1.0, 0.0, 0.0])
else:
U = na([0.0, 1.0, 0.0])
vz = np.cross(vx, U)
vz = vz / np.linalg.norm(vz)
vy = np.cross(vx, vz)
vy = vy / np.linalg.norm(vy)
vx = bv(vx[0], vx[1], vx[2])
vy = bv(vy[0], vy[1], vy[2])
vz = bv(vz[0], vz[1], vz[2])
return vx, vy, vz
def angle_btw_vectors(v1=[1,0,0], v2=[0,1,0]):
"""
Return the angle, in radians, between 2 vectors
"""
v1 = np.asarray(v1)
v2 = np.asarray(v2)
mags = np.linalg.norm(v1) * np.linalg.norm(v2)
return np.arccos(np.dot(v1, v2) / mags)
edgeLen = 4.0 # length of square side
dirVec = na([4,4,4]) # direction of given vector
pos = na([0.0, 0.0, 0.0]) # starting point of given vector
x = pos[0]
y = pos[1]
z = pos[2]
Ux = na([1,0,0]) # Unit basis vectors for static frame
Uy = na([0,1,0])
Uz = na([0,0,1])
Tx, Ty, Tz = rotated_axes(dirVec) # Unit basis vectors for rotated frame
# where Tx = dirVec / |dirVec|
nodeLine = np.cross(Uz, Tz) # Node line - xy intersect XY
alpha = angle_btw_vectors(Ux, nodeLine) #Euler angles
beta = angle_btw_vectors(Uz, Tz)
gamma = angle_btw_vectors(nodeLine, Tx)
Rzxz = rotation_ZXZ(alpha, beta, gamma) # Rotation matrix
print '--------------------------------------'
print 'Tx: ', Tx
print 'Ty: ', Ty
print 'Tz: ', Tz
print 'Node line: ', nodeLine
print 'Tx.dirVec: ', np.dot(Tx, (dirVec / np.linalg.norm(dirVec)))
print 'Ty.dirVec: ', np.dot(Ty, dirVec)
print 'Tz.dirVec: ', np.dot(Tz, dirVec)
print '(Node Line).Tx: ', np.dot(Tx, nodeLine)
print 'alpha: ', alpha * 180 / np.pi
print 'beta: ', beta * 180 / np.pi
print 'gamma: ', gamma * 180 / np.pi
#print 'Rzxz: ', Rxzx
# corner vertices of square in yz plane
verts = na([[0, edgeLen/2.0, edgeLen/2.0],
[0, edgeLen/2.0, -edgeLen/2.0],
[0, -edgeLen/2.0, -edgeLen/2.0],
[0, -edgeLen/2.0, edgeLen/2.0]])
rotVerts = na([0,0,0])
for v in verts:
temp = np.dot(Rzxz, v)
temp = na([temp[0]+x, temp[1]+y, temp[2]+z])
rotVerts = np.vstack((rotVerts, temp))
rotVerts = np.delete(rotVerts, rotVerts[0], axis=0)
# plot
# oringinal square
ax.scatter(verts[:,0], verts[:,1], verts[:,2], s=10, c='g', marker='o')
ax.plot([verts[0,0], verts[1,0]], [verts[0,1], verts[1,1]], [verts[0,2], verts[1,2]], color='g', linewidth=1.0)
ax.plot([verts[1,0], verts[2,0]], [verts[1,1], verts[2,1]], [verts[1,2], verts[2,2]], color='g', linewidth=1.0)
ax.plot([verts[2,0], verts[3,0]], [verts[2,1], verts[3,1]], [verts[2,2], verts[3,2]], color='g', linewidth=1.0)
ax.plot([verts[0,0], verts[3,0]], [verts[0,1], verts[3,1]], [verts[0,2], verts[3,2]], color='g', linewidth=1.0)
# rotated & translated square
ax.scatter(rotVerts[:,0], rotVerts[:,1], rotVerts[:,2], s=10, c='b', marker='o')
ax.plot([rotVerts[0,0], rotVerts[1,0]], [rotVerts[0,1], rotVerts[1,1]], [rotVerts[0,2], rotVerts[1,2]], color='b', linewidth=1.0)
ax.plot([rotVerts[1,0], rotVerts[2,0]], [rotVerts[1,1], rotVerts[2,1]], [rotVerts[1,2], rotVerts[2,2]], color='b', linewidth=1.0)
ax.plot([rotVerts[2,0], rotVerts[3,0]], [rotVerts[2,1], rotVerts[3,1]], [rotVerts[2,2], rotVerts[3,2]], color='b', linewidth=1.0)
ax.plot([rotVerts[0,0], rotVerts[3,0]], [rotVerts[0,1], rotVerts[3,1]], [rotVerts[0,2], rotVerts[3,2]], color='b', linewidth=1.0)
# Rotated reference coordinate system
ax.plot([pos[0], pos[0]+Tx[0]], [pos[1], pos[1]+Tx[1]], [pos[2], pos[2]+Tx[2]], color='r', linewidth=1.0)
ax.plot([pos[0], pos[0]+Ty[0]], [pos[1], pos[1]+Ty[1]], [pos[1], pos[2]+Ty[2]], color='b', linewidth=1.0)
ax.plot([pos[0], pos[0]+Tz[0]], [pos[1], pos[1]+Tz[1]], [pos[1], pos[2]+Tz[2]], color='g', linewidth=1.0)
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')
Here's a solution that I came up with - it should work, although there wasn't a whole lot of testing. The solution is somewhat more general, as it would work for any 2D object of any orientation, the only thing you have to adjust are the vertices stored in obj (this could be done better but here I just created a list of points by hand).
Note, that I defined mObj as the "center" of the object - this does not change the functionality but is the anchor point of the normal vector that is displayed.
Here's some explanation for the math:
What we need to do is to find the right rotation axis and angle, such that we only need one matrix multiplication (in principle you could use the Euler angles which would be an equivalent solution). The angle is easy, since it is given by the dot-product:
dot(a, b) = |a| |b| * cos(theta)
where theta is the angle between the vector a and b. To find the rotation axis, we can use the normal vector of the plane spanned by a and b, i.e. use the cross product and normalize it:
rotAxis = cross(a, b) / |cross(a, b)|
Note that this vector is orthogonal to a and b, hence the axis we are looking for.
Hope this helps.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
def rotateVector3D(v, theta, axis):
""" Takes a three-dimensional vector v and rotates it by the angle theta around the specified axis.
"""
return np.dot(rotationMatrix3D(theta, axis), v)
def rotationMatrix3D(theta, axis):
""" Return the rotation matrix associated with counterclockwise rotation about
the given axis by theta radians.
"""
axis = np.asarray(axis) / np.sqrt(np.dot(axis, axis))
a = np.cos(theta/2.0)
b, c, d = -axis*np.sin(theta/2.0)
aa, bb, cc, dd = a**2, b**2, c**2, d**2
bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
return np.array([[aa+bb-cc-dd, 2*(bc+ad), 2*(bd-ac)],
[2*(bc-ad), aa+cc-bb-dd, 2*(cd+ab)],
[2*(bd+ac), 2*(cd-ab), aa+dd-bb-cc]])
def drawObject(ax, pts, color="red"):
""" Draws an object on a specified 3D axis with points and lines between consecutive points.
"""
map(lambda pt: ax.scatter(*pt, s=10, color=color), pts)
for k in range(len(pts)-1):
x, y, z = zip(*pts[k:k+2])
ax.plot(x, y, z, color=color, linewidth=1.0)
x, y, z = zip(*[pts[-1],pts[0]])
ax.plot(x, y, z, color=color, linewidth=1.0)
def normalVector(obj):
""" Takes a set of points, assumed to be flat, and returns a normal vector with unit length.
"""
n = np.cross(np.array(obj[1])-np.array(obj[0]), np.array(obj[2])-np.array(obj[0]))
return n/np.sqrt(np.dot(n,n))
# Set the original object (can be any set of points)
obj = [(2, 0, 2), (2, 0, 4), (4, 0, 4), (4, 0, 2)]
mObj = (3, 0, 3)
nVecObj = normalVector(obj)
# Given vector.
vec = (6, 6, 6)
# Find rotation axis and angle.
rotAxis = normalVector([(0,0,0), nVecObj, vec])
angle = np.arccos(np.dot(nVecObj, vec) / (np.sqrt(np.dot(vec, vec)) * np.sqrt(np.dot(nVecObj, nVecObj))))
print "Rotation angle: {:.2f} degrees".format(angle/np.pi*180)
# Generate the rotated object.
rotObj = map(lambda pt: rotateVector3D(pt, angle, rotAxis), obj)
mRotObj = rotateVector3D(mObj, angle, rotAxis)
nVecRotObj = normalVector(rotObj)
# Set up Plot.
fig = plt.figure()
fig.set_size_inches(18,18)
ax = fig.add_subplot(111, projection='3d')
# Draw.
drawObject(ax, [[0,0,0], np.array(vec)/np.sqrt(np.dot(vec,vec))], color="gray")
drawObject(ax, [mObj, mObj+nVecObj], color="red")
drawObject(ax, obj, color="red")
drawObject(ax, [mRotObj, mRotObj + nVecRotObj], color="green")
drawObject(ax, rotObj, color="green")
# Plot cosmetics.
ax.set_xlabel('X axis')
ax.set_ylabel('Y axis')
ax.set_zlabel('Z axis')
# Check if the given vector and the normal of the rotated object are parallel (cross product should be zero).
print np.round(np.sum(np.cross(vec, nVecRotObj)**2), 5)

Categories