Correct rotations even when improper (#225)

* correct treatment of WCS angle, flip, scale (fixes #218); fixed tests of renderer and moments

* make Moments fluent

* angle and scale get astropy units

* fixed pixels size comparison

* removed units from angle and scale because not supported by jax;

one could use unxt, but with several new dependencies

* moved set_spectra_to_match to measure.forced_photometry

* removed set_spectra_to_match

* corrected plotting centers and boxes for WCS

* create dummy WCS if not set

* minor fixes

* added check for sub-pixel shift

* added moment normalization

Author: pmelchior

docs/0-quickstart.ipynb

@@ -340,7 +340,6 @@
    "outputs": [],
    "source": [
     "maxiter = 1000\n",
-    "scene.set_spectra_to_match(obs, parameters)\n",
     "scene_ = scene.fit(obs, parameters, max_iter=maxiter, e_rel=1e-4, progress_bar=True)"
    ]
   },

docs/howto/multiresolution.ipynb

@@ -296,7 +296,7 @@
    "id": "81bb89a6412b8ce8",
    "metadata": {},
    "source": [
-    " But the initial linear solver for the spectrum amplitudes and the fitting method receive lists of observations now:"
+    " The fitting method now receives a list of observations:"
    ]
   },
   {
@@ -308,7 +308,6 @@
    },
    "outputs": [],
    "source": [
-    "scene.set_spectra_to_match([obs_hsc, obs_hst], parameters)\n",
     "scene_ = scene.fit([obs_hsc, obs_hst], parameters, max_iter=50, progress_bar=True)"
    ]
   },

docs/howto/priors.ipynb

@@ -197,7 +197,6 @@
    "source": [
     "maxiter = 1000\n",
     "print(\"Initial likelihood:\", obs.log_likelihood(scene()))\n",
-    "scene.set_spectra_to_match(obs, parameters)\n",
     "scene_ = scene.fit(obs, parameters, max_iter=maxiter, e_rel=1e-4, progress_bar=True)\n",
     "print(\"Optimized likelihood:\", obs.log_likelihood(scene_()))"
    ]

src/scarlet2/frame.py

@@ -2,6 +2,7 @@
 import astropy.wcs
 import equinox as eqx
 import jax.numpy as jnp
+import numpy as np
 from astropy.coordinates import SkyCoord
 
 from .bbox import Box
@@ -27,10 +28,14 @@ class Frame(eqx.Module):
     def __init__(self, bbox, psf=None, wcs=None, channels=None):
         self.bbox = bbox
         self.psf = psf
+        if wcs is None:
+            wcs = astropy.wcs.WCS(naxis=2)  # Dummy WCS
+            wcs._naxis = bbox.spatial.shape[::-1]
+            wcs.wcs.ctype = ["RA---TAN", "DEC--TAN"]
         self.wcs = wcs
+
         if channels is None:
             channels = list(range(bbox.shape[0]))
-
         self.channels = channels
 
     def __hash__(self):
@@ -47,14 +52,10 @@ def pixel_size(self):
 
         Returns
         -------
-        float
-            Pixel size in arcsec, averaged over x and y direction
+        float, astropy.units.quantity.Quantity
+            Pixel size in units of the WCS sky coordinates
         """
-        if self.wcs is not None:
-            # return get_pixel_size(get_affine(self.wcs)) * 60 * 60  # in arcsec
-            return get_scale(self.wcs).mean() * 60 * 60  # in arcsec
-        else:
-            return 1
+        return get_scale(self.wcs)
 
     def get_pixel(self, pos):
         """Get the sky coordinates from a world coordinate
@@ -69,7 +70,6 @@ def get_pixel(self, pos):
         pixel coordinates in the model frame
         """
         if isinstance(pos, SkyCoord):
-            assert self.wcs is not None, "SkyCoord can only be converted with valid WCS"
             wcs_ = self.wcs.celestial  # only use celestial portion
             pixel = jnp.asarray(pos.to_pixel(wcs_), dtype="float32").T
             return pixel[..., ::-1]
@@ -85,14 +85,12 @@ def get_sky_coord(self, pos):
 
         Returns
         ----------
-        astropy.coordinates.SkyCoord if WCS is set, otherwise pos
+        astropy.coordinates.SkyCoord
         """
