Stable starlets (#249)

* remove proximal thresholds and add proper edge treatment in bspline_convolve

* added scales to starelet_reconstruction; docs  call for constraint/prior

Author: pmelchior

src/scarlet2/morphology.py

@@ -211,51 +211,53 @@ def f(self, r2):
         return jnp.exp(-bn * (r2 ** (0.5 / n) - 1))
 
 
-prox_plus = lambda x: jnp.maximum(x, 0)  # noqa: E731
-prox_soft = lambda x, thresh: jnp.sign(x) * prox_plus(jnp.abs(x) - thresh)  # noqa: E731
-prox_soft_plus = lambda x, thresh: prox_plus(prox_soft(x, thresh))  # noqa: E731
-
-
 class StarletMorphology(Morphology):
     """Morphology in the starlet basis
 
+    Notes
+    -----
+    The starlet basis is overcomplete, which means it can exactly represent the same image in multiple ways.
+    If used without constraints or priors on the starlet coefficients, this morphology model is functionally
+    indistinguishable from a 2D pixel array, while using more operations.
+
     See Also
     --------
     scarlet2.wavelets.Starlet
     """
 
     coeffs: jnp.ndarray
     """Starlet coefficients"""
-    l1_thresh: float = eqx.field(default=0)
-    """L1 threshold for coefficient to create sparse representation"""
-    positive: bool = eqx.field(default=True)
-    """Whether the coefficients are restricted to non-negative values"""
 
     def __call__(self, **kwargs):
         """Evaluate the model"""
-        f = prox_soft_plus if self.positive else prox_soft
-        return starlet_reconstruction(f(self.coeffs, self.l1_thresh))
+        return starlet_reconstruction(self.coeffs)
 
     @property
     def shape(self):
         """Shape (2D) of the morphology model"""
         return self.coeffs.shape[-2:]  # wavelet coeffs: scales x n1 x n2
 
     @staticmethod
-    def from_image(image, **kwargs):
+    def from_image(image, min_value=None, max_value=None):
         """Create starlet morphology from `image`
 
         Parameters
         ----------
         image: array
             2D image array to determine coefficients from.
-        kwargs: dict
-            Additional arguments for `__init__`
+        min_value: (float, None):
+            Minimum value threshold for coefficients
+        max_value: (float, None):
+            Minimum value threshold for coefficients
 
         Returns
         -------
         StarletMorphology
         """
         # Starlet transform of image (n1,n2) into coefficient with 3 dimensions: (scales+1,n1,n2)
         coeffs = starlet_transform(image)
-        return StarletMorphology(coeffs, **kwargs)
+        if min_value is not None:
+            coeffs = coeffs.at[coeffs < min_value].set(min_value)
+        if max_value is not None:
+            coeffs = coeffs.at[coeffs > max_value].set(max_value)
+        return StarletMorphology(coeffs)

src/scarlet2/wavelets.py

@@ -90,28 +90,37 @@ def bspline_convolve(image, scale):
         spacing between adjacent pixels with the spline.
 
     """
-    # Filter for the scarlet transform. Here bspline
     h1d = jnp.array([1.0 / 16, 1.0 / 4, 3.0 / 8, 1.0 / 4, 1.0 / 16])
-    j = scale
-
-    slice0 = slice(None, -(2 ** (j + 1)))
-    slice1 = slice(None, -(2**j))
-    slice3 = slice(2**j, None)
-    slice4 = slice(2 ** (j + 1), None)
-
-    # row
-    col = image * h1d[2]
-    col = col.at[slice4].add(image[slice0] * h1d[0])
-    col = col.at[slice3].add(image[slice1] * h1d[1])
-    col = col.at[slice1].add(image[slice3] * h1d[3])
-    col = col.at[slice0].add(image[slice4] * h1d[4])
-
-    # column
-    result = col * h1d[2]
-    result = result.at[:, slice4].add(col[:, slice0] * h1d[0])
-    result = result.at[:, slice3].add(col[:, slice1] * h1d[1])
-    result = result.at[:, slice1].add(col[:, slice3] * h1d[3])
-    result = result.at[:, slice0].add(col[:, slice4] * h1d[4])
+    step = 2**scale
+    ny, nx = image.shape
+
+    row_idx = jnp.arange(ny)
+    col_idx = jnp.arange(nx)
+
+    def reflect(idx, size):
+        # reflect indices at boundaries into [0, size-1]
+        idx = jnp.abs(idx)
+        # Map into [0, 2*size - 2] period, then fold back
+        idx = idx % (2 * size - 2)
+        return jnp.where(idx >= size, 2 * size - 2 - idx, idx)
+
+        # a simpler version: clamp pixels beyond the edge to the edge pixels
+        #         return jnp.clip(idx, 0, size - 1)  # clamp — or use true reflection below
+
+    # Row convolution
+    col = jnp.zeros_like(image)
+    for k, offset in enumerate([-2 * step, -step, 0, step, 2 * step]):
+        reflected = reflect(row_idx + offset, ny)
+        shifted = jnp.take(image, reflected, axis=0)
+        col += shifted * h1d[k]
+
+    # Column convolution
+    result = jnp.zeros_like(col)
+    for k, offset in enumerate([-2 * step, -step, 0, step, 2 * step]):
+        reflected = reflect(col_idx + offset, nx)
+        shifted = jnp.take(col, reflected, axis=1)
+        result += shifted * h1d[k]
+
     return result
 
 
@@ -164,7 +173,7 @@ def starlet_transform(image, scales=None, generation=2, convolve2d=None):
     return starlet
 
 
-def starlet_reconstruction(starlets, generation=2, convolve2d=None):
+def starlet_reconstruction(starlets, generation=2, convolve2d=None, scales=None):
     """Reconstruct an image from a dictionary of starlets
 
     Parameters
@@ -177,6 +186,8 @@ def starlet_reconstruction(starlets, generation=2, convolve2d=None):
     convolve2d: function
         The filter function to use to convolve the image
         with starlets in 2D.
+    scales: list of int
+        The scales to include in the reconstruction (0 being the smallest)
 
     Returns
     -------
@@ -187,11 +198,17 @@ def starlet_reconstruction(starlets, generation=2, convolve2d=None):
         return jnp.sum(starlets, axis=0)
     if convolve2d is None:
         convolve2d = bspline_convolve
-    scales = len(starlets) - 1
 
-    c = starlets[-1]
-    for i in range(1, scales + 1):
-        j = scales - i
+    # scales sorted in reverse order: from largest to smallest
+    max_scale = len(starlets) - 1
+    if scales is None:
+        scales = tuple(max_scale - i for i in range(1, max_scale + 1))
+    else:
+        scales = sorted(tuple(scale for scale in scales if scale <= max_scale), reverse=True)
+
+    # reconstruct: initialize from largest, go to smallest
+    c = starlets[scales[0]]
+    for j in scales[1:]:
         cj = convolve2d(c, j)
         c = cj + starlets[j]
     return c