Fix BatchLM vmap+jacfwd failures caused by in-place tensor mutations

astrophot/backend_obj.py

@@ -90,7 +90,6 @@ def setup_torch(self):
         self.detach = lambda x: x.detach()
         self.fill_at_indices = self._fill_at_indices_torch
         self.add_at_indices = self._add_at_indices_torch
-        self.and_at_indices = self._and_at_indices_torch
 
     def setup_jax(self):
         self.jax = importlib.import_module("jax")
@@ -135,7 +134,6 @@ def setup_jax(self):
         self.detach = lambda x: x
         self.fill_at_indices = self._fill_at_indices_jax
         self.add_at_indices = self._add_at_indices_jax
-        self.and_at_indices = self._and_at_indices_jax
 
     @property
     def array_type(self):
@@ -367,28 +365,28 @@ def _vmap_jax(self, *args, in_dims=0, **kwargs):
         return self.jax.vmap(*args, in_axes=in_dims, **kwargs)
 
     def _fill_at_indices_torch(self, array, indices, values):
+        if isinstance(indices, self.module.Tensor) and indices.dtype != self.module.bool:
+            # Long (integer) tensor indices: use index_put for vmap+jacfwd compatibility
+            return array.index_put((indices,), values)
+        # Bool tensor or tuple/slice indices: use clone + in-place
+        array = array.clone()
         array[indices] = values
         return array
 
     def _fill_at_indices_jax(self, array, indices, values):
-        array = array.at[indices].set(values)
-        return array
+        return array.at[indices].set(values)
 
     def _add_at_indices_torch(self, array, indices, values):
+        if isinstance(indices, self.module.Tensor) and indices.dtype != self.module.bool:
+            # Long (integer) tensor indices: use index_put for vmap+jacfwd compatibility
+            return array.index_put((indices,), values, accumulate=True)
+        # Bool tensor or tuple/slice indices: use clone + in-place
+        array = array.clone()
         array[indices] += values
         return array
 
     def _add_at_indices_jax(self, array, indices, values):
-        array = array.at[indices].add(values)
-        return array
-
-    def _and_at_indices_torch(self, array, indices, values):
-        array[indices] &= values
-        return array
-
-    def _and_at_indices_jax(self, array, indices, values):
-        array = array.at[indices].set(array[indices] & values)
-        return array
+        return array.at[indices].add(values)
 
     def _flatten_torch(self, array, start_dim=0, end_dim=-1):
         return array.flatten(start_dim, end_dim)

astrophot/models/func/spline.py

