geo.py 17 KB

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