FlowRegSuite · phflot · Mar 30, 2026 · Mar 31, 2026 · Mar 31, 2026 · Apr 30, 2026
diff --git a/docs/api/core.md b/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}

diff --git a/docs/api/session.md b/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,
 }

diff --git a/docs/theory/parameters.md b/docs/theory/parameters.md
@@ -38,6 +38,11 @@ The sublinear value follows best practices from {cite}`sun2010secrets` for handl
 - **Edge preservation**: Sublinear diffusion allows the model to handle brightness discontinuities at cell boundaries more gracefully
 - **Empirical validation**: This value has been validated across diverse 2-photon imaging datasets {cite}`flotho2022flow`
 
+Optional GNC staging can be enabled with `gnc_schedule`, for example
+`(0.0, 0.5, 1.0)`. This reruns the pyramid from quadratic to robust stages
+while keeping the final `a_data` and `a_smooth` values unchanged. Leaving
+`gnc_schedule=None` preserves the legacy solver path.
+
 ## Spatial-Temporal Filtering: `sigma`
 
 The `sigma` parameter controls Gaussian filtering applied to the video data before optical flow computation. It is specified as `[σx, σy, σt]` for each channel, where:

diff --git a/docs/user_guide/configuration.md b/docs/user_guide/configuration.md
@@ -47,9 +47,24 @@ options = OFOptions(
     # Nonlinear diffusion parameters
     a_smooth=1.0,  # Smoothness diffusion parameter
     a_data=0.45,  # Data term diffusion parameter
+
+    # Optional solver-level GNC stages for sublinear penalties
+    gnc_schedule=(0.0, 0.5, 1.0),
+
+    # 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.
+
+Set `gnc_schedule=None` to keep the default solver path. When provided,
+PyFlowReg reruns the pyramid once per stage, warm-starting each stage from
+the previous result.
+
 ### Alpha (Smoothness Weight)
 
 Controls the tradeoff between fitting the data and enforcing smooth flow fields:

diff --git a/docs/user_guide/multi_session.md b/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:

diff --git a/examples/session_config.toml b/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"
diff --git a/examples/session_config.yml b/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"
diff --git a/src/pyflowreg/core/__init__.py b/src/pyflowreg/core/__init__.py
@@ -145,6 +145,7 @@ def torch_level_solver(
             a_smooth,
             hx,
             hy,
+            gnc_beta=None,
         ):
             # Convert to tensors
             dtype_map = {"float64": torch.float64, "float32": torch.float32}
@@ -171,6 +172,8 @@ def to_tensor(a):
                 a_smooth,
                 hx,
                 hy,
+                update_lag_semantics="matlab" if gnc_beta is not None else "torch",
+                gnc_beta=gnc_beta,
             )
 
             return du.cpu().numpy(), dv.cpu().numpy()
