-
Notifications
You must be signed in to change notification settings - Fork 281
/
plot.py
2046 lines (1670 loc) · 69.8 KB
/
plot.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
# Copyright Iris contributors
#
# This file is part of Iris and is released under the LGPL license.
# See COPYING and COPYING.LESSER in the root of the repository for full
# licensing details.
"""
Iris-specific extensions to matplotlib, mimicking the :mod:`matplotlib.pyplot`
interface.
See also: :ref:`matplotlib <matplotlib:users-guide-index>`.
"""
import collections
import datetime
import warnings
import cartopy.crs as ccrs
from cartopy.geodesic import Geodesic
import cartopy.mpl.geoaxes
import cftime
import matplotlib.animation as animation
import matplotlib.axes
import matplotlib.collections as mpl_collections
import matplotlib.dates as mpl_dates
from matplotlib.offsetbox import AnchoredText
import matplotlib.pyplot as plt
import matplotlib.ticker as mpl_ticker
import matplotlib.transforms as mpl_transforms
import numpy as np
import numpy.ma as ma
import iris.analysis.cartography as cartography
import iris.coord_systems
import iris.coords
import iris.cube
from iris.exceptions import IrisError, IrisUnsupportedPlottingWarning
# Importing iris.palette to register the brewer palettes.
import iris.palette
from iris.util import _meshgrid
# Cynthia Brewer citation text.
BREWER_CITE = "Colours based on ColorBrewer.org"
PlotDefn = collections.namedtuple("PlotDefn", ("coords", "transpose"))
def _get_plot_defn_custom_coords_picked(cube, coords, mode, ndims=2):
def names(coords):
result = []
for coord in coords:
if isinstance(coord, int):
result.append("dim={}".format(coord))
else:
result.append(coord.name())
return ", ".join(result)
def as_coord(coord):
if isinstance(coord, int):
# Pass through valid dimension indexes.
if coord >= ndims:
emsg = (
"The data dimension ({}) is out of range for "
"the dimensionality of the required plot ({})"
)
raise IndexError(emsg.format(coord, ndims))
else:
coord = cube.coord(coord)
return coord
coords = list(map(as_coord, coords))
# Check that we were given the right number of coordinates/dimensions.
if len(coords) != ndims:
raise ValueError(
"The list of coordinates given (%s) should have the"
" same length (%s) as the dimensionality of the"
" required plot (%s)" % (names(coords), len(coords), ndims)
)
# Check which dimensions are spanned by each coordinate.
def get_span(coord):
if isinstance(coord, int):
span = set([coord])
else:
span = set(cube.coord_dims(coord))
return span
spans = list(map(get_span, coords))
for span, coord in zip(spans, coords):
if not span:
msg = "The coordinate {!r} doesn't span a data dimension."
raise ValueError(msg.format(coord.name()))
if mode == iris.coords.BOUND_MODE and len(span) not in [1, 2]:
raise ValueError(
"The coordinate {!r} has {} dimensions."
"Cell-based plotting is only supported for"
"coordinates with one or two dimensions.".format(
coord.name(), len(span)
)
)
# Check the combination of coordinates spans enough (ndims) data
# dimensions.
total_span = set().union(*spans)
if len(total_span) != ndims:
raise ValueError(
"The given coordinates ({}) don't span the {} data"
" dimensions.".format(names(coords), ndims)
)
# If we have 2-dimensional data, and one or more 1-dimensional
# coordinates, check if we need to transpose.
transpose = False
if ndims == 2 and min(map(len, spans)) == 1:
for i, span in enumerate(spans):
if len(span) == 1:
if list(span)[0] == i:
transpose = True
break
# Note the use of `reversed` to convert from the X-then-Y
# convention of the end-user API to the V-then-U convention used by
# the plotting routines.
plot_coords = list(reversed(coords))
return PlotDefn(plot_coords, transpose)
def _valid_bound_dim_coord(coord):
result = None
if coord and coord.ndim == 1 and coord.nbounds:
result = coord
return result
def _get_plot_defn(cube, mode, ndims=2):
"""
Return data and plot-axis coords given a cube & a mode of either
POINT_MODE or BOUND_MODE.
"""
if cube.ndim != ndims:
msg = "Cube must be %s-dimensional. Got %s dimensions."
raise ValueError(msg % (ndims, cube.ndim))
# Start by taking the DimCoords from each dimension.
coords = [None] * ndims
for dim_coord in cube.dim_coords:
dim = cube.coord_dims(dim_coord)[0]
coords[dim] = dim_coord
# When appropriate, restrict to 1D with bounds.
if mode == iris.coords.BOUND_MODE:
coords = list(map(_valid_bound_dim_coord, coords))
def guess_axis(coord):
axis = None
if coord is not None:
axis = iris.util.guess_coord_axis(coord)
return axis
# Allow DimCoords in aux_coords to fill in for missing dim_coords.
for dim, coord in enumerate(coords):
if coord is None:
aux_coords = cube.coords(dimensions=dim)
aux_coords = [
coord
for coord in aux_coords
if isinstance(coord, iris.coords.DimCoord)
]
if aux_coords:
aux_coords.sort(key=lambda coord: coord.metadata)
coords[dim] = aux_coords[0]
# If plotting a 2 dimensional plot, check for 2d coordinates
if ndims == 2:
missing_dims = [
dim for dim, coord in enumerate(coords) if coord is None
]
if missing_dims:
# Note that this only picks up coordinates that span the dims
two_dim_coords = cube.coords(dimensions=missing_dims)
two_dim_coords = [
coord for coord in two_dim_coords if coord.ndim == 2
]
if len(two_dim_coords) >= 2:
two_dim_coords.sort(key=lambda coord: coord.metadata)
coords = two_dim_coords[:2]
if mode == iris.coords.POINT_MODE:
# Allow multi-dimensional aux_coords to override the dim_coords
# along the Z axis. This results in a preference for using the
# derived altitude over model_level_number or level_height.
# Limit to Z axis to avoid preferring latitude over grid_latitude etc.
axes = list(map(guess_axis, coords))
axis = "Z"
if axis in axes:
for coord in cube.coords(dim_coords=False):
if (
max(coord.shape) > 1
and iris.util.guess_coord_axis(coord) == axis
):
coords[axes.index(axis)] = coord
# Re-order the coordinates to achieve the preferred
# horizontal/vertical associations. If we can't associate
# an axis to order the coordinates, fall back to using the cube dimension
# followed by the name of the coordinate.
def sort_key(coord):
order = {"X": 2, "T": 1, "Y": -1, "Z": -2}
axis = guess_axis(coord)
return (
order.get(axis, 0),
coords.index(coord),
coord and coord.name(),
)
sorted_coords = sorted(coords, key=sort_key)
transpose = sorted_coords != coords
return PlotDefn(sorted_coords, transpose)
def _can_draw_map(coords):
std_names = [
c and c.standard_name
for c in coords
if isinstance(c, iris.coords.Coord)
]
valid_std_names = [
["latitude", "longitude"],
["grid_latitude", "grid_longitude"],
["projection_y_coordinate", "projection_x_coordinate"],
]
return std_names in valid_std_names
def _broadcast_2d(u, v):
# Matplotlib needs the U and V coordinates to have the same
# dimensionality (either both 1D, or both 2D). So we simply
# broadcast both to 2D to be on the safe side.
u = np.atleast_2d(u)
v = np.atleast_2d(v.T).T
u, v = np.broadcast_arrays(u, v)
return u, v
def _string_coord_axis_tick_labels(string_axes, axes=None):
"""Apply tick labels for string coordinates."""
ax = axes if axes else plt.gca()
for axis, ticks in string_axes.items():
# Define a tick formatter. This will assign a label to all ticks
# located precisely on an integer in range(len(ticks)) and assign
# an empty string to any other ticks.
def ticker_func(tick_location, _):
tick_locations = range(len(ticks))
labels = ticks
label_dict = dict(zip(tick_locations, labels))
label = label_dict.get(tick_location, "")
return label
formatter = mpl_ticker.FuncFormatter(ticker_func)
locator = mpl_ticker.MaxNLocator(integer=True)
this_axis = getattr(ax, axis)
this_axis.set_major_formatter(formatter)
this_axis.set_major_locator(locator)
def _invert_yaxis(v_coord, axes=None):
"""
Inverts the y-axis of the current plot based on conditions:
* If the y-axis is already inverted we don't want to re-invert it.
* If v_coord is None then it will not have any attributes.
* If neither of the above are true then invert y if v_coord has
attribute 'positive' set to 'down'.
Args:
* v_coord - the coord to be plotted on the y-axis
"""
axes = axes if axes else plt.gca()
yaxis_is_inverted = axes.yaxis_inverted()
if not yaxis_is_inverted and isinstance(v_coord, iris.coords.Coord):
attr_pve = v_coord.attributes.get("positive")
if attr_pve is not None and attr_pve.lower() == "down":
axes.invert_yaxis()
def _check_bounds_contiguity_and_mask(coord, data, atol=None, rtol=None):
"""
Checks that any discontiguities in the bounds of the given coordinate only
occur where the data is masked.
Where a discontinuity occurs the grid created for plotting will not be
correct. This does not matter if the data is masked in that location as
this is not plotted.
If a discontiguity occurs where the data is *not* masked, an error is
raised.
Args:
coord: (iris.coord.Coord)
Coordinate the bounds of which will be checked for contiguity
data: (array)
Data of the the cube we are plotting
atol:
Absolute tolerance when checking the contiguity. Defaults to None.
If an absolute tolerance is not set, 1D coords are not checked (so
as to not introduce a breaking change without a major release) but
2D coords are always checked, by calling
:meth:`iris.coords.Coord._discontiguity_in_bounds` with its default
tolerance.
"""
kwargs = {}
data_is_masked = hasattr(data, "mask")
if data_is_masked:
# When checking the location of the discontiguities, we check against
# the opposite of the mask, which is True where data exists.
mask_invert = np.logical_not(data.mask)
if coord.ndim == 1:
# 1D coords are only checked if an absolute tolerance is set, to avoid
# introducing a breaking change.
if atol:
contiguous, diffs = coord._discontiguity_in_bounds(atol=atol)
if not contiguous and data_is_masked:
not_masked_at_discontiguity = np.any(
np.logical_and(mask_invert[:-1], diffs)
)
else:
return
elif coord.ndim == 2:
if atol:
kwargs["atol"] = atol
if rtol:
kwargs["rtol"] = rtol
contiguous, diffs = coord._discontiguity_in_bounds(**kwargs)
if not contiguous and data_is_masked:
diffs_along_x, diffs_along_y = diffs
# Check along both dimensions that any discontiguous
# points are correctly masked.
not_masked_at_discontiguity_along_x = np.any(
np.logical_and(mask_invert[:, :-1], diffs_along_x)
)
not_masked_at_discontiguity_along_y = np.any(
np.logical_and(mask_invert[:-1], diffs_along_y)
)
not_masked_at_discontiguity = (
not_masked_at_discontiguity_along_x
or not_masked_at_discontiguity_along_y
)
# If any discontiguity occurs where the data is not masked the grid will be
# created incorrectly, so raise an error.
if not contiguous:
if not data_is_masked:
raise ValueError(
"The bounds of the {} coordinate are not "
"contiguous. Not able to create a suitable grid"
"to plot. You can use "
"iris.util.find_discontiguities() to identify "
"discontiguities in your x and y coordinate "
"bounds arrays.".format(coord.name())
)
if not_masked_at_discontiguity:
raise ValueError(
"The bounds of the {} coordinate are not "
"contiguous and data is not masked where the "
"discontiguity occurs. Not able to create a "
"suitable grid to plot. You can use "
"iris.util.find_discontiguities() to identify "
"discontiguities in your x and y coordinate "
"bounds arrays, and then mask them with "
"iris.util.mask_cube()"
"".format(coord.name())
)
def _draw_2d_from_bounds(draw_method_name, cube, *args, **kwargs):
# NB. In the interests of clarity we use "u" and "v" to refer to the
# horizontal and vertical axes on the matplotlib plot.
mode = iris.coords.BOUND_MODE
# Get & remove the coords entry from kwargs.
coords = kwargs.pop("coords", None)
if coords is not None:
plot_defn = _get_plot_defn_custom_coords_picked(
cube, coords, mode, ndims=2
)
else:
plot_defn = _get_plot_defn(cube, mode, ndims=2)
contig_tol = kwargs.pop("contiguity_tolerance", None)
for coord in plot_defn.coords:
if hasattr(coord, "has_bounds") and coord.has_bounds():
_check_bounds_contiguity_and_mask(
coord, data=cube.data, atol=contig_tol
)
if _can_draw_map(plot_defn.coords):
result = _map_common(
draw_method_name,
None,
iris.coords.BOUND_MODE,
cube,
plot_defn,
*args,
**kwargs,
)
else:
# Obtain data array.
data = cube.data
if plot_defn.transpose:
data = data.T
# Obtain U and V coordinates
v_coord, u_coord = plot_defn.coords
# Track numpy arrays to use for the actual plotting.
plot_arrays = []
# Map axis name to associated values.
string_axes = {}
for coord, axis_name, data_dim in zip(
[u_coord, v_coord], ["xaxis", "yaxis"], [1, 0]
):
if coord is None:
values = np.arange(data.shape[data_dim] + 1)
elif isinstance(coord, int):
dim = 1 - coord if plot_defn.transpose else coord
values = np.arange(data.shape[dim] + 1)
else:
if coord.points.dtype.char in "SU":
if coord.points.ndim != 1:
msg = "Coord {!r} must be one-dimensional."
raise ValueError(msg.format(coord))
if coord.bounds is not None:
msg = "Cannot plot bounded string coordinate."
raise ValueError(msg)
string_axes[axis_name] = coord.points
values = np.arange(data.shape[data_dim] + 1) - 0.5
else:
values = coord.contiguous_bounds()
values = _fixup_dates(coord, values)
if values.dtype == np.dtype(object) and isinstance(
values[0], datetime.datetime
):
values = mpl_dates.date2num(values)
plot_arrays.append(values)
u, v = plot_arrays
# If the data is transposed, 2D coordinates will also need to be
# transposed.
if plot_defn.transpose is True:
u, v = [coord.T if coord.ndim == 2 else coord for coord in [u, v]]
if u.ndim == v.ndim == 1:
u, v = _broadcast_2d(u, v)
axes = kwargs.pop("axes", None)
draw_method = getattr(axes if axes else plt, draw_method_name)
result = draw_method(u, v, data, *args, **kwargs)
# Apply tick labels for string coordinates.
_string_coord_axis_tick_labels(string_axes, axes)
# Invert y-axis if necessary.
_invert_yaxis(v_coord, axes)
return result
def _draw_2d_from_points(draw_method_name, arg_func, cube, *args, **kwargs):
# NB. In the interests of clarity we use "u" and "v" to refer to the
# horizontal and vertical axes on the matplotlib plot.
mode = iris.coords.POINT_MODE
# Get & remove the coords entry from kwargs.
coords = kwargs.pop("coords", None)
if coords is not None:
plot_defn = _get_plot_defn_custom_coords_picked(cube, coords, mode)
else:
plot_defn = _get_plot_defn(cube, mode, ndims=2)
if _can_draw_map(plot_defn.coords):
result = _map_common(
draw_method_name,
arg_func,
iris.coords.POINT_MODE,
cube,
plot_defn,
*args,
**kwargs,
)
else:
# Obtain data array.
data = cube.data
if plot_defn.transpose:
data = data.T
# Also transpose the scatter marker color array,
# as now mpl 2.x does not do this for free.
if draw_method_name == "scatter" and "c" in kwargs:
c = kwargs["c"]
if hasattr(c, "T") and cube.data.shape == c.shape:
kwargs["c"] = c.T
# Obtain U and V coordinates
v_coord, u_coord = plot_defn.coords
if u_coord is None:
u = np.arange(data.shape[1])
elif isinstance(u_coord, int):
dim = 1 - u_coord if plot_defn.transpose else u_coord
u = np.arange(data.shape[dim])
else:
u = u_coord.points
u = _fixup_dates(u_coord, u)
if v_coord is None:
v = np.arange(data.shape[0])
elif isinstance(v_coord, int):
dim = 1 - v_coord if plot_defn.transpose else v_coord
v = np.arange(data.shape[dim])
else:
v = v_coord.points
v = _fixup_dates(v_coord, v)
if plot_defn.transpose:
u = u.T
v = v.T
# Track numpy arrays to use for the actual plotting.
plot_arrays = []
# Map axis name to associated values.
string_axes = {}
for values, axis_name in zip([u, v], ["xaxis", "yaxis"]):
# Replace any string coordinates with "index" coordinates.
if values.dtype.char in "SU":
if values.ndim != 1:
raise ValueError(
"Multi-dimensional string coordinates "
"not supported."
)
plot_arrays.append(np.arange(values.size))
string_axes[axis_name] = values
elif values.dtype == np.dtype(object) and isinstance(
values[0], datetime.datetime
):
plot_arrays.append(mpl_dates.date2num(values))
else:
plot_arrays.append(values)
u, v = plot_arrays
u, v = _broadcast_2d(u, v)
axes = kwargs.pop("axes", None)
draw_method = getattr(axes if axes else plt, draw_method_name)
if arg_func is not None:
args, kwargs = arg_func(u, v, data, *args, **kwargs)
result = draw_method(*args, **kwargs)
else:
result = draw_method(u, v, data, *args, **kwargs)
# Apply tick labels for string coordinates.
_string_coord_axis_tick_labels(string_axes, axes)
# Invert y-axis if necessary.
_invert_yaxis(v_coord, axes)
return result
def _fixup_dates(coord, values):
if coord.units.calendar is not None and values.ndim == 1:
# Convert coordinate values into tuples of
# (year, month, day, hour, min, sec)
dates = [coord.units.num2date(val).timetuple()[0:6] for val in values]
if coord.units.calendar == "standard":
r = [datetime.datetime(*date) for date in dates]
else:
try:
import nc_time_axis # noqa: F401
except ImportError:
msg = (
"Cannot plot against time in a non-standard "
'calendar, because "nc_time_axis" is not available : '
"Install the package from "
"https://github.com/SciTools/nc-time-axis to enable "
"this usage."
)
raise IrisError(msg)
r = [
cftime.datetime(*date, calendar=coord.units.calendar)
for date in dates
]
values = np.empty(len(r), dtype=object)
values[:] = r
return values
def _data_from_coord_or_cube(c):
if isinstance(c, iris.cube.Cube):
data = c.data
elif isinstance(c, iris.coords.Coord):
data = _fixup_dates(c, c.points)
else:
raise TypeError("Plot arguments must be cubes or coordinates.")
return data
def _uv_from_u_object_v_object(u_object, v_object):
ndim_msg = "Cube or coordinate must be 1-dimensional. Got {} dimensions."
if u_object is not None and u_object.ndim > 1:
raise ValueError(ndim_msg.format(u_object.ndim))
if v_object.ndim > 1:
raise ValueError(ndim_msg.format(v_object.ndim))
v = _data_from_coord_or_cube(v_object)
if u_object is None:
u = np.arange(v.shape[0])
else:
u = _data_from_coord_or_cube(u_object)
return u, v
def _u_object_from_v_object(v_object):
u_object = None
if isinstance(v_object, iris.cube.Cube):
plot_defn = _get_plot_defn(v_object, iris.coords.POINT_MODE, ndims=1)
(u_object,) = plot_defn.coords
return u_object
def _get_plot_objects(args):
if len(args) > 2 and isinstance(
args[2], (iris.cube.Cube, iris.coords.Coord)
):
# three arguments
u_object, v_object1, v_object2 = args[:3]
u1, v1 = _uv_from_u_object_v_object(u_object, v_object1)
_, v2 = _uv_from_u_object_v_object(u_object, v_object2)
args = args[3:]
if u1.size != v1.size or u1.size != v2.size:
msg = "The x and y-axis objects are not all compatible. They should have equal sizes but got ({}: {}), ({}: {}) and ({}: {})"
raise ValueError(
msg.format(
u_object.name(),
u1.size,
v_object1.name(),
v1.size,
v_object2.name(),
v2.size,
)
)
u = u1
v = (v1, v2)
v_object = (v_object1, v_object2)
elif len(args) > 1 and isinstance(
args[1], (iris.cube.Cube, iris.coords.Coord)
):
# two arguments
u_object, v_object = args[:2]
u, v = _uv_from_u_object_v_object(u_object, v_object)
args = args[2:]
if u.size != v.size:
msg = (
"The x and y-axis objects are not compatible. They should "
"have equal sizes but got ({}: {}) and ({}: {})."
)
raise ValueError(
msg.format(u_object.name(), u.size, v_object.name(), v.size)
)
else:
# single argument
v_object = args[0]
u_object = _u_object_from_v_object(v_object)
u, v = _uv_from_u_object_v_object(u_object, args[0])
# If a single cube argument, and the associated dimension coordinate
# is vertical-like, put the coordinate on the y axis, and the data o
# the x.
if (
isinstance(v_object, iris.cube.Cube)
and isinstance(u_object, iris.coords.Coord)
and iris.util.guess_coord_axis(u_object) == "Z"
):
u_object, v_object = v_object, u_object
u, v = v, u
args = args[1:]
return u_object, v_object, u, v, args
def _get_geodesic_params(globe):
# Derive the semimajor axis and flattening values for a given globe from
# its attributes. If the values are under specified, raise a ValueError
flattening = globe.flattening
semimajor = globe.semimajor_axis
try:
if semimajor is None:
# Has semiminor or raises error
if flattening is None:
# Has inverse flattening or raises error
flattening = 1.0 / globe.inverse_flattening
semimajor = globe.semiminor_axis / (1.0 - flattening)
elif flattening is None:
if globe.semiminor_axis is not None:
flattening = (semimajor - globe.semiminor_axis) / float(
semimajor
)
else:
# Has inverse flattening or raises error
flattening = 1.0 / globe.inverse_flattening
except TypeError:
# One of the required attributes was None
raise ValueError("The globe was underspecified.")
return semimajor, flattening
def _shift_plot_sections(u_object, u, v):
"""
Shifts subsections of u by multiples of 360 degrees within ranges
defined by the points where the line should cross over the 0/360 degree
longitude boundary.
e.g. [ 300, 100, 200, 300, 100, 300 ] => [ 300, 460, 560, 660, 820, 660 ]
"""
# Convert coordinates to true lat-lon
src_crs = (
u_object.coord_system.as_cartopy_crs()
if u_object.coord_system is not None
else ccrs.Geodetic()
)
tgt_crs = ccrs.Geodetic(globe=src_crs.globe)
tgt_proj = ccrs.PlateCarree(globe=src_crs.globe)
points = tgt_crs.transform_points(src_crs, u, v)
startpoints = points[:-1, :2]
endpoints = points[1:, :2]
proj_x, proj_y, _ = tgt_proj.transform_points(src_crs, u, v).T
# Calculate the inverse geodesic for each pair of points in turn, and
# convert the start point's azimuth into a vector in the source coordinate
# system.
try:
radius, flattening = _get_geodesic_params(src_crs.globe)
geodesic = Geodesic(radius, flattening)
except ValueError:
geodesic = Geodesic()
dists, azms, _ = geodesic.inverse(startpoints, endpoints).T
azms_lon = np.sin(np.deg2rad(azms))
azms_lat = np.cos(np.deg2rad(azms))
azms_u, _ = src_crs.transform_vectors(
tgt_proj, proj_x[:-1], proj_y[:-1], azms_lon, azms_lat
)
# Use the grid longitude values and the geodesic azimuth to determine
# the points where the line should cross the 0/360 degree boundary, and
# in which direction
lwraps = np.logical_and(u[1:] > u[:-1], azms_u < 0)
rwraps = np.logical_and(u[1:] < u[:-1], azms_u > 0)
shifts = np.where(rwraps, 1, 0) - np.where(lwraps, 1, 0)
shift_vals = shifts.cumsum() * u_object.units.modulus
new_u = np.empty_like(u)
new_u[0] = u[0]
new_u[1:] = u[1:] + shift_vals
return new_u
def _draw_1d_from_points(draw_method_name, arg_func, *args, **kwargs):
# NB. In the interests of clarity we use "u" to refer to the horizontal
# axes on the matplotlib plot and "v" for the vertical axes.
# retrieve the objects that are plotted on the horizontal and vertical
# axes (cubes or coordinates) and their respective values, along with the
# argument tuple with these objects removed
u_object, v_object, u, v, args = _get_plot_objects(args)
# Track numpy arrays to use for the actual plotting.
plot_arrays = []
# Map axis name to associated values.
string_axes = {}
for values, axis_name in zip([u, v], ["xaxis", "yaxis"]):
# Replace any string coordinates with "index" coordinates.
if values.dtype.char in "SU":
if values.ndim != 1:
msg = "Multi-dimensional string coordinates are not supported."
raise ValueError(msg)
plot_arrays.append(np.arange(values.size))
string_axes[axis_name] = values
else:
plot_arrays.append(values)
u, v = plot_arrays
# if both u_object and v_object are coordinates then check if a map
# should be drawn
if (
isinstance(u_object, iris.coords.Coord)
and isinstance(v_object, iris.coords.Coord)
and _can_draw_map([v_object, u_object])
):
# Replace non-cartopy subplot/axes with a cartopy alternative and set
# the transform keyword.
kwargs = _ensure_cartopy_axes_and_determine_kwargs(
u_object, v_object, kwargs
)
if draw_method_name == "plot" and u_object.standard_name not in (
"projection_x_coordinate",
"projection_y_coordinate",
):
u = _shift_plot_sections(u_object, u, v)
axes = kwargs.pop("axes", None)
draw_method = getattr(axes if axes else plt, draw_method_name)
if arg_func is not None:
args, kwargs = arg_func(u, v, *args, **kwargs)
result = draw_method(*args, **kwargs)
else:
result = draw_method(u, v, *args, **kwargs)
# Apply tick labels for string coordinates.
_string_coord_axis_tick_labels(string_axes, axes)
# Invert y-axis if necessary.
_invert_yaxis(v_object, axes)
return result
def _draw_two_1d_from_points(draw_method_name, arg_func, *args, **kwargs):
"""
This function is equivalend to _draw_two_1d_from_points but expects two
y-axis variables rather than one (such as is required for .fill_between). It
can't be used where the y-axis variables are string coordinates. The y-axis
variable provided first has precedence where the two differ on whether the
axis should be inverted or whether a map should be drawn.
"""
# NB. In the interests of clarity we use "u" to refer to the horizontal
# axes on the matplotlib plot and "v" for the vertical axes.
# retrieve the objects that are plotted on the horizontal and vertical
# axes (cubes or coordinates) and their respective values, along with the
# argument tuple with these objects removed
u_object, v_objects, u, vs, args = _get_plot_objects(args)
v_object1, _ = v_objects
v1, v2 = vs
# if both u_object and v_object are coordinates then check if a map
# should be drawn
if (
isinstance(u_object, iris.coords.Coord)
and isinstance(v_object1, iris.coords.Coord)
and _can_draw_map([v_object1, u_object])
):
# Replace non-cartopy subplot/axes with a cartopy alternative and set
# the transform keyword.
kwargs = _ensure_cartopy_axes_and_determine_kwargs(
u_object, v_object1, kwargs
)
axes = kwargs.pop("axes", None)
draw_method = getattr(axes if axes else plt, draw_method_name)
if arg_func is not None:
args, kwargs = arg_func(u, v1, v2, *args, **kwargs)
result = draw_method(*args, **kwargs)
else:
result = draw_method(u, v1, v2, *args, **kwargs)
# Invert y-axis if necessary.
_invert_yaxis(v_object1, axes)
return result
def _replace_axes_with_cartopy_axes(cartopy_proj):
"""
Replace non-cartopy subplot/axes with a cartopy alternative
based on the provided projection. If the current axes are already an
instance of :class:`cartopy.mpl.geoaxes.GeoAxes` then no action is taken.
"""
ax = plt.gca()
if not isinstance(ax, cartopy.mpl.geoaxes.GeoAxes):
fig = ax.get_figure()
if isinstance(ax, matplotlib.axes.SubplotBase):
_ = fig.add_subplot(
ax.get_subplotspec(),
projection=cartopy_proj,
title=ax.get_title(),
xlabel=ax.get_xlabel(),
ylabel=ax.get_ylabel(),
)
else:
position = ax.get_position(original=True)
_ = fig.add_axes(
position,
projection=cartopy_proj,
title=ax.get_title(),
xlabel=ax.get_xlabel(),
ylabel=ax.get_ylabel(),
)
# delete the axes which didn't have a cartopy projection
fig.delaxes(ax)
def _ensure_cartopy_axes_and_determine_kwargs(x_coord, y_coord, kwargs):
"""
Replace the current non-cartopy axes with
:class:`cartopy.mpl.geoaxes.GeoAxes` and return the appropriate kwargs dict
based on the provided coordinates and kwargs.
"""
# Determine projection.
if x_coord.coord_system != y_coord.coord_system:
raise ValueError(
"The X and Y coordinates must have equal coordinate" " systems."
)
cs = x_coord.coord_system
if cs is not None:
cartopy_proj = cs.as_cartopy_projection()
else:
cartopy_proj = ccrs.PlateCarree()
# Ensure the current axes are a cartopy.mpl.geoaxes.GeoAxes instance.
axes = kwargs.get("axes")
if axes is None:
if (
isinstance(cs, iris.coord_systems.RotatedGeogCS)
and x_coord.points.max() > 180
and x_coord.points.max() < 360
and x_coord.points.min() > 0
):
# The RotatedGeogCS has 0 - 360 extent, different from the
# assumptions made by Cartopy: rebase longitudes for the map axes
# to set the datum longitude to the International Date Line.
cs_kwargs = cs._ccrs_kwargs()
cs_kwargs["central_rotated_longitude"] = 180.0
adapted_cartopy_proj = ccrs.RotatedPole(**cs_kwargs)
_replace_axes_with_cartopy_axes(adapted_cartopy_proj)
else:
_replace_axes_with_cartopy_axes(cartopy_proj)
elif axes and not isinstance(axes, cartopy.mpl.geoaxes.GeoAxes):
raise TypeError(
"The supplied axes instance must be a cartopy " "GeoAxes instance."
)
# Set the "from transform" keyword.
if "transform" in kwargs:
raise ValueError(
"The 'transform' keyword is not allowed as it "
"automatically determined from the coordinate "
"metadata."
)
new_kwargs = kwargs.copy()
new_kwargs["transform"] = cartopy_proj
return new_kwargs
def _check_geostationary_coords_and_convert(x, y, kwargs):
# Geostationary stores projected coordinates as scanning angles (
# radians), in line with CF definition (this behaviour is unique to
# Geostationary). Before plotting, must be converted by multiplying by
# satellite height.
x, y = (i.copy() for i in (x, y))
transform = kwargs.get("transform")
if isinstance(transform, cartopy.crs.Geostationary):
satellite_height = transform.proj4_params["h"]
for i in (x, y):
i *= satellite_height
return x, y
def _map_common(
draw_method_name, arg_func, mode, cube, plot_defn, *args, **kwargs