Improve wcs metadata handling

[lowderchris]

Creating a pull request for a conversation on WCS handling. This might be useful before fully merging to use for generating synthetic L3 data for SDAC testing.

Before fully merging:

• Remove any comment blocks no longer needed.
• Resolve conflicts with test fits files
• improve matching of celestial and helio WCSes

work log:

• @lowderchris : 3 hours
• @jmbhughes : 3.5 hours

[codecov]

Codecov Report

Attention: 85 lines in your changes are missing coverage. Please review.

Comparison is base (633dd70) 94.32% compared to head (2367e84) 94.67%.
Report is 269 commits behind head on develop.

Files	Patch %	Lines
punchbowl/data.py	89.09%	46 Missing ⚠️
punchbowl/level3/f_corona_model.py	84.41%	12 Missing ⚠️
punchbowl/level1/deficient_pixel.py	89.04%	8 Missing ⚠️
punchbowl/level1/psf.py	64.28%	5 Missing ⚠️
punchbowl/util.py	75.00%	5 Missing ⚠️
punchbowl/level1/quartic_fit.py	88.00%	3 Missing ⚠️
punchbowl/tests/test_data.py	99.17%	3 Missing ⚠️
punchbowl/level0/decode_sqrt.py	98.63%	1 Missing ⚠️
punchbowl/level1/flow.py	96.42%	1 Missing ⚠️
punchbowl/level3/tests/test_f_corona_model.py	99.07%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           develop      #62      +/-   ##
===========================================
+ Coverage    94.32%   94.67%   +0.35%     
===========================================
  Files           28       28              
  Lines         1585     1729     +144     
===========================================
+ Hits          1495     1637     +142     
- Misses          90       92       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[jmbhughes]

Can you make codecov happy first? I see you've already noted the conflict to resolve.

[jmbhughes]

We're getting there!

[jmbhughes]

What's causing the one test to fail?

needs a new review

[jmbhughes]

We're still getting WCSes that differ by multiple degrees. We should figure out how to resolve that and tighten the tests for this.

[jmbhughes]

See this SunPy example https://docs.sunpy.org/en/stable/generated/gallery/units_and_coordinates/radec_to_hpc_map.html#sphx-glr-generated-gallery-units-and-coordinates-radec-to-hpc-map-py

A good plan would be to convert the celestial WCS to helio according to it. That might just be good enough for us.

[jmbhughes]

I experimented with a solution. This works pretty well and only introduces a few pixels error in the test.

from astropy.time import Time
from astropy.coordinates import get_sun
import astropy.units as u
from astropy.coordinates import EarthLocation, SkyCoord, GCRS
from astropy.wcs import WCS

import numpy as np

import sunpy.map
from sunpy.coordinates import frames, sun

import matplotlib.pyplot as plt


def calculate_pc_matrix(crota, cdelt):
    return np.array([[np.cos(crota), -np.sin(crota) * (cdelt[1]/cdelt[0])],
                     [np.sin(crota)*(cdelt[0]/cdelt[1]), np.cos(crota)]])


