Repository: gammapy/gammapy (python)

Issue:
TSMapEstimator fails for some values of sigma with using a GaussianSpatialModel
**Gammapy version**
Both 1.2 and  dev

**Bug description**
Hi, I am having a strange error with the `TSMapEstimator`. I am running it on a dataset using a `GaussianSpatialModel`. If I set sigma=0.15º, it works, but 0.16º (or higher) does not. 

The issue arises both in 1.2 and the dev version, although the error message is slightly different. 

Because the error message (pasted below) is related to broadcasting of the kernel, I suspected a numpy/astropy version problem. But on my dev environment I have numpy>2 and astropy>7 and the error persists.


Since my dataset has private data, I made a toy example with a dataset that attempts to get as close to mine as possible. I clearly did not manage completely because while I reproduce the error, the boundary is now between 0.2 and 0.21º degree. I pasted it below.

**To Reproduce**
```
from gammapy.irf import PSFMap, EDispKernelMap
from gammapy.maps import MapAxis,WcsGeom
import numpy as np
import astropy.units as u
from astropy.coordinates import SkyCoord
from gammapy.estimators import TSMapEstimator
from gammapy.modeling.models import GaussianSpatialModel, PowerLawSpectralModel, SkyModel
from gammapy.datasets import MapDataset


energy_axis = MapAxis.from_energy_bounds(0.8*u.TeV,80*u.TeV, nbin=14, name='energy')


energy_axis_true = MapAxis.from_energy_bounds("0.1 TeV", "100 TeV", nbin=20, per_decade=True, name="energy_true")

center = SkyCoord(l = 6.773950 ,b= -4.789072 , unit='deg', frame='galactic')

geom = WcsGeom.create(binsz = 0.01,
                        skydir =center,
                        width = 3*u.deg,
                        frame = 'galactic',
                        axes = [energy_axis])

test = MapDataset.create(geom, energy_axis_true)
test.exposure.data = 1
test.background.data = 2
test.counts.data = 3
test.mask_safe.data = True
test.psf = PSFMap.from_gauss(energy_axis_true, rad_axis =test.psf.psf_map.geom.axes["rad"], geom = test.psf.psf_map.geom.to_image())
test.edisp = EDispKernelMap.from_gauss(energy_axis,energy_axis_true, sigma=0.1, bias =0,  geom = test.edisp.edisp_map.geom.to_image())


# 0.2 will work
ts_estimator1 = TSMapEstimator(SkyModel(spatial_model=GaussianSpatialModel(sigma=0.2*u.deg), spectral_model=PowerLawSpectralModel(index=2)), 
                      selection_optional = [], energy_edges =[0.8, 5,22]*u.TeV, n_jobs=6)

result_1 = ts_estimator1.run(test)
print(result_1)



# 0.21 will not
ts_estimator2 = TSMapEstimator(SkyModel(spatial_model=GaussianSpatialModel(sigma=0.21*u.deg), spectral_model=PowerLawSpectralModel(index=2)), 
                      selection_optional = [], energy_edges =[0.8, 5,22]*u.TeV, n_jobs=6)

result_2 = ts_estimator2.run(test)
print(result_2)
```

