123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510 |
- import numpy as np
- import matplotlib as mpl
- from matplotlib import _api
- 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 = round(np.rad2deg(x) / self._round_to) * self._round_to
- return f"{degrees:0.0f}\N{DEGREE SIGN}"
- RESOLUTION = 75
- def _init_axis(self):
- self.xaxis = maxis.XAxis(self, clear=False)
- self.yaxis = maxis.YAxis(self, clear=False)
- self.spines['geo'].register_axis(self.yaxis)
- def clear(self):
- # docstring inherited
- super().clear()
- 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(mpl.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((np.pi, 0))
- _, yscale = transform.transform((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'):
- _api.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'):
- _api.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):
- """Not supported. Please consider using Cartopy."""
- raise TypeError("Changing axes limits of a geographic projection is "
- "not supported. Please consider using Cartopy.")
- set_ylim = set_xlim
- def format_coord(self, lon, lat):
- """Return a format string formatting the coordinate."""
- lon, lat = np.rad2deg([lon, lat])
- ns = 'N' if lat >= 0.0 else 'S'
- ew = 'E' if lon >= 0.0 else '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 whether this Axes supports the zoom box button functionality.
- This Axes object does not support interactive zoom box.
- """
- return False
- def can_pan(self):
- """
- Return whether 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 = output_dims = 2
- 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.
- """
- super().__init__()
- self._resolution = resolution
- def __str__(self):
- return f"{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."""
- @_api.rename_parameter("3.8", "ll", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- longitude, latitude = values.T
- # Pre-compute some values
- half_long = longitude / 2.0
- cos_latitude = np.cos(latitude)
- alpha = np.arccos(cos_latitude * np.cos(half_long))
- sinc_alpha = np.sinc(alpha / np.pi) # np.sinc is sin(pi*x)/(pi*x).
- x = (cos_latitude * np.sin(half_long)) / sinc_alpha
- y = np.sin(latitude) / sinc_alpha
- return np.column_stack([x, y])
- def inverted(self):
- # docstring inherited
- return AitoffAxes.InvertedAitoffTransform(self._resolution)
- class InvertedAitoffTransform(_GeoTransform):
- @_api.rename_parameter("3.8", "xy", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- # MGDTODO: Math is hard ;(
- return np.full_like(values, np.nan)
- def inverted(self):
- # docstring inherited
- return AitoffAxes.AitoffTransform(self._resolution)
- def __init__(self, *args, **kwargs):
- self._longitude_cap = np.pi / 2.0
- super().__init__(*args, **kwargs)
- self.set_aspect(0.5, adjustable='box', anchor='C')
- self.clear()
- def _get_core_transform(self, resolution):
- return self.AitoffTransform(resolution)
- class HammerAxes(GeoAxes):
- name = 'hammer'
- class HammerTransform(_GeoTransform):
- """The base Hammer transform."""
- @_api.rename_parameter("3.8", "ll", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- longitude, latitude = values.T
- 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.column_stack([x, y])
- def inverted(self):
- # docstring inherited
- return HammerAxes.InvertedHammerTransform(self._resolution)
- class InvertedHammerTransform(_GeoTransform):
- @_api.rename_parameter("3.8", "xy", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- x, y = values.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
- super().__init__(*args, **kwargs)
- self.set_aspect(0.5, adjustable='box', anchor='C')
- self.clear()
- def _get_core_transform(self, resolution):
- return self.HammerTransform(resolution)
- class MollweideAxes(GeoAxes):
- name = 'mollweide'
- class MollweideTransform(_GeoTransform):
- """The base Mollweide transform."""
- @_api.rename_parameter("3.8", "ll", "values")
- def transform_non_affine(self, values):
- # 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, latitude = values.T
- 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(values.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):
- @_api.rename_parameter("3.8", "xy", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- x, y = values.T
- # from Equations (7, 8) of
- # https://mathworld.wolfram.com/MollweideProjection.html
- theta = np.arcsin(y / np.sqrt(2))
- longitude = (np.pi / (2 * np.sqrt(2))) * x / np.cos(theta)
- latitude = np.arcsin((2 * theta + np.sin(2 * theta)) / np.pi)
- return np.column_stack([longitude, latitude])
- def inverted(self):
- # docstring inherited
- return MollweideAxes.MollweideTransform(self._resolution)
- def __init__(self, *args, **kwargs):
- self._longitude_cap = np.pi / 2.0
- super().__init__(*args, **kwargs)
- self.set_aspect(0.5, adjustable='box', anchor='C')
- self.clear()
- 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
- @_api.rename_parameter("3.8", "ll", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- longitude, latitude = values.T
- 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.column_stack([x, y])
- 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
- @_api.rename_parameter("3.8", "xy", "values")
- def transform_non_affine(self, values):
- # docstring inherited
- x, y = values.T
- 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)
- latitude = np.arcsin(cos_c*np.sin(clat) +
- ((y*sin_c*np.cos(clat)) / p))
- longitude = clong + np.arctan(
- (x*sin_c) / (p*np.cos(clat)*cos_c - y*np.sin(clat)*sin_c))
- return np.column_stack([longitude, latitude])
- 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
- super().__init__(*args, **kwargs)
- self.set_aspect('equal', adjustable='box', anchor='C')
- self.clear()
- def clear(self):
- # docstring inherited
- super().clear()
- 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)
|