From 278a31beab4213cd785fe6279054c7d61cd6c593 Mon Sep 17 00:00:00 2001
From: Thibault Barnouin <thibault.barnouin@astro.unistra.fr>
Date: Thu, 10 Apr 2025 18:14:38 +0200
Subject: [PATCH] fix WCS computation when importing fits

---
 package/lib/fits.py  | 42 ++++++++++++++++++++++-------------------
 package/lib/utils.py | 45 +++++++++++++++++++++++++++++++-------------
 2 files changed, 55 insertions(+), 32 deletions(-)

diff --git a/package/lib/fits.py b/package/lib/fits.py
index 4e5aef7..0c82afd 100755
--- a/package/lib/fits.py
+++ b/package/lib/fits.py
@@ -16,7 +16,7 @@ from astropy.io import fits
 from astropy.wcs import WCS
 
 from .convex_hull import clean_ROI
-from .utils import wcs_PA
+from .utils import wcs_CD_to_PC, wcs_PA
 
 
 def get_obs_data(infiles, data_folder="", compute_flux=False):
@@ -61,26 +61,30 @@ def get_obs_data(infiles, data_folder="", compute_flux=False):
     for i in range(len(data_array)):
         data_array[i][data_array[i] < 0.0] = 0.0
 
-    # force WCS to convention PCi_ja unitary, cdelt in deg
+    # Compute CDELT, ORIENTAT from header
     for wcs, header in zip(wcs_array, headers):
         new_wcs = wcs.deepcopy()
-        if new_wcs.wcs.has_cd() or (new_wcs.wcs.cdelt[:2] == np.array([1.0, 1.0])).all():
-            # Update WCS with relevant information
-            if new_wcs.wcs.has_cd():
-                del new_wcs.wcs.cd
-                keys = list(new_wcs.to_header().keys()) + ["CD1_1", "CD1_2", "CD1_3", "CD2_1", "CD2_2", "CD2_3", "CD3_1", "CD3_2", "CD3_3"]
-                for key in keys:
-                    header.remove(key, ignore_missing=True)
-                new_cdelt = np.linalg.eigvals(wcs.wcs.cd)
-                # new_cdelt.sort()
-            new_wcs.wcs.pc = wcs.wcs.cd.dot(np.diag(1.0 / new_cdelt))
-            new_wcs.wcs.cdelt = new_cdelt
-            for key, val in new_wcs.to_header().items():
-                header[key] = val
-        try:
-            header["ORIENTAT"] = float(header["ORIENTAT"])
-        except KeyError:
-            header["ORIENTAT"] = wcs_PA(new_wcs.wcs.pc[1, 0], np.diag(new_wcs.wcs.pc).mean())
+        if new_wcs.wcs.has_cd():
+            del new_wcs.wcs.cd
+            keys = list(new_wcs.to_header().keys()) + ["CD1_1", "CD1_2", "CD1_3", "CD2_1", "CD2_2", "CD2_3", "CD3_1", "CD3_2", "CD3_3"]
+            for key in keys:
+                header.remove(key, ignore_missing=True)
+            cdelt, orient = wcs_CD_to_PC(wcs.wcs.cd)
+            new_wcs.wcs.pc = wcs.wcs.cd.dot(np.diag(1.0 / cdelt))
+            new_wcs.wcs.cdelt = cdelt
+            try:
+                header["ORIENTAT"] = float(header["ORIENTAT"])
+            except KeyError:
+                header["ORIENTAT"] = -orient
+        elif new_wcs.wcs.has_pc() and (new_wcs.wcs.cdelt[:2] != np.array([1.0, 1.0])).all():
+            try:
+                header["ORIENTAT"] = float(header["ORIENTAT"])
+            except KeyError:
+                header["ORIENTAT"] = -wcs_PA(new_wcs.wcs.pc, new_wcs.wcs.cdelt)
+        else:
+            print("Could not compute ORIENTAT or CDELT from WCS")
+        for key, val in new_wcs.to_header().items():
+            header[key] = val
 
     # force WCS for POL60 to have same pixel size as POL0 and POL120
     is_pol60 = np.array([head["filtnam1"].lower() == "pol60" for head in headers], dtype=bool)
diff --git a/package/lib/utils.py b/package/lib/utils.py
index ee9fc5b..583f44d 100755
--- a/package/lib/utils.py
+++ b/package/lib/utils.py
@@ -154,27 +154,46 @@ def sci_not(v, err, rnd=1, out=str):
         return *output[1:], -power
 
 
-def wcs_PA(PC21, PC22):
+def wcs_CD_to_PC(CD):
     """
     Return the position angle in degrees to the North direction of a wcs
     from the values of coefficient of its transformation matrix.
     ----------
     Inputs:
-    PC21 : float
-        Value of the WCS matric PC[1,0]
-    PC22 : float
-        Value of the WCS matric PC[1,1]
+    CD : np.ndarray
+        Value of the WCS matrix CD
+    ----------
+    Returns:
+    cdelt : (float, float)
+        Scaling factor in arcsec between pixel in X and Y directions.
+    orient : float
+        Angle in degrees between the North direction and the Up direction of the WCS.
+    """
+    det = CD[0, 0] * CD[1, 1] - CD[0, 1] * CD[1, 0]
+    sgn = -1.0 if det < 0 else 1.0
+    cdelt = np.array([sgn, 1.0]) * np.sqrt(np.sum(CD**2, axis=1))
+    rot = np.arctan2(-CD[1, 0], sgn * CD[0, 0])
+    rot2 = np.arctan2(sgn * CD[0, 1], CD[1, 1])
+    orient = 0.5 * (rot + rot2) * 180.0 / np.pi
+    return cdelt, orient
+
+
+def wcs_PA(PC, cdelt):
+    """
+    Return the position angle in degrees to the North direction of a wcs
+    from the values of coefficient of its transformation matrix.
+    ----------
+    Inputs:
+    PC : np.ndarray
+        Value of the WCS matrix PC
+    cdelt : (float, float)
+        Scaling factor in arcsec between pixel in X and Y directions.
     ----------
     Returns:
     orient : float
         Angle in degrees between the North direction and the Up direction of the WCS.
     """
-    if (abs(PC21) > abs(PC22)) and (PC21 >= 0):
-        orient = -np.arccos(PC22) * 180.0 / np.pi
-    elif (abs(PC21) > abs(PC22)) and (PC21 < 0):
-        orient = np.arccos(PC22) * 180.0 / np.pi
-    elif (abs(PC21) < abs(PC22)) and (PC22 >= 0):
-        orient = np.arccos(PC22) * 180.0 / np.pi
-    elif (abs(PC21) < abs(PC22)) and (PC22 < 0):
-        orient = -np.arccos(PC22) * 180.0 / np.pi
+    rot = np.pi / 2.0 - np.arctan2(-cdelt[1] * PC[1, 0], abs(cdelt[0]) * PC[0, 0])
+    rot2 = np.pi / 2.0 - np.arctan2(abs(cdelt[0]) * PC[0, 1], cdelt[1] * PC[1, 1])
+    orient = 0.5 * (rot + rot2) * 180.0 / np.pi
     return orient