Skip to content

Adjust Time.argsort() for slightly faster performance. #11907

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

Merged
taldcroft merged 2 commits into from
Jul 19, 2021
Merged

Adjust Time.argsort() for slightly faster performance. #11907

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
30f678e
Adjust Time.argsort() for slightly faster performance.
mhvk Jun 30, 2021
888deb9
Add hypothesis tests for argmin, argmax, argsort
mhvk Jul 1, 2021
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
27 changes: 14 additions & 13 deletions astropy/time/core.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1195,10 +1195,9 @@ def argmin(self, axis=None, out=None):
that the full precision given by the two doubles ``jd1`` and ``jd2``
is used. See :func:`~numpy.argmin` for detailed documentation.
"""
# first get the minimum at normal precision.
jd = self.jd1 + self.jd2

approx = np.min(jd, axis, keepdims=True)
# First get the minimum at normal precision.
jd1, jd2 = self.jd1, self.jd2
approx = np.min(jd1 + jd2, axis, keepdims=True)

# Approx is very close to the true minimum, and by subtracting it at
# full precision, all numbers near 0 can be represented correctly,
Expand All @@ -1208,7 +1207,7 @@ def argmin(self, axis=None, out=None):
# which translates to:
# approx_jd1, approx_jd2 = day_frac(approx, 0.)
# dt = (self.jd1 - approx_jd1) + (self.jd2 - approx_jd2)
dt = (self.jd1 - approx) + self.jd2
dt = (jd1 - approx) + jd2

return dt.argmin(axis, out)

Expand All @@ -1220,11 +1219,10 @@ def argmax(self, axis=None, out=None):
is used. See :func:`~numpy.argmax` for detailed documentation.
"""
# For procedure, see comment on argmin.
jd = self.jd1 + self.jd2

approx = np.max(jd, axis, keepdims=True)
jd1, jd2 = self.jd1, self.jd2
approx = np.max(jd1 + jd2, axis, keepdims=True)

dt = (self.jd1 - approx) + self.jd2
dt = (jd1 - approx) + jd2

return dt.argmax(axis, out)

Expand All @@ -1236,12 +1234,15 @@ def argsort(self, axis=-1):
is used, and that corresponding attributes are copied. Internally,
it uses :func:`~numpy.lexsort`, and hence no sort method can be chosen.
"""
jd_approx = self.jd
jd_remainder = (self - self.__class__(jd_approx, format='jd', scale=self.scale)).jd

@taldcroft taldcroft Jul 1, 2021

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Do you recall why it was written this way? That is likely the performance bottleneck here. I can't think of why it should be done this way, but sometimes the original coder (even when it was yourself) understood something that has since been forgotten. 😄

@mhvk mhvk Jul 1, 2021

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Indeed! This line has been here from the start (#3681), and argmin and argmax were treated differently already too. I have no recollection why at the time I did not think the same logic that is written down for argmin might not hold for argsort. Perhaps it is that at the time we did not as strictly guarantee that jd1 would be exact (being integer + 0.5)? In any case, because jd1 is exactly representable, this means there is no point going through initializing ourselves and doing the subtraction that way; one can just write it out like for argmin and realize that jd1-approx cannot loose bits, and the subsequent +jd2 cannot loose bits either. Indeed, if we completely trusted that jd1, jd2 had absolutely always gone through day_frac, we should be able to sort on jd1, jd2. But I'm just a bit wary that rounding on +/-0.5 can sometimes throw a spanner in the works.

Anyway, it is right to be worried, so I wrote an extra test for test_precision using hypothesis, which will warn us if things fail (the existing tests are actually quite good, I think, but fine to be safe here!)

# For procedure, see comment on argmin.
jd1, jd2 = self.jd1, self.jd2
approx = jd1 + jd2
remainder = (jd1 - approx) + jd2

if axis is None:
return np.lexsort((jd_remainder.ravel(), jd_approx.ravel()))
return np.lexsort((remainder.ravel(), approx.ravel()))
else:
return np.lexsort(keys=(jd_remainder, jd_approx), axis=axis)
return np.lexsort(keys=(remainder, approx), axis=axis)

def min(self, axis=None, out=None, keepdims=False):
"""Minimum along a given axis.
Expand Down
22 changes: 22 additions & 0 deletions astropy/time/tests/test_precision.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -498,6 +498,28 @@ def test_jd_add_subtract_round_trip(scale, jds, delta):
raise


@given(scale=sampled_from(STANDARD_TIME_SCALES),
jds=reasonable_jd(),
delta=floats(-3*tiny, 3*tiny))
@example(scale='tai', jds=(0.0, 3.5762786865234384), delta=2.220446049250313e-16)
def test_time_argminmaxsort(scale, jds, delta):
jd1, jd2 = jds
t = Time(jd1, jd2+np.array([0, delta]), scale=scale, format="jd")
imin = t.argmin()
imax = t.argmax()
isort = t.argsort()
diff = (t.jd1[1]-t.jd1[0]) + (t.jd2[1]-t.jd2[0])
if diff < 0: # item 1 smaller
assert delta < 0
assert imin == 1 and imax == 0 and np.all(isort == [1, 0])
elif diff == 0: # identical within precision
assert abs(delta) <= tiny
assert imin == 0 and imax == 0 and np.all(isort == [0, 1])
else:
assert delta > 0
assert imin == 0 and imax == 1 and np.all(isort == [0, 1])


@given(sampled_from(STANDARD_TIME_SCALES), unreasonable_jd(), unreasonable_jd())
@example(scale='utc',
jds_a=(2455000.0, 0.0),
Expand Down