@@ -59,9 +59,9 @@ def spline(R: ArrayLike, profR: ArrayLike, profI: ArrayLike, extend: str = "zero
     """
     I = cubic_spline_torch(profR, profI, R.flatten()).reshape(*R.shape)
     if extend == "zeros":
-        backend.fill_at_indices(I, R > profR[-1], 0)
+        I = backend.where(R > profR[-1], backend.zeros_like(I), I)
     elif extend == "const":
-        backend.fill_at_indices(I, R > profR[-1], profI[-1])
+        I = backend.where(R > profR[-1], profI[-1] * backend.ones_like(I), I)
     else:
         raise ValueError(f"Unknown extend option: {extend}. Use 'zeros' or 'const'.")
     return I

astrophot/models/mixins/brightness.py

@@ -68,7 +68,7 @@ def polar_model(self, R: ArrayLike, T: ArrayLike) -> ArrayLike:
         v = w * np.arange(self.segments)
         for s in range(self.segments):
             indices = (angles >= v[s]) & (angles < (v[s] + w))
-            model = backend.add_at_indices(model, indices, self.iradial_model(s, R[indices]))
+            model = model + backend.where(indices, self.iradial_model(s, R), backend.zeros_like(R))
         return model
 
     def brightness(self, x: Tensor, y: Tensor) -> Tensor:
@@ -116,11 +116,15 @@ def polar_model(self, R: ArrayLike, T: ArrayLike) -> ArrayLike:
         for s in range(self.segments):
             angles = (T + cycle / 2 - v[s]) % cycle - cycle / 2
             indices = (angles >= -w) & (angles < w)
-            weights = (backend.cos(angles[indices] * self.segments) + 1) / 2
-            model = backend.add_at_indices(
-                model, indices, weights * self.iradial_model(s, R[indices])
+            # Weights and iradial_model are evaluated for all pixels (not just indices)
+            # to keep static tensor shapes compatible with vmap+jacfwd used in BatchLM.
+            weights = backend.where(
+                indices,
+                (backend.cos(angles * self.segments) + 1) / 2,
+                backend.zeros_like(angles),
             )
-            weight = backend.add_at_indices(weight, indices, weights)
+            model = model + weights * self.iradial_model(s, R)
+            weight = weight + weights
         return model / weight
 
     def brightness(self, x: ArrayLike, y: ArrayLike) -> ArrayLike:

tests/test_fit.py

@@ -9,6 +9,126 @@
 ######################################################################
 
 
+def _make_batch_lm_setup(
+    integrate_mode="none", sampling_mode="quad:3", n_images=3, pixelscale=1.0, cd_angles=None
+):
+    """Helper to create a BatchLM fitting setup with the given parameters."""
+    np.random.seed(42)
+    if cd_angles is None:
+        cd_angles = [np.pi / 3, np.pi / 6, np.pi / 16]
+
+    true_centers = [(32, 40), (15, 15), (45, 20)]
+    base_target = ap.TargetImage(data=np.zeros((64, 64)), pixelscale=pixelscale)
+    model_gen = ap.Model(
+        model_type="sersic galaxy model",
+        name="batch_gen",
+        center=true_centers[0],
+        q=0.6,
+        PA=np.pi / 3,
+        n=1,
+        Re=10,
+        Ie=1,
+        target=base_target,
+        integrate_mode="none",
+        sampling_mode="quad:3",
+    )
+    model_gen.initialize()
+
+    images = []
+    for k in range(n_images):
+        model_gen.center = true_centers[k % len(true_centers)]
+        kwargs = dict(
+            name=f"target{k}",
+            data=ap.backend.to_numpy(model_gen().data) + np.random.normal(scale=0.5, size=(64, 64)),
+            pixelscale=pixelscale,
+        )
+        if cd_angles is not None:
+            kwargs["CD"] = ap.utils.initialize.R(cd_angles[k % len(cd_angles)])
+        images.append(ap.TargetImage(**kwargs))
+
+    batch_target = ap.TargetImageBatch(images)
+
+    model = ap.Model(
+        model_type="sersic galaxy model",
+        name="batch_fit",
+        center=true_centers[0],
+        q=0.6,
+        PA=np.pi / 3,
+        n=1,
+        Re=10,
+        Ie=1,
+        target=batch_target.images[0],
+        integrate_mode=integrate_mode,
+        sampling_mode=sampling_mode,
+    )
+    model.initialize()
+
+    init_centers = [(30, 42), (16, 16), (46, 21)]
+    model.center = tuple(init_centers[k % len(init_centers)] for k in range(n_images))
+    model.q = tuple(0.6 for _ in range(n_images))
+    model.PA = tuple(np.pi / 3 for _ in range(n_images))
+    model.n = tuple(1.1 if k % 2 == 0 else 0.9 for k in range(n_images))
+    model.Re = tuple(11 for _ in range(n_images))
+    model.Ie = tuple(1.0 for _ in range(n_images))
+
+    return model, batch_target
+
+
+@pytest.mark.parametrize("integrate_mode", ["none", "bright", "curvature"])
+def test_batch_lm_integrate_modes(integrate_mode):
+    """BatchLM must converge without error for all integrate_mode values."""
+    model, batch_target = _make_batch_lm_setup(integrate_mode=integrate_mode)
+    res = ap.fit.BatchLM(model, batch_target, batch_target.window, max_iter=5).fit()
+    assert len(res.loss_history) >= 2, f"BatchLM ({integrate_mode}) must take at least one step"
+    assert np.all(
+        np.isfinite(res.loss_history[-1])
+    ), f"BatchLM ({integrate_mode}) final loss should be finite"
+
+
+@pytest.mark.parametrize("sampling_mode", ["midpoint", "simpsons", "quad:3"])
+def test_batch_lm_sampling_modes(sampling_mode):
+    """BatchLM must converge without error for common sampling_mode values."""
+    model, batch_target = _make_batch_lm_setup(sampling_mode=sampling_mode)
+    res = ap.fit.BatchLM(model, batch_target, batch_target.window, max_iter=5).fit()
+    assert len(res.loss_history) >= 2, f"BatchLM ({sampling_mode}) must take at least one step"
+    assert np.all(
+        np.isfinite(res.loss_history[-1])
+    ), f"BatchLM ({sampling_mode}) final loss should be finite"
+
+
+def test_batch_lm_rotated_images():
+    """BatchLM should work on images with non-trivial (rotated) CD matrices."""
+    model, batch_target = _make_batch_lm_setup(
+        cd_angles=[np.pi / 3, np.pi / 6, np.pi / 16],
+    )
+    res = ap.fit.BatchLM(model, batch_target, batch_target.window, max_iter=5).fit()
+    assert len(res.loss_history) >= 2, "BatchLM should take at least one step on rotated images"
+    assert np.all(
+        np.isfinite(res.loss_history[-1])
+    ), "BatchLM final loss should be finite for rotated images"
+
+
+def test_batch_lm_poisson_likelihood():
+    """BatchLM should work with Poisson likelihood."""
+    model, batch_target = _make_batch_lm_setup()
+    res = ap.fit.BatchLM(
+        model, batch_target, batch_target.window, max_iter=5, likelihood="poisson"
+    ).fit()
+    assert len(res.loss_history) >= 2, "BatchLM (poisson) must take at least one step"
+    assert np.all(
+        np.isfinite(res.loss_history[-1])
+    ), "BatchLM (poisson) final loss should be finite"
+
+
+def test_batch_lm_bright_integrate_improves():
+    """BatchLM with integrate_mode='bright' should improve the loss (chi^2/ndf)."""
+    model, batch_target = _make_batch_lm_setup(integrate_mode="bright")
+    res = ap.fit.BatchLM(model, batch_target, batch_target.window, max_iter=10).fit()
+    assert np.all(
+        res.loss_history[-1] <= res.loss_history[0]
+    ), "BatchLM with integrate_mode='bright' should not worsen the loss"
+
+
 @pytest.mark.parametrize("center", [[20.01, 20.02], [25.1, 17.324567]])
 @pytest.mark.parametrize("PA", [0, 60 * np.pi / 180])
 @pytest.mark.parametrize("q", [0.4, 0.8])