geo.py 17 KB

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