Save matplotlib pcolormesh output to file - matplotlib

How to save pcolormesh output (a QuadMesh) to file?
My figure looks like this.
And the axes's children are the following.
[<matplotlib.lines.Line2D at 0x7fca1e4a7b20>,
<matplotlib.lines.Line2D at 0x7fca1e4a7d30>,
<matplotlib.lines.Line2D at 0x7fca1e4d0b20>,
<matplotlib.collections.QuadMesh at 0x7fca1e063f40>,
<matplotlib.spines.Spine at 0x7fca1e486890>,
<matplotlib.spines.Spine at 0x7fca1e4869b0>,
<matplotlib.spines.Spine at 0x7fca1e486ad0>,
<matplotlib.spines.Spine at 0x7fca1e486bf0>,
<matplotlib.axis.XAxis at 0x7fca1e486830>,
<matplotlib.axis.YAxis at 0x7fca1e4870d0>,
Text(0.5, 1.0, ''),
Text(0.0, 1.0, ''),
Text(1.0, 1.0, ''),
<matplotlib.legend.Legend at 0x7fca1e487310>,
<matplotlib.patches.Rectangle at 0x7fca1e4a5840>]
I want to save <matplotlib.collections.QuadMesh at 0x7fca1e063f40> to file.
Hacking it through
quad_mesh = ax.get_children()[3]
face_colors = quad_mesh.get_facecolors() # 90601 x 4 numpy array
These 90k (301^2) quad_mesh._facecolors values are put in place following quad_mesh._coordinates.
So, I just need to reshape them correctly and convert them into an image.
face_colors_reshaped = face_colors.reshape((301, 301, 4))
face_colors_reshaped_8 = (face_colors_reshaped * 255).astype(np.uint8)
im = Image.fromarray(face_colors_reshaped_8).transpose(Image.FLIP_TOP_BOTTOM)
im.save('quad_mesh.png')
So, I did manageā€¦ I just wonder if there was now a better way.
Moreover, the code written above needs to be executed after the figure has been rendered, or nothing will come out of it.