@@ -232,6 +235,7 @@ def cuda_level_solver(
             a_smooth,
             hx,
             hy,
+            gnc_beta=None,
         ):
             # CUDA solver handles numpy/cupy conversion internally
             return level_solver_rbgs_cuda(
@@ -251,6 +255,8 @@ def cuda_level_solver(
                 a_smooth,
                 hx,
                 hy,
+                update_lag_semantics="matlab" if gnc_beta is not None else "torch",
+                gnc_beta=gnc_beta,
             )
 
         # Return a partial function with the custom level solver

diff --git a/src/pyflowreg/core/level_solver.py b/src/pyflowreg/core/level_solver.py
@@ -367,3 +367,189 @@ def compute_flow(
     flow[:, :, 0] = du
     flow[:, :, 1] = dv
     return flow
+
+
+@njit(fastmath=True, cache=True)
+def compute_flow_gnc(
+    J11,
+    J22,
+    J33,
+    J12,
+    J13,
+    J23,
+    weight,
+    u,
+    v,
+    alpha_x,
+    alpha_y,
+    iterations,
+    update_lag,
+    a_data,
+    a_smooth,
+    hx,
+    hy,
+    gnc_beta,
+):
+    """
+    Iterative solver for one fixed GNC stage at a single pyramid level.
+
+    This variant keeps the baseline solver unchanged and mixes the quadratic
+    and robust penalties with a fixed stage weight ``gnc_beta`` in ``[0, 1]``.
+    """
+    m, n, n_channels = J11.shape
+    du = np.zeros((m, n))
+    dv = np.zeros((m, n))
+    psi = np.ones((m, n, n_channels))
+    psi_smooth = np.ones((m, n))
+
+    OMEGA = 1.95
+    alpha = np.array([alpha_x, alpha_y], dtype=np.float64)
+    mix_quadratic = 1.0 - gnc_beta
+
+    for iteration_counter in range(iterations):
+        if (iteration_counter + 1) % update_lag == 0:
+            for k in range(n_channels):
+                a_k = a_data[k]
+                use_gnc_data = 0.0 < a_k < 1.0 and gnc_beta < 1.0
+                for i in range(n):
+                    for j in range(m):
+                        val = (
+                            J11[j, i, k] * du[j, i] * du[j, i]
+                            + J22[j, i, k] * dv[j, i] * dv[j, i]
+                            + J23[j, i, k] * dv[j, i]
+                            + 2.0 * J12[j, i, k] * du[j, i] * dv[j, i]
+                            + 2.0 * J13[j, i, k] * du[j, i]
+                            + J23[j, i, k] * dv[j, i]
+                            + J33[j, i, k]
+                        )
+                        if val < 0.0:
+                            val = 0.0
+                        robust = a_k * (val + 0.00001) ** (a_k - 1.0)
+                        if use_gnc_data:
+                            psi[j, i, k] = mix_quadratic + gnc_beta * robust
+                        else:
+                            psi[j, i, k] = robust
+
+            if a_smooth != 1.0:
+                nonlinearity_smoothness_2d(
+                    psi_smooth, u, du, v, dv, m, n, a_smooth, hx, hy
+                )
+                if 0.0 < a_smooth < 1.0 and gnc_beta < 1.0:
+                    for i in range(n):
+                        for j in range(m):
+                            psi_smooth[j, i] = (
+                                mix_quadratic + gnc_beta * psi_smooth[j, i]
+                            )
+            else:
+                for i in range(n):
+                    for j in range(m):
+                        psi_smooth[j, i] = 1.0
+
+        set_boundary_2d(du)
+        set_boundary_2d(dv)
+
+        for i in range(1, n - 1):
+            for j in range(1, m - 1):
+                denom_u = 0.0
+                denom_v = 0.0
+                num_u = 0.0
+                num_v = 0.0
+
+                left = (j, i - 1)
+                right = (j, i + 1)
+                down = (j + 1, i)
+                up = (j - 1, i)
+
+                if a_smooth != 1.0:
+                    tmp = (
+                        0.5
+                        * (psi_smooth[j, i] + psi_smooth[left])
+                        * (alpha[0] / (hx * hx))
+                    )
+                    num_u += tmp * (u[left] + du[left] - u[j, i])
+                    num_v += tmp * (v[left] + dv[left] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                    tmp = (
+                        0.5
+                        * (psi_smooth[j, i] + psi_smooth[right])
+                        * (alpha[0] / (hx * hx))
+                    )
+                    num_u += tmp * (u[right] + du[right] - u[j, i])
+                    num_v += tmp * (v[right] + dv[right] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                    tmp = (
+                        0.5
+                        * (psi_smooth[j, i] + psi_smooth[down])
+                        * (alpha[1] / (hy * hy))
+                    )
+                    num_u += tmp * (u[down] + du[down] - u[j, i])
+                    num_v += tmp * (v[down] + dv[down] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                    tmp = (
+                        0.5
+                        * (psi_smooth[j, i] + psi_smooth[up])
+                        * (alpha[1] / (hy * hy))
+                    )
+                    num_u += tmp * (u[up] + du[up] - u[j, i])
+                    num_v += tmp * (v[up] + dv[up] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+                else:
+                    tmp = alpha[0] / (hx * hx)
+                    num_u += tmp * (u[left] + du[left] - u[j, i])
+                    num_v += tmp * (v[left] + dv[left] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                    tmp = alpha[0] / (hx * hx)
+                    num_u += tmp * (u[right] + du[right] - u[j, i])
+                    num_v += tmp * (v[right] + dv[right] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                    tmp = alpha[1] / (hy * hy)
+                    num_u += tmp * (u[down] + du[down] - u[j, i])
+                    num_v += tmp * (v[down] + dv[down] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                    tmp = alpha[1] / (hy * hy)
+                    num_u += tmp * (u[up] + du[up] - u[j, i])
+                    num_v += tmp * (v[up] + dv[up] - v[j, i])
+                    denom_u += tmp
+                    denom_v += tmp
+
+                for k in range(n_channels):
+                    val_u = (
+                        weight[j, i, k]
+                        * psi[j, i, k]
+                        * (J13[j, i, k] + J12[j, i, k] * dv[j, i])
+                    )
+                    num_u -= val_u
+                    denom_u += weight[j, i, k] * psi[j, i, k] * J11[j, i, k]
+                    denom_v += weight[j, i, k] * psi[j, i, k] * J22[j, i, k]
+
+                du_kp1 = num_u / denom_u if denom_u != 0.0 else 0.0
+                du[j, i] = (1.0 - OMEGA) * du[j, i] + OMEGA * du_kp1
+
+                num_v2 = num_v
+                for k in range(n_channels):
+                    num_v2 -= (
+                        weight[j, i, k]
+                        * psi[j, i, k]
+                        * (J23[j, i, k] + J12[j, i, k] * du[j, i])
+                    )
+
+                dv_kp1 = num_v2 / denom_v if denom_v != 0.0 else 0.0
+                dv[j, i] = (1.0 - OMEGA) * dv[j, i] + OMEGA * dv_kp1
+
+    flow = np.zeros((m, n, 2))
+    flow[:, :, 0] = du
+    flow[:, :, 1] = dv
+    return flow