Wire census constancy dataterm through flow and session pipelines

docs/api/core.md

@@ -30,6 +30,14 @@ Low-level optical flow computation engine implementing variational optical flow
 .. autofunction:: pyflowreg.core.optical_flow.get_motion_tensor_gc
 ```
 
+```{eval-rst}
+.. autofunction:: pyflowreg.core.optical_flow.get_motion_tensor_gray
+```
+
+```{eval-rst}
+.. autofunction:: pyflowreg.core.optical_flow.get_motion_tensor_cs
+```
+
 ### Boundary Handling
 
 ```{eval-rst}

docs/api/session.md

@@ -37,6 +37,7 @@ from pyflowreg.session.config import SessionConfig
             - sigma_smooth
             - alpha_between
             - iterations_between
+            - stage2_constancy_assumption
 
 **Configuration File Support**
 
@@ -78,6 +79,7 @@ cc_upsample = 4
 sigma_smooth = 6.0
 alpha_between = 25.0
 iterations_between = 100
+stage2_constancy_assumption = "gc"  # Options: "gc", "gray", "cs"
 ```
 
 ## Stage 1: Per-Recording Compensation
@@ -321,13 +323,15 @@ config = SessionConfig(
     cc_upsample=4,
     sigma_smooth=6.0,
     alpha_between=25.0,
-    iterations_between=100
+    iterations_between=100,
+    stage2_constancy_assumption="gc",
 )
 
 # Stage 1: Motion correct each recording
 print("Running Stage 1...")
 config.flow_options = {
     "quality_setting": "balanced",
+    "constancy_assumption": "gc",
     "save_valid_idx": True,
     "save_w": False,
 }

docs/user_guide/configuration.md

@@ -47,9 +47,17 @@ options = OFOptions(
     # Nonlinear diffusion parameters
     a_smooth=1.0,  # Smoothness diffusion parameter
     a_data=0.45,  # Data term diffusion parameter
+
+    # Data term, default preserves MATLAB Flow-Registration behavior
+    constancy_assumption="gc",  # Options: "gc", "gray", "cs"
 )
 ```
 
+`constancy_assumption="gc"` is the default gradient constancy data term used by
+the MATLAB Flow-Registration reference. `"gray"` selects gray-value constancy,
+and `"cs"` selects census constancy. These data terms are implemented by the
+native `flowreg` backend; the `diso` backend rejects non-default values.
+
 ### Alpha (Smoothness Weight)
 
 Controls the tradeoff between fitting the data and enforcing smooth flow fields:

docs/user_guide/multi_session.md

@@ -62,6 +62,7 @@ cc_upsample = 4        # Subpixel accuracy
 sigma_smooth = 6.0     # Gaussian smoothing
 alpha_between = 25.0   # Regularization
 iterations_between = 100
+stage2_constancy_assumption = "gc"  # Options: "gc", "gray", "cs"
 ```
 
 ### 3. Run Processing
@@ -118,6 +119,9 @@ The `pyflowreg.session` pipeline always runs the same three deterministic stages
 
 - Temporal averages are reloaded from disk and the reference recording (center) is selected automatically or from `SessionConfig.center`.
 - `compute_between_displacement()` smooths both averages, applies phase cross-correlation for a rigid guess, then refines with the configured flow backend (`src/pyflowreg/session/stage2_between_avgs.py`).
+- `stage2_constancy_assumption` controls the Stage 2 data term. The default
+  `"gc"` preserves MATLAB Flow-Registration behavior; `"cs"` enables the
+  census term for the native `flowreg` backend.
 - Results are written to `w_to_reference.npz` (separate `u`/`v` arrays) so MATLAB users can load them directly.
 
 **Outputs:** `w_to_reference.npz`, per-recording `status.json` updates, and `middle_idx` (0-based) pointing to the reference average.
@@ -170,6 +174,7 @@ quality_setting = "fast"     # Options: fast, balanced, quality
 buffer_size = 1000           # Frames per batch
 save_w = false               # Don't save displacement fields
 save_valid_idx = true        # Required for Stage 3
+# constancy_assumption = "cs"  # Optional Stage 1 data term override
 ```
 
 Alternatively, point to a saved MATLAB/Python options file:

examples/session_config.toml

@@ -42,6 +42,7 @@ sigma_smooth = 6.0                    # Sigma for Gaussian filter
 # Optical flow refinement
 alpha_between = 25.0                  # Regularization strength (higher = smoother)
 iterations_between = 100              # Solver iterations (higher = more accurate)
