Speed up HGC > HCI transformation

[dstansby]

In pfsspy field lines are often traced in a Heliographic Carrington frame of reference. In sunkit-pyvsita, the plotting is all done in a Heliocentric Inertial frame, necessitating a transformation between lots (sometimes 180 x 360 ~ 65k) of SkyCoord objects that range in length (but are typically around 50 points long). This transformation is relatively slow - the example below with 1000 dummy field lines takes 40 seconds to run on my laptop.

It would be good to speed this transformation up. The profile is:

39.253 <module>  test.py:1
└─ 39.241 transform_to  astropy/coordinates/sky_coordinate.py:591
   ├─ 36.125 __call__  astropy/coordinates/transformations.py:1463
   │  └─ 35.795 __call__  astropy/coordinates/transformations.py:999
   │     └─ 35.774 wrapped_func  sunpy/coordinates/transformations.py:239
   │        ├─ 26.366 hgc_to_hgs  sunpy/coordinates/transformations.py:403
   │        │  ├─ 24.166 _rotation_matrix_hgs_to_hgc  sunpy/coordinates/transformations.py:358
   │        │  │  ├─ 21.701 L0  sunpy/coordinates/sun.py:536
   │        │  │  │  ├─ 9.884 _detilt_lon  sunpy/coordinates/sun.py:517
   │        │  │  │  │  ├─ 8.810 __getattr__  astropy/coordinates/sky_coordinate.py:823
   │        │  │  │  │  │  └─ 8.386 transform_to  astropy/coordinates/sky_coordinate.py:591
   │        │  │  │  │  │     ├─ 5.366 __call__  astropy/coordinates/transformations.py:1463
   │        │  │  │  │  │     │  └─ 5.277 __call__  astropy/coordinates/transformations.py:999
   │        │  │  │  │  │     │     └─ 5.260 wrapped_func  sunpy/coordinates/transformations.py:239
   │        │  │  │  │  │     │        └─ 5.248 hgs_to_hcrs  sunpy/coordinates/transformations.py:710
   │        │  │  │  │  │     │           ├─ 3.731 _affine_params_hcrs_to_hgs  sunpy/coordinates/transformations.py:652
   │        │  │  │  │  │     │           │  ├─ 1.496 _rotation_matrix_reprs_to_xz_about_z  sunpy/coordinates/transformations.py:612
   │        │  │  │  │  │     │           │  │  └─ 1.023 _rotation_matrix_reprs_to_reprs  sunpy/coordinates/transformations.py:591
   │        │  │  │  │  │     │           │  ├─ 1.073 get_body_barycentric  astropy/coordinates/solar_system.py:344
   │        │  │  │  │  │     │           │  │  └─ 1.073 _get_body_barycentric_posvel  astropy/coordinates/solar_system.py:182
   │        │  │  │  │  │     │           │  │     └─ 0.616 get_jd12  astropy/coordinates/builtin_frames/utils.py:96
   │        │  │  │  │  │     │           │  │        └─ 0.527 __getattr__  astropy/time/core.py:1324
   │        │  │  │  │  │     │           │  │           └─ 0.411 _set_scale  astropy/time/core.py:538
   │        │  │  │  │  │     │           │  └─ 0.473 transform  astropy/coordinates/representation.py:1353
   │        │  │  │  │  │     │           └─ 0.448 transform  astropy/coordinates/representation.py:1353
   │        │  │  │  │  │     ├─ 2.169 __init__  astropy/coordinates/sky_coordinate.py:286
   │        │  │  │  │  │     │  ├─ 1.379 _parse_coordinate_data  astropy/coordinates/sky_coordinate_parsers.py:215
   │        │  │  │  │  │     │  │  └─ 1.350 _parse_coordinate_arg  astropy/coordinates/sky_coordinate_parsers.py:363
   │        │  │  │  │  │     │  │     └─ 0.769 represent_as  astropy/coordinates/representation.py:833
   │        │  │  │  │  │     │  │        └─ 0.764 from_cartesian  astropy/coordinates/representation.py:2016
   │        │  │  │  │  │     │  │           └─ 0.558 __init__  astropy/coordinates/representation.py:1918
   │        │  │  │  │  │     │  │              └─ 0.404 __init__  astropy/coordinates/representation.py:659
   │        │  │  │  │  │     │  └─ 0.654 __init__  astropy/coordinates/baseframe.py:318
   │        │  │  │  │  │     │     └─ 0.524 _infer_data  astropy/coordinates/baseframe.py:425
   │        │  │  │  │  │     │        └─ 0.472 __init__  astropy/coordinates/representation.py:1918
   │        │  │  │  │  │     └─ 0.418 __getattr__  astropy/coordinates/sky_coordinate.py:823
   │        │  │  │  │  └─ 0.640 represent_as  astropy/coordinates/representation.py:833
   │        │  │  │  │     └─ 0.635 from_cartesian  astropy/coordinates/representation.py:2016
   │        │  │  │  │        └─ 0.562 __init__  astropy/coordinates/representation.py:1918
   │        │  │  │  ├─ 8.376 get_earth  sunpy/coordinates/ephemeris.py:143
   │        │  │  │  │  ├─ 5.868 get_body_heliographic_stonyhurst  sunpy/coordinates/ephemeris.py:34
   │        │  │  │  │  │  └─ 5.384 transform_to  astropy/coordinates/baseframe.py:1177
   │        │  │  │  │  │     └─ 5.351 __call__  astropy/coordinates/transformations.py:1463
   │        │  │  │  │  │        ├─ 4.087 __call__  astropy/coordinates/transformations.py:999
   │        │  │  │  │  │        │  └─ 4.075 wrapped_func  sunpy/coordinates/transformations.py:239
   │        │  │  │  │  │        │     └─ 4.062 hcrs_to_hgs  sunpy/coordinates/transformations.py:690
   │        │  │  │  │  │        │        └─ 3.173 _affine_params_hcrs_to_hgs  sunpy/coordinates/transformations.py:652
   │        │  │  │  │  │        │           ├─ 1.252 _rotation_matrix_reprs_to_xz_about_z  sunpy/coordinates/transformations.py:612
   │        │  │  │  │  │        │           │  └─ 0.975 _rotation_matrix_reprs_to_reprs  sunpy/coordinates/transformations.py:591
   │        │  │  │  │  │        │           ├─ 0.549 get_body_barycentric  astropy/coordinates/solar_system.py:344
   │        │  │  │  │  │        │           │  └─ 0.545 _get_body_barycentric_posvel  astropy/coordinates/solar_system.py:182
   │        │  │  │  │  │        │           ├─ 0.495 transform  astropy/coordinates/representation.py:1353
   │        │  │  │  │  │        │           └─ 0.478 __mul__  astropy/coordinates/representation.py:442
   │        │  │  │  │  │        │              └─ 0.472 _scale_operation  astropy/coordinates/representation.py:1040
   │        │  │  │  │  │        └─ 0.942 __call__  astropy/coordinates/transformations.py:1252
   │        │  │  │  │  │           └─ 0.569 _affine_params  astropy/coordinates/transformations.py:1307
   │        │  │  │  │  │              └─ 0.565 icrs_to_hcrs  astropy/coordinates/builtin_frames/icrs_cirs_transforms.py:227
   │        │  │  │  │  │                 └─ 0.546 get_offset_sun_from_barycenter  astropy/coordinates/builtin_frames/utils.py:386
   │        │  │  │  │  └─ 1.488 __getattr__  astropy/coordinates/baseframe.py:1624
   │        │  │  │  │     └─ 1.446 represent_as  sunpy/coordinates/frames.py:142
   │        │  │  │  │        └─ 1.213 represent_as  astropy/coordinates/baseframe.py:995
   │        │  │  │  │           └─ 0.777 represent_as  astropy/coordinates/representation.py:833
   │        │  │  │  │              └─ 0.769 from_cartesian  astropy/coordinates/representation.py:2016
   │        │  │  │  │                 └─ 0.608 __init__  astropy/coordinates/representation.py:1918
   │        │  │  │  │                    └─ 0.440 __init__  astropy/coordinates/representation.py:659
   │        │  │  │  │                       └─ 0.427 __init__  astropy/coordinates/representation.py:182
   │        │  │  │  ├─ 1.381 __sub__  astropy/time/core.py:2127
   │        │  │  │  │  └─ 0.493 __init__  astropy/time/core.py:2309
   │        │  │  │  │     └─ 0.485 _init_from_vals  astropy/time/core.py:345
   │        │  │  │  │        └─ 0.437 _get_time_fmt  astropy/time/core.py:402
   │        │  │  │  │           └─ 0.424 __init__  astropy/time/formats.py:106
   │        │  │  │  ├─ 0.667 __getattr__  astropy/coordinates/sky_coordinate.py:823
   │        │  │  │  │  └─ 0.652 __getattr__  astropy/coordinates/baseframe.py:1624
   │        │  │  │  │     └─ 0.627 represent_as  sunpy/coordinates/frames.py:142
   │        │  │  │  │        └─ 0.440 represent_as  astropy/coordinates/baseframe.py:995
   │        │  │  │  └─ 0.425 __mul__  astropy/time/core.py:2426
   │        │  │  └─ 1.729 earth_distance  sunpy/coordinates/sun.py:642
   │        │  │     └─ 1.128 get_body_barycentric  astropy/coordinates/solar_system.py:344
   │        │  │        └─ 1.123 _get_body_barycentric_posvel  astropy/coordinates/solar_system.py:182
   │        │  │           └─ 0.615 get_jd12  astropy/coordinates/builtin_frames/utils.py:96
   │        │  │              └─ 0.509 __getattr__  astropy/time/core.py:1324
   │        │  │                 └─ 0.407 _set_scale  astropy/time/core.py:538
   │        │  ├─ 0.635 __getattr__  astropy/coordinates/baseframe.py:1624
   │        │  │  └─ 0.624 represent_as  sunpy/coordinates/frames.py:142
   │        │  │     └─ 0.506 represent_as  astropy/coordinates/baseframe.py:995
   │        │  ├─ 0.483 _transform_obstime  sunpy/coordinates/transformations.py:340
   │        │  │  └─ 0.474 _times_are_equal  sunpy/coordinates/transformations.py:324
   │        │  └─ 0.449 cartesian  astropy/coordinates/baseframe.py:1794
   │        │     └─ 0.446 represent_as  sunpy/coordinates/frames.py:142
   │        │        └─ 0.440 represent_as  astropy/coordinates/baseframe.py:995
   │        └─ 9.396 hgs_to_hci  sunpy/coordinates/transformations.py:962
   │           └─ 8.632 _rotation_matrix_hgs_to_hci  sunpy/coordinates/transformations.py:942
   │              ├─ 6.484 transform_to  astropy/coordinates/baseframe.py:1177
   │              │  └─ 6.449 __call__  astropy/coordinates/transformations.py:1463
   │              │     ├─ 4.010 __call__  astropy/coordinates/transformations.py:999
   │              │     │  └─ 4.004 wrapped_func  sunpy/coordinates/transformations.py:239
   │              │     │     └─ 3.996 hcrs_to_hgs  sunpy/coordinates/transformations.py:690
   │              │     │        ├─ 2.965 _affine_params_hcrs_to_hgs  sunpy/coordinates/transformations.py:652
   │              │     │        │  ├─ 1.327 _rotation_matrix_reprs_to_xz_about_z  sunpy/coordinates/transformations.py:612
   │              │     │        │  │  └─ 1.005 _rotation_matrix_reprs_to_reprs  sunpy/coordinates/transformations.py:591
   │              │     │        │  ├─ 0.476 transform  astropy/coordinates/representation.py:1353
   │              │     │        │  └─ 0.463 get_body_barycentric  astropy/coordinates/solar_system.py:344
   │              │     │        │     └─ 0.462 _get_body_barycentric_posvel  astropy/coordinates/solar_system.py:182
   │              │     │        └─ 0.420 __add__  astropy/coordinates/representation.py:467
   │              │     │           └─ 0.418 _combine_operation  astropy/coordinates/representation.py:1399
   │              │     └─ 2.045 __call__  astropy/coordinates/transformations.py:1252
   │              │        ├─ 1.070 _affine_params  astropy/coordinates/transformations.py:1307
   │              │        │  ├─ 0.543 helioecliptic_to_icrs  astropy/coordinates/builtin_frames/ecliptic_transforms.py:123
   │              │        │  └─ 0.523 icrs_to_hcrs  astropy/coordinates/builtin_frames/icrs_cirs_transforms.py:227
   │              │        │     └─ 0.513 get_offset_sun_from_barycenter  astropy/coordinates/builtin_frames/utils.py:386
   │              │        └─ 0.824 _apply_transform  astropy/coordinates/transformations.py:1097
   │              │           └─ 0.544 __add__  astropy/coordinates/representation.py:467
   │              │              └─ 0.541 _combine_operation  astropy/coordinates/representation.py:1399
   │              └─ 1.320 _rotation_matrix_reprs_to_xz_about_z  sunpy/coordinates/transformations.py:612
   │                 └─ 0.925 _rotation_matrix_reprs_to_reprs  sunpy/coordinates/transformations.py:591
   ├─ 2.187 __init__  astropy/coordinates/sky_coordinate.py:286
   │  ├─ 1.308 _parse_coordinate_data  astropy/coordinates/sky_coordinate_parsers.py:215
   │  │  └─ 1.271 _parse_coordinate_arg  astropy/coordinates/sky_coordinate_parsers.py:363
   │  │     └─ 0.729 represent_as  astropy/coordinates/representation.py:833
   │  │        └─ 0.721 from_cartesian  astropy/coordinates/representation.py:2016
   │  │           └─ 0.606 __init__  astropy/coordinates/representation.py:1918
   │  └─ 0.789 __init__  sunpy/coordinates/frames.py:127
   │     └─ 0.693 __init__  astropy/coordinates/baseframe.py:318
   │        └─ 0.573 _infer_data  astropy/coordinates/baseframe.py:425
   │           └─ 0.521 __init__  astropy/coordinates/representation.py:1918
   └─ 0.528 __getattr__  astropy/coordinates/sky_coordinate.py:823