-        if self.wcs is not None:
-            pixels = pos.reshape(-1, 2)
-            wcs = self.wcs.celestial  # only use celestial portion
-            sky_coord = SkyCoord.from_pixel(pixels[:, 1], pixels[:, 0], wcs)
-            return sky_coord
-        return pos
+        pixels = pos.reshape(-1, 2)
+        wcs = self.wcs.celestial  # only use celestial portion
+        sky_coord = SkyCoord.from_pixel(pixels[:, 1], pixels[:, 0], wcs)
+        return sky_coord
 
     def convert_pixel_to(self, target, pixel=None):
         """Converts pixel coordinates from this frame to `target` frame
@@ -130,12 +128,11 @@ def u_to_pixel(self, distance):
         float
             size in pixels
         """
-        assert u.get_physical_type(distance) == "angle"
-
-        # first computer the pixel size
-        pixel_size = get_pixel_size(self.wcs.celestial) * 60 * 60  # in arcsec/pixel
-
-        return distance.to(u.arcsec).value / pixel_size
+        if self.wcs is not None:
+            pixel_size = get_pixel_size(self.wcs)
+            return distance / pixel_size
+        else:
+            return distance
 
     def pixel_to_angle(self, size):
         """Converts pixel size to celestial distance according to this frame WCS
@@ -148,12 +145,10 @@ def pixel_to_angle(self, size):
         Returns
         -------
         distance: :py:class:`astropy.units.Quantity`
-            Physical size, must be `PhysicalType("angle")`
         """
-        # first computer the pixel size
-        pixel_size = get_pixel_size(self.wcs.celestial) * 60 * 60  # in arcsec/pixel
 
-        distance = size * pixel_size * u.arcsec
+        pixel_size = get_pixel_size(self.wcs)
+        distance = size * pixel_size
         return distance
 
     @staticmethod
