Make STIXFrame work nicely with sunpy Maps (#90)

* Make STIXFrame work nicely with `Map`. * Apply suggestions from code review * Fix docs and STIX map axis labels * Plot axis --------- Co-authored-by: DanRyanIrish <[email protected]>
TCDSolar · Apr 30, 2024 · d44a1c8 · d44a1c8
1 parent 73f8d02
commit d44a1c8
Show file tree

Hide file tree

Showing 15 changed files with 337 additions and 149 deletions.
diff --git a/changelog/90.feature.rst b/changelog/90.feature.rst
@@ -0,0 +1,2 @@
+Add a custom coordinate frame :class:`stixpy.coordinate.frames.STIXImaging` and the associated transformations (`~stixpy.coordinates.transforms.stixim_to_hpc`, `~stixpy.coordinates.transforms.hpc_to_stixim`) to and from `~sunpy.coordinates.frames.Helioprojective`.
+Also add a `stixpy.map.stix.STIXMap` for `~sunpy.map.Map` source which properly displays the :class:`stixpy.coordinate.frames.STIXImaging` coordinate frame.
diff --git a/docs/reference/coordinates.rst b/docs/reference/coordinates.rst
@@ -0,0 +1,12 @@
+.. _coordinates:
+
+Coordinates ('stixpy.coordinates')
+**********************************
+
+The ``coordinates`` submodule contains STIX specific coordinate frames and transforms.
+
+.. automodapi:: stixpy.coordinates
+
+.. automodapi:: stixpy.coordinates.frames
+
+.. automodapi:: stixpy.coordinates.transforms
diff --git a/docs/reference/frames.rst b/docs/reference/frames.rst
diff --git a/docs/reference/index.rst b/docs/reference/index.rst
@@ -7,12 +7,13 @@ Reference
 .. toctree::
    :maxdepth: 2
 
-   data
    calibration
-   timeseries
-   science
+   coordinates
+   data
+   map
    product
+   science
+   timeseries
    visualisation
-   frames
 
    ../whatsnew/index
diff --git a/docs/reference/map.rst b/docs/reference/map.rst
@@ -0,0 +1,10 @@
+.. _map:
+
+Map ('stixpy.map')
+******************
+
+The ``map`` submodule contains stix specific map source for `~sunpy.map.Map`
+
+.. automodapi:: stixpy.map
+
+.. automodapi:: stixpy.map.stix
diff --git a/examples/imaging_demo.py b/examples/imaging_demo.py
@@ -27,16 +27,9 @@
     create_visibility,
     get_uv_points_data,
 )
-# from stixpy.frames import get_hpc_info
 from stixpy.imaging.em import em
 from stixpy.product import Product
 
-# from astropy.coordinates import SkyCoord
-# from astropy.time import Time
-# from sunpy.map import make_fitswcs_header, Map
-
-
-
 logger = logging.getLogger(__name__)
 logger.setLevel("DEBUG")
 

diff --git a/setup.cfg b/setup.cfg
@@ -68,7 +68,7 @@ filterwarnings =
 
 [isort]
 balanced_wrapping = true
-skip =
+extend_skip =
     docs/conf.py,
     stixpy/__init__.py,
 default_section = THIRDPARTY