Error message with gammapy 1.2 (the specific numbers in the error depend on sigma)
```
(gammapy-1.2) [nieto@lfs1 ~]$ python test_kernel_error.py 
FluxMaps
--------

  geom                   : WcsGeom
  axes                   : ['lon', 'lat', 'energy']
  shape                  : (300, 300, 2)
  quantities             : ['ts', 'norm', 'niter', 'norm_err', 'npred', 'npred_excess', 'stat', 'stat_null', 'success']
  ref. model             : pl
  n_sigma                : 1
  n_sigma_ul             : 2
  sqrt_ts_threshold_ul   : 2
  sed type init          : likelihood

multiprocessing.pool.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/multiprocessing/pool.py", line 51, in starmapstar
    return list(itertools.starmap(args[0], args[1]))
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/site-packages/gammapy/estimators/map/ts.py", line 837, in _ts_value
    dataset = SimpleMapDataset.from_arrays(
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/site-packages/gammapy/estimators/map/ts.py", line 595, in from_arrays
    model=kernel * exposure_cutout,
ValueError: operands could not be broadcast together with shapes (1,303,303) (1,0,0) 
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/nfs/us0/nieto/test_kernel_error.py", line 46, in <module>
    result_2 = ts_estimator2.run(test)
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/site-packages/gammapy/estimators/map/ts.py", line 510, in run
    result = self.estimate_flux_map(dataset_sliced)
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/site-packages/gammapy/estimators/map/ts.py", line 433, in estimate_flux_map
    results = parallel.run_multiprocessing(
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/site-packages/gammapy/utils/parallel.py", line 283, in run_multiprocessing
    results = pool_func(
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/site-packages/gammapy/utils/parallel.py", line 313, in run_pool_star_map
    return pool.starmap(func, progress_bar(inputs, desc=task_name), **method_kwargs)
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/multiprocessing/pool.py", line 375, in starmap
    return self._map_async(func, iterable, starmapstar, chunksize).get()
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-1.2/lib/python3.10/multiprocessing/pool.py", line 774, in get
    raise self._value
ValueError: operands could not be broadcast together with shapes (1,303,303) (1,0,0) 
```
Error message with the dev version
```
(gammapy-dev) [nieto@lfs1 ~]$ python test_kernel_error.py 
FluxMaps
--------

  geom                   : WcsGeom
  axes                   : ['lon', 'lat', 'energy']
  shape                  : (np.int64(300), np.int64(300), 2)
  quantities             : ['ts', 'norm', 'niter', 'norm_err', 'npred', 'npred_excess', 'stat', 'stat_null', 'success']
  ref. model             : pl
  n_sigma                : 1
  n_sigma_ul             : 2
  sqrt_ts_threshold_ul   : 2
  sed type init          : likelihood

multiprocessing.pool.RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-dev/lib/python3.11/multiprocessing/pool.py", line 125, in worker
    result = (True, func(*args, **kwds))
                    ^^^^^^^^^^^^^^^^^^^
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-dev/lib/python3.11/multiprocessing/pool.py", line 51, in starmapstar
    return list(itertools.starmap(args[0], args[1]))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/d22/hin/nieto/gammapy/gammapy/estimators/map/ts.py", line 1068, in _ts_value
    SimpleMapDataset.from_arrays(
  File "/nfs/d22/hin/nieto/gammapy/gammapy/estimators/map/ts.py", line 711, in from_arrays
    kernel = (kernel * weights).sum(axis=0, keepdims=True)
              ~~~~~~~^~~~~~~~~
ValueError: operands could not be broadcast together with shapes (6,303,303) (6,0,0) 
"""

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/nfs/us0/nieto/test_kernel_error.py", line 46, in <module>
    result_2 = ts_estimator2.run(test)
               ^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/d22/hin/nieto/gammapy/gammapy/estimators/map/ts.py", line 626, in run
    result = self.estimate_flux_map(datasets_sliced)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/d22/hin/nieto/gammapy/gammapy/estimators/map/ts.py", line 507, in estimate_flux_map
    results = parallel.run_multiprocessing(
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/d22/hin/nieto/gammapy/gammapy/utils/parallel.py", line 288, in run_multiprocessing
    results = pool_func(
              ^^^^^^^^^^
  File "/nfs/d22/hin/nieto/gammapy/gammapy/utils/parallel.py", line 318, in run_pool_star_map
    return pool.starmap(func, progress_bar(inputs, desc=task_name), **method_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-dev/lib/python3.11/multiprocessing/pool.py", line 375, in starmap
    return self._map_async(func, iterable, starmapstar, chunksize).get()
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/nfs/scratch/nieto/miniforge3/envs/gammapy-dev/lib/python3.11/multiprocessing/pool.py", line 774, in get
    raise self._value
ValueError: operands could not be broadcast together with shapes (6,303,303) (6,0,0)
``

**Other information**
Any other information you think will be useful for us to fix the issue can go here.


Pull request description:
Fix TSmapEstimator pad_width
Resolves #5950
This issue with the TSmapEstimator seems to come from an inconsistancy as pad_width is determined using estimate_kernel computed on the original dataset while later on the estimate_kernel is computed from the padded dataset.
Adding one pixel to the padding fix it.

Could you check it work also for your real use case @LauraOlivera


Patch:
diff --git a/gammapy/estimators/map/ts.py b/gammapy/estimators/map/ts.py
index 8c85e580b..b05c94dfd 100644
--- a/gammapy/estimators/map/ts.py
+++ b/gammapy/estimators/map/ts.py
@@ -2,6 +2,7 @@
 """Functions to compute test statistic images."""
 
 import warnings
+import astropy.units as u
 from itertools import repeat
 import numpy as np
 import scipy.optimize
@@ -594,9 +595,15 @@ class TSMapEstimator(Estimator, parallel.ParallelMixin):
         datasets_models = datasets.models
 
         pad_width = (0, 0)
+        kernel_width = 0 * u.deg
         for dataset in datasets:
             pad_width_dataset = self.estimate_pad_width(dataset=dataset)
+            kernel_width_dataset = np.max(self.estimate_kernel(dataset).geom.width)
             pad_width = tuple(np.maximum(pad_width, pad_width_dataset))
+            kernel_width = np.maximum(kernel_width, kernel_width_dataset)
+
+        if self.kernel_width is None:
+            self.kernel_width = kernel_width
 
         datasets_padded = Datasets()
         for dataset in datasets:


Test patch:
diff --git a/gammapy/estimators/map/tests/test_ts.py b/gammapy/estimators/map/tests/test_ts.py
index 25c701dd7..a29ffe51b 100644
--- a/gammapy/estimators/map/tests/test_ts.py
+++ b/gammapy/estimators/map/tests/test_ts.py
@@ -589,7 +589,7 @@ def test_with_TemplateSpatialModel():
     )
 
     result = estimator.run(dataset)
-    assert_allclose(result["sqrt_ts"].data[0, 12, 16], 22.932, rtol=1e-3)
+    assert_allclose(result["sqrt_ts"].data[0, 12, 16], 22.3658, rtol=1e-3)
 
 
 def test_joint_ts_map(fake_dataset):
@@ -606,13 +606,13 @@ def test_joint_ts_map(fake_dataset):
     assert estimator.sum_over_energy_groups
 
     result = estimator.run(fake_dataset)
-    assert_allclose(result["npred_excess"].data.sum(), 902.403647, rtol=1e-3)
-    assert_allclose(result["sqrt_ts"].data[0, 10, 10], 1.360219, rtol=1e-3)
+    assert_allclose(result["npred_excess"].data.sum(), 912.92157, rtol=1e-3)
+    assert_allclose(result["sqrt_ts"].data[0, 10, 10], 1.462328, rtol=1e-3)
 
     result = get_combined_significance_maps(estimator, [fake_dataset, fake_dataset2])
 
-    assert_allclose(result["npred_excess"].data.sum(), 2 * 902.403647, rtol=1e-3)
-    assert_allclose(result["significance"].data[10, 10], 1.414529, rtol=1e-3)
+    assert_allclose(result["npred_excess"].data.sum(), 2 * 912.92157, rtol=1e-3)
+    assert_allclose(result["significance"].data[10, 10], 1.563892, rtol=1e-3)
     assert_allclose(
         result["df"].data, 2 * (~np.isnan(result["significance"].data)), rtol=1e-3
     )
@@ -621,7 +621,7 @@ def test_joint_ts_map(fake_dataset):
         model=model, threshold=1, selection_optional="all", sum_over_energy_groups=True
     )
     result = estimator.run([fake_dataset, fake_dataset2])
-    assert_allclose(result["sqrt_ts"].data[0, 10, 10], 1.92364, rtol=1e-3)
+    assert_allclose(result["sqrt_ts"].data[0, 10, 10], 2.063912, rtol=1e-3)
 
 
 @requires_data()
@@ -684,3 +684,70 @@ def test_joint_ts_map_hawc():
     result = estimator.run(datasets)
     assert_allclose(result["norm_sensitivity"].data[0, 59, 59], 0.04897, rtol=1e-3)
     assert_allclose(result["flux_sensitivity"].data[0, 59, 59], 4.881527e-14, rtol=1e-3)
+
+
+def test_tsmap_extended_source():
+    energy_axis = MapAxis.from_energy_bounds(
+        0.8 * u.TeV, 80 * u.TeV, nbin=2, name="energy"
+    )
+    energy_axis_true = MapAxis.from_energy_bounds(
+        "0.1 TeV", "100 TeV", nbin=2, per_decade=True, name="energy_true"
+    )
+
+    center = SkyCoord(l=6.773950, b=-4.789072, unit="deg", frame="galactic")
+
+    geom = WcsGeom.create(
+        binsz=0.1, skydir=center, width=3 * u.deg, frame="galactic", axes=[energy_axis]
+    )
+
+    test = MapDataset.create(geom, energy_axis_true)
+    test.exposure.data = 1
+    test.background.data = 2
+    test.counts.data = 3
+    test.mask_safe.data = True
+    test.psf = PSFMap.from_gauss(
+        energy_axis_true,
+        rad_axis=test.psf.psf_map.geom.axes["rad"],
+        geom=test.psf.psf_map.geom.to_image(),
+    )
+    test.edisp = EDispKernelMap.from_gauss(
+        energy_axis,
+        energy_axis_true,
+        sigma=0.1,
+        bias=0,
+        geom=test.edisp.edisp_map.geom.to_image(),
+    )
+
+    sigmas = np.arange(0.11, 0.6, 0.05) * u.deg
+    for sigma in sigmas:
+        ts_estimator = TSMapEstimator(
+            SkyModel(
+                spatial_model=GaussianSpatialModel(sigma=sigma),
+                spectral_model=PowerLawSpectralModel(index=2),
+            ),
+            selection_optional=[],
+            energy_edges=[0.8, 22] * u.TeV,
+            n_jobs=1,
+            kernel_width=10 * u.deg,
+        )
+
+        result = ts_estimator.run(test)
+
+        ts_estimator2 = TSMapEstimator(
+            SkyModel(
+                spatial_model=GaussianSpatialModel(sigma=sigma),
+                spectral_model=PowerLawSpectralModel(index=2),
+            ),
+            selection_optional=[],
+            energy_edges=[0.8, 22] * u.TeV,
+            n_jobs=1,
+            kernel_width=None,
+        )
+
+        result2 = ts_estimator2.run(test)
+
+        assert_allclose(
+            result2["sqrt_ts"].data[0, 12, 16],
+            result["sqrt_ts"].data[0, 12, 16],
+            rtol=5e-2,
+        )