With ax.pcolormesh you supply C, which maps to the colors via norm and cmap (see https://matplotlib.org/stable/tutorials/colors/colormapnorms.html).
Here is an example to do this manually:
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
np.random.seed(42)
C = np.random.uniform(-1, 1, size=(301, 301))
norm = mpl.colors.Normalize(vmin=-1, vmax=1)
C_norm = norm(C)
viridis = plt.get_cmap('viridis')
C_color = (viridis(C_norm) * 255).astype(np.uint8)
fig, ax = plt.subplots()
ax.imshow(C_color)

Related

Matplotlib- Colour scheme

I have a (256,256) numpy array. Each element has a value of 1 or 2 or 3.
I would like to plot this array as an image by:`plt.plot(A)'
What I would like to do is assigning this colour scheme to every element;
1: red, 2:blue and 3: gray.
I have no idea how I can do that, Can someone please help me?
Thanks
You could use plt.imshow() and a ListedColormap. The example below uses yellow instead of gray to see the different colors a bit better ('lightgray' probably looks better than 'gray').
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import numpy as np
from scipy.ndimage import gaussian_filter
fig = plt.figure(figsize=(18, 14))
A = (gaussian_filter(np.random.randn(256, 256), 20) * 200).astype(int) % 3 + 1
cmap = ListedColormap(['crimson', 'dodgerblue', 'yellow'])
plt.imshow(A, origin='lower', cmap=cmap)
cbar = plt.colorbar(ticks=np.linspace(1, 3, 7)[1::2])
cbar.set_ticklabels([1, 2, 3])
plt.show()

Extracting BCI Geodetic and ECI coordinates of an orbit

I am using astropy to define a Tundra orbit around Earth and subsequently, I would like to extract the ECI and geodetic coordinates as the object propagates in time. I was able to get something but it does not agree with what I would expect (ECI coordinates extracted from another SW). The two orbits are not even on the same plane, which is clearly wrong.
Can anybody tell me if I am doing something obviously wrong?
The plot below shows the two results. Orange is with Astropy.
import astropy
from astropy import units as u
from poliastro.bodies import Earth
from astropy.coordinates import CartesianRepresentation
from poliastro.twobody import Orbit
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
epoch = astropy.time.Time('2020-01-01T00:00:00.000', scale='tt')
# Tundra
tundra1 = Orbit.from_classical(attractor=Earth,
a = 42164 *u.km,
ecc = 0.2684 * u.one,
inc = 63.4 * u.deg,
raan = 25 * u.deg,
argp = 270 * u.deg,
nu = 50 * u.deg,
# epoch=epoch
)
def plot_orb(orb, start_t, end_t, step_t, ax, c='k'):
orb_list = []
for t in np.arange(start_t, end_t, step_t):
single_orb = orb.propagate(t*u.min)
orb_list = orb_list + [single_orb]
xyz = orb.sample().xyz
ax.plot(*xyz,'r')
s_xyz_ar = np.zeros((len(orb_list), 3))
for i, s_orb in enumerate(orb_list):
s_xyz = s_orb.represent_as(CartesianRepresentation).xyz
s_xyz_ar[i, :] = s_xyz
ax.scatter(s_xyz_ar[:, 0], s_xyz_ar[:, 1], s_xyz_ar[:, 2], c)
return s_xyz_ar, t
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
s_xyz_ar1, t1 = plot_orb(orb=tundra1, start_t=0, end_t=1440, step_t=10, ax=ax, c='k')
When I wrote that you can do this more efficiently I was under the mistaken assumption that Orbit.propagate can be called directly on an array of time steps like:
>>> tt = np.arange(0, 1440, 10) * u.min
>>> orb = tundra1.propagate(tt)
While this "works" in that it returns a new orbit with an array of epochs, it appears Orbit is not really designed to work with an array of epochs and trying to do something like orb.represent_as just returns a value for the first epoch in the array. This would be a nice possible enhancement to poliastro.
However, the code you wrote for the scatter plot can still be significantly simplified to something like this:
>>> tt = np.arange(0, 1440, 10) * u.min
>>> xyz = np.vstack([tundra1.propagate(t).represent_as(CartesianRepresentation).xyz for t in tt])
>>> fig = plt.figure()
>>> ax = fig.add_subplot(111, projection='3d')
>>> ax.scatter(*xyz.T)
>>> fig.show()
Result:
Ideally you should be able to do this without the np.vstack and instead just call tundra1.propagate(tt).represent_as(CartesianRepresentation).xyz without a for loop. But as the above demonstrates you can still simplify a lot by using np.vstack to make an array from a list of (x, y, z) triplets.
I'm not sure this really answers your original question though, which it seems you found the answer to that wasn't really related to the code. Still, I hope this helps!

Animate a point moving along path between two points

I want to animate a point moving along a path from one location to another on the map.
For example, I drawn a path from New York to New Delhi, using Geodetic transform. Eg. taken from docs Adding data to the map
plt.plot([ny_lon, delhi_lon], [ny_lat, delhi_lat],
color='blue', linewidth=2, marker='o',
transform=ccrs.Geodetic(),
)
Now i want to move a point along this path.
My idea was to somehow get some (say 50) points, along the path and plot a marker on each point for each frame. But I am not able to find a way to get the points on the path.
I found a function transform_points under classCRS, but I am unable to use this, as this gives me the same number of points i have, not the points in between.
Thanks in advance!
There are a couple of approaches to this.
The matplotlib approach
I'll start with perhaps the most basic if you are familiar with matplotlib, but this approach suffers from indirectly using cartopy's functionality, and is therefore harder to configure/extend.
There is a private _get_transformed_path method on a Line2D object (the thing that is returned from plt.plot). The resulting TransformedPath object has a get_transformed_path_and_affine method, which basically will give us the projected line (in the coordinate system of the Axes being drawn).
In [1]: import cartopy.crs as ccrs
In [3]: import matplotlib.pyplot as plt
In [4]: ax = plt.axes(projection=ccrs.Robinson())
In [6]: ny_lon, ny_lat = -75, 43
In [7]: delhi_lon, delhi_lat = 77.23, 28.61
In [8]: [line] = plt.plot([ny_lon, delhi_lon], [ny_lat, delhi_lat],
...: color='blue', linewidth=2, marker='o',
...: transform=ccrs.Geodetic(),
...: )
In [9]: t_path = line._get_transformed_path()
In [10]: path_in_data_coords, _ = t_path.get_transformed_path_and_affine()
In [11]: path_in_data_coords.vertices
Out[11]:
array([[-6425061.82215208, 4594257.92617961],
[-5808923.84969279, 5250795.00604155],
[-5206753.88613758, 5777772.51828996],
[-4554622.94040482, 6244967.03723341],
[-3887558.58343227, 6627927.97123701],
[-3200922.19194864, 6932398.19937816],
[-2480001.76507805, 7165675.95095855],
[-1702269.5101901 , 7332885.72276795],
[ -859899.12295981, 7431215.78426759],
[ 23837.23431173, 7453455.61302756],
[ 889905.10635756, 7397128.77301289],
[ 1695586.66856764, 7268519.87627204],
[ 2434052.81300274, 7073912.54130764],
[ 3122221.22299409, 6812894.40443648],
[ 3782033.80448001, 6478364.28561403],
[ 4425266.18173684, 6062312.15662039],
[ 5049148.25986903, 5563097.6328901 ],
[ 5616318.74912886, 5008293.21452795],
[ 6213232.98764984, 4307186.23400115],
[ 6720608.93929235, 3584542.06839575],
[ 7034261.06659143, 3059873.62740856]])
We can pull this together with matplotlib's animation functionality to do as requested:
import cartopy.crs as ccrs
import matplotlib.animation as animation
import matplotlib.pyplot as plt
ax = plt.axes(projection=ccrs.Robinson())
ax.stock_img()
ny_lon, ny_lat = -75, 43
delhi_lon, delhi_lat = 77.23, 28.61
[line] = plt.plot([ny_lon, delhi_lon], [ny_lat, delhi_lat],
color='blue', linewidth=2, marker='o',
transform=ccrs.Geodetic(),
)
t_path = line._get_transformed_path()
path_in_data_coords, _ = t_path.get_transformed_path_and_affine()
# Draw the point that we want to animate.
[point] = plt.plot(ny_lon, ny_lat, marker='o', transform=ax.projection)
def animate_point(i):
verts = path_in_data_coords.vertices
i = i % verts.shape[0]
# Set the coordinates of the line to the coordinate of the path.
point.set_data(verts[i, 0], verts[i, 1])
ani = animation.FuncAnimation(
ax.figure, animate_point,
frames= path_in_data_coords.vertices.shape[0],
interval=125, repeat=True)
ani.save('point_ani.gif', writer='imagemagick')
plt.show()
The cartopy approach
Under the hood, cartopy's matplotlib implementation (as used above), is calling the project_geometry method. We may as well make use of this directly as it is often more convenient to be using Shapely geometries than it is matplotlib Paths.
With this approach, we simply define a shapely geometry, and then construct the source and target coordinate reference systems that we want to convert the geometry from/to:
target_cs.project_geometry(geometry, source_cs)
The only thing we have to watch out for is that the result can be a MultiLineString (or more generally, any Multi- geometry type). However, in our simple case, we don't need to deal with that (incidentally, the same was true of the simple Path returned in the first example).
The code to produce a similar plot to above:
import cartopy.crs as ccrs
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import shapely.geometry as sgeom
ax = plt.axes(projection=ccrs.Robinson())
ax.stock_img()
ny_lon, ny_lat = -75, 43
delhi_lon, delhi_lat = 77.23, 28.61
line = sgeom.LineString([[ny_lon, ny_lat], [delhi_lon, delhi_lat]])
projected_line = ccrs.PlateCarree().project_geometry(line, ccrs.Geodetic())
# We only animate along one of the projected lines.
if isinstance(projected_line, sgeom.MultiLineString):
projected_line = projected_line.geoms[0]
ax.add_geometries(
[projected_line], ccrs.PlateCarree(),
edgecolor='blue', facecolor='none')
[point] = plt.plot(ny_lon, ny_lat, marker='o', transform=ccrs.PlateCarree())
def animate_point(i):
verts = np.array(projected_line.coords)
i = i % verts.shape[0]
# Set the coordinates of the line to the coordinate of the path.
point.set_data(verts[i, 0], verts[i, 1])
ani = animation.FuncAnimation(
ax.figure, animate_point,
frames=len(projected_line.coords),
interval=125, repeat=True)
ani.save('projected_line_ani.gif', writer='imagemagick')
plt.show()
Final remaaaaarrrrrrks....
The approach naturally generalises to animating any type of matplotlib Arrrrtist.... in this case, I took a bit more control over the great circle resolution, and I animated an image along the great circle:
import cartopy.crs as ccrs
import matplotlib.animation as animation
import matplotlib.pyplot as plt
import numpy as np
import shapely.geometry as sgeom
ax = plt.axes(projection=ccrs.Mercator())
ax.stock_img()
line = sgeom.LineString([[-5.9845, 37.3891], [-82.3666, 23.1136]])
# Higher resolution version of Mercator. Same workaround as found in
# https://github.com/SciTools/cartopy/issues/8#issuecomment-326987465.
class HighRes(ax.projection.__class__):
#property
def threshold(self):
return super(HighRes, self).threshold / 100
projected_line = HighRes().project_geometry(line, ccrs.Geodetic())
# We only animate along one of the projected lines.
if isinstance(projected_line, sgeom.MultiLineString):
projected_line = projected_line.geoms[0]
# Add the projected line to the map.
ax.add_geometries(
[projected_line], ax.projection,
edgecolor='blue', facecolor='none')
def ll_to_extent(x, y, ax_size=(4000000, 4000000)):
"""
Return an image extent in centered on the given
point with the given width and height.
"""
return [x - ax_size[0] / 2, x + ax_size[0] / 2,
y - ax_size[1] / 2, y + ax_size[1] / 2]
# Image from https://pixabay.com/en/sailing-ship-boat-sail-pirate-28930/.
pirate = plt.imread('pirates.png')
img = ax.imshow(pirate, extent=ll_to_extent(0, 0), transform=ax.projection, origin='upper')
ax.set_global()
def animate_ship(i):
verts = np.array(projected_line.coords)
i = i % verts.shape[0]
# Set the extent of the image to the coordinate of the path.
img.set_extent(ll_to_extent(verts[i, 0], verts[i, 1]))
ani = animation.FuncAnimation(
ax.figure, animate_ship,
frames=len(projected_line.coords),
interval=125, repeat=False)
ani.save('arrrr.gif', writer='imagemagick')
plt.show()
All code and images for this answer can be found at https://gist.github.com/pelson/618a5f4ca003e56f06d43815b21848f6.

Using perceptually uniform colormaps in Mayavi volumetric visualization

AFAIK Mayavi does not come with any perceptually uniform colormaps. I tried naively to just pass it one of Matplotlib's colormaps but it failed:
from mayavi import mlab
import multiprocessing
import matplotlib.pyplot as plt
plasma = plt.get_cmap('plasma')
...
mlab.pipeline.volume(..., colormap=plasma)
TraitError: Cannot set the undefined 'colormap' attribute of a 'VolumeFactory' object.
Edit: I found a guide to convert Matplotlib colormaps to Mayavi colormaps. However, it unfortunately doesn't work since I am trying to use a volume using a perceptually uniform colormap.
from matplotlib.cm import get_cmap
import numpy as np
from mayavi import mlab
values = np.linspace(0., 1., 256)
lut_dict = {}
lut_dict['plasma'] = get_cmap('plasma')(values.copy())
x, y, z = np.ogrid[-10:10:20j, -10:10:20j, -10:10:20j]
s = np.sin(x*y*z)/(x*y*z)
mlab.pipeline.volume(mlab.pipeline.scalar_field(s), vmin=0, vmax=0.8, colormap=lut_dict['plasma']) # still getting the same error
mlab.axes()
mlab.show()
...
Instead of setting it as the colormap argument, if you set it as the ColorTransferFunction of the volume, it works as expected.
import numpy as np
from mayavi import mlab
from tvtk.util import ctf
from matplotlib.pyplot import cm
values = np.linspace(0., 1., 256)
x, y, z = np.ogrid[-10:10:20j, -10:10:20j, -10:10:20j]
s = np.sin(x*y*z)/(x*y*z)
volume = mlab.pipeline.volume(mlab.pipeline.scalar_field(s), vmin=0, vmax=0.8)
# save the existing colormap
c = ctf.save_ctfs(volume._volume_property)
# change it with the colors of the new colormap
# in this case 'plasma'
c['rgb']=cm.get_cmap('plasma')(values.copy())
# load the color transfer function to the volume
ctf.load_ctfs(c, volume._volume_property)
# signal for update
volume.update_ctf = True
mlab.show()
While the previous answer by like444 helped me partially with a similar problem, it leads to incorrect translation between colormaps. This is because the format in which matplotlib and tvtk store color information is slightly different: Matplotlib uses RGBA, while ColorTransferFunction uses VRGB, where V is the value in the shown data that this part of the colormap is assigned to. So by doing a 1-to-1 copy, green becomes red, blue becomes green and alpha becomes blue. The following code snippet fixes that:
def cmap_to_ctf(cmap_name):
values = list(np.linspace(0, 1, 256))
cmap = cm.get_cmap(cmap_name)(values)
transfer_function = ctf.ColorTransferFunction()
for i, v in enumerate(values):
transfer_function.add_rgb_point(v, cmap[i, 0], cmap[i, 1], cmap[i, 2])
return transfer_function

Plotting masked numpy array leads to incorrect colorbar

I'm trying to create a custom color bar for a matplotlib PolyCollection. Everything seems ok until I attempt to plot a masked array. The color bar no longer shows the correct colors even though the plot does. Is there a different procedure for plotting masked arrays?
I'm using matplotlib 1.4.0 and numpy 1.8.
Here's my plotting code:
import numpy
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.collections import PolyCollection
vertices = numpy.load('vertices.npy')
array = numpy.load('array.npy')
# Take 2d slice out of 3D array
slice_ = array[:, :, 0:1].flatten(order='F')
fig, ax = plt.subplots()
poly = PolyCollection(vertices, array=slice_, edgecolors='black', linewidth=.25)
cm = mpl.colors.ListedColormap([(1.0, 0.0, 0.0), (.2, .5, .2)])
poly.set_cmap(cm)
bounds = [.1, .4, .6]
norm = mpl.colors.BoundaryNorm(bounds, cm.N)
fig.colorbar(poly, ax=ax, orientation='vertical', boundaries=bounds, norm=norm)
ax.add_collection(poly, autolim=True)
ax.autoscale_view()
plt.show()
Here's what the plot looks like:
However, when I plot a masked array with the following change before the slicing:
array = numpy.ma.array(array, mask=array > .5)
I get a color bar that now shows only a single color. Even though both colors are (correctly) still shown in the plot.
Is there some trick to keeping a colobar consistent when plotting a masked array? I know I can use cm.set_bad to change the color of masked values, but that's not quite what I'm looking for. I want the color bar to show up the same between these two plots since both colors and the color bar itself should remain unchanged.
Pass the BoundaryNorm to the PolyCollection, poly. Otherwise, poly.norm gets set to a matplotlib.colors.Normalize instance by default:
In [119]: poly.norm
Out[119]: <matplotlib.colors.Normalize at 0x7faac4dc8210>
I have not stepped through the source code sufficiently to explain exactly what is happening in the code you posted, but I speculate that the interaction of this Normalize instance and the BoundaryNorm make the range of values seen by the fig.colorbar different than what you expected.
In any case, if you pass norm=norm to PolyCollection, then the result looks correct:
import numpy
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.collections as mcoll
import matplotlib.colors as mcolors
numpy.random.seed(4)
N, M = 3, 3
vertices = numpy.random.random((N, M, 2))
array = numpy.random.random((1, N, 2))
# vertices = numpy.load('vertices.npy')
# array = numpy.load('array.npy')
array = numpy.ma.array(array, mask=array > .5)
# Take 2d slice out of 3D array
slice_ = array[:, :, 0:1].flatten(order='F')
fig, ax = plt.subplots()
bounds = [.1, .4, .6]
cm = mpl.colors.ListedColormap([(1.0, 0.0, 0.0), (.2, .5, .2)])
norm = mpl.colors.BoundaryNorm(bounds, cm.N)
poly = mcoll.PolyCollection(
vertices,
array=slice_,
edgecolors='black', linewidth=.25, norm=norm)
poly.set_cmap(cm)
fig.colorbar(poly, ax=ax, orientation='vertical')
ax.add_collection(poly, autolim=True)
ax.autoscale_view()
plt.show()