diff --git a/stixpy/coordinates/__init__.py b/stixpy/coordinates/__init__.py
diff --git a/stixpy/coordinates/frames.py b/stixpy/coordinates/frames.py
@@ -0,0 +1,164 @@
+import astropy
+import astropy.coordinates as coord
+import astropy.units as u
+from astropy.coordinates import QuantityAttribute
+from astropy.wcs import WCS
+from sunpy.coordinates.frameattributes import ObserverCoordinateAttribute
+from sunpy.coordinates.frames import HeliographicStonyhurst, SunPyBaseCoordinateFrame
+from sunpy.sun.constants import radius as _RSUN
+
+from stixpy.net.client import STIXClient  # noqa
+from stixpy.utils.logging import get_logger
+
+logger = get_logger(__name__, "DEBUG")
+
+__all__ = ["STIXImaging"]
+
+STIX_X_SHIFT = 26.1 * u.arcsec  # fall back to this when non sas solution available
+STIX_Y_SHIFT = 58.2 * u.arcsec  # fall back to this when non sas solution available
+STIX_X_OFFSET = 60.0 * u.arcsec  # remaining offset after SAS solution
+STIX_Y_OFFSET = 8.0 * u.arcsec  # remaining offset after SAS solution
+
+STIX_X_CTYPE = "SXLN-TAN"
+STIX_Y_CTYPE = "SXLT-TAN"
+
+
+class STIXImaging(SunPyBaseCoordinateFrame):
+    r"""
+    STIX Imaging Frame
+
+    - ``Tx`` (aka "theta_x") along the pixel rows (e.g. 0, 1, 2, 3; 4, 5, 6, 8).
+    - ``Ty`` (aka "theta_y") along the pixel columns (e.g. A, B, C, D).
+    - ``distance`` is the Sun-observer distance.
+
+    Aligned with STIX 'pixels' +X corresponds direction along pixel row toward pixels
+    0, 4 and +Y corresponds direction along columns towards pixels 0, 1, 2, 3.
+
+    .. code-block:: text
+
+      Col\Row
+            ---------
+           | 3  | 7  |
+        D  |   _|_   |
+           |  |11 |  |
+            ---------
+           | 2  | 6  |
+        C  |   _|_   |
+           |  |10 |  |
+            ---------
+           | 1  | 5  |
+        B  |   _|_   |
+           |  | 9 |  |
+            ---------
+           | 0  | 4  |
+        A  |   _|_   |
+           |  | 8 |  |
+            ---------
+
+    """
+    observer = ObserverCoordinateAttribute(HeliographicStonyhurst)
+    rsun = QuantityAttribute(default=_RSUN, unit=u.km)
+
+    frame_specific_representation_info = {
+        coord.SphericalRepresentation: [
+            coord.RepresentationMapping("lon", "Tx", u.arcsec),
+            coord.RepresentationMapping("lat", "Ty", u.arcsec),
+            coord.RepresentationMapping("distance", "distance"),
+        ],
+        coord.UnitSphericalRepresentation: [coord.RepresentationMapping('lon', 'Tx', u.arcsec),
+                                            coord.RepresentationMapping('lat', 'Ty', u.arcsec)]
+    }
+
+
+def stix_wcs_to_frame(wcs):
+    r"""
+    This function registers the coordinate frames to their FITS-WCS coordinate
+    type values in the `astropy.wcs.utils.wcs_to_celestial_frame` registry.
+
+    Parameters
+    ----------
+    wcs : `astropy.wcs.WCS`
+
+    Returns
+    -------
+    `astropy.coordinates.BaseCoordinateFrame`
+    """
+    if hasattr(wcs, "coordinate_frame"):
+        return wcs.coordinate_frame
+
+    # Not a STIX wcs bail out early
+    if set(wcs.wcs.ctype) != {STIX_X_CTYPE, STIX_Y_CTYPE}:
+        return None
+
+    dateobs = wcs.wcs.dateobs
+
+    rsun = wcs.wcs.aux.rsun_ref
+    if rsun is not None:
+        rsun *= u.m
+
+    hgs_longitude = wcs.wcs.aux.hgln_obs
+    hgs_latitude = wcs.wcs.aux.hglt_obs
+    hgs_distance = wcs.wcs.aux.dsun_obs
+
+    observer = HeliographicStonyhurst(lat=hgs_latitude * u.deg,
+                                      lon=hgs_longitude * u.deg,
+                                      radius=hgs_distance * u.m,
+                                      obstime=dateobs,
+                                      rsun=rsun)
+
+    frame_args = {'obstime': dateobs,
+                  'observer': observer,
+                  'rsun': rsun}
+
+    return STIXImaging(**frame_args)
+
+
+def stix_frame_to_wcs(frame, projection='TAN'):
+    r"""
+    For a given frame, this function returns the corresponding WCS object.
+
+    It registers the WCS coordinates types from their associated frame in the
+    `astropy.wcs.utils.celestial_frame_to_wcs` registry.
+
+    Parameters
+    ----------
+    frame : `astropy.coordinates.BaseCoordinateFrame`
+    projection : `str`, optional
+
+    Returns
+    -------
+    `astropy.wcs.WCS`
+    """
+    # Bail out early if not STIXImaging frame
+    if not isinstance(frame, STIXImaging):
+        return None
+
+    wcs = WCS(naxis=2)
+    wcs.wcs.aux.rsun_ref = frame.rsun.to_value(u.m)
+
+    # Sometimes obs_coord can be a SkyCoord, so convert down to a frame
+    obs_frame = frame.observer
+    if hasattr(obs_frame, 'frame'):
+        obs_frame = frame.observer.frame
+
+    wcs.wcs.aux.hgln_obs = obs_frame.lon.to_value(u.deg)
+    wcs.wcs.aux.hglt_obs = obs_frame.lat.to_value(u.deg)
+    wcs.wcs.aux.dsun_obs = obs_frame.radius.to_value(u.m)
+
+    wcs.wcs.dateobs = frame.obstime.utc.iso
+    wcs.wcs.cunit = ['arcsec', 'arcsec']
+    wcs.wcs.ctype = [STIX_X_CTYPE, STIX_Y_CTYPE]
+
+    return wcs
+
+
+# Remove once min version of sunpy has https://github.com/sunpy/sunpy/pull/7594
+astropy.wcs.utils.WCS_FRAME_MAPPINGS.insert(1, [stix_wcs_to_frame])
+astropy.wcs.utils.FRAME_WCS_MAPPINGS.insert(1, [stix_frame_to_wcs])
+
+STIX_CTYPE_TO_UCD1 = {
+    "SXLT": "custom:pos.stiximaging.lat",
+    "SXLN": "custom:pos.stiximaging.lon",
+    "SXRZ": "custom:pos.stiximaging.z"
+}
+astropy.wcs.wcsapi.fitswcs.CTYPE_TO_UCD1.update(STIX_CTYPE_TO_UCD1)
diff --git a/stixpy/coordinates/tests/__init__.py b/stixpy/coordinates/tests/__init__.py
diff --git a/stixpy/coordinates/tests/test_frames.py b/stixpy/coordinates/tests/test_frames.py
@@ -0,0 +1,100 @@
+import astropy.units as u
+import numpy as np
+import pytest
+from astropy.coordinates import SkyCoord
+from astropy.wcs import WCS
+from sunpy.coordinates import get_horizons_coord
+from sunpy.coordinates.frames import HeliographicStonyhurst, Helioprojective
+from sunpy.map import Map, make_fitswcs_header
+from sunpy.map.mapbase import SpatialPair
+
+from stixpy.coordinates.frames import STIXImaging, stix_frame_to_wcs, stix_wcs_to_frame
+from stixpy.map.stix import STIXMap
+
+
+@pytest.fixture
+def stix_wcs():
+    w = WCS(naxis=2)
+
+    w.wcs.dateobs = '2024-01-01 00:00:00.000'
+    w.wcs.crpix = [10, 20]
+    w.wcs.cdelt = np.array([2, 2])
+    w.wcs.crval = [0, 0]
+    w.wcs.ctype = ["SXLN-TAN", "SXLT-TAN"]
+
+    w.wcs.aux.hgln_obs = 10
+    w.wcs.aux.hglt_obs = 20
+    w.wcs.aux.dsun_obs = 1.5e11
+
+    return w
+
+
+@pytest.fixture
+def stix_frame():
+    obstime = '2024-01-01'
+    observer = HeliographicStonyhurst(10 * u.deg,
+                                      20 * u.deg,
+                                      1.5e11 * u.m,
+                                      obstime=obstime)
+
+    frame_args = {'obstime': obstime,
+                  'observer': observer,
+                  'rsun': 695_700_000 * u.m}
+
+    frame = STIXImaging(**frame_args)
+    return frame
+
+
+def test_stix_wcs_to_frame(stix_wcs):
+    frame = stix_wcs_to_frame(stix_wcs)
+    assert isinstance(frame, STIXImaging)
+
+    assert frame.obstime.isot == '2024-01-01T00:00:00.000'
+    assert frame.rsun == 695700 * u.km
+    assert frame.observer == HeliographicStonyhurst(10 * u.deg,
+                                                    20 * u.deg,
+                                                    1.5e11 * u.m,
+                                                    obstime='2024-01-01T00:00:00.000')
+
+
+def test_stix_wcs_to_frame_none():
+    w = WCS(naxis=2)
+    w.wcs.ctype = ['ham', 'cheese']
+    frame = stix_wcs_to_frame(w)
+
+    assert frame is None
+
+
+def test_stix_frame_to_wcs(stix_frame):
+    wcs = stix_frame_to_wcs(stix_frame)
+
+    assert isinstance(wcs, WCS)
+    assert wcs.wcs.ctype[0] == 'SXLN-TAN'
+    assert wcs.wcs.cunit[0] == 'arcsec'
+    assert wcs.wcs.dateobs == '2024-01-01 00:00:00.000'
+
+    assert wcs.wcs.aux.rsun_ref == stix_frame.rsun.to_value(u.m)
+    assert wcs.wcs.aux.dsun_obs == 1.5e11
+    assert wcs.wcs.aux.hgln_obs == 10
+    assert wcs.wcs.aux.hglt_obs == 20
+
+
+def test_stix_frame_to_wcs_none():
+    wcs = stix_frame_to_wcs(Helioprojective())
+    assert wcs is None
+
+@pytest.mark.remote_data
+def test_stix_frame_map():
+    data = np.random.rand(512, 512)
+    obstime = '2023-01-01 12:00:00'
+    solo = get_horizons_coord('solo', time=obstime)
+    coord = SkyCoord(0 * u.arcsec, 0 * u.arcsec, obstime=obstime, observer=solo, frame=STIXImaging)
+    header = make_fitswcs_header(data, coord, scale=[8, 8] * u.arcsec / u.pix, telescope='STIX',
+                                 instrument='STIX', observatory='Solar Orbiter')
+    smap = Map((data, header))
+    assert isinstance(smap, STIXMap)
+    assert smap.coordinate_system == SpatialPair(axis1='SXLN-TAN', axis2='SXLT-TAN')
+    assert isinstance(smap.coordinate_frame, STIXImaging)
+    smap.plot()
+    smap.draw_limb()
+    smap.draw_grid()
diff --git a/stixpy/tests/test_frames.py → stixpy/coordinates/tests/test_transforms.py b/stixpy/tests/test_frames.py → stixpy/coordinates/tests/test_transforms.py
@@ -5,9 +5,9 @@
 import pytest
 from astropy.tests.helper import assert_quantity_allclose
 from astropy.time import Time