if __name__ == "__main__":
    empty_data = np.ones((4096, 4096))

    # just set the date_obs and fig
    date_obs = Time("2024-01-01T00:00:00", format='isot', scale='utc')

    # set up a celestial coordinate system centered on the Sun
    sun_radec = get_sun(date_obs)
    wcs_celestial = WCS(naxis=2)
    wcs_celestial.wcs.crpix = [2047.5, 2047.5]
    wcs_celestial.wcs.crval = [sun_radec.ra.to(u.deg).value, sun_radec.dec.to(u.deg).value]
    wcs_celestial.wcs.cdelt = np.array([-0.025, 0.025])
    wcs_celestial.wcs.ctype = ["RA---TAN", "DEC--TAN"]
    wcs_celestial.wcs.cunit = ['deg', 'deg']
    crota = 84 * u.deg
    wcs_celestial.wcs.pc = calculate_pc_matrix(crota, wcs_celestial.wcs.cdelt)

    # we're at the center of the Earth so let's try that
    test_loc = EarthLocation.from_geocentric(0, 0, 0, unit=u.m)
    test_gcrs = SkyCoord(test_loc.get_gcrs(date_obs))

    # follow the SunPy tutorial from here
    # https://docs.sunpy.org/en/stable/generated/gallery/units_and_coordinates/radec_to_hpc_map.html#sphx-glr-generated-gallery-units-and-coordinates-radec-to-hpc-map-py
    reference_coord = SkyCoord(wcs_celestial.wcs.crval[0] * u.Unit(wcs_celestial.wcs.cunit[0]),
                               wcs_celestial.wcs.crval[1] * u.Unit(wcs_celestial.wcs.cunit[1]),
                               frame='gcrs',
                               obstime=date_obs,
                               obsgeoloc=test_gcrs.cartesian,
                               obsgeovel=test_gcrs.velocity.to_cartesian(),
                               distance=test_gcrs.hcrs.distance
                               )

    reference_coord_arcsec = reference_coord.transform_to(frames.Helioprojective(observer=test_gcrs))

    cdelt1 = (np.abs(wcs_celestial.wcs.cdelt[0]) * u.deg).to(u.arcsec)
    cdelt2 = (np.abs(wcs_celestial.wcs.cdelt[1]) * u.deg).to(u.arcsec)

    p_angle = sun.P(date_obs)
    print("p_angle", p_angle)

    rotation_matrix = np.array([[np.cos(p_angle), -np.sin(p_angle)], [np.sin(p_angle), np.cos(p_angle)]])
    pc_helio = rotation_matrix @ wcs_celestial.wcs.pc

    new_header = sunpy.map.make_fitswcs_header(empty_data, reference_coord_arcsec,
                                               reference_pixel=u.Quantity([wcs_celestial.wcs.crpix[0] - 1,
                                                                           wcs_celestial.wcs.crpix[1] - 1] * u.pixel),
                                               scale=u.Quantity([cdelt1, cdelt2] * u.arcsec / u.pix),
                                               rotation_angle=-p_angle-crota,
                                               #rotation_matrix=pc_helio,
                                               observatory='PUNCH')

    wcs_helio = WCS(new_header)
    print(wcs_helio)

    # run a version of Chris's tests to see how it goes
    npoints = 50
    output_coords = []
    input_coords = np.stack([
                              (np.random.random(npoints) * 4095).astype(int),
                              (np.random.random(npoints) * 4095).astype(int)], axis=1)

    points_celestial = wcs_celestial.all_pix2world(input_coords, 0)
    points_helio = wcs_helio.all_pix2world(input_coords, 0)

    print("input coord \t celestial \t helio \t output coord")
    for c_pix, c_celestial, c_helio in zip(input_coords, points_celestial, points_helio):
        skycoord_helio = SkyCoord(c_helio[0] * u.deg, c_helio[1] * u.deg,
                                  frame=frames.Helioprojective,
                                  obstime=date_obs,
                                  observer=test_gcrs)
        skycoord_celestial = SkyCoord(c_celestial[0] * u.deg, c_celestial[1] * u.deg,
                                      frame=GCRS,
                                      obstime=date_obs,
                                      observer=test_gcrs,
                                      obsgeoloc=test_gcrs.cartesian,
                                      obsgeovel=test_gcrs.velocity.to_cartesian(),
                                      distance=test_gcrs.hcrs.distance
                                      )

        intermediate = skycoord_celestial.transform_to(frames.Helioprojective)

        output_coords.append(wcs_helio.all_world2pix(intermediate.data.lon.to(u.deg).value,
                                                               intermediate.data.lat.to(u.deg).value, 0))
        print(*c_pix, "\t", *c_celestial, "\t", *c_helio, "\t", *output_coords[-1])

    output_coords = np.array(output_coords)

    fig, ax = plt.subplots()
    ax.plot(input_coords[:, 0], output_coords[:, 0], 'bo')
    plt.show()
    #
    fig, ax = plt.subplots()
    ax.plot(input_coords[:, 1], output_coords[:, 1], 'ro')
    plt.show()

    print(wcs_helio)
    print(wcs_celestial)
        #wcs_helio.all_world2pix(intermediate.make_3d()[0], intermediate.make_3d[1], 0)
        # including this causes a warning
        # print("\t", skycoord_celestial.separation(skycoord_helio))

        # # including this causes an error
        # with frames.Helioprojective.assume_spherical_screen(skycoord_helio.observer):
        #     print("\t", skycoord_celestial.separation(skycoord_helio))