import numpy as np from matplotlib import cbook, rcParams from matplotlib.axes import Axes import matplotlib.axis as maxis from matplotlib.patches import Circle from matplotlib.path import Path import matplotlib.spines as mspines from matplotlib.ticker import ( Formatter, NullLocator, FixedLocator, NullFormatter) from matplotlib.transforms import Affine2D, BboxTransformTo, Transform class GeoAxes(Axes): """An abstract base class for geographic projections.""" class ThetaFormatter(Formatter): """ Used to format the theta tick labels. Converts the native unit of radians into 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 / np.pi) * 180.0 degrees = np.round(degrees / self._round_to) * self._round_to if rcParams['text.usetex'] and not rcParams['text.latex.unicode']: return r"$%0.0f^\circ$" % degrees else: return "%0.0f\N{DEGREE SIGN}" % degrees RESOLUTION = 75 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 GeoAxes.xaxis.cla() works. # self.spines['geo'].register_axis(self.yaxis) self._update_transScale() def cla(self): Axes.cla(self) self.set_longitude_grid(30) self.set_latitude_grid(15) self.set_longitude_grid_ends(75) self.xaxis.set_minor_locator(NullLocator()) self.yaxis.set_minor_locator(NullLocator()) self.xaxis.set_ticks_position('none') self.yaxis.set_ticks_position('none') self.yaxis.set_tick_params(label1On=True) # Why do we need to turn on yaxis tick labels, but # xaxis tick labels are already on? self.grid(rcParams['axes.grid']) Axes.set_xlim(self, -np.pi, np.pi) Axes.set_ylim(self, -np.pi / 2.0, np.pi / 2.0) def _set_lim_and_transforms(self): # A (possibly non-linear) projection on the (already scaled) data self.transProjection = self._get_core_transform(self.RESOLUTION) self.transAffine = self._get_affine_transform() self.transAxes = BboxTransformTo(self.bbox) # The complete data transformation stack -- from data all the # way to display coordinates self.transData = \ self.transProjection + \ self.transAffine + \ self.transAxes # This is the transform for longitude ticks. self._xaxis_pretransform = \ Affine2D() \ .scale(1, self._longitude_cap * 2) \ .translate(0, -self._longitude_cap) self._xaxis_transform = \ self._xaxis_pretransform + \ self.transData self._xaxis_text1_transform = \ Affine2D().scale(1, 0) + \ self.transData + \ Affine2D().translate(0, 4) self._xaxis_text2_transform = \ Affine2D().scale(1, 0) + \ self.transData + \ Affine2D().translate(0, -4) # This is the transform for latitude ticks. yaxis_stretch = Affine2D().scale(np.pi * 2, 1).translate(-np.pi, 0) yaxis_space = Affine2D().scale(1, 1.1) 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) self._yaxis_text2_transform = \ yaxis_text_base + \ Affine2D().translate(8, 0) def _get_affine_transform(self): transform = self._get_core_transform(1) xscale, _ = transform.transform_point((np.pi, 0)) _, yscale = transform.transform_point((0, np.pi / 2)) return Affine2D() \ .scale(0.5 / xscale, 0.5 / yscale) \ .translate(0.5, 0.5) def get_xaxis_transform(self, which='grid'): cbook._check_in_list(['tick1', 'tick2', 'grid'], which=which) return self._xaxis_transform def get_xaxis_text1_transform(self, pad): return self._xaxis_text1_transform, 'bottom', 'center' def get_xaxis_text2_transform(self, pad): return self._xaxis_text2_transform, 'top', 'center' def get_yaxis_transform(self, which='grid'): cbook._check_in_list(['tick1', 'tick2', 'grid'], which=which) return self._yaxis_transform def get_yaxis_text1_transform(self, pad): return self._yaxis_text1_transform, 'center', 'right' def get_yaxis_text2_transform(self, pad): return self._yaxis_text2_transform, 'center', 'left' def _gen_axes_patch(self): return Circle((0.5, 0.5), 0.5) def _gen_axes_spines(self): return {'geo': mspines.Spine.circular_spine(self, (0.5, 0.5), 0.5)} def set_yscale(self, *args, **kwargs): if args[0] != 'linear': raise NotImplementedError set_xscale = set_yscale def set_xlim(self, *args, **kwargs): raise TypeError("It is not possible to change axes limits " "for geographic projections. Please consider " "using Basemap or Cartopy.") set_ylim = set_xlim def format_coord(self, lon, lat): 'return a format string formatting the coordinate' lon, lat = np.rad2deg([lon, lat]) if lat >= 0.0: ns = 'N' else: ns = 'S' if lon >= 0.0: ew = 'E' else: ew = 'W' return ('%f\N{DEGREE SIGN}%s, %f\N{DEGREE SIGN}%s' % (abs(lat), ns, abs(lon), ew)) def set_longitude_grid(self, degrees): """ Set the number of degrees between each longitude grid. """ # Skip -180 and 180, which are the fixed limits. grid = np.arange(-180 + degrees, 180, degrees) self.xaxis.set_major_locator(FixedLocator(np.deg2rad(grid))) self.xaxis.set_major_formatter(self.ThetaFormatter(degrees)) def set_latitude_grid(self, degrees): """ Set the number of degrees between each latitude grid. """ # Skip -90 and 90, which are the fixed limits. grid = np.arange(-90 + degrees, 90, degrees) self.yaxis.set_major_locator(FixedLocator(np.deg2rad(grid))) self.yaxis.set_major_formatter(self.ThetaFormatter(degrees)) def set_longitude_grid_ends(self, degrees): """ Set the latitude(s) at which to stop drawing the longitude grids. """ self._longitude_cap = np.deg2rad(degrees) self._xaxis_pretransform \ .clear() \ .scale(1.0, self._longitude_cap * 2.0) \ .translate(0.0, -self._longitude_cap) def get_data_ratio(self): ''' Return the aspect ratio of the data itself. ''' return 1.0 ### Interactive panning def can_zoom(self): """ Return *True* if this axes supports the zoom box button functionality. This axes object does not support interactive zoom box. """ return False def can_pan(self) : """ Return *True* if this axes supports the pan/zoom button functionality. This axes object does not support interactive pan/zoom. """ 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 _GeoTransform(Transform): # Factoring out some common functionality. input_dims = 2 output_dims = 2 is_separable = False def __init__(self, resolution): """ Create a new geographical transform. Resolution is the number of steps to interpolate between each input line segment to approximate its path in curved space. """ Transform.__init__(self) self._resolution = resolution def __str__(self): return "{}({})".format(type(self).__name__, self._resolution) def transform_path_non_affine(self, path): # docstring inherited ipath = path.interpolated(self._resolution) return Path(self.transform(ipath.vertices), ipath.codes) class AitoffAxes(GeoAxes): name = 'aitoff' class AitoffTransform(_GeoTransform): """The base Aitoff transform.""" def transform_non_affine(self, ll): # docstring inherited longitude = ll[:, 0] latitude = ll[:, 1] # Pre-compute some values half_long = longitude / 2.0 cos_latitude = np.cos(latitude) alpha = np.arccos(cos_latitude * np.cos(half_long)) # Avoid divide-by-zero errors using same method as NumPy. alpha[alpha == 0.0] = 1e-20 # We want unnormalized sinc. numpy.sinc gives us normalized sinc_alpha = np.sin(alpha) / alpha xy = np.empty_like(ll, float) xy[:, 0] = (cos_latitude * np.sin(half_long)) / sinc_alpha xy[:, 1] = np.sin(latitude) / sinc_alpha return xy def inverted(self): # docstring inherited return AitoffAxes.InvertedAitoffTransform(self._resolution) class InvertedAitoffTransform(_GeoTransform): def transform_non_affine(self, xy): # docstring inherited # MGDTODO: Math is hard ;( return xy def inverted(self): # docstring inherited return AitoffAxes.AitoffTransform(self._resolution) def __init__(self, *args, **kwargs): self._longitude_cap = np.pi / 2.0 GeoAxes.__init__(self, *args, **kwargs) self.set_aspect(0.5, adjustable='box', anchor='C') self.cla() def _get_core_transform(self, resolution): return self.AitoffTransform(resolution) class HammerAxes(GeoAxes): name = 'hammer' class HammerTransform(_GeoTransform): """The base Hammer transform.""" def transform_non_affine(self, ll): # docstring inherited longitude = ll[:, 0:1] latitude = ll[:, 1:2] # Pre-compute some values half_long = longitude / 2.0 cos_latitude = np.cos(latitude) sqrt2 = np.sqrt(2.0) alpha = np.sqrt(1.0 + cos_latitude * np.cos(half_long)) x = (2.0 * sqrt2) * (cos_latitude * np.sin(half_long)) / alpha y = (sqrt2 * np.sin(latitude)) / alpha return np.concatenate((x, y), 1) def inverted(self): # docstring inherited return HammerAxes.InvertedHammerTransform(self._resolution) class InvertedHammerTransform(_GeoTransform): def transform_non_affine(self, xy): # docstring inherited x, y = xy.T z = np.sqrt(1 - (x / 4) ** 2 - (y / 2) ** 2) longitude = 2 * np.arctan((z * x) / (2 * (2 * z ** 2 - 1))) latitude = np.arcsin(y*z) return np.column_stack([longitude, latitude]) def inverted(self): # docstring inherited return HammerAxes.HammerTransform(self._resolution) def __init__(self, *args, **kwargs): self._longitude_cap = np.pi / 2.0 GeoAxes.__init__(self, *args, **kwargs) self.set_aspect(0.5, adjustable='box', anchor='C') self.cla() def _get_core_transform(self, resolution): return self.HammerTransform(resolution) class MollweideAxes(GeoAxes): name = 'mollweide' class MollweideTransform(_GeoTransform): """The base Mollweide transform.""" def transform_non_affine(self, ll): # docstring inherited def d(theta): delta = (-(theta + np.sin(theta) - pi_sin_l) / (1 + np.cos(theta))) return delta, np.abs(delta) > 0.001 longitude = ll[:, 0] latitude = ll[:, 1] clat = np.pi/2 - np.abs(latitude) ihigh = clat < 0.087 # within 5 degrees of the poles ilow = ~ihigh aux = np.empty(latitude.shape, dtype=float) if ilow.any(): # Newton-Raphson iteration pi_sin_l = np.pi * np.sin(latitude[ilow]) theta = 2.0 * latitude[ilow] delta, large_delta = d(theta) while np.any(large_delta): theta[large_delta] += delta[large_delta] delta, large_delta = d(theta) aux[ilow] = theta / 2 if ihigh.any(): # Taylor series-based approx. solution e = clat[ihigh] d = 0.5 * (3 * np.pi * e**2) ** (1.0/3) aux[ihigh] = (np.pi/2 - d) * np.sign(latitude[ihigh]) xy = np.empty(ll.shape, dtype=float) xy[:, 0] = (2.0 * np.sqrt(2.0) / np.pi) * longitude * np.cos(aux) xy[:, 1] = np.sqrt(2.0) * np.sin(aux) return xy def inverted(self): # docstring inherited return MollweideAxes.InvertedMollweideTransform(self._resolution) class InvertedMollweideTransform(_GeoTransform): def transform_non_affine(self, xy): # docstring inherited x = xy[:, 0:1] y = xy[:, 1:2] # from Equations (7, 8) of # http://mathworld.wolfram.com/MollweideProjection.html theta = np.arcsin(y / np.sqrt(2)) lon = (np.pi / (2 * np.sqrt(2))) * x / np.cos(theta) lat = np.arcsin((2 * theta + np.sin(2 * theta)) / np.pi) return np.concatenate((lon, lat), 1) def inverted(self): # docstring inherited return MollweideAxes.MollweideTransform(self._resolution) def __init__(self, *args, **kwargs): self._longitude_cap = np.pi / 2.0 GeoAxes.__init__(self, *args, **kwargs) self.set_aspect(0.5, adjustable='box', anchor='C') self.cla() def _get_core_transform(self, resolution): return self.MollweideTransform(resolution) class LambertAxes(GeoAxes): name = 'lambert' class LambertTransform(_GeoTransform): """The base Lambert transform.""" def __init__(self, center_longitude, center_latitude, resolution): """ Create a new Lambert transform. Resolution is the number of steps to interpolate between each input line segment to approximate its path in curved Lambert space. """ _GeoTransform.__init__(self, resolution) self._center_longitude = center_longitude self._center_latitude = center_latitude def transform_non_affine(self, ll): # docstring inherited longitude = ll[:, 0:1] latitude = ll[:, 1:2] clong = self._center_longitude clat = self._center_latitude cos_lat = np.cos(latitude) sin_lat = np.sin(latitude) diff_long = longitude - clong cos_diff_long = np.cos(diff_long) inner_k = np.maximum( # Prevent divide-by-zero problems 1 + np.sin(clat)*sin_lat + np.cos(clat)*cos_lat*cos_diff_long, 1e-15) k = np.sqrt(2 / inner_k) x = k * cos_lat*np.sin(diff_long) y = k * (np.cos(clat)*sin_lat - np.sin(clat)*cos_lat*cos_diff_long) return np.concatenate((x, y), 1) def inverted(self): # docstring inherited return LambertAxes.InvertedLambertTransform( self._center_longitude, self._center_latitude, self._resolution) class InvertedLambertTransform(_GeoTransform): def __init__(self, center_longitude, center_latitude, resolution): _GeoTransform.__init__(self, resolution) self._center_longitude = center_longitude self._center_latitude = center_latitude def transform_non_affine(self, xy): # docstring inherited x = xy[:, 0:1] y = xy[:, 1:2] clong = self._center_longitude clat = self._center_latitude p = np.maximum(np.hypot(x, y), 1e-9) c = 2 * np.arcsin(0.5 * p) sin_c = np.sin(c) cos_c = np.cos(c) lat = np.arcsin(cos_c*np.sin(clat) + ((y*sin_c*np.cos(clat)) / p)) lon = clong + np.arctan( (x*sin_c) / (p*np.cos(clat)*cos_c - y*np.sin(clat)*sin_c)) return np.concatenate((lon, lat), 1) def inverted(self): # docstring inherited return LambertAxes.LambertTransform( self._center_longitude, self._center_latitude, self._resolution) def __init__(self, *args, center_longitude=0, center_latitude=0, **kwargs): self._longitude_cap = np.pi / 2 self._center_longitude = center_longitude self._center_latitude = center_latitude GeoAxes.__init__(self, *args, **kwargs) self.set_aspect('equal', adjustable='box', anchor='C') self.cla() def cla(self): GeoAxes.cla(self) self.yaxis.set_major_formatter(NullFormatter()) def _get_core_transform(self, resolution): return self.LambertTransform( self._center_longitude, self._center_latitude, resolution) def _get_affine_transform(self): return Affine2D() \ .scale(0.25) \ .translate(0.5, 0.5)