How to animate a BuiltinSurface in Mayavi mlab? - python

I'm attempting to animate the Earth rotating using Mayavi mlab. I've succeeded in the past by just rotating the camera around a BuiltinSurface representation of the Earth, but this becomes inconvenient when I need to plot many other objects (spacecraft, stars, etc) in the frame as well. The code below seems to "almost" work: on my Windows 10 machine, it runs for 8 iterations and then the animation freezes. How can I fix this code, or is there a better way to animate a BuiltinSurface in general?
import numpy as np
from mayavi import mlab
from mayavi.sources.builtin_surface import BuiltinSurface
from mayavi.modules.surface import Surface
from mayavi.filters.transform_data import TransformData
def rotMat3D(axis, angle, tol=1e-12):
"""Return the rotation matrix for 3D rotation by angle `angle` degrees about an
arbitrary axis `axis`.
"""
t = np.radians(angle)
x, y, z = axis
R = (np.cos(t))*np.eye(3) +\
(1-np.cos(t))*np.matrix(((x**2,x*y,x*z),(x*y,y**2,y*z),(z*x,z*y,z**2))) + \
np.sin(t)*np.matrix(((0,-z,y),(z,0,-x),(-y,x,0)))
R[np.abs(R)<tol]=0.0
return R
#mlab.show
#mlab.animate(delay=200)
def anim():
fig = mlab.figure()
engine = mlab.get_engine()
# Add a cylinder builtin source
cylinder_src = BuiltinSurface()
engine.add_source(cylinder_src)
cylinder_src.source = 'earth'
# Add transformation filter to rotate cylinder about an axis
transform_data_filter = TransformData()
engine.add_filter(transform_data_filter, cylinder_src)
Rt = np.eye(4)
Rt[0:3,0:3] = rotMat3D((0,0,1), 0) # in homogeneous coordinates
Rtl = list(Rt.flatten()) # transform the rotation matrix into a list
transform_data_filter.transform.matrix.__setstate__({'elements': Rtl})
transform_data_filter.widget.set_transform(transform_data_filter.transform)
transform_data_filter.filter.update()
transform_data_filter.widget.enabled = False # disable the rotation control further.
# Add surface module to the cylinder source
cyl_surface = Surface()
engine.add_filter(cyl_surface, transform_data_filter)
#add color property
#cyl_surface.actor.property.color = (1.0, 0.0, 0.0)
ind=1
while ind<90:
print ind
Rt[0:3,0:3] = rotMat3D((0,0,1), ind) # in homogeneous coordinates
Rtl = list(Rt.flatten()) # transform the rotation matrix into a list
transform_data_filter.transform.matrix.__setstate__({'elements': Rtl})
transform_data_filter.widget.set_transform(transform_data_filter.transform)
transform_data_filter.filter.update()
transform_data_filter.widget.enabled = False # disable the rotation control further.
# Add surface module to the cylinder source
cyl_surface = Surface()
engine.add_filter(cyl_surface, transform_data_filter)
# add color property
#cyl_surface.actor.property.color = (1.0, 0.0, 0.0)
yield
ind+=1
anim()

I haven't been able to figure out a way to use Mayavi to make this happen. However, Vpython appears to be much better suited to accomplish this task. I've posted an example section of code below to make a revolving Earth, along with a few other features.
from visual import *
def destroy():
for obj in scene.objects:
obj.visible = False
del obj
R = 6378. # radius of sphere
angle=0.
scene.range = 10000.
SunDirection=vector(.77,.77,0)
# scene.fov = 0.5
scene.center = (0,0,0)
scene.forward = (-1,0,-1)
scene.up = (0,0,1)
scene.lights=[distant_light(direction=SunDirection, color=color.gray(0.8)),
distant_light(direction=-SunDirection, color=color.gray(0.3))]
x=0
y=0
while True:
rate(10)
angle=angle+1.*pi/180.
destroy()
s = sphere(pos=(x,y,0), radius=R, material=materials.BlueMarble)
s.rotate(angle=90.*pi/180.,axis=(1,0,0)) # Always include this to rotate Earth into correct ECI x y z frame
s.rotate(angle=90.*pi/180.,axis=(0,0,1)) # Always include this to rotate Earth into correct ECI x y z frame
s.rotate(angle=angle, axis=(0,0,1)) # This rotation causes Earth to spin on its axis
xaxis = arrow(pos=(0,0,0), axis=vector(1,0,0)*7000, shaftwidth=100, color=color.red)
yaxis = arrow(pos=(0,0,0), axis=vector(0,1,0)*7000, shaftwidth=100, color=color.green)
zaxis = arrow(pos=(0,0,0), axis=vector(0,0,1)*7000, shaftwidth=100, color=color.blue)
ST = cone(pos=(0,8000,0),axis=(0,700,0),radius=700*tan(10*pi/180),color=color.blue,opacity=1)

Related

Manim objects separated in random motion

I am trying to animate a particle and a vector that is attached to its center as a random motion is applied to the particle. The particle behaves as intended, but the vector always has a offset from the particle. I tried to set the random seed of each scene, but it didn't work.
from manim import *
import numpy as np
class Particles(ThreeDScene):
def construct(self):
self.camera.background_color = WHITE
coordinate = np.array([ (0,0,0) ])
angle = 90
radius = 1.25
#Polar coordinates
def polar2cart(theta_degrees, rho):
theta = theta_degrees*(np.pi/180)
x = rho * np.cos(theta)
y = rho * np.sin(theta)
return(x, y, 0)
start_arrow = np.array(coordinate) - np.array( [polar2cart(angle,radius )] )
end_arrow = np.array(coordinate) + np.array( [polar2cart(angle,radius )] )
arrow = Arrow( start= start_arrow, end = end_arrow, color = '#ff0000')
dot = Dot(point=coordinate, radius=0.2, color = '#000000')
#Display
self.add(arrow, dot)
##############################################################################################################################
#Random motion
valor = ValueTracker(0)
dot.add_updater(lambda m: m.move_to( m.get_center() + (np.random.normal(0,0.05),np.random.normal(0,0.05),np.random.normal(0,0.05)) ) )
arrow.add_updater(lambda m: m.move_to( dot.get_center() + (np.random.normal(0,0.05),np.random.normal(0,0.05),np.random.normal(0,0.05)) ) )
self.play(valor.animate.set_value(10),rate_func=smooth, run_time=5)
self.wait()
Particle and vector
How can I make that the vector and the particle remain fixed in all frames?
Here is some things that you need to know:
The order in which updaters attached to mobjects are executed is the order in which their mobjects have been added to the scene. In your code, the updater attached to arrow runs before the updater attached to dot (which will make the arrow always lag behind).
Even if you fix the seed, subsequent calls to np.random.norm will not yield the same numbers (fortunately, otherwise your dot would move in a straight line out of the scene). If you want to fix the arrow to the center of the dot, then there is no need to add the random offset to that as well.
Here are possible solutions:
Add the objects in the other way around, self.add(dot, arrow), and remove the random offset from the updater attached to the arrow, arrow.add_updater(lambda m: m.move_to(dot.get_center())). If the order of mobjects on the screen is important to you, you can also run dot.set_z_index(1).
Alternatively, you could also simply create a group of your two objects and attach an updater to that; particle_group = VGroup(arrow, dot) followed by particle_group.add_updater(lambda m: m.shift(...)).
I'd personally go for the group, but it really depends on what else you intend to do in this scene.

