geo.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511
  1. import numpy as np
  2. import matplotlib as mpl
  3. from matplotlib import _api
  4. from matplotlib.axes import Axes
  5. import matplotlib.axis as maxis
  6. from matplotlib.patches import Circle
  7. from matplotlib.path import Path
  8. import matplotlib.spines as mspines
  9. from matplotlib.ticker import (
  10. Formatter, NullLocator, FixedLocator, NullFormatter)
  11. from matplotlib.transforms import Affine2D, BboxTransformTo, Transform
  12. class GeoAxes(Axes):
  13. """An abstract base class for geographic projections."""
  14. class ThetaFormatter(Formatter):
  15. """
  16. Used to format the theta tick labels. Converts the native
  17. unit of radians into degrees and adds a degree symbol.
  18. """
  19. def __init__(self, round_to=1.0):
  20. self._round_to = round_to
  21. def __call__(self, x, pos=None):
  22. degrees = round(np.rad2deg(x) / self._round_to) * self._round_to
  23. return f"{degrees:0.0f}\N{DEGREE SIGN}"
  24. RESOLUTION = 75
  25. def _init_axis(self):
  26. self.xaxis = maxis.XAxis(self, clear=False)
  27. self.yaxis = maxis.YAxis(self, clear=False)
  28. self.spines['geo'].register_axis(self.yaxis)
  29. def clear(self):
  30. # docstring inherited
  31. super().clear()
  32. self.set_longitude_grid(30)
  33. self.set_latitude_grid(15)
  34. self.set_longitude_grid_ends(75)
  35. self.xaxis.set_minor_locator(NullLocator())
  36. self.yaxis.set_minor_locator(NullLocator())
  37. self.xaxis.set_ticks_position('none')
  38. self.yaxis.set_ticks_position('none')
  39. self.yaxis.set_tick_params(label1On=True)
  40. # Why do we need to turn on yaxis tick labels, but
  41. # xaxis tick labels are already on?
  42. self.grid(mpl.rcParams['axes.grid'])
  43. Axes.set_xlim(self, -np.pi, np.pi)
  44. Axes.set_ylim(self, -np.pi / 2.0, np.pi / 2.0)
  45. def _set_lim_and_transforms(self):
  46. # A (possibly non-linear) projection on the (already scaled) data
  47. self.transProjection = self._get_core_transform(self.RESOLUTION)
  48. self.transAffine = self._get_affine_transform()
  49. self.transAxes = BboxTransformTo(self.bbox)
  50. # The complete data transformation stack -- from data all the
  51. # way to display coordinates
  52. self.transData = \
  53. self.transProjection + \
  54. self.transAffine + \
  55. self.transAxes
  56. # This is the transform for longitude ticks.
  57. self._xaxis_pretransform = \
  58. Affine2D() \
  59. .scale(1, self._longitude_cap * 2) \
  60. .translate(0, -self._longitude_cap)
  61. self._xaxis_transform = \
  62. self._xaxis_pretransform + \
  63. self.transData
  64. self._xaxis_text1_transform = \
  65. Affine2D().scale(1, 0) + \
  66. self.transData + \
  67. Affine2D().translate(0, 4)
  68. self._xaxis_text2_transform = \
  69. Affine2D().scale(1, 0) + \
  70. self.transData + \
  71. Affine2D().translate(0, -4)
  72. # This is the transform for latitude ticks.
  73. yaxis_stretch = Affine2D().scale(np.pi * 2, 1).translate(-np.pi, 0)
  74. yaxis_space = Affine2D().scale(1, 1.1)
  75. self._yaxis_transform = \
  76. yaxis_stretch + \
  77. self.transData
  78. yaxis_text_base = \
  79. yaxis_stretch + \
  80. self.transProjection + \
  81. (yaxis_space +
  82. self.transAffine +
  83. self.transAxes)
  84. self._yaxis_text1_transform = \
  85. yaxis_text_base + \
  86. Affine2D().translate(-8, 0)
  87. self._yaxis_text2_transform = \
  88. yaxis_text_base + \
  89. Affine2D().translate(8, 0)
  90. def _get_affine_transform(self):
  91. transform = self._get_core_transform(1)
  92. xscale, _ = transform.transform((np.pi, 0))
  93. _, yscale = transform.transform((0, np.pi/2))
  94. return Affine2D() \
  95. .scale(0.5 / xscale, 0.5 / yscale) \
  96. .translate(0.5, 0.5)
  97. def get_xaxis_transform(self, which='grid'):
  98. _api.check_in_list(['tick1', 'tick2', 'grid'], which=which)
  99. return self._xaxis_transform
  100. def get_xaxis_text1_transform(self, pad):
  101. return self._xaxis_text1_transform, 'bottom', 'center'
  102. def get_xaxis_text2_transform(self, pad):
  103. return self._xaxis_text2_transform, 'top', 'center'
  104. def get_yaxis_transform(self, which='grid'):
  105. _api.check_in_list(['tick1', 'tick2', 'grid'], which=which)
  106. return self._yaxis_transform
  107. def get_yaxis_text1_transform(self, pad):
  108. return self._yaxis_text1_transform, 'center', 'right'
  109. def get_yaxis_text2_transform(self, pad):
  110. return self._yaxis_text2_transform, 'center', 'left'
  111. def _gen_axes_patch(self):
  112. return Circle((0.5, 0.5), 0.5)
  113. def _gen_axes_spines(self):
  114. return {'geo': mspines.Spine.circular_spine(self, (0.5, 0.5), 0.5)}
  115. def set_yscale(self, *args, **kwargs):
  116. if args[0] != 'linear':
  117. raise NotImplementedError
  118. set_xscale = set_yscale
  119. def set_xlim(self, *args, **kwargs):
  120. """Not supported. Please consider using Cartopy."""
  121. raise TypeError("Changing axes limits of a geographic projection is "
  122. "not supported. Please consider using Cartopy.")
  123. set_ylim = set_xlim
  124. set_xbound = set_xlim
  125. set_ybound = set_ylim
  126. def invert_xaxis(self):
  127. """Not supported. Please consider using Cartopy."""
  128. raise TypeError("Changing axes limits of a geographic projection is "
  129. "not supported. Please consider using Cartopy.")
  130. invert_yaxis = invert_xaxis
  131. def format_coord(self, lon, lat):
  132. """Return a format string formatting the coordinate."""
  133. lon, lat = np.rad2deg([lon, lat])
  134. ns = 'N' if lat >= 0.0 else 'S'
  135. ew = 'E' if lon >= 0.0 else 'W'
  136. return ('%f\N{DEGREE SIGN}%s, %f\N{DEGREE SIGN}%s'
  137. % (abs(lat), ns, abs(lon), ew))
  138. def set_longitude_grid(self, degrees):
  139. """
  140. Set the number of degrees between each longitude grid.
  141. """
  142. # Skip -180 and 180, which are the fixed limits.
  143. grid = np.arange(-180 + degrees, 180, degrees)
  144. self.xaxis.set_major_locator(FixedLocator(np.deg2rad(grid)))
  145. self.xaxis.set_major_formatter(self.ThetaFormatter(degrees))
  146. def set_latitude_grid(self, degrees):
  147. """
  148. Set the number of degrees between each latitude grid.
  149. """
  150. # Skip -90 and 90, which are the fixed limits.
  151. grid = np.arange(-90 + degrees, 90, degrees)
  152. self.yaxis.set_major_locator(FixedLocator(np.deg2rad(grid)))
  153. self.yaxis.set_major_formatter(self.ThetaFormatter(degrees))
  154. def set_longitude_grid_ends(self, degrees):
  155. """
  156. Set the latitude(s) at which to stop drawing the longitude grids.
  157. """
  158. self._longitude_cap = np.deg2rad(degrees)
  159. self._xaxis_pretransform \
  160. .clear() \
  161. .scale(1.0, self._longitude_cap * 2.0) \
  162. .translate(0.0, -self._longitude_cap)
  163. def get_data_ratio(self):
  164. """Return the aspect ratio of the data itself."""
  165. return 1.0
  166. ### Interactive panning
  167. def can_zoom(self):
  168. """
  169. Return whether this Axes supports the zoom box button functionality.
  170. This Axes object does not support interactive zoom box.
  171. """
  172. return False
  173. def can_pan(self):
  174. """
  175. Return whether this Axes supports the pan/zoom button functionality.
  176. This Axes object does not support interactive pan/zoom.
  177. """
  178. return False
  179. def start_pan(self, x, y, button):
  180. pass
  181. def end_pan(self):
  182. pass
  183. def drag_pan(self, button, key, x, y):
  184. pass
  185. class _GeoTransform(Transform):
  186. # Factoring out some common functionality.
  187. input_dims = output_dims = 2
  188. def __init__(self, resolution):
  189. """
  190. Create a new geographical transform.
  191. Resolution is the number of steps to interpolate between each input
  192. line segment to approximate its path in curved space.
  193. """
  194. super().__init__()
  195. self._resolution = resolution
  196. def __str__(self):
  197. return f"{type(self).__name__}({self._resolution})"
  198. def transform_path_non_affine(self, path):
  199. # docstring inherited
  200. ipath = path.interpolated(self._resolution)
  201. return Path(self.transform(ipath.vertices), ipath.codes)
  202. class AitoffAxes(GeoAxes):
  203. name = 'aitoff'
  204. class AitoffTransform(_GeoTransform):
  205. """The base Aitoff transform."""
  206. def transform_non_affine(self, values):
  207. # docstring inherited
  208. longitude, latitude = values.T
  209. # Pre-compute some values
  210. half_long = longitude / 2.0
  211. cos_latitude = np.cos(latitude)
  212. alpha = np.arccos(cos_latitude * np.cos(half_long))
  213. sinc_alpha = np.sinc(alpha / np.pi) # np.sinc is sin(pi*x)/(pi*x).
  214. x = (cos_latitude * np.sin(half_long)) / sinc_alpha
  215. y = np.sin(latitude) / sinc_alpha
  216. return np.column_stack([x, y])
  217. def inverted(self):
  218. # docstring inherited
  219. return AitoffAxes.InvertedAitoffTransform(self._resolution)
  220. class InvertedAitoffTransform(_GeoTransform):
  221. def transform_non_affine(self, values):
  222. # docstring inherited
  223. # MGDTODO: Math is hard ;(
  224. return np.full_like(values, np.nan)
  225. def inverted(self):
  226. # docstring inherited
  227. return AitoffAxes.AitoffTransform(self._resolution)
  228. def __init__(self, *args, **kwargs):
  229. self._longitude_cap = np.pi / 2.0
  230. super().__init__(*args, **kwargs)
  231. self.set_aspect(0.5, adjustable='box', anchor='C')
  232. self.clear()
  233. def _get_core_transform(self, resolution):
  234. return self.AitoffTransform(resolution)
  235. class HammerAxes(GeoAxes):
  236. name = 'hammer'
  237. class HammerTransform(_GeoTransform):
  238. """The base Hammer transform."""
  239. def transform_non_affine(self, values):
  240. # docstring inherited
  241. longitude, latitude = values.T
  242. half_long = longitude / 2.0
  243. cos_latitude = np.cos(latitude)
  244. sqrt2 = np.sqrt(2.0)
  245. alpha = np.sqrt(1.0 + cos_latitude * np.cos(half_long))
  246. x = (2.0 * sqrt2) * (cos_latitude * np.sin(half_long)) / alpha
  247. y = (sqrt2 * np.sin(latitude)) / alpha
  248. return np.column_stack([x, y])
  249. def inverted(self):
  250. # docstring inherited
  251. return HammerAxes.InvertedHammerTransform(self._resolution)
  252. class InvertedHammerTransform(_GeoTransform):
  253. def transform_non_affine(self, values):
  254. # docstring inherited
  255. x, y = values.T
  256. z = np.sqrt(1 - (x / 4) ** 2 - (y / 2) ** 2)
  257. longitude = 2 * np.arctan((z * x) / (2 * (2 * z ** 2 - 1)))
  258. latitude = np.arcsin(y*z)
  259. return np.column_stack([longitude, latitude])
  260. def inverted(self):
  261. # docstring inherited
  262. return HammerAxes.HammerTransform(self._resolution)
  263. def __init__(self, *args, **kwargs):
  264. self._longitude_cap = np.pi / 2.0
  265. super().__init__(*args, **kwargs)
  266. self.set_aspect(0.5, adjustable='box', anchor='C')
  267. self.clear()
  268. def _get_core_transform(self, resolution):
  269. return self.HammerTransform(resolution)
  270. class MollweideAxes(GeoAxes):
  271. name = 'mollweide'
  272. class MollweideTransform(_GeoTransform):
  273. """The base Mollweide transform."""
  274. def transform_non_affine(self, values):
  275. # docstring inherited
  276. def d(theta):
  277. delta = (-(theta + np.sin(theta) - pi_sin_l)
  278. / (1 + np.cos(theta)))
  279. return delta, np.abs(delta) > 0.001
  280. longitude, latitude = values.T
  281. clat = np.pi/2 - np.abs(latitude)
  282. ihigh = clat < 0.087 # within 5 degrees of the poles
  283. ilow = ~ihigh
  284. aux = np.empty(latitude.shape, dtype=float)
  285. if ilow.any(): # Newton-Raphson iteration
  286. pi_sin_l = np.pi * np.sin(latitude[ilow])
  287. theta = 2.0 * latitude[ilow]
  288. delta, large_delta = d(theta)
  289. while np.any(large_delta):
  290. theta[large_delta] += delta[large_delta]
  291. delta, large_delta = d(theta)
  292. aux[ilow] = theta / 2
  293. if ihigh.any(): # Taylor series-based approx. solution
  294. e = clat[ihigh]
  295. d = 0.5 * (3 * np.pi * e**2) ** (1.0/3)
  296. aux[ihigh] = (np.pi/2 - d) * np.sign(latitude[ihigh])
  297. xy = np.empty(values.shape, dtype=float)
  298. xy[:, 0] = (2.0 * np.sqrt(2.0) / np.pi) * longitude * np.cos(aux)
  299. xy[:, 1] = np.sqrt(2.0) * np.sin(aux)
  300. return xy
  301. def inverted(self):
  302. # docstring inherited
  303. return MollweideAxes.InvertedMollweideTransform(self._resolution)
  304. class InvertedMollweideTransform(_GeoTransform):
  305. def transform_non_affine(self, values):
  306. # docstring inherited
  307. x, y = values.T
  308. # from Equations (7, 8) of
  309. # https://mathworld.wolfram.com/MollweideProjection.html
  310. theta = np.arcsin(y / np.sqrt(2))
  311. longitude = (np.pi / (2 * np.sqrt(2))) * x / np.cos(theta)
  312. latitude = np.arcsin((2 * theta + np.sin(2 * theta)) / np.pi)
  313. return np.column_stack([longitude, latitude])
  314. def inverted(self):
  315. # docstring inherited
  316. return MollweideAxes.MollweideTransform(self._resolution)
  317. def __init__(self, *args, **kwargs):
  318. self._longitude_cap = np.pi / 2.0
  319. super().__init__(*args, **kwargs)
  320. self.set_aspect(0.5, adjustable='box', anchor='C')
  321. self.clear()
  322. def _get_core_transform(self, resolution):
  323. return self.MollweideTransform(resolution)
  324. class LambertAxes(GeoAxes):
  325. name = 'lambert'
  326. class LambertTransform(_GeoTransform):
  327. """The base Lambert transform."""
  328. def __init__(self, center_longitude, center_latitude, resolution):
  329. """
  330. Create a new Lambert transform. Resolution is the number of steps
  331. to interpolate between each input line segment to approximate its
  332. path in curved Lambert space.
  333. """
  334. _GeoTransform.__init__(self, resolution)
  335. self._center_longitude = center_longitude
  336. self._center_latitude = center_latitude
  337. def transform_non_affine(self, values):
  338. # docstring inherited
  339. longitude, latitude = values.T
  340. clong = self._center_longitude
  341. clat = self._center_latitude
  342. cos_lat = np.cos(latitude)
  343. sin_lat = np.sin(latitude)
  344. diff_long = longitude - clong
  345. cos_diff_long = np.cos(diff_long)
  346. inner_k = np.maximum( # Prevent divide-by-zero problems
  347. 1 + np.sin(clat)*sin_lat + np.cos(clat)*cos_lat*cos_diff_long,
  348. 1e-15)
  349. k = np.sqrt(2 / inner_k)
  350. x = k * cos_lat*np.sin(diff_long)
  351. y = k * (np.cos(clat)*sin_lat - np.sin(clat)*cos_lat*cos_diff_long)
  352. return np.column_stack([x, y])
  353. def inverted(self):
  354. # docstring inherited
  355. return LambertAxes.InvertedLambertTransform(
  356. self._center_longitude,
  357. self._center_latitude,
  358. self._resolution)
  359. class InvertedLambertTransform(_GeoTransform):
  360. def __init__(self, center_longitude, center_latitude, resolution):
  361. _GeoTransform.__init__(self, resolution)
  362. self._center_longitude = center_longitude
  363. self._center_latitude = center_latitude
  364. def transform_non_affine(self, values):
  365. # docstring inherited
  366. x, y = values.T
  367. clong = self._center_longitude
  368. clat = self._center_latitude
  369. p = np.maximum(np.hypot(x, y), 1e-9)
  370. c = 2 * np.arcsin(0.5 * p)
  371. sin_c = np.sin(c)
  372. cos_c = np.cos(c)
  373. latitude = np.arcsin(cos_c*np.sin(clat) +
  374. ((y*sin_c*np.cos(clat)) / p))
  375. longitude = clong + np.arctan(
  376. (x*sin_c) / (p*np.cos(clat)*cos_c - y*np.sin(clat)*sin_c))
  377. return np.column_stack([longitude, latitude])
  378. def inverted(self):
  379. # docstring inherited
  380. return LambertAxes.LambertTransform(
  381. self._center_longitude,
  382. self._center_latitude,
  383. self._resolution)
  384. def __init__(self, *args, center_longitude=0, center_latitude=0, **kwargs):
  385. self._longitude_cap = np.pi / 2
  386. self._center_longitude = center_longitude
  387. self._center_latitude = center_latitude
  388. super().__init__(*args, **kwargs)
  389. self.set_aspect('equal', adjustable='box', anchor='C')
  390. self.clear()
  391. def clear(self):
  392. # docstring inherited
  393. super().clear()
  394. self.yaxis.set_major_formatter(NullFormatter())
  395. def _get_core_transform(self, resolution):
  396. return self.LambertTransform(
  397. self._center_longitude,
  398. self._center_latitude,
  399. resolution)
  400. def _get_affine_transform(self):
  401. return Affine2D() \
  402. .scale(0.25) \
  403. .translate(0.5, 0.5)