-from sunpy.coordinates.frames import HeliographicStonyhurst, Helioprojective
+from sunpy.coordinates import HeliographicStonyhurst, Helioprojective
 
-from stixpy.frames import STIXImaging
+from stixpy.coordinates.frames import STIXImaging
 
 
 @pytest.mark.skip(reason="Test data maybe incorrect")
@@ -28,7 +28,7 @@ def test_transforms_with_know_values():
     assert_quantity_allclose(earth_coord.Ty, stix_to_earth.Ty)
 
 
-@mock.patch("stixpy.frames._get_aux_data")
+@mock.patch("stixpy.coordinates.transforms._get_aux_data")
 def test_hpc_to_stx(mock):
     # fake some data to make check easier
     obstime = Time("2021-05-22T02:52:00")
@@ -39,11 +39,11 @@ def test_hpc_to_stx(mock):
     stix_frame = STIXImaging(0 * u.arcsec, 0 * u.arcsec, obstime=obstime)
     stix_coord = solo_hpc_coord.transform_to(stix_frame)
     # should match the offset -8, 60
-    assert_quantity_allclose(stix_coord.x, -8 * u.arcsec)
-    assert_quantity_allclose(stix_coord.y, 60 * u.arcsec)
+    assert_quantity_allclose(stix_coord.Tx, -8 * u.arcsec)
+    assert_quantity_allclose(stix_coord.Ty, 60 * u.arcsec)
 
 
-@mock.patch("stixpy.frames._get_aux_data")
+@mock.patch("stixpy.coordinates.transforms._get_aux_data")
 def test_hpc_to_stx_no_sas(mock):
     # fake some data to make check easier
     obstime = Time("2021-05-22T02:52:00")
@@ -55,5 +55,5 @@ def test_hpc_to_stx_no_sas(mock):
     with pytest.warns(Warning, match="SAS solution not"):
         stix_coord = solo_hpc_coord.transform_to(stix_frame)
         # should match the offset -8, 60 added to yaw and pitch 10, 10
-        assert_quantity_allclose(stix_coord.x, (10 - 8) * u.arcsec)
-        assert_quantity_allclose(stix_coord.y, (10 + 60) * u.arcsec)
+        assert_quantity_allclose(stix_coord.Tx, (10 - 8) * u.arcsec)
+        assert_quantity_allclose(stix_coord.Ty, (10 + 60) * u.arcsec)
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1,2 @@
		Add a custom coordinate frame :class:`stixpy.coordinate.frames.STIXImaging` and the associated transformations (`~stixpy.coordinates.transforms.stixim_to_hpc`, `~stixpy.coordinates.transforms.hpc_to_stixim`) to and from `~sunpy.coordinates.frames.Helioprojective`.
		Also add a `stixpy.map.stix.STIXMap` for `~sunpy.map.Map` source which properly displays the :class:`stixpy.coordinate.frames.STIXImaging` coordinate frame.