import numpy as np
from astropy.coordinates import SkyCoord
import astropy.units as u
from sunpy.coordinates import HeliocentricInertial

nflines = 1000
coords = [SkyCoord(np.random.rand(50) * u.deg,
                   np.random.rand(50) * u.deg,
                   np.random.rand(50) * u.au,
                   frame='heliographic_carrington',
                   observer='self',
                   obstime='2020-01-01') for _ in range(nflines)]

to_frame = HeliocentricInertial()
from pyinstrument import Profiler

profiler = Profiler()
profiler.start()

for coord in coords:
    coord.transform_to(to_frame)

profiler.stop()
profiler.print()

[ayshih]

The HeliographicCarrington->HeliocentricInertial transformation can certainly be sped up. For historical reasons, that transformation goes through HeliographicStonyhurst (HGS), which is detrimental to performance because the orientation of the HGS axes depends on the location of Earth. At least half of the runtime is wasted on figuring out the location of Earth, which has no bearing on the final result.

However, the reason that this is even noticeable is because your test script is highly inefficient and amplifies the above waste by three orders of magnitude. You perform 1000 independent SkyCoord transformations, but all 1000 are actually the same rotation. Thus, ~80% of the runtime is wasted on redundantly calculating the same rotation matrices. Assuming that your test code is representative of the real situation, can you instead re-jigger sunkit-pyvista to combine everything into a single SkyCoord before performing the transformation?

[dstansby]

Yes, there's probably a way to do that. I realised that I could also set the coordinate frame in sunkit-pyvista to be HeliographicCarrington, which would avoid the transformation altogether... probably still worth leaving the issue open though if you think there's a way to speed the transformation up.