Is it possible to fill in a circular graph with a solid colour and save it as svg in matplotlib?

I wrote some code that creates randomised patches from graphs in matplotlib. Basically how it works is that you create a graph from nodes taken from a circle using the parametric equation for a circle and then you randomly displace the nodes along the vector of (0,0) to the node point on the circumference of the circle. That way you can be certain to avoid lines from crossing each other once the circle is drawn. In the end you just append the first (x,y) coordinate to the list of coordinates to close the circle.
What I want to do next is to find a way to fill that circular graph with a solid colour so that I can create a "stamp" that can be used to make randomised patches on a canvas that hopefully will not create crossing edges. I want to use this to make procedural risk maps in svg format, because a lot of those are uploaded with terrible edges using raster image formats using jpeg.
I am pretty sure that my information of the nodes should be sufficient to make that happen but I have no idea how to implement that. Can anyone help?
import numpy as np
import matplotlib.pyplot as plt
def node_circle(r=0.5,res=100):
# Create arrays (x and y coordinates) for the nodes on the circumference of a circle. Use parametric equation.
# x = r cos(t) y = r sin(t)
t = np.linspace(0,2*np.pi,res)
x = r*np.cos(t)
y = r*np.sin(t)
return t,x,y
def sgn(x,x_shift=-0.5,y_shift=1):
# A shifted sign function to use as a switching function
# in order to avoid shifts lower than -0.5 which is
# the radius of the circle.
return -0.5*(np.abs(x -x_shift)/(x -x_shift)) +y_shift
def displacer(x,y,low=-0.5,high=0.5,maxrad=0.5):
# Displaces the node points of the circle
shift = 0
shift_increment = 0
for i in range(len(x)):
shift_increment = np.random.uniform(low,high)
shift += shift_increment*sgn(maxrad)
x[i] += x[i]*shift
y[i] += y[i]*shift
x = np.append(x,x[0])
y = np.append(y,y[0])
return x,y
def plot():
# Actually visualises everything
fig, ax = plt.subplots(figsize=(4,4))
# np.random.seed(1)
ax.axis('off')
t,x,y = node_circle(res=100)
a = 0
x,y = displacer(x,y,low=-0.15,high=0.15)
ax.plot(x,y,'r-')
# ax.scatter(x,y,)
plt.show()
plot()
got it: the answer is to use matplotlib.Patches.Polygon
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon
def node_circle(r=0.5,res=100):
# Create arrays (x and y coordinates) for the nodes on the circumference of a circle. Use parametric equation.
# x = r cos(t) y = r sin(t)
t = np.linspace(0,2*np.pi,res)
x = r*np.cos(t)
y = r*np.sin(t)
return x,y
def sgn(x,x_shift=-0.5,y_shift=1):
# A shifted sign function to use as a switching function
# in order to avoid shifts lower than -0.5 which is
# the radius of the circle.
return -0.5*(np.abs(x -x_shift)/(x -x_shift)) +y_shift
def displacer(x,y,low=-0.5,high=0.5,maxrad=0.5):
# Displaces the node points of the circle
shift = 0
shift_increment = 0
for i in range(len(x)):
shift_increment = np.random.uniform(low,high)
shift += shift_increment*sgn(maxrad)
x[i] += x[i]*shift
y[i] += y[i]*shift
x = np.append(x,x[0])
y = np.append(y,y[0])
return x,y
def patch_distributor(M,N,res,grid='square'):
# Distribute Patches based on a specified pattern/grid.
if grid == 'square':
data = np.zeros(shape=(M,N,2,res+1))
for i in range(M):
for j in range(N):
x,y = displacer(*node_circle(res=res),low=-0.2,high=0.2)
data[i,j,0,:] = x
data[i,j,1,:] = y
return data
def plot(res):
# Actually visualises everything
fig, ax = plt.subplots(figsize=(4,4))
# np.random.seed(1)
ax.axis('off')
# x,y = node_circle(res=res)
# x,y = displacer(x,y,low=-0.15,high=0.15)
# xy = np.zeros((len(x),2))
# xy[:,0] = x
# xy[:,1] = y
patch_data = patch_distributor(10,10,res)
for i in range(patch_data.shape[0]):
for j in range(patch_data.shape[1]):
x,y = patch_data[i,j]
x += i*0.5
y += j*0.5
xy = np.zeros((len(x),2))
xy[:,0] = x
xy[:,1] = y
patch = Polygon(xy,fc='w',ec='k',lw=2,zorder=np.random.randint(2),antialiased=False)
ax.add_patch(patch)
ax.autoscale_view()
# ax.plot(x,y,'r-')
# ax.scatter(x,y,)
plt.savefig('lol.png')
plot(res=40)
# Displace circle along the line of (0,0) -> (cos(t),sin(t))
# Make the previous step influence the next to avoid jaggedness
# limit displacement level to an acceptable amount
# Random displaced cubic grid as placing points for stamps.

Test whether points are inside ellipses, without using Matplotlib?

