|
1
|
|
|
# Copyright (c) 2008-2015 MetPy Developers. |
|
2
|
|
|
# Distributed under the terms of the BSD 3-Clause License. |
|
3
|
|
|
# SPDX-License-Identifier: BSD-3-Clause |
|
4
|
|
|
"""Utilities for use in making plots.""" |
|
5
|
|
|
|
|
6
|
|
|
import functools |
|
7
|
|
|
|
|
8
|
|
|
from matplotlib.collections import LineCollection |
|
9
|
|
|
import numpy as np |
|
10
|
|
|
|
|
11
|
|
|
from ..units import concatenate |
|
12
|
|
|
|
|
13
|
|
|
|
|
14
|
|
|
# Not part of public API |
|
15
|
|
|
def colored_line(x, y, c, **kwargs): |
|
16
|
|
|
"""Create a multi-colored line. |
|
17
|
|
|
|
|
18
|
|
|
Takes a set of points and turns them into a collection of lines colored by another array. |
|
19
|
|
|
|
|
20
|
|
|
Parameters |
|
21
|
|
|
---------- |
|
22
|
|
|
x : array-like |
|
23
|
|
|
x-axis coordinates |
|
24
|
|
|
y : array-like |
|
25
|
|
|
y-axis coordinates |
|
26
|
|
|
c : array-like |
|
27
|
|
|
values used for color-mapping |
|
28
|
|
|
kwargs : dict |
|
29
|
|
|
Other keyword arguments passed to :class:`matplotlib.collections.LineCollection` |
|
30
|
|
|
|
|
31
|
|
|
Returns |
|
32
|
|
|
------- |
|
33
|
|
|
The created :class:`matplotlib.collections.LineCollection` instance. |
|
34
|
|
|
|
|
35
|
|
|
""" |
|
36
|
|
|
# Mask out any NaN values |
|
37
|
|
|
nan_mask = ~(np.isnan(x) | np.isnan(y) | np.isnan(c)) |
|
38
|
|
|
x = x[nan_mask] |
|
39
|
|
|
y = y[nan_mask] |
|
40
|
|
|
c = c[nan_mask] |
|
41
|
|
|
|
|
42
|
|
|
# Paste values end to end |
|
43
|
|
|
points = concatenate([x, y]) |
|
44
|
|
|
|
|
45
|
|
|
# Exploit numpy's strides to present a view of these points without copying. |
|
46
|
|
|
# Dimensions are (segment, start/end, x/y). Since x and y are concatenated back to back, |
|
47
|
|
|
# moving between segments only moves one item; moving start to end is only an item; |
|
48
|
|
|
# The move between x any moves from one half of the array to the other |
|
49
|
|
|
num_pts = points.size // 2 |
|
50
|
|
|
final_shape = (num_pts - 1, 2, 2) |
|
51
|
|
|
final_strides = (points.itemsize, points.itemsize, num_pts * points.itemsize) |
|
52
|
|
|
segments = np.lib.stride_tricks.as_strided(points, shape=final_shape, |
|
53
|
|
|
strides=final_strides) |
|
54
|
|
|
|
|
55
|
|
|
# Create a LineCollection from the segments and set it to colormap based on c |
|
56
|
|
|
lc = LineCollection(segments, **kwargs) |
|
57
|
|
|
lc.set_array(c) |
|
58
|
|
|
return lc |
|
59
|
|
|
|
|
60
|
|
|
|
|
61
|
|
|
def delete_masked_points(*arrs): |
|
62
|
|
|
"""Delete masked points from arrays. |
|
63
|
|
|
|
|
64
|
|
|
Takes arrays and removes masked points to help with calculations and plotting. |
|
65
|
|
|
|
|
66
|
|
|
Parameters |
|
67
|
|
|
---------- |
|
68
|
|
|
arrs : one or more array-like |
|
69
|
|
|
source arrays |
|
70
|
|
|
|
|
71
|
|
|
Returns |
|
72
|
|
|
------- |
|
73
|
|
|
arrs : one or more array-like |
|
74
|
|
|
arrays with masked elements removed |
|
75
|
|
|
|
|
76
|
|
|
""" |
|
77
|
|
|
if any(hasattr(a, 'mask') for a in arrs): |
|
78
|
|
|
keep = ~functools.reduce(np.logical_or, (np.ma.getmaskarray(a) for a in arrs)) |
|
79
|
|
|
return tuple(a[keep] for a in arrs) |
|
80
|
|
|
else: |
|
81
|
|
|
return arrs |
|
82
|
|
|
|
Unidata /
MetPy
