Repository: gammapy/gammapy

Issue:
WCSGeom.get_coord() does not work with TimeMapAxis
```
center = SkyCoord("0d", "0d", frame="icrs")

t_ref = Time(temporal_model.t_ref.value, format="mjd")

energy_axis = MapAxis.from_energy_bounds(
        "1 TeV", "10 TeV", nbin=3, name="energy_true"
    )

time_min = t_ref + [1, 3, 5, 7] * u.day
time_max = t_ref + [2, 4, 6, 8] * u.day

time_axis = TimeMapAxis.from_time_edges(time_min=time_min, time_max=time_max)

wcs_geom = WcsGeom.create(
        width=[1, 1.2],
        binsz=0.05,
        skydir=center,
        axes=[energy_axis, time_axis]
    )
wcs_geom.get_coord(sparse=False)
```

creates the error:
```
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[51], line 1
----> 1 wcs_geom.get_coord(sparse=False)

File ~/Gammapy-dev/gammapy/gammapy/maps/wcs/geom.py:628, in WcsGeom.get_coord(self, idx, mode, frame, sparse, axis_name)
    624     frame = self.frame
    626 pix = self._get_pix_all(idx=idx, mode=mode, sparse=sparse, axis_name=axis_name)
--> 628 data = self.pix_to_coord(pix)
    630 coords = MapCoord.create(
    631     data=data, frame=self.frame, axis_names=self.axes.names
    632 )
    633 return coords.to_frame(frame)

File ~/Gammapy-dev/gammapy/gammapy/maps/wcs/geom.py:669, in WcsGeom.pix_to_coord(self, pix)
    662     coords = self._wcs.wcs_pix2world(pix[0], pix[1], 0)
    664 coords = (
    665     u.Quantity(coords[0], unit="deg", copy=False),
    666     u.Quantity(coords[1], unit="deg", copy=False),
    667 )
--> 669 coords += self.axes.pix_to_coord(pix[self._slice_non_spatial_axes])
    670 return coords

File ~/Gammapy-dev/gammapy/gammapy/maps/axes.py:1810, in MapAxes.pix_to_coord(self, pix)
   1797 def pix_to_coord(self, pix):
   1798     """Convert pixel coordinates to map coordinates.
   1799 
   1800     Parameters
   (...)
   1808         Tuple of map coordinates.
   1809     """
-> 1810     return tuple([ax.pix_to_coord(p) for ax, p in zip(self, pix)])

File ~/Gammapy-dev/gammapy/gammapy/maps/axes.py:1810, in <listcomp>(.0)
   1797 def pix_to_coord(self, pix):
   1798     """Convert pixel coordinates to map coordinates.
   1799 
   1800     Parameters
   (...)
   1808         Tuple of map coordinates.
   1809     """
-> 1810     return tuple([ax.pix_to_coord(p) for ax, p in zip(self, pix)])

File ~/Gammapy-dev/gammapy/gammapy/maps/axes.py:2556, in TimeMapAxis.pix_to_coord(self, pix)
   2552 idx_valid = np.where(valid)
   2553 idx_invalid = np.where(~valid)
   2555 coords[idx_valid] = (
-> 2556     frac[idx_valid] * self.time_delta[idx_valid] + self.edges_min[idx_valid]
   2557 ).jd
   2558 coords = coords * self.unit + self.reference_time
   2559 coords[idx_invalid] = Time(INVALID_VALUE.time, scale=self.reference_time.scale)

File ~/anaconda3/envs/gammapy-dev/lib/python3.9/site-packages/astropy/utils/shapes.py:225, in ShapedLikeNDArray.__getitem__(self, item)
    223 def __getitem__(self, item):
    224     try:
--> 225         return self._apply("__getitem__", item)
    226     except IndexError:
    227         if self.isscalar:

File ~/anaconda3/envs/gammapy-dev/lib/python3.9/site-packages/astropy/time/core.py:1273, in TimeBase._apply(self, method, format, cls, *args, **kwargs)
   1271 jd1, jd2 = self._time.jd1, self._time.jd2
   1272 if apply_method:
-> 1273     jd1 = apply_method(jd1)
   1274     jd2 = apply_method(jd2)
   1276 # Get a new instance of our class and set its attributes directly.

IndexError: index 4 is out of bounds for axis 0 with size 4
```

This arises because `TimeMapAxis.pix_to_coord` does not handle multiple dimensions correctly (eg, when a slice is passed). This is because of the raveling to a 1-d array and then reshaping back into the original shape

https://github.com/gammapy/gammapy/blob/c5c1a5b66f9f8574bf10bacb7e3d4623995b264a/gammapy/maps/axes.py#L2548


Discussion:

Resolving patch:
diff --git a/gammapy/maps/axes.py b/gammapy/maps/axes.py
--- a/gammapy/maps/axes.py
+++ b/gammapy/maps/axes.py
@@ -2545,16 +2545,17 @@ def pix_to_coord(self, pix):
             Array of axis coordinate values.
         """
         shape = np.shape(pix)
-        pix = np.atleast_1d(pix).ravel()
-        coords = np.zeros(len(pix))
+        pix = np.atleast_1d(pix)
+        coords = np.zeros_like(pix)
         frac, idx = np.modf(pix)
+        idx1 = idx.astype(int)
         valid = np.logical_and(idx >= 0, idx < self.nbin, np.isfinite(idx))
         idx_valid = np.where(valid)
         idx_invalid = np.where(~valid)
 
         coords[idx_valid] = (
-            frac[idx_valid] * self.time_delta[idx_valid] + self.edges_min[idx_valid]
-        ).jd
+            frac[idx_valid] * self.time_delta[idx1[valid]] + self.edges_min[idx1[valid]]
+        ).value
         coords = coords * self.unit + self.reference_time
         coords[idx_invalid] = Time(INVALID_VALUE.time, scale=self.reference_time.scale)
         return coords.reshape(shape)


Test patch:
diff --git a/gammapy/maps/tests/test_axes.py b/gammapy/maps/tests/test_axes.py
--- a/gammapy/maps/tests/test_axes.py
+++ b/gammapy/maps/tests/test_axes.py
@@ -544,7 +544,7 @@ def test_pix_to_coord_time_axis(time_intervals):
     pixels = [1.3, 3.2, 5.4, 7, 15.33, 17.21, 19.11]
     coords = axis.pix_to_coord(pixels)
     assert_allclose(
-        coords[0:3].mjd, [58927.0125, 58927.534649, 58928.069298], rtol=1e-5
+        coords[0:3].mjd, [58927.538816, 58928.587281, 58929.648246], rtol=1e-5
     )
 
     # test with nan indices
@@ -552,7 +552,7 @@ def test_pix_to_coord_time_axis(time_intervals):
     coords = axis.pix_to_coord(pixels)
     assert_allclose(
         coords[-3:].mjd,
-        [58929.64032894737, 58930.162478070175, -3.725000e-04],
+        [58935.9561, 58937.0046, -3.72500000e-04],
         rtol=1e-5,
     )
 
@@ -560,13 +560,13 @@ def test_pix_to_coord_time_axis(time_intervals):
     coords = axis.pix_to_coord([[-1.2, 0.6], [1.5, 24.7]])
     assert_allclose(
         coords.mjd,
-        [[-3.725000e-04, 58927.551315789475], [58928.07346491228, -3.725000e-04]],
+        [[-3.72500000e-04, 5.89270250e04], [5.89275471e04, -3.72500000e-04]],
         rtol=1e-5,
     )
 
     # test with one value
     coords = axis.pix_to_coord(3)
-    assert_allclose(coords.mjd, 58927.0, rtol=1e-5)
+    assert_allclose(coords.mjd, axis.time_min[3].mjd, rtol=1e-5)
 
 
 def test_slice_time_axis(time_intervals):
diff --git a/gammapy/maps/wcs/tests/test_geom.py b/gammapy/maps/wcs/tests/test_geom.py
--- a/gammapy/maps/wcs/tests/test_geom.py
+++ b/gammapy/maps/wcs/tests/test_geom.py
@@ -5,8 +5,9 @@
 import astropy.units as u
 from astropy.coordinates import Angle, SkyCoord
 from astropy.io import fits
+from astropy.time import Time
 from regions import CircleSkyRegion
-from gammapy.maps import Map, MapAxis, WcsGeom
+from gammapy.maps import Map, MapAxis, TimeMapAxis, WcsGeom
 from gammapy.maps.utils import _check_binsz, _check_width
 
 axes1 = [MapAxis(np.logspace(0.0, 3.0, 3), interp="log", name="energy")]
@@ -637,3 +638,30 @@ def test_wcs_geom_no_zero_shape_cut():
 
     assert geom_cut.data_shape != (1, 0)
     assert geom_cut.data_shape == (1, 1)
+
+
+def test_wcs_geom_with_timeaxis():
+    center = SkyCoord("0d", "0d", frame="icrs")
+
+    t_ref = Time(55555, format="mjd")
+
+    energy_axis = MapAxis.from_energy_bounds(
+        "1 TeV", "10 TeV", nbin=3, name="energy_true"
+    )
+
+    time_min = t_ref + [1, 3, 5, 7] * u.day
+    time_max = t_ref + [2, 4, 6, 8] * u.day
+
+    time_axis = TimeMapAxis.from_time_edges(time_min=time_min, time_max=time_max)
+
+    wcs_geom = WcsGeom.create(
+        width=[1, 1.2], binsz=0.2, skydir=center, axes=[energy_axis, time_axis]
+    )
+
+    coords = wcs_geom.get_coord()
+    assert coords.shape == (4, 3, 6, 5)
+    assert_allclose(
+        coords[0][0][0][0:-1:5].value,
+        [[4.000e-01, 2.000e-01, 0.000e00, 3.598e02, 3.596e02]],
+        rtol=1e-5,
+    )