+stage2_constancy_assumption = "gc"     # Options: "gc", "gray", "cs"
 
 # === Stage 1 Flow Options (Optional) ===
 # Provide inline overrides passed to OFOptions
@@ -51,6 +52,7 @@ buffer_size = 1000                  # Frames per batch
 save_w = false                      # Save displacement fields
 save_valid_idx = true               # Required for Stage 3
 save_meta_info = true               # Save statistics
+# constancy_assumption = "gc"       # Options: "gc", "gray", "cs"
 
 # Alternatively reference a saved JSON file:
 # flow_options = "./saved_options/session_stage1.json"

examples/session_config.yml

@@ -31,6 +31,7 @@ cc_upsample: 4                    # Subpixel accuracy (higher = more precise)
 sigma_smooth: 6.0                 # Sigma for Gaussian filter
 alpha_between: 25.0               # Regularization strength (higher = smoother)
 iterations_between: 100           # Solver iterations (higher = more accurate)
+stage2_constancy_assumption: gc    # Options: gc, gray, cs
 
 # Stage 1 Flow Options (Optional)
 # Provide inline overrides passed to OFOptions
@@ -40,6 +41,7 @@ flow_options:
   save_w: false                  # Save displacement fields
   save_valid_idx: true           # Required for Stage 3
   save_meta_info: true           # Save statistics
+  # constancy_assumption: gc      # Options: gc, gray, cs
 
 # Alternatively, reference a JSON file saved via OF_options:
 # flow_options: "./saved_options/session_stage1.json"

src/pyflowreg/core/optical_flow.py

@@ -14,6 +14,10 @@
     Main API function for computing optical flow between two frames
 get_motion_tensor_gc
     Compute motion tensor components for gradient constancy
+get_motion_tensor_gray
+    Compute motion tensor components for gray-value constancy
+get_motion_tensor_cs
+    Compute motion tensor components for census constancy
 imregister_wrapper
     Warp an image using computed displacement fields
 warpingDepth
@@ -222,7 +226,7 @@ def get_motion_tensor_gray(f1, f2, hy, hx):
     return J11, J22, J33, J12, J13, J23
 
 
-def get_motion_tensor_cs(f1, f2, hy, hx):
+def get_motion_tensor_cs(f1, f2, hy, hx, eps=None):
     """
     Compute motion tensor components for census-based constancy assumption.
 
@@ -242,6 +246,14 @@ def get_motion_tensor_cs(f1, f2, hy, hx):
         Spatial grid spacing in y-direction.
     hx : float
         Spatial grid spacing in x-direction.
+    eps : float, optional
+        Smoothing width for the smoothed Heaviside function applied to
+        directional differences ``r = (neighbor - center) / dist``. If None,
+        uses ``0.1 / 255.0``, matching the Hafner/Demetz/Weickert
+        ``epsilon = 0.1`` convention for images scaled from ``[0, 255]`` to
+        approximately ``[0, 1]``. When ``hx`` or ``hy`` are physical units
+        rather than pixel-like units, callers may need to scale ``eps``
+        consistently.
 
     Returns
     -------
@@ -252,26 +264,33 @@ def get_motion_tensor_cs(f1, f2, hy, hx):
 
     Notes
     -----
-    The census transform is invariant to monotonically increasing grey-level
-    changes, improving robustness to illumination variations compared to
-    gray-value or gradient constancy. A smoothed Heaviside approximation
-    controlled by `eps` stabilizes the derivatives while preserving the
-    ordering information that drives the census constraint.
-
-    Symmetric padding enforces Neumann boundary conditions, and the border
-    is zeroed after aggregation to avoid wrap-around effects from the
-    cyclic shifts used during neighbor comparisons.
+    The hard census transform is invariant to global monotonically increasing
+    grey-value transforms because it depends only on ordering. This
+    implementation uses finite differences, Gaussian-preprocessed inputs, a
+    smoothed Heaviside function, and a linearized motion tensor, so invariance
+    is approximate.
+
+    Additive offsets cancel exactly in neighbor-center differences. Positive
+    multiplicative changes are approximately handled only when ``eps`` is small
+    relative to the directional-difference scale, or when ``eps`` is scaled
+    consistently with image intensity scale. Gamma and other nonlinear
+    monotone transforms preserve hard ordering but not exact smoothed
+    Heaviside values.
 
     References
     ----------
     .. [1] Hafner, D., Demetz, O., and Weickert, J. "Why is the Census
        Transform Good for Robust Optic Flow Computation?", SSVM 2013.
     """
