Skip to content

Improve speed in projections/geo.py #22677

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Improve speed in projections/geo.py #22677

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 33 additions & 28 deletions lib/matplotlib/projections/geo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import math
import numpy as np

from matplotlib import _api, rcParams
Expand Down Expand Up @@ -52,8 +53,8 @@ def cla(self):

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)
Axes.set_xlim(self, -math.pi, math.pi)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can just use np.pi in most of these places, because set_xlim/etc. will directly convert everything to numpy scalars anyways, obviating any speedup. (Using python scalars is only useful if you do some computations with them, and even then, probably the gain of writing math.pi*2 below is obviated by the additional builtin->numpy conversion.)

Axes.set_ylim(self, -math.pi / 2.0, math.pi / 2.0)

def _set_lim_and_transforms(self):
# A (possibly non-linear) projection on the (already scaled) data
Expand Down Expand Up @@ -88,7 +89,7 @@ def _set_lim_and_transforms(self):
Affine2D().translate(0, -4)

# This is the transform for latitude ticks.
yaxis_stretch = Affine2D().scale(np.pi * 2, 1).translate(-np.pi, 0)
yaxis_stretch = Affine2D().scale(math.pi * 2, 1).translate(-math.pi, 0)
yaxis_space = Affine2D().scale(1, 1.1)
self._yaxis_transform = \
yaxis_stretch + \
Expand All @@ -108,8 +109,8 @@ def _set_lim_and_transforms(self):

def _get_affine_transform(self):
transform = self._get_core_transform(1)
xscale, _ = transform.transform((np.pi, 0))
_, yscale = transform.transform((0, np.pi/2))
xscale, _ = transform.transform((math.pi, 0))
_, yscale = transform.transform((0, math.pi / 2))
return Affine2D() \
.scale(0.5 / xscale, 0.5 / yscale) \
.translate(0.5, 0.5)
Expand Down Expand Up @@ -287,7 +288,7 @@ def inverted(self):
return AitoffAxes.AitoffTransform(self._resolution)

def __init__(self, *args, **kwargs):
self._longitude_cap = np.pi / 2.0
self._longitude_cap = math.pi / 2
super().__init__(*args, **kwargs)
self.set_aspect(0.5, adjustable='box', anchor='C')
self.cla()
Expand All @@ -305,11 +306,11 @@ class HammerTransform(_GeoTransform):
def transform_non_affine(self, ll):
# docstring inherited
longitude, latitude = ll.T
half_long = longitude / 2.0
half_long = longitude / 2
cos_latitude = np.cos(latitude)
sqrt2 = np.sqrt(2.0)
sqrt2 = 2 ** (1/2)
alpha = np.sqrt(1.0 + cos_latitude * np.cos(half_long))
x = (2.0 * sqrt2) * (cos_latitude * np.sin(half_long)) / alpha
x = (2 * sqrt2) * (cos_latitude * np.sin(half_long)) / alpha
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd just write 2**(3/2) here and 2**(1/2) below and drop the sqrt(2) variable, as noticed elsewhere this will get inlined anyways.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I follow the argument in #22678 (comment) that 2**0.5 reads better than 2**(1/2).

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That's fine with me too.

y = (sqrt2 * np.sin(latitude)) / alpha
return np.column_stack([x, y])

Expand All @@ -324,15 +325,15 @@ def transform_non_affine(self, xy):
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)
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
self._longitude_cap = math.pi / 2
super().__init__(*args, **kwargs)
self.set_aspect(0.5, adjustable='box', anchor='C')
self.cla()
Expand All @@ -351,18 +352,18 @@ def transform_non_affine(self, ll):
# docstring inherited
def d(theta):
delta = (-(theta + np.sin(theta) - pi_sin_l)
/ (1 + np.cos(theta)))
/ (1.0 + np.cos(theta)))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's stick with 1 (and likewise 2 for 2.0 below) unless it matters significantly (I doubt so...); it reads better IMO (e.g. it matches the math formula).

return delta, np.abs(delta) > 0.001

longitude, latitude = ll.T

clat = np.pi/2 - np.abs(latitude)
pi_half = math.pi / 2
clat = pi_half - 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])
pi_sin_l = math.pi * np.sin(latitude[ilow])
theta = 2.0 * latitude[ilow]
delta, large_delta = d(theta)
while np.any(large_delta):
Expand All @@ -372,12 +373,13 @@ def d(theta):

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])
d = np.cbrt(3.0 * math.pi * e ** 2) / 2
aux[ihigh] = (pi_half - 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)
sqrt2 = 2 ** (1 / 2)
xy[:, 0] = (sqrt2 / pi_half) * longitude * np.cos(aux)
xy[:, 1] = sqrt2 * np.sin(aux)

return xy

Expand All @@ -392,17 +394,18 @@ def transform_non_affine(self, xy):
x, y = xy.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)
sqrt2 = 2 ** (1 / 2)
theta = np.arcsin(y / sqrt2)
longitude = (math.pi / (2.0 * sqrt2)) * x / np.cos(theta)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

again sqrt2 doesn't warrant being a separate variable; the compiler will inline 2**(1/2) in theta and 2**(3/2) in longitude; and again 2.0 -> 2

latitude = np.arcsin((2.0 * theta + np.sin(2.0 * theta)) / math.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
self._longitude_cap = math.pi / 2
super().__init__(*args, **kwargs)
self.set_aspect(0.5, adjustable='box', anchor='C')
self.cla()
Expand Down Expand Up @@ -434,15 +437,17 @@ def transform_non_affine(self, ll):
clat = self._center_latitude
cos_lat = np.cos(latitude)
sin_lat = np.sin(latitude)
cos_clat = np.cos(clat)
sin_clat = np.sin(clat)
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,
1.0 + sin_clat*sin_lat + cos_clat*cos_lat*cos_diff_long,
1e-15)
k = np.sqrt(2 / inner_k)
k = np.sqrt(2.0 / 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)
y = k * (cos_clat*sin_lat - sin_clat*cos_lat*cos_diff_long)

return np.column_stack([x, y])

Expand All @@ -466,7 +471,7 @@ def transform_non_affine(self, xy):
clong = self._center_longitude
clat = self._center_latitude
p = np.maximum(np.hypot(x, y), 1e-9)
c = 2 * np.arcsin(0.5 * p)
c = 2.0 * np.arcsin(0.5 * p)
sin_c = np.sin(c)
cos_c = np.cos(c)

Expand All @@ -485,7 +490,7 @@ def inverted(self):
self._resolution)

def __init__(self, *args, center_longitude=0, center_latitude=0, **kwargs):
self._longitude_cap = np.pi / 2
self._longitude_cap = math.pi / 2
self._center_longitude = center_longitude
self._center_latitude = center_latitude
super().__init__(*args, **kwargs)
Expand Down