@@ -189,7 +184,6 @@ def from_observations(observations, model_psf=None, model_wcs=None, obs_id=None,
         # Array of pixel sizes for each observation
         pix_tab = []
         # Array of psf size for each psf of each observation
-        fat_psf_size = None
         small_psf_size = None
         channels = []
         # Create frame channels and find smallest and largest psf
@@ -199,14 +193,14 @@ def from_observations(observations, model_psf=None, model_wcs=None, obs_id=None,
 
             # concatenate all pixel sizes
             h_temp = get_pixel_size(obs.frame.wcs)
-
+            if isinstance(h_temp, u.Quantity):
+                h_temp = h_temp.to(u.arcsec).value  # standardize pixel sizes, using simple scalars below
             pix_tab.append(h_temp)
-            # Looking for the sharpest and the fatest psf
+
+            # Looking for the sharpest PSF
             psf = obs.frame.psf.morphology
             for psf_channel in psf:
                 psf_size = get_psf_size(psf_channel) * h_temp
-                if (fat_psf_size is None) or (psf_size > fat_psf_size):
-                    fat_psf_size = psf_size
                 if (
                     model_psf is None
                     and ((obs_id is None) or (c == obs_id))
@@ -225,23 +219,22 @@ def from_observations(observations, model_psf=None, model_wcs=None, obs_id=None,
 
         # Reference wcs
         if model_wcs is None:
-            model_wcs = obs_ref.frame.wcs
+            model_wcs = obs_ref.frame.wcs.deepcopy()
 
-        # Scale of the smallest pixel
+        # Scale of the model pixel
         h = get_pixel_size(model_wcs)
 
         # If needed and psf is not provided: interpolate psf to smallest pixel
         if model_psf is None:
             # create Gaussian PSF with a sigma smaller than the smallest observed PSF
             sigma = 0.7
-            assert (
-                small_psf_size / h > sigma
-            ), f"Default model PSF width ({sigma} pixel) too large for best-seeing observation"
+            assert small_psf_size / h > sigma, (
+                f"Default model PSF width ({sigma} pixel) too large for best-seeing observation"
+            )
             model_psf = GaussianPSF(sigma=sigma)
 
         # Dummy frame for WCS computations
         model_shape = (len(channels), 0, 0)
-
         model_frame = Frame(Box(model_shape), channels=channels, psf=model_psf, wcs=model_wcs)
 
         # Determine overlap of all observations in pixel coordinates of the model frame
@@ -262,11 +255,14 @@ def from_observations(observations, model_psf=None, model_wcs=None, obs_id=None,
                 else:
                     model_box &= this_box
 
+        # update model_wcs to change NAXIS1/2 and CRPIX1/2, but don't change frame_origin!
+        model_wcs._naxis = list(model_wcs._naxis)
+        model_wcs._naxis[:2] = model_box.shape[::-1]  # x/y needed here
+        model_wcs.wcs.crpix[:2] -= model_box.origin[::-1]  # x/y needed here
+
+        # frame_origin = (0,) + model_box.origin
         frame_shape = (len(channels),) + model_box.shape
-        frame_origin = (0,) + model_box.origin
-        model_frame = Frame(
-            Box(shape=frame_shape, origin=frame_origin), channels=channels, psf=model_psf, wcs=model_wcs
-        )
+        model_frame = Frame(Box(shape=frame_shape), channels=channels, psf=model_psf, wcs=model_wcs)
 
         # Match observations to this frame
         for obs in observations:
@@ -311,68 +307,145 @@ def get_psf_size(psf):
 
 def get_affine(wcs):
     """Return the WCS transformation matrix"""
+    if wcs is None:
+        return jnp.diag(jnp.ones(2))
+    wcs_ = wcs.celestial
     try:
-        model_affine = wcs.wcs.pc
+        model_affine = wcs_.wcs.pc
     except AttributeError:
         try:
-            model_affine = wcs.cd
+            model_affine = wcs_.cd
         except AttributeError:
-            model_affine = wcs.wcs.cd
-    return model_affine
+            model_affine = wcs_.wcs.cd
+    return model_affine[:2, :2]
 
 
-def get_pixel_size(wcs):
-    """Extracts the pixel size from a wcs, and returns it in deg/pixel"""
-    if wcs is None:
-        return 1
-    model_affine = get_affine(wcs)
-    pix = jnp.sqrt(
-        jnp.abs(model_affine[0, 0]) * jnp.abs(model_affine[1, 1] - model_affine[0, 1] * model_affine[1, 0])
-    )
-    return pix
+# for WCS linear matrix calculations:
+# rotation matrix for counter-clockwise rotation from positive x-axis
+# uses (x,y) coordinates and phi in radian!!
+def _rot_matrix(phi):
+    sinphi, cosphi = jnp.sin(phi), jnp.cos(phi)
+    return jnp.array([[cosphi, -sinphi], [sinphi, cosphi]])
+
+
+# flip in y!!!
+# uses (x,y) coordinates!
+_flip_matrix = lambda flip: jnp.diag(jnp.array((1.0, flip)))
 
+# 2x2 matrix determinant
+_det = lambda m: m[0, 0] * m[1, 1] - m[0, 1] * m[1, 0]
+
+
+def get_scale_angle_flip(trans):
+    """Return, scale, angle, flip from the WCS transformation matrix
+
+    Parameters
+    ----------
+    trans: (`astropy.wcs.WCS`, array)
+        WCS or WCS transformation matrix
 
-def get_scale(wcs):
+    Returns
+    -------
+    scale: `float`
+    angle: `float`, in radian
+    flip: -1 or 1
     """
-    Return WCS axis scales in deg/pixel
+    if isinstance(trans, (np.ndarray, jnp.ndarray)):  # noqa: SIM108
+        M = trans  # noqa: N806
+    else:
+        M = get_affine(trans)  # noqa: N806
+
+    det = _det(M)
+    # this requires pixels to be square
+    # if not, use scale = jnp.linalg.svd(M, compute_uv=False)
+    # but be careful with rotations as anisotropic stretch and rotation do not commute
+    scale = jnp.sqrt(jnp.abs(det)).item(0)
+
+    # if rotation is improper: need to apply y-flip to M to get pure rotation matrix (and unique angle)
+    improper = det < 0
+    flip = -1 if improper else 1
+    F = _flip_matrix(flip)  # noqa: N806, flip in y, is identity if flip = 1!!!
+    M_ = F @ M  # noqa: N806, flip = inverse flip
+    angle = jnp.arctan2(M_[1, 0], M_[0, 0]).item()
+
+    return scale, angle, flip
+
+
+def get_pixel_size(wcs):
+    """Extracts the pixel size from a wcs, and returns it in deg/pixel
+
+    Parameters
+    ----------
+    wcs: `astropy.wcs.WCS`
+        WCS structure or transformation matrix
+
+    Returns
+    -------
+    pixel_size: `float`
     """
-    if wcs is None:
-        return 1
-    model_affine = get_affine(wcs)
-    c1 = (model_affine[0, :2] ** 2).sum() ** 0.5
-    c2 = (model_affine[1, :2] ** 2).sum() ** 0.5
-    return jnp.array([c1, c2])
+    scale, angle, flip = get_scale_angle_flip(wcs)
+    return scale
 
 
-def get_angle(wcs):
+def get_scale(wcs, separate=False):
     """
-    Return WCS rotation angle in rad
+    Get WCS axis scales in deg/pixel
+
+    Parameters
+    ----------
+    wcs: `astropy.wcs.WCS`
+        WCS structure or transformation matrix
+    separate: `bool`
+          Compute separate axis scales
+
+    Returns
+    -------
+    float
     """
-    if wcs is None:
-        return 0
-    model_affine = get_affine(wcs)
-    c = get_scale(wcs)
-    c = c.reshape([c.shape[-1], 1])
-    r = model_affine[:2, :2] / c  # removing the scaling factors from the pc
-
-    if r[0, 0] == 0.0:
-        return jnp.arcsin(r[0, 1])
+    if separate:
+        M = get_affine(wcs)  # noqa: N806
+        c1 = (M[0, :] ** 2).sum() ** 0.5
+        c2 = (M[1, :] ** 2).sum() ** 0.5
+        return jnp.array([c1, c2])
     else:
-        return jnp.arctan(r[0, 1] / r[0, 0])
+        scale, angle, flip = get_scale_angle_flip(wcs)
+        return scale
 
 
-def get_sign(wcs):
+def get_angle(wcs):
+    """
+    Get WCS rotation angle
+
+    The angle is computed counter-clockwise from the positive x-axis, in radians.
+
+    Parameters
+    ----------
+    wcs: `astropy.wcs.WCS`
+        WCS structure or transformation matrix
+
+    Returns
+    -------
+    `astropy.units.quantity.Quantity`, unit = u.rad
     """
-    Return WCS flip signs
+    scale, angle, flip = get_scale_angle_flip(wcs)
+    return angle
+
+
+def get_flip(wcs):
     """
-    model_affine = get_affine(wcs)
-    c = get_scale(wcs)
-    c = c.reshape([c.shape[-1], 1])
-    r = model_affine[:2, :2] / c  # removing the absolute scaling factors from the pc
+    Return WCS sign convention
 
-    phi = jnp.arcsin(r[0, 1]) if r[0, 0] == 0.0 else jnp.arctan(r[0, 1] / r[0, 0])
+    A negative sign means that the rotation is improper and requires a flip.
+    By convention, we define this to be a flip in the y-axis.
 
-    r_inv = jnp.array([[jnp.cos(phi), -jnp.sin(phi)], [jnp.sin(phi), jnp.cos(phi)]])
+    Parameters
+    ----------
+    wcs: `astropy.wcs.WCS`
+        WCS structure or transformation matrix
 
-    r = r_inv @ r
-    return jnp.round(jnp.diag(r))
+    Returns
+    -------
+    -1 or 1
+    """
+    scale, angle, flip = get_scale_angle_flip(wcs)
+    return flip