Skip to content

Wrap erfa functions with ufuncs - #7484

Closed
mhvk wants to merge 7 commits into
astropy:masterfrom
mhvk:erfa-wrapped-by-ufuncs
Closed

mhvk wants to merge 7 commits into
astropy:masterfrom
mhvk:erfa-wrapped-by-ufuncs

Conversation

@mhvk

@mhvk mhvk commented May 21, 2018 •

Copy link
Copy Markdown
Contributor

This implements the suggestion I made in #6865, of wrapping the ERFA functions inside ufuncs instead of using custom wrappers. This has the advantage that casting to arrays with the right dtype and iteration is taken care of automatically, in numpy C code (well, except that it turned out it needed a custom type resolver...). Another advantage is that classes can now override the ufuncs using __array_ufunc__ - this is important both for ensuring Quantity input is disallowed (a small bug, though one with zero impact, was in fact found this way; see first commit), and that the new Distribution class proposed by @eteq can work with the erfa functions, a first step towards making it work with SkyCoord and Time.

@jwoillez - now that I've done this, I admire your initial implementation even more!

Notes

The ufunc are still lightly wrapped by python functions as the ufunc implementation requires that, e.g., a vector uses a structured dtype, so that its 3 elements are considered one element of the array, and broadcasting can be done properly. This wrapping may need adjusting for the use in Distribution, or perhaps it makes more sense to use those structured dtypes throughout. But I think that makes more sense in a follow-up PR. Here, though, I do expose the ufunc themselves inside _erfa, so it is easier to experiment (unlike the original C wrappers, which made no sense on their own); I also think some of the small ERFA matrix and vector routines may well be faster than the numpy equivalents, given that the number of elements is always small (the numpy ones use blas under the hood, which may well introduce an overhead for small dimensions).

UPDATE: with the old implementation, functions that needed to be passed in bodies for light deflection calculations did not work correctly. Here, they are wrapped by generalized ufuncs, which works.

Initial timing

Because this removes casts to and creation of outout arrays in python, the implementation is sped up, especially for small arrays. E.g.,

import numpy as np
from astropy import _erfa as erfa
utc1 = np.array(2400000); utc2 = np.array(50000.)
%timeit erfa.utctai(utc1, utc2)
# 10000 loops, best of 3: 13 µs per loop
# with this PR: 5.48 µs
%timeit erfa.apcg13(utc1, utc2)
# 10000 loops, best of 3: 111 µs per loop
# with this PR: 70.6 µs
utc2 = np.arange(50000., 55000., 0.1)
%timeit erfa.utctai(utc1, utc2)
# 10 loops, best of 3: 15.8 ms per loop
# with this PR: 11.1 ms
%timeit erfa.apcg13(utc1, utc2)
# 1 loop, best of 3: 5.11 s per loop
# with this PR: 3.78 s

For actual conversions, the effect is more modest, but still noticeable for Time at least:

from astropy.time import Time
t = Time('2012-01-01')
%timeit del t.cache; t.tdb
# 1000 loops, best of 3: 791 µs per loop
# with this PR: 714 µs

from astropy.coordinates import SkyCoord
sc = SkyCoord(0., 10., unit='deg')
%timeit sc.transform_to('fk4')
# 100 loops, best of 3: 13.8 ms per loop
# with this PR: no change.

@mhvk mhvk added this to the v3.1 milestone May 21, 2018
@mhvk
mhvk requested review from astrofrog, eteq and jwoillez May 21, 2018 20:15
@astropy-bot

astropy-bot Bot commented May 21, 2018 •

Copy link
Copy Markdown

Hi there @mhvk 👋 - thanks for the pull request! I'm just a friendly 🤖 that checks for issues related to the changelog and making sure that this pull request is milestoned and labeled correctly. This is mainly intended for the maintainers, so if you are not a maintainer you can ignore this, and a maintainer will let you know if any action is required on your part 😃.

Everything looks good from my point of view! 👍

If there are any issues with this message, please report them here.

@astrofrog

Copy link
Copy Markdown
Member

Very nice! I can help review this soon but too jet lagged for this level right now ;) In the mean time, a simple question - how hard do you think it would be to add support for OpenMP at the C level? (now that astropy-helpers supports adding OpenMP flags). Just thought I'd ask since this is being refactored anyway, we might as well think about this.

@mhvk

mhvk commented May 21, 2018

Copy link
Copy Markdown
Contributor Author

@astrofrog - OpenMP at C level is a bit beyond me... But overall I would think that it should be possible to have it be done automatically to the extent that numpy can be compiled with OpenMP; my guess would be that it mostly needs adjustment in setup_package, to be sure compilation is done with the same flags that are used by the compilation of numpy. But I'm not sure how to actually do that! (I do think the inner loops will already be called within an enviroment that is inside being/end allow threads, etc.)

@mhvk
mhvk force-pushed the erfa-wrapped-by-ufuncs branch from d5931c2 to 5d19076 Compare May 22, 2018 18:45
@mhvk

mhvk commented May 22, 2018

Copy link
Copy Markdown
Contributor Author

Fixed the failure on 32-bit (which was very real, due to a poorly defined dtype) and rebased. Think all is OK now.

@mhvk

mhvk commented May 23, 2018

Copy link
Copy Markdown
Contributor Author

Note that the new circle-ci 32-bit failure is a false positive; see #7486 (comment)

@astrofrog

Copy link
Copy Markdown
Member

@mhvk - to help with reviewing this, could you paste the generated C code in a gist?

@mhvk

mhvk commented May 23, 2018

Copy link
Copy Markdown
Contributor Author

@mhvk

mhvk commented May 23, 2018

Copy link
Copy Markdown
Contributor Author

p.s. While you are looking at it, maybe also consider something I'm still wondering about: should all ufuncs that act on vectors/matrices just be gufuncs, so that one doesn't have to cast to the temporary structured dtypes. If this is done, the python wrappers become even thinner, just checking the status values. (Note that since it won't change the python interface, this does not have to be decided now; it is partially inspired by a numpy pull request which would bring fixed axis dimensions to the gufunc signature; see numpy/numpy#11132. (For now, I'd have to do a dimension check, but that's not too hard.)

mhvk added 5 commits May 23, 2018 14:57
Using ufuncs instead of C routines and a custom iterator has the
advantage that casting to arrays with the right dtype and
iteration is taken care of automatically, in numpy C code.
For small arrays, this means substantially faster execution.

Another advantage is that classes can now override the ufuncs
using __array_ufunc__ - this is important both for ensuring
Quantity input is disallowed (or at least treated correctly
rather than just cast to array), and that a new Distribution
class can work with the erfa functions.
Before, contents of astrom and ldbody were kept as well.
@mhvk
mhvk force-pushed the erfa-wrapped-by-ufuncs branch from 98e11c4 to 7d97bb0 Compare May 23, 2018 18:57
@mhvk

mhvk commented May 30, 2018

Copy link
Copy Markdown
Contributor Author

I've become convinced that the wrapping with gufuncs is way better, so am closing this in favour of #7502.

@pllim pllim closed this May 30, 2018
@mhvk
mhvk deleted the erfa-wrapped-by-ufuncs branch May 30, 2018 21:21
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants