Merge pull request #176 from punch-mission/fix-NFI_imax-distortion

Fix nfi imax distortion

Author: s0larish

docs/source/example.ipynb

@@ -13,8 +13,8 @@
     "solpolpy is invoked using a simple interface:`output = sp.resolve(input_data, out_system)` where the `input_data` parameter is an appropriate NDCollection (a collection of [SunPy NDCubes](https://docs.sunpy.org/projects/ndcube/en/stable/)), such as a triplet of polarized images. `out_system` is the desired output polarization state (MZP, BtBr, Stokes, BpB, Bp3, or Bthp). It has a utility function to load data into the expected NDCollection format called `load_data`. There is also a plotting utility to plot any collection called `plot_collection`. \n",
     "\n",
     "## Polarization systems handled by solpolpy\n",
-    "\n",
-    "- Minus, Zero, Plus (MZP): Triplet of images taken at -60°, 0°, and +60° polarizing angles.\n",
+    "- Minus, Zero, Plus in solar frame (mzpsolar): Triplet of images taken at -60°, 0°, and +60° polarizing angles with a reference angle set to solar frame.\n",
+    "- Minus, Zero, Plus in instrument frame (mzpinstru): Triplet of images taken at -60°, 0°, and +60° polarizing angles with a reference angle set to instrument frame.\n",
     "- Tangential and Radial Brightness (BtBr): Pair of images with polarization along the tangential and radial direction with respect to the Sun respectively.\n",
     "- Stokes (Stokes): Total (unpolarized) brightness (I), polarized brightness along vertical and horizontal axes (Q) and polarized brightness along ±45° (U) .\n",
     "- Brightness, polarized Brightness (BpB): Total (unpolarized) brightness and ‘excess polarized’ brightness images pair respectively.\n",
@@ -150,7 +150,9 @@
    "metadata": {
     "collapsed": false
    },
-   "source": "In this case, we can see they are angle measurements (in degrees). If we want to know more, we can look at the appropriate cube:"
+   "source": [
+    "In this case, we can see they are angle measurements (in degrees). If we want to know more, we can look at the appropriate cube:"
+   ]
   },
   {
    "cell_type": "code",
@@ -160,7 +162,9 @@
      "start_time": "2024-08-06T14:26:55.020121Z"
     }
    },
-   "source": "stereo_collection['240.0 deg']",
+   "source": [
+    "stereo_collection['240.0 deg']"
+   ],
    "outputs": [
     {
      "data": {
@@ -197,7 +201,9 @@
      "start_time": "2024-08-06T14:26:55.023786Z"
     }
    },
-   "source": "stereo_collection['240.0 deg'].meta['POLAR']",
+   "source": [
+    "stereo_collection['240.0 deg'].meta['POLAR']"
+   ],
    "outputs": [
     {
      "data": {
@@ -350,7 +356,9 @@
      "start_time": "2024-08-06T14:26:56.857685Z"
     }
    },
-   "source": "output_bpb = resolve(stereo_collection, 'bpb')",
+   "source": [
+    "output_bpb = resolve(stereo_collection, 'bpb')"
+   ],
    "outputs": [],
    "execution_count": 8
   },
@@ -521,7 +529,9 @@
      "start_time": "2024-08-06T14:26:58.564677Z"
     }
    },
-   "source": "output_stokes = resolve(output_bpb, 'stokes')",
+   "source": [
+    "output_stokes = resolve(output_bpb, 'stokes')"
+   ],
    "outputs": [],
    "execution_count": 13
   },

solpolpy/core.py

@@ -1,4 +1,5 @@
 """Core transformation functions for solpolpy."""
+import copy
 import typing as t
 
 import astropy.units as u
@@ -12,7 +13,7 @@
 from solpolpy.errors import UnsupportedTransformationError
 from solpolpy.instruments import load_data
 from solpolpy.transforms import SYSTEM_REQUIRED_KEYS, System, transform_graph
-from solpolpy.util import extract_crota_from_wcs
+from solpolpy.util import apply_distortion_shift, extract_crota_from_wcs
 
 
 @u.quantity_input
@@ -33,8 +34,9 @@ def resolve(input_data: list[str] | NDCollection,
         The polarization state you want to convert your input dataframes to.
         Must be one of the following strings:
 
-        - "mzp": Triplet of images taken at -60°, 0°, and +60° polarizing angles.
-        - "btbr": Pair of images with polarization along the tangential and radial direction with respect to the Sun respectively.
+        - "mzpsolar": Triplet of images taken at -60°, 0°, and +60° polarizing angles with a reference angle set to solar frame.
+        - "mzpinstru": Triplet of images taken at -60°, 0°, and +60° polarizing angles with a reference angle set to instrument frame.
+        - "btbr": A Pair of images with polarization along the tangential and radial direction with respect to the Sun respectively.
         - "stokes": Total brightness ("I"), polarized brightness along vertical and horizontal axes (Q) and polarized brightness along ±45° (U) .
         - "bpb": Total brightness and ‘excess polarized’ brightness images pair respectively.
         - "bp3": Analogous to Stokes I, Q and U, but rotates around the Sun instead of a fixed frame of reference of the instrument.
@@ -74,7 +76,7 @@ def resolve(input_data: list[str] | NDCollection,
         raise ValueError("Out angles must be specified for this transform.")
 
     if imax_effect:
-        if input_kind == "mzp":
+        if input_kind in ('mzpsolar', 'mzpinstru'):
             input_data = resolve_imax_effect(input_data)
         else:
             msg = "IMAX effect applies only for transformations starting from MZP."
@@ -125,6 +127,9 @@ def determine_input_kind(input_data: NDCollection) -> System:
         raise ValueError(msg)
 
     for valid_kind, param_set in SYSTEM_REQUIRED_KEYS.items():
+        if valid_kind in [System.mzpinstru, System.mzpsolar] and param_set == input_keys:
+            polarref_value = input_data['Z'].meta.get("POLARREF", "solar").lower()
+            return System.mzpinstru if polarref_value == "instrument" else System.mzpsolar
         if valid_kind != System.npol and param_set == input_keys:
             return valid_kind
     try:
@@ -225,20 +230,25 @@ def generate_imax_matrix(array_shape: (int, int), cumulative_offset: u.deg, wcs:
         raise ValueError(msg)
 
     thmzp = [-60, 0, 60] * u.degree + cumulative_offset
+    # Define the A matrix
+    mat_a = np.empty((array_shape[0], array_shape[1], 3, 3))
 
-    long_extent, lat_extent = wcs.wcs.cdelt[0]*array_shape[0], wcs.wcs.cdelt[1]*array_shape[1]
-    long_arr, lat_arr = np.meshgrid(np.linspace(-long_extent/2, long_extent/2, array_shape[0]),
-                                    np.linspace(-lat_extent/2, lat_extent, array_shape[1]))
+    long_extent, lat_extent = wcs.wcs.cdelt[0] * array_shape[0], wcs.wcs.cdelt[1] * array_shape[1]
+    long_arr, lat_arr = np.meshgrid(np.linspace(-long_extent / 2, long_extent / 2, array_shape[0]),
+                                    np.linspace(-lat_extent / 2, lat_extent, array_shape[1]))
 
     # Foreshortening (IMAX) effect on polarizer angle
     phi_m = np.arctan2(np.tan(thmzp[0]) * np.cos(long_arr * u.degree), np.cos(lat_arr * u.degree)).to(u.degree)
     phi_z = np.arctan2(np.tan(thmzp[1]) * np.cos(long_arr * u.degree), np.cos(lat_arr * u.degree)).to(u.degree)
     phi_p = np.arctan2(np.tan(thmzp[2]) * np.cos(long_arr * u.degree), np.cos(lat_arr * u.degree)).to(u.degree)
 
-    phi = np.stack([phi_m, phi_z, phi_p])
+    # Apply distortion to IMAX matrix
+    if wcs.has_distortion:
+        phi_m = apply_distortion_shift(phi_m, wcs)
+        phi_z = apply_distortion_shift(phi_z, wcs)
+        phi_p = apply_distortion_shift(phi_p, wcs)
 
-    # Define the A matrix
-    mat_a = np.empty((array_shape[0], array_shape[1], 3, 3))
+    phi = np.stack([phi_m, phi_z, phi_p])
 
     for i in range(3):
         for j in range(3):
@@ -263,10 +273,15 @@ def resolve_imax_effect(input_data: NDCollection) -> NDCollection:
     """
     data_shape = _determine_image_shape(input_data)
     data_mzp_camera = np.zeros([data_shape[0], data_shape[1], 3, 1])
+    input_keys = list(input_data.keys())
 
-    satellite_orientation = extract_crota_from_wcs(input_data['M'])
-    polarizer_difference = [input_data[k].meta['POLAROFF'] if 'POLAROFF' in input_data[k].meta else 0
-                            for k in ['M', 'Z', 'P']] * u.degree
+    if input_data['M'].meta['POLARREF'].lower() == 'instrument':
+        satellite_orientation = extract_crota_from_wcs(input_data['M'])
+        polarizer_difference = [input_data[k].meta['POLAROFF'] if 'POLAROFF' in input_data[k].meta else 0
+                                for k in ['M', 'Z', 'P']] * u.degree
+    else:
+        satellite_orientation = 0 * u.degree
+        polarizer_difference = [0, 0, 0] * u.degree
 
     for i, key in enumerate(["M", "Z", "P"]):
         data_mzp_camera[:, :, i, 0] = input_data[key].data
@@ -283,10 +298,21 @@ def resolve_imax_effect(input_data: NDCollection) -> NDCollection:
 
     data_mzp_solar = np.matmul(imax_matrix_inv, data_mzp_camera)
 
-    for i, key in enumerate(input_data.keys()):
-        input_data[key].data[:, :] = data_mzp_solar[:, :, i, 0]
+    metaM, metaZ, metaP = (copy.copy(input_data[key].meta) for key in ["M", "Z", "P"])
+    for meta in [metaM, metaZ, metaP]:
+        meta.update(POLARREF='Solar')
 
-    return input_data
+    cube_list = [
+        (key, NDCube(data_mzp_solar[:, :, i, 0], wcs=input_data[key].wcs, meta=meta))
+        for i, (key, meta) in enumerate(zip(["M", "Z", "P"], [metaM, metaZ, metaP]))
+    ]
+
+    for p_angle in input_keys:
+        if p_angle.lower() == "alpha":
+            cube_list.append(("alpha", NDCube(input_data["alpha"].data * u.radian,
+                                              wcs=input_data[input_keys[0]].wcs,
+                                              meta=input_data["alpha"].meta)))
+    return NDCollection(cube_list, meta={}, aligned_axes="all")
 
 
 def _determine_image_shape(input_collection: NDCollection) -> tuple[int, int]:
@@ -354,8 +380,10 @@ def _compose2(f: t.Callable, g: t.Callable) -> t.Callable:
         composed function
 
     """
+
     def out(*a, **kw):
         return f(g(*a, **kw), **kw)
+
     out.uses_alpha = getattr(f, "uses_alpha", False) or getattr(g, "uses_alpha", False)
     out.uses_out_angles = getattr(f, "uses_out_angles", False) or getattr(g, "uses_out_angles", False)