I'm working on a Python-based data analysis. I have some x-y data points, and some ellipses, and I want to determine whether points are inside any of the ellipses. The way that I've been doing this works, but it's kludgy. As I think about distributing my software to other people, I find myself wanting a cleaner way.
Right now, I'm using matplotlib.patches.Ellipse objects. Matplotlib Ellipses have a useful method called contains_point(). You can work in data coordinates on a Matplotlib Axes object by calling Axes.transData.transform().
The catch is that I have to create a Figure and an Axes object to hold the Ellipses. And when my program runs, an annoying Matplotlib Figure object will get rendered, showing the Ellipses, which I don't actually need to see. I have tried several methods to suppress this output. I have succeeded in deleting the Ellipses from the Axes, using Axes.clear(), resulting in an empty graph. But I can't get Matplolib's pyplot.close(fig_number) to delete the Figure itself before calling pyplot.show().
Any advice is appreciated, thanks!
Inspired by how a carpenter draws an ellipse using two nails and a piece of string, here is a numpy-friendly implementation to test whether points lie inside given ellipses.
One of the definitions of an ellipse, is that the sum of the distances to the two foci is constant, equal to the width (or height if it would be larger) of the ellipse. The distance between the center and the foci is sqrt(a*a - b*b), where a and b are half of the width and height. Using that distance and rotation by the desired angle finds the locations of the foci. numpy.linalg.norm can be used to calculate the distances using numpy's efficient array operations.
After the calculations, a plot is generated to visually check whether everything went correct.
import numpy as np
from numpy.linalg import norm # calculate the length of a vector
x = np.random.uniform(0, 40, 20000)
y = np.random.uniform(0, 20, 20000)
xy = np.dstack((x, y))
el_cent = np.array([20, 10])
el_width = 28
el_height = 17
el_angle = 20
# distance between the center and the foci
foc_dist = np.sqrt(np.abs(el_height * el_height - el_width * el_width) / 4)
# vector from center to one of the foci
foc_vect = np.array([foc_dist * np.cos(el_angle * np.pi / 180), foc_dist * np.sin(el_angle * np.pi / 180)])
# the two foci
el_foc1 = el_cent + foc_vect
el_foc2 = el_cent - foc_vect
# for each x,y: calculate z as the sum of the distances to the foci;
# np.ravel is needed to change the array of arrays (of 1 element) into a single array
z = np.ravel(norm(xy - el_foc1, axis=-1) + norm(xy - el_foc2, axis=-1) )
# points are exactly on the ellipse when the sum of distances is equal to the width
# z = np.where(z <= max(el_width, el_height), 1, 0)
# now create a plot to check whether everything makes sense
from matplotlib import pyplot as plt
from matplotlib import patches as mpatches
fig, ax = plt.subplots()
# show the foci as red dots
plt.plot(*el_foc1, 'ro')
plt.plot(*el_foc2, 'ro')
# create a filter to separate the points inside the ellipse
filter = z <= max(el_width, el_height)
# draw all the points inside the ellipse with the plasma colormap
ax.scatter(x[filter], y[filter], s=5, c=z[filter], cmap='plasma')
# draw all the points outside with the cool colormap
ax.scatter(x[~filter], y[~filter], s=5, c=z[~filter], cmap='cool')
# add the original ellipse to verify that the boundaries match
ellipse = mpatches.Ellipse(xy=el_cent, width=el_width, height=el_height, angle=el_angle,
facecolor='None', edgecolor='black', linewidth=2,
transform=ax.transData)
ax.add_patch(ellipse)
ax.set_aspect('equal', 'box')
ax.autoscale(enable=True, axis='both', tight=True)
plt.show()
The simplest solution here is to use shapely.
If you have an array of shape Nx2 containing a set of vertices (xy) then it is trivial to construct the appropriate shapely.geometry.polygon object and check if an arbitrary point or set of points (points) is contained within -
import shapely.geometry as geom
ellipse = geom.Polygon(xy)
for p in points:
if ellipse.contains(geom.Point(p)):
# ...
Alternatively, if the ellipses are defined by their parameters (i.e. rotation angle, semimajor and semiminor axis) then the array containing the vertices must be constructed and then the same process applied. I would recommend using the polar form relative to center as this is the most compatible with how shapely constructs the polygons.
import shapely.geometry as geom
from shapely import affinity
n = 360
a = 2
b = 1
angle = 45
theta = np.linspace(0, np.pi*2, n)
r = a * b / np.sqrt((b * np.cos(theta))**2 + (a * np.sin(theta))**2)
xy = np.stack([r * np.cos(theta), r * np.sin(theta)], 1)
ellipse = affinity.rotate(geom.Polygon(xy), angle, 'center')
for p in points:
if ellipse.contains(geom.Point(p)):
# ...
This method is advantageous because it supports any properly defined polygons - not just ellipses, it doesn't rely on matplotlib methods to perform the containment checking, and it produces a very readable code (which is often important when "distributing [one's] software to other people").
Here is a complete example (with added plotting to show it working)
import shapely.geometry as geom
from shapely import affinity
import matplotlib.pyplot as plt
import numpy as np
n = 360
theta = np.linspace(0, np.pi*2, n)
a = 2
b = 1
angle = 45.0
r = a * b / np.sqrt((b * np.cos(theta))**2 + (a * np.sin(theta))**2)
xy = np.stack([r * np.cos(theta), r * np.sin(theta)], 1)
ellipse = affinity.rotate(geom.Polygon(xy), angle, 'center')
x, y = ellipse.exterior.xy
# Create a Nx2 array of points at grid coordinates throughout
# the ellipse extent
rnd = np.array([[i,j] for i in np.linspace(min(x),max(x),50)
for j in np.linspace(min(y),max(y),50)])
# Filter for points which are contained in the ellipse
res = np.array([p for p in rnd if ellipse.contains(geom.Point(p))])
plt.plot(x, y, lw = 1, color='k')
plt.scatter(rnd[:,0], rnd[:,1], s = 50, color=(0.68, 0.78, 0.91)
plt.scatter(res[:,0], res[:,1], s = 15, color=(0.12, 0.67, 0.71))
plt.show()

Matplotlib: Custom functions, that are called each time a figure is drawn

I want to create a matplotlib plot containing arrows, whose head's shape is independent from the data coordinates. This is similar to FancyArrowPatch, but when the arrow length is smaller than the head length is shrank to fit the length of the arrow.
Currently, I solve this by setting the length of the arrow head by transforming the width to display coordinates, calculating the head length in display coordinates and transform it back into data coordinates.
This approach works well as long the axes' dimensions do not change, which can happen due to set_xlim(), set_ylim() or tight_layout() for example.
I want to cover these cases, by redrawing the arrow whenever the plot's dimensions do change. At the moment I handle this by registering a function on_draw(event) via
axes.get_figure().canvas.mpl_connect("resize_event", on_draw)
but this does only work for interactive backends. I also need a solution for cases, where I save the plot as image file. Is there any other place, where I can register my callback function?
EDIT: Here is the code, I am currently using:
def draw_adaptive_arrow(axes, x, y, dx, dy,
tail_width, head_width, head_ratio, draw_head=True,
shape="full", **kwargs):
from matplotlib.patches import FancyArrow
from matplotlib.transforms import Bbox
arrow = None
def on_draw(event=None):
"""
Callback function that is called, every time the figure is resized
Removes the current arrow and replaces it with an arrow with
recalcualted head
"""
nonlocal tail_width
nonlocal head_width
nonlocal arrow
if arrow is not None:
arrow.remove()
# Create a head that looks equal, independent of the aspect
# ratio
# Hence, a transformation into display coordinates has to be
# performed to fix the head width to length ratio
# In this transformation only the height and width are
# interesting, absolute coordinates are not needed
# -> box origin at (0,0)
arrow_box = Bbox([(0,0),(0,head_width)])
arrow_box_display = axes.transData.transform_bbox(arrow_box)
head_length_display = np.abs(arrow_box_display.height * head_ratio)
arrow_box_display.x1 = arrow_box_display.x0 + head_length_display
# Transfrom back to data coordinates for plotting
arrow_box = axes.transData.inverted().transform_bbox(arrow_box_display)
head_length = arrow_box.width
if head_length > np.abs(dx):
# If the head would be longer than the entire arrow,
# only draw the arrow head with reduced length
head_length = np.abs(dx)
if not draw_head:
head_length = 0
head_width = tail_width
arrow = FancyArrow(
x, y, dx, dy,
width=tail_width, head_width=head_width, head_length=head_length,
length_includes_head=True, **kwargs)
axes.add_patch(arrow)
axes.get_figure().canvas.mpl_connect("resize_event", on_draw)
# Some place in the user code...
fig = plt.figure(figsize=(8.0, 3.0))
ax = fig.add_subplot(1,1,1)
# 90 degree tip
draw_adaptive_arrow(
ax, 0, 0, 4, 0, tail_width=0.4, head_width=0.8, head_ratio=0.5
)
# Still 90 degree tip
draw_adaptive_arrow(
ax, 5, 0, 2, 0, tail_width=0.4, head_width=0.8, head_ratio=0.5
)
# Smaller head, since otherwise head would be longer than entire arrow
draw_adaptive_arrow(
ax, 8, 0, 0.5, 0, tail_width=0.4, head_width=0.8, head_ratio=0.5
)
ax.set_xlim(0,10)
ax.set_ylim(-1,1)
# Does not work in non-interactive backend
plt.savefig("test.pdf")
# But works in interactive backend
plt.show()
Here is a solution without callback. I took over mostly the algorithm from the question, because I'm not sure I understand the requirements for the arrow. I'm pretty sure that can be simplified, but that's also beyond the point of the question.
So here we subclass FancyArrow and let it add itself to the axes. We then override the draw method to calculate the needed parameters and then - which is somehow unusual and may in other cases fail - call __init__ again inside the draw method.
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import FancyArrow
from matplotlib.transforms import Bbox
class MyArrow(FancyArrow):
def __init__(self, *args, **kwargs):
self.ax = args[0]
self.args = args[1:]
self.kw = kwargs
self.head_ratio = self.kw.pop("head_ratio", 1)
self.draw_head = self.kw.pop("draw_head", True)
self.kw.update(length_includes_head=True)
super().__init__(*self.args,**self.kw)
self.ax.add_patch(self)
self.trans = self.get_transform()
def draw(self, renderer):
self.kw.update(transform = self.trans)
arrow_box = Bbox([(0,0),(0,self.kw["head_width"])])
arrow_box_display = self.ax.transData.transform_bbox(arrow_box)
head_length_display = np.abs(arrow_box_display.height * self.head_ratio)
arrow_box_display.x1 = arrow_box_display.x0 + head_length_display
# Transfrom back to data coordinates for plotting
arrow_box = self.ax.transData.inverted().transform_bbox(arrow_box_display)
self.kw["head_length"] = arrow_box.width
if self.kw["head_length"] > np.abs(self.args[2]):
# If the head would be longer than the entire arrow,
# only draw the arrow head with reduced length
self.kw["head_length"] = np.abs(self.args[2])
if not self.draw_head:
self.kw["head_length"] = 0
self.kw["head_width"] = self.kw["width"]
super().__init__(*self.args,**self.kw)
self.set_clip_path(self.ax.patch)
self.ax._update_patch_limits(self)
super().draw(renderer)
fig = plt.figure(figsize=(8.0, 3.0))
ax = fig.add_subplot(1,1,1)
# 90 degree tip
MyArrow( ax, 0, 0, 4, 0, width=0.4, head_width=0.8, head_ratio=0.5 )
MyArrow( ax, 5, 0, 2, 0, width=0.4, head_width=0.8, head_ratio=0.5 )
# Smaller head, since otherwise head would be longer than entire arrow
MyArrow( ax, 8, 0, 0.5, 0, width=0.4, head_width=0.8, head_ratio=0.5 )
ax.set_xlim(0,10)
ax.set_ylim(-1,1)
# Does not work in non-interactive backend
plt.savefig("test.pdf")
# But works in interactive backend
plt.show()
I found a solution to the problem, however, it is not very elegant.
The only callback function, I found, that is called in non-interactive backends, is the draw_path() method of AbstractPathEffect subclasses.
I created an AbstractPathEffect subclass that updates the vertices of the arrow head
in its draw_path() method.
I am still open for other probably more straight forward solutions to my problem.
import numpy as np
from numpy.linalg import norm
from matplotlib.patches import FancyArrow
from matplotlib.patheffects import AbstractPathEffect
class AdaptiveFancyArrow(FancyArrow):
"""
A `FancyArrow` with fixed head shape.
The length of the head is proportional to the width the head
in display coordinates.
If the head length is longer than the length of the entire
arrow, the head length is limited to the arrow length.
"""
def __init__(self, x, y, dx, dy,
tail_width, head_width, head_ratio, draw_head=True,
shape="full", **kwargs):
if not draw_head:
head_width = tail_width
super().__init__(
x, y, dx, dy,
width=tail_width, head_width=head_width,
overhang=0, shape=shape,
length_includes_head=True, **kwargs
)
self.set_path_effects(
[_ArrowHeadCorrect(self, head_ratio, draw_head)]
)
class _ArrowHeadCorrect(AbstractPathEffect):
"""
Updates the arrow head length every time the arrow is rendered
"""
def __init__(self, arrow, head_ratio, draw_head):
self._arrow = arrow
self._head_ratio = head_ratio
self._draw_head = draw_head
def draw_path(self, renderer, gc, tpath, affine, rgbFace=None):
# Indices to certain vertices in the arrow
TIP = 0
HEAD_OUTER_1 = 1
HEAD_INNER_1 = 2
TAIL_1 = 3
TAIL_2 = 4
HEAD_INNER_2 = 5
HEAD_OUTER_2 = 6
transform = self._arrow.axes.transData
vert = tpath.vertices
# Transform data coordiantes to display coordinates
vert = transform.transform(vert)
# The direction vector alnog the arrow
arrow_vec = vert[TIP] - (vert[TAIL_1] + vert[TAIL_2]) / 2
tail_width = norm(vert[TAIL_2] - vert[TAIL_1])
# Calculate head length from head width
head_width = norm(vert[HEAD_OUTER_2] - vert[HEAD_OUTER_1])
head_length = head_width * self._head_ratio
if head_length > norm(arrow_vec):
# If the head would be longer than the entire arrow,
# only draw the arrow head with reduced length
head_length = norm(arrow_vec)
# The new head start vector; is on the arrow vector
if self._draw_head:
head_start = \
vert[TIP] - head_length * arrow_vec/norm(arrow_vec)
else:
head_start = vert[TIP]
# vector that is orthogonal to the arrow vector
arrow_vec_ortho = vert[TAIL_2] - vert[TAIL_1]
# Make unit vector
arrow_vec_ortho = arrow_vec_ortho / norm(arrow_vec_ortho)
# Adjust vertices of the arrow head
vert[HEAD_OUTER_1] = head_start - arrow_vec_ortho * head_width/2
vert[HEAD_OUTER_2] = head_start + arrow_vec_ortho * head_width/2
vert[HEAD_INNER_1] = head_start - arrow_vec_ortho * tail_width/2
vert[HEAD_INNER_2] = head_start + arrow_vec_ortho * tail_width/2
# Transform back to data coordinates
# and modify path with manipulated vertices
tpath.vertices = transform.inverted().transform(vert)
renderer.draw_path(gc, tpath, affine, rgbFace)

Matplotlib custom projection: How to transform points

I am working with a custom projection of Matplotlib and don't understand how to do vector transformations within the projection (Note: The custom projection is a Lambert azimuthal equal-area projection with equatorial aspect).
In my example I want to transform a point that is dipping 30° to the North (meaning that the point is 60°N of the equator) into a point that dips 30° East (meaning that is lies 60° east of the prime meridian). I want to do this with the help of a vector transformation matrix, in order to do more complicated calculations with the program in the future. But I don't really understand how to get the length of the transformed vector right (or getting the correct longitude and latitude of that point).
I am also studying this example, but it uses a slightly different approach for the transformations:
https://github.com/joferkington/mplstereonet/blob/master/mplstereonet/stereonet_math.py
Testfile:
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from numpy import pi, sin, cos, sqrt, tan, arctan2, arccos
#Internal imports
import projection
def transformVector(geom, raxis, rot):
"""
Input:
geom: single point geometry (vector)
raxis: rotation axis as a vector (vector)
([0][1][2]) = (x,y,z) = (Longitude, Latitude, Down)
rot: rotation in radian
Returns:
Array: a vector that has been transformed
"""
sr = sin(rot)
cr = cos(rot)
omcr = 1.0 - cr
tf = np.array([
[cr + raxis[0]**2 * omcr,
-raxis[2] * sr + raxis[0] * raxis[1] * omcr,
raxis[1] * sr + raxis[0] * raxis[2] * omcr],
[raxis[2] * sr + raxis[1] * raxis[0] * omcr,
cr + raxis[1]**2 * omcr,
-raxis[0] * sr + raxis[1] * raxis[2] * omcr],
[-raxis[1] * sr + raxis[2] * raxis[0] * omcr,
raxis[0] * sr + raxis[2] * raxis[1] * omcr,
cr + raxis[2]**2 * omcr]])
ar = np.dot(geom, tf)
return ar
def sphericalToVector(inp_ar):
"""
Convert a spherical measurement into a vector in cartesian space
[0] = x (+) east (-) west
[1] = y (+) north (-) south
[2] = z (+) down
"""
ar = np.array([0.0, 0.0, 0.0])
ar[0] = sin(inp_ar[0]) * cos(inp_ar[1])
ar[1] = cos(inp_ar[0]) * cos(inp_ar[1])
ar[2] = sin(inp_ar[1])
return ar
def vectorToGeogr(vect):
"""
Returns:
Array with the components [0] longitude, [1] latitude
"""
ar = np.array([0.0, 0.0])
ar[0] = np.arctan2(vect[0], vect[2])
ar[1] = np.arctan2(vect[1], vect[2])
ar = ar * pi/2
return ar
def plotPoint(dip):
"""
Testfunction for converting, transforming and plotting a point
"""
plt.subplot(111, projection="lmbrt_equ_area_equ_aspect")
#Convert to radians
dip_rad = np.radians(dip)
#Set rotation to azimuth and convert dip to latitude on north-south axis
rot = dip_rad[0]
dip_lat = pi/2 - dip_rad[1]
plt.plot(0, dip_lat, "ro")
print(dip_lat, rot)
#Convert the dip into a vector along the north-south axis
#x = 0, y = dip
vect = sphericalToVector([0, dip_lat])
print(vect, np.linalg.norm(vect))
#Transfrom the dip to its proper azimuth
tvect = transformVector(vect, [0,0,1], rot)
print(tvect, np.linalg.norm(tvect))
#Transform the vector back to geographic coordinates
geo = vectorToGeogr(tvect)
print(geo)
plt.plot(geo[0], geo[1], "bo")
plt.grid(True)
plt.show()
datapoint = np.array([090.0,30])
plotPoint(datapoint)
Custom projection:
import matplotlib
from matplotlib.axes import Axes
from matplotlib.patches import Circle
from matplotlib.path import Path
from matplotlib.ticker import NullLocator, Formatter, FixedLocator
from matplotlib.transforms import Affine2D, BboxTransformTo, Transform
from matplotlib.projections import register_projection
import matplotlib.spines as mspines
import matplotlib.axis as maxis
import matplotlib.pyplot as plt
import numpy as np
from numpy import pi, sin, cos, sqrt, arctan2
# This example projection class is rather long, but it is designed to
# illustrate many features, not all of which will be used every time.
# It is also common to factor out a lot of these methods into common
# code used by a number of projections with similar characteristics
# (see geo.py).
class LambertAxes(Axes):
"""
A custom class for the Lambert azimuthal equal-area projection
with equatorial aspect. In geosciences this is also referre to
as a "Schmidt plot". For more information see:
http://pubs.er.usgs.gov/publication/pp1395
"""
# The projection must specify a name. This will be used be the
# user to select the projection, i.e. ``subplot(111,
# projection='lmbrt_equ_area_equ_aspect')``.
name = 'lmbrt_equ_area_equ_aspect'
def __init__(self, *args, **kwargs):
Axes.__init__(self, *args, **kwargs)
self.set_aspect(1, adjustable='box', anchor='C')
self.cla()
def _init_axis(self):
self.xaxis = maxis.XAxis(self)
self.yaxis = maxis.YAxis(self)
# Do not register xaxis or yaxis with spines -- as done in
# Axes._init_axis() -- until LambertAxes.xaxis.cla() works.
# self.spines['hammer'].register_axis(self.yaxis)
self._update_transScale()
def cla(self):
"""
Override to set up some reasonable defaults.
"""
# Don't forget to call the base class
Axes.cla(self)
# Set up a default grid spacing
self.set_longitude_grid(10)
self.set_latitude_grid(10)
self.set_longitude_grid_ends(80)
# Turn off minor ticking altogether
self.xaxis.set_minor_locator(NullLocator())
self.yaxis.set_minor_locator(NullLocator())
# Do not display ticks -- we only want gridlines and text
self.xaxis.set_ticks_position('none')
self.yaxis.set_ticks_position('none')
# The limits on this projection are fixed -- they are not to
# be changed by the user. This makes the math in the
# transformation itself easier, and since this is a toy
# example, the easier, the better.
Axes.set_xlim(self, -pi/2, pi/2)
Axes.set_ylim(self, -pi, pi)
def _set_lim_and_transforms(self):
"""
This is called once when the plot is created to set up all the
transforms for the data, text and grids.
"""
# There are three important coordinate spaces going on here:
#
# 1. Data space: The space of the data itself
#
# 2. Axes space: The unit rectangle (0, 0) to (1, 1)
# covering the entire plot area.
#
# 3. Display space: The coordinates of the resulting image,
# often in pixels or dpi/inch.
# This function makes heavy use of the Transform classes in
# ``lib/matplotlib/transforms.py.`` For more information, see
# the inline documentation there.
# The goal of the first two transformations is to get from the
# data space (in this case longitude and latitude) to axes
# space. It is separated into a non-affine and affine part so
# that the non-affine part does not have to be recomputed when
# a simple affine change to the figure has been made (such as
# resizing the window or changing the dpi).
# 1) The core transformation from data space into
# rectilinear space defined in the LambertEqualAreaTransform class.
self.transProjection = self.LambertEqualAreaTransform()
# 2) The above has an output range that is not in the unit
# rectangle, so scale and translate it so it fits correctly
# within the axes. The peculiar calculations of xscale and
# yscale are specific to a Aitoff-Hammer projection, so don't
# worry about them too much.
xscale = sqrt(2.0) * sin(0.5 * pi)
yscale = sqrt(2.0) * sin(0.5 * pi)
self.transAffine = Affine2D() \
.scale(0.5 / xscale, 0.5 / yscale) \
.translate(0.5, 0.5)
# 3) This is the transformation from axes space to display
# space.
self.transAxes = BboxTransformTo(self.bbox)
# Now put these 3 transforms together -- from data all the way
# to display coordinates. Using the '+' operator, these
# transforms will be applied "in order". The transforms are
# automatically simplified, if possible, by the underlying
# transformation framework.
self.transData = \
self.transProjection + \
self.transAffine + \
self.transAxes
# The main data transformation is set up. Now deal with
# gridlines and tick labels.
# Longitude gridlines and ticklabels. The input to these
# transforms are in display space in x and axes space in y.
# Therefore, the input values will be in range (-xmin, 0),
# (xmax, 1). The goal of these transforms is to go from that
# space to display space. The tick labels will be offset 4
# pixels from the equator.
self._xaxis_pretransform = \
Affine2D() \
.scale(1.0, pi) \
.translate(0.0, -pi)
self._xaxis_transform = \
self._xaxis_pretransform + \
self.transData
self._xaxis_text1_transform = \
Affine2D().scale(1.0, 0.0) + \
self.transData + \
Affine2D().translate(0.0, 4.0)
self._xaxis_text2_transform = \
Affine2D().scale(1.0, 0.0) + \
self.transData + \
Affine2D().translate(0.0, -4.0)
# Now set up the transforms for the latitude ticks. The input to
# these transforms are in axes space in x and display space in
# y. Therefore, the input values will be in range (0, -ymin),
# (1, ymax). The goal of these transforms is to go from that
# space to display space. The tick labels will be offset 4
# pixels from the edge of the axes ellipse.
yaxis_stretch = Affine2D().scale(pi * 2.0, 1.0).translate(-pi, 0.0)
yaxis_space = Affine2D().scale(1.0, 1.0)
self._yaxis_transform = \
yaxis_stretch + \
self.transData
yaxis_text_base = \
yaxis_stretch + \
self.transProjection + \
(yaxis_space + \
self.transAffine + \
self.transAxes)
self._yaxis_text1_transform = \
yaxis_text_base + \
Affine2D().translate(-8.0, 0.0)
self._yaxis_text2_transform = \
yaxis_text_base + \
Affine2D().translate(8.0, 0.0)
def get_xaxis_transform(self,which='grid'):
"""
Override this method to provide a transformation for the
x-axis grid and ticks.
"""
assert which in ['tick1','tick2','grid']
return self._xaxis_transform
def get_xaxis_text1_transform(self, pixelPad):
"""
Override this method to provide a transformation for the
x-axis tick labels.
Returns a tuple of the form (transform, valign, halign)
"""
return self._xaxis_text1_transform, 'bottom', 'center'
def get_xaxis_text2_transform(self, pixelPad):
"""
Override this method to provide a transformation for the
secondary x-axis tick labels.
Returns a tuple of the form (transform, valign, halign)
"""
return self._xaxis_text2_transform, 'top', 'center'
def get_yaxis_transform(self,which='grid'):
"""
Override this method to provide a transformation for the
y-axis grid and ticks.
"""
assert which in ['tick1','tick2','grid']
return self._yaxis_transform
def get_yaxis_text1_transform(self, pixelPad):
"""
Override this method to provide a transformation for the
y-axis tick labels.
Returns a tuple of the form (transform, valign, halign)
"""
return self._yaxis_text1_transform, 'center', 'right'
def get_yaxis_text2_transform(self, pixelPad):
"""
Override this method to provide a transformation for the
secondary y-axis tick labels.
Returns a tuple of the form (transform, valign, halign)
"""
return self._yaxis_text2_transform, 'center', 'left'
def _gen_axes_patch(self):
"""
Override this method to define the shape that is used for the
background of the plot. It should be a subclass of Patch.
In this case, it is a Circle (that may be warped by the axes
transform into an ellipse). Any data and gridlines will be
clipped to this shape.
"""
return Circle((0.5, 0.5), 0.5)
def _gen_axes_spines(self):
return {'lmbrt_equ_area_equ_aspect':mspines.Spine.circular_spine(self,
(0.5, 0.5), 0.5)}
# Prevent the user from applying scales to one or both of the
# axes. In this particular case, scaling the axes wouldn't make
# sense, so we don't allow it.
def set_xscale(self, *args, **kwargs):
if args[0] != 'linear':
raise NotImplementedError
Axes.set_xscale(self, *args, **kwargs)
def set_yscale(self, *args, **kwargs):
if args[0] != 'linear':
raise NotImplementedError
Axes.set_yscale(self, *args, **kwargs)
# Prevent the user from changing the axes limits. In our case, we
# want to display the whole sphere all the time, so we override
# set_xlim and set_ylim to ignore any input. This also applies to
# interactive panning and zooming in the GUI interfaces.
def set_xlim(self, *args, **kwargs):
Axes.set_xlim(self, -pi, pi)
Axes.set_ylim(self, -pi, pi)
set_ylim = set_xlim
def format_coord(self, lon, lat):
"""
Override this method to change how the values are displayed in
the status bar.
In this case, we want them to be displayed in degrees N/S/E/W.
"""
lon = np.degrees(lon)
lat = np.degrees(lat)
#if lat >= 0.0:
# ns = 'N'
#else:
# ns = 'S'
#if lon >= 0.0:
# ew = 'E'
#else:
# ew = 'W'
return "{0} / {1}".format(round(lon,1), round(lat,1))
class DegreeFormatter(Formatter):
"""
This is a custom formatter that converts the native unit of
radians into (truncated) degrees and adds a degree symbol.
"""
def __init__(self, round_to=1.0):
self._round_to = round_to
def __call__(self, x, pos=None):
degrees = (x / pi) * 180.0
degrees = round(degrees / self._round_to) * self._round_to
return "%d\u00b0" % degrees
def set_longitude_grid(self, degrees):
"""
Set the number of degrees between each longitude grid.
This is an example method that is specific to this projection
class -- it provides a more convenient interface to set the
ticking than set_xticks would.
"""
# Set up a FixedLocator at each of the points, evenly spaced
# by degrees.
number = (360.0 / degrees) + 1
self.xaxis.set_major_locator(
plt.FixedLocator(
np.linspace(-pi, pi, number, True)[1:-1]))
# Set the formatter to display the tick labels in degrees,
# rather than radians.
self.xaxis.set_major_formatter(self.DegreeFormatter(degrees))
def set_latitude_grid(self, degrees):
"""
Set the number of degrees between each longitude grid.
This is an example method that is specific to this projection
class -- it provides a more convenient interface than
set_yticks would.
"""
# Set up a FixedLocator at each of the points, evenly spaced
# by degrees.
number = (180.0 / degrees) + 1
self.yaxis.set_major_locator(
FixedLocator(
np.linspace(-pi / 2.0, pi / 2.0, number, True)[1:-1]))
# Set the formatter to display the tick labels in degrees,
# rather than radians.
self.yaxis.set_major_formatter(self.DegreeFormatter(degrees))
def set_longitude_grid_ends(self, degrees):
"""
Set the latitude(s) at which to stop drawing the longitude grids.
Often, in geographic projections, you wouldn't want to draw
longitude gridlines near the poles. This allows the user to
specify the degree at which to stop drawing longitude grids.
This is an example method that is specific to this projection
class -- it provides an interface to something that has no
analogy in the base Axes class.
"""
longitude_cap = degrees * (pi / 180.0)
# Change the xaxis gridlines transform so that it draws from
# -degrees to degrees, rather than -pi to pi.
self._xaxis_pretransform \
.clear() \
.scale(1.0, longitude_cap * 2.0) \
.translate(0.0, -longitude_cap)
def get_data_ratio(self):
"""
Return the aspect ratio of the data itself.
This method should be overridden by any Axes that have a
fixed data ratio.
"""
return 1.0
# Interactive panning and zooming is not supported with this projection,
# so we override all of the following methods to disable it.
def can_zoom(self):
"""
Return True if this axes support the zoom box
"""
return False
def start_pan(self, x, y, button):
pass
def end_pan(self):
pass
def drag_pan(self, button, key, x, y):
pass
class LambertEqualAreaTransform(Transform):
"""
The basic transformation class.
"""
input_dims = 2
output_dims = 2
is_separable = False
def transform_non_affine(self, ll):
"""
Override the transform_non_affine method to implement the custom
transform.
The input and output are Nx2 numpy arrays.
"""
xi = ll[:, 0:1]
yi = ll[:, 1:2]
k = 1 + np.absolute(cos(yi) * cos(xi))
k = 2 / k
if np.isposinf(k[0]) == True:
k[0] = 1e+15
if np.isneginf(k[0]) == True:
k[0] = -1e+15
if k[0] == 0:
k[0] = 1e-15
k = sqrt(k)
x = k * cos(yi) * sin(xi)
y = k * sin(yi)
return np.concatenate((x, y), 1)
# This is where things get interesting. With this projection,
# straight lines in data space become curves in display space.
# This is done by interpolating new values between the input
# values of the data. Since ``transform`` must not return a
# differently-sized array, any transform that requires
# changing the length of the data array must happen within
# ``transform_path``.
def transform_path_non_affine(self, path):
ipath = path.interpolated(path._interpolation_steps)
return Path(self.transform(ipath.vertices), ipath.codes)
transform_path_non_affine.__doc__ = \
Transform.transform_path_non_affine.__doc__
if matplotlib.__version__ < '1.2':
# Note: For compatibility with matplotlib v1.1 and older, you'll
# need to explicitly implement a ``transform`` method as well.
# Otherwise a ``NotImplementedError`` will be raised. This isn't
# necessary for v1.2 and newer, however.
transform = transform_non_affine
# Similarly, we need to explicitly override ``transform_path`` if
# compatibility with older matplotlib versions is needed. With v1.2
# and newer, only overriding the ``transform_path_non_affine``
# method is sufficient.
transform_path = transform_path_non_affine
transform_path.__doc__ = Transform.transform_path.__doc__
def inverted(self):
return LambertAxes.InvertedLambertEqualAreaTransform()
inverted.__doc__ = Transform.inverted.__doc__
class InvertedLambertEqualAreaTransform(Transform):
#This is not working yet !!!
input_dims = 2
output_dims = 2
is_separable = False
def transform_non_affine(self, xy):
x = xy[:, 0:1]
y = xy[:, 1:2]
#quarter_x = 0.25 * x
#half_y = 0.5 * y
#z = sqrt(1.0 - quarter_x*quarter_x - half_y*half_y)
#longitude = 2 * np.arctan((z*x) / (2.0 * (2.0*z*z - 1.0)))
r = sqrt(2)
p = sqrt(x**2 * y**2)
c = 2 * np.arcsin(p / (2 * r))
phi1 = pi/2
lbd0 = 0
#print(x,y)
if y[0] == 0:
lat = 0
else:
lat = np.arcsin(cos(c) * sin(phi1) + (y * sin(c) * cos(phi1 / p)))
#if phi == phi1:
# lon = lbd0 + np.arctan(x / (-y))
#elif phi == -phi1:
# lon = lbd0 + np.arctan(x / y)
#else:
# lon = lbd0 + np.arctan(x * sin(c) / (p * cos(phi1) * cos(c) - y * sin(phi1) * sin(c)))
if x[0] == 0:
lon = 0
else:
lon = lbd0 + np.arctan(x * sin(c) / (p * cos(phi1) * cos(c) - y * sin(phi1) * sin(c)))
return np.concatenate((lon, lat), 1)
transform_non_affine.__doc__ = Transform.transform_non_affine.__doc__
# As before, we need to implement the "transform" method for
# compatibility with matplotlib v1.1 and older.
if matplotlib.__version__ < '1.2':
transform = transform_non_affine
def inverted(self):
# The inverse of the inverse is the original transform... ;)
return LambertAxes.LambertEqualAreaTransform()
inverted.__doc__ = Transform.inverted.__doc__
# Now register the projection with matplotlib so the user can select
# it.
register_projection(LambertAxes)
It appears that the problem is in both your vectorToGeogr and spherical2vector functions. Based on the comments in those and the pole that you were rotating around, it looks (?) like you were intending to have the following relationship:
x : east-west (east-positive)
y : north-south (north-positive)
z : up-down (down-positive)
However, you had mixed in math in places that assumed mathematical coordinates:
x : towards the equator/prime-meridian intersection
y : towards the equator/90 intersection
z : towards the north pole
A quick-but-not-foolproof test is to try "round-tripping" any coordinate conversion functions. It doesn't guarantee that what you're doing is correct, but it guarantees that it's internally consistent. Your current version of things fails this test:
for lat in range(-90, 100, 10):
for lon in range(-180, 190, 10):
point = np.radians([lon, lat])
round_trip = vectorToGeogr(sphericalToVector(point))
assert np.allclose(point, round_trip)
As an aside, I highly reccomend getting at least a few tests up and running and using a test runner of some sort (py.test is my favorite). It will save you a lot of pain in the long run!
Quick side note:
Personally I prefer to separate "real-world" Cartesian space from the Cartesian space used in a stereonet.
It makes the math simpler, and converting between real-world and "stereonet" space is straight-forward (e.g. see the mplstereonet.xyz2stereonet and mplstereonet.stereonet2xyz functions. They're both in the file you linked to.). The examples in stereonet_math.py all use the second set of conventions. When you need to deal with "real" vectors, (e.g. the contour_normal_vectors.py example) they can be converted over with either xyz2stereonet (outputs lon, lat) or one of the various normal2<foo> functions (outputs plunge/bearing, strike/dip, etc).
However, if you do want to use "real-world" Cartesian coordinates internally, you'll need to change your conversion functions.
Your original sphericalToVector function:
def sphericalToVector(inp_ar):
ar = np.array([0.0, 0.0, 0.0])
ar[0] = sin(inp_ar[0]) * cos(inp_ar[1])
ar[1] = cos(inp_ar[0]) * cos(inp_ar[1])
ar[2] = sin(inp_ar[1])
return ar
Should be changed to:
def sphericalToVector(inp_ar):
ar = np.array([0.0, 0.0, 0.0])
ar[0] = -sin(inp_ar[1])
ar[1] = sin(inp_ar[0]) * cos(inp_ar[1])
ar[2] = cos(inp_ar[0]) * cos(inp_ar[1])
return ar
And your original vectorToGeogr function:
def vectorToGeogr(vect):
ar = np.array([0.0, 0.0])
ar[0] = np.arctan2(vect[0], vect[2])
ar[1] = np.arctan2(vect[1], vect[2])
ar = ar * pi/2
return ar
Should be changed to:
def vectorToGeogr(vect):
ar = np.array([0.0, 0.0])
ar[0] = np.arctan2(vect[1], vect[2])
ar[1] = np.arcsin(-vect[0] / np.linalg.norm(vect))
return ar
The modified version of your example is here: https://gist.github.com/joferkington/ddd90715421720033066 The only things changed are the functions above in test.py. As an example of the result:

Categories