-    eps = 0.1
+    if eps is None:
+        eps = 0.1 / 255.0
     eps2 = eps * eps
 
+    H, W = f1.shape
     f1p = np.pad(f1, ((1, 1), (1, 1)), mode="symmetric")
     f2p = np.pad(f2, ((1, 1), (1, 1)), mode="symmetric")
+    center1 = f1p[1:-1, 1:-1]
+    center2 = f2p[1:-1, 1:-1]
 
     offsets = [
         (dy, dx) for dy in (-1, 0, 1) for dx in (-1, 0, 1) if not (dy == 0 and dx == 0)
@@ -289,22 +308,22 @@ def get_motion_tensor_cs(f1, f2, hy, hx):
     for dy, dx in offsets:
         dist = float(np.sqrt((hy * dy) * (hy * dy) + (hx * dx) * (hx * dx)))
 
-        r1 = (np.roll(f1p, shift=(-dy, -dx), axis=(0, 1)) - f1p) / dist
-        r2 = (np.roll(f2p, shift=(-dy, -dx), axis=(0, 1)) - f2p) / dist
+        neigh1 = f1p[1 + dy : 1 + dy + H, 1 + dx : 1 + dx + W]
+        neigh2 = f2p[1 + dy : 1 + dy + H, 1 + dx : 1 + dx + W]
+
+        r1_core = (neigh1 - center1) / dist
+        r2_core = (neigh2 - center2) / dist
 
-        s1 = 0.5 * (1.0 + r1 / np.sqrt(r1 * r1 + eps2))
-        s2 = 0.5 * (1.0 + r2 / np.sqrt(r2 * r2 + eps2))
+        s1_core = 0.5 * (1.0 + r1_core / np.sqrt(r1_core * r1_core + eps2))
+        s2_core = 0.5 * (1.0 + r2_core / np.sqrt(r2_core * r2_core + eps2))
 
-        s1[0, :] = s1[1, :]
-        s1[-1, :] = s1[-2, :]
-        s1[:, 0] = s1[:, 1]
-        s1[:, -1] = s1[:, -2]
-        s2[0, :] = s2[1, :]
-        s2[-1, :] = s2[-2, :]
-        s2[:, 0] = s2[:, 1]
-        s2[:, -1] = s2[:, -2]
+        s1 = np.pad(s1_core, 1, mode="edge")
+        s2 = np.pad(s2_core, 1, mode="edge")
 
-        sy, sx = np.gradient(s1, hy, hx)
+        sy1, sx1 = np.gradient(s1, hy, hx)
+        sy2, sx2 = np.gradient(s2, hy, hx)
+        sx = 0.5 * (sx1 + sx2)
+        sy = 0.5 * (sy1 + sy2)
         st = s2 - s1
 
         J11 += sx * sx
@@ -367,6 +386,36 @@ def level_solver(
     return du, dv
 
 
+def _resolve_motion_tensor_func(const_assumption):
+    """
+    Resolve a constancy-assumption selector to a motion tensor function.
+
+    The default ``gc`` path is the MATLAB Flow-Registration behavior. Census
+    and gray-value constancy are explicit opt-in alternatives.
+    """
+    if hasattr(const_assumption, "value"):
+        const_assumption = const_assumption.value
+
+    key = str(const_assumption).strip().lower()
+    tensor_funcs = {
+        "gc": get_motion_tensor_gc,
+        "gradient": get_motion_tensor_gc,
+        "gray": get_motion_tensor_gray,
+        "brightness": get_motion_tensor_gray,
+        "cs": get_motion_tensor_cs,
+        "census": get_motion_tensor_cs,
+    }
+
+    try:
+        return tensor_funcs[key]
+    except KeyError as e:
+        supported = "', '".join(sorted(tensor_funcs))
+        raise ValueError(
+            f"Unknown constancy assumption: '{const_assumption}'. "
+            f"Supported values are: '{supported}'."
+        ) from e
+
+
 def get_displacement(
     fixed,
     moving,
@@ -426,7 +475,9 @@ def get_displacement(
         - a = 0.5: linear (L1) penalty
         - a = 0.45: sublinear, robust to noisy microscopy data
     const_assumption : str, default='gc'
-        Constancy assumption: 'gc' for gradient constancy (only option implemented)
+        Constancy assumption. Supported values are 'gc'/'gradient' for gradient
+        constancy, 'gray'/'brightness' for gray-value constancy, and
+        'cs'/'census' for census constancy.
     uv : np.ndarray, optional
         Initial displacement field (H, W, 2) with [u, v] components to initialize
         the coarsest (highest) pyramid level. If None, initializes with zeros.
@@ -477,6 +528,7 @@ def get_displacement(
     assert (
         fixed.ndim == moving.ndim
     ), f"Fixed and moving must have same dimensions: fixed.shape={fixed.shape}, moving.shape={moving.shape}"
+    motion_tensor_func = _resolve_motion_tensor_func(const_assumption)
     fixed = fixed.astype(np.float64)
     moving = moving.astype(np.float64)
     if fixed.ndim == 3:
@@ -578,7 +630,7 @@ def get_displacement(
         J13 = np.zeros(J_size, dtype=np.float64)
         J23 = np.zeros(J_size, dtype=np.float64)
         for ch in range(n_channels):
-            J11_ch, J22_ch, J33_ch, J12_ch, J13_ch, J23_ch = get_motion_tensor_gc(
+            J11_ch, J22_ch, J33_ch, J12_ch, J13_ch, J23_ch = motion_tensor_func(
                 f1_level[:, :, ch], tmp[:, :, ch], current_hx, current_hy
             )
             J11[:, :, ch] = J11_ch

src/pyflowreg/motion_correction/OF_options.py

@@ -78,6 +78,22 @@ class InterpolationMethod(str, Enum):
 class ConstancyAssumption(str, Enum):
     GRAY = "gray"
     GRADIENT = "gc"
+    CENSUS = "cs"
+
+
+def _normalize_constancy_assumption_value(v):
+    """Normalize constancy assumption aliases to serialized option values."""
+    if hasattr(v, "value"):
+        v = v.value
+    if isinstance(v, str):
+        aliases = {
+            "gradient": ConstancyAssumption.GRADIENT.value,
+            "brightness": ConstancyAssumption.GRAY.value,
+            "census": ConstancyAssumption.CENSUS.value,
+        }
+        key = v.strip().lower()
+        return aliases.get(key, key)
+    return v
 
 
 class NamingConvention(str, Enum):
@@ -176,7 +192,8 @@ class OFOptions(BaseModel):
         NamingConvention.DEFAULT, description="Output filename style"
     )
     constancy_assumption: ConstancyAssumption = Field(
-        ConstancyAssumption.GRADIENT, description="Constancy assumption"
+        ConstancyAssumption.GRADIENT,
+        description="Optical-flow data term: 'gc', 'gray', or 'cs'",
     )
 
     # Backend configuration
@@ -278,6 +295,12 @@ def normalize_sigma(cls, v):
             raise ValueError("Sigma must be [sx,sy,st] or (n_channels, 3)")
         return v
 
+    @field_validator("constancy_assumption", mode="before")
+    @classmethod
+    def normalize_constancy_assumption(cls, v):
+        """Normalize constancy assumption aliases to serialized option values."""
+        return _normalize_constancy_assumption_value(v)
+
     @model_validator(mode="after")
     def validate_and_normalize(self) -> "OFOptions":
         """Normalize fields and maintain MATLAB parity."""
@@ -691,6 +714,16 @@ def resolve_get_displacement(self) -> Callable:
         # Priority 3: Registry backend
         from pyflowreg.core.backend_registry import get_backend
 
+        constancy_assumption = _normalize_constancy_assumption_value(
+            self.constancy_assumption
+        )
+        if self.flow_backend == "diso" and constancy_assumption != "gc":
+            raise ValueError(
+                "The 'diso' backend does not support variational constancy "
+                f"assumption '{constancy_assumption}'. Use "
+                "flow_backend='flowreg' for 'gray' or 'cs'."
+            )
+
         factory = get_backend(self.flow_backend)
         return factory(**self.backend_params)
 
@@ -706,7 +739,9 @@ def to_dict(self) -> dict:
             "update_lag": self.update_lag,
             "a_data": self.a_data,
             "a_smooth": self.a_smooth,
-            "const_assumption": self.constancy_assumption.value,  # Fixed: use const_assumption for API compatibility
+            "const_assumption": _normalize_constancy_assumption_value(
+                self.constancy_assumption
+            ),
         }
 
     def __repr__(self) -> str: