Feat: added optional tau/mu to Segment

[mrava87]

This PR adds two new parameters to Segment, namely tau and mu, which were previously hard-coded internally.

Note that this PR fixes also a bug since L=8 was chosen independent of the dimensionality of y. Now it is selected to be 8 for 2d and 12 for 3d

[mrava87]

@claude review

[claude]

Claude finished @mrava87's task in 3m 46s —— View job

---

Code Review: Segment tau/mu parameters

• Gather context (PR diff, repo guidelines)
• Understand the request (code review of Segment tau/mu change)
• Review pyproximal/optimization/segmentation.py changes
• Post review feedback

Overall this is a nice improvement: exposing tau/mu and fixing the Lipschitz constant L = 4 * ndims (previously hardcoded 8.0 regardless of y's dimensionality) is correct — for ndims=2 it still gives L=8, matching the old default exactly, and L=12 for 3D is consistent with the block-diagonal forward-difference gradient's operator norm. A few things worth addressing before merge:

1. Likely mypy --strict failure (lines 109–142)

if tau is None and mu is None:
    ...
    raise ValueError(msg)
...
L = 4 * ndims
tau = 1.0 / (mu * L) if tau is None else tau   # line 141
mu = 1.0 / (tau * L) if mu is None else mu     # line 142

mu is still typed float | None at line 141 — mypy cannot infer from the compound A is None and B is None guard that mu is non-None when tau is None. There's an identical pattern already in pyproximal/optimization/cls_primal.py:115-121 (if x0 is None and z0 is None: raise) that needed a # type: ignore[union-attr] on z0.copy() for exactly this reason — strong precedent that mu * L here will trip the same strict-mode check.

Suggested fix — restructure into per-variable nested checks so mypy can narrow without needing an ignore:

if tau is None:
    if mu is None:
        msg = "Either tau or mu must be provided."
        raise ValueError(msg)
    tau = 1.0 / (mu * L)
elif mu is None:
    mu = 1.0 / (tau * L)

Fix this →

2. Backward-compatible positional args (signature, lines 13–26)

tau and mu are inserted before the existing clsigmas/z/niter/... parameters:

def Segment(y, cl, sigma, alpha, tau=1.0, mu=None, clsigmas=None, z=None, ...):

Any external caller currently invoking Segment(y, cl, sigma, alpha, my_clsigmas) positionally will silently have my_clsigmas bound to tau instead, which is a behavior-breaking change (not just a TypeError, since tau legitimately accepts arbitrary values in early arithmetic). None of the in-repo callers (tutorials/brainmri.py, tutorials/segmentation.py) pass more than 4 positional args, so this repo is unaffected, but it's a public function. Consider appending tau/mu after the existing optional params instead, or making them keyword-only (*, tau=1.0, mu=None).

3. Minor

• No tests currently exist for Segment in pytests/ (pre-existing gap, not introduced by this PR) — might be worth adding a small case exercising the new tau/mu/ValueError paths per CLAUDE.md's guidance to update tests when public APIs change.
• If a caller supplies both tau and mu explicitly, they're used as-is with no check that tau * mu * L <= 1 (the convergence condition for Chambolle-Pock); likely fine to leave as user responsibility, just flagging.
  (branch)

[codacy-production]

Up to standards ✅

🟢 Issues 0 issues

Results:
0 new issues

View in Codacy

🟢 Metrics 0 complexity · 0 duplication

Metric	Results
Complexity	0
Duplication	0

View in Codacy

_{NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.}

[mrava87]

@claude review again. Note that I do not care about 2 as this library is still in v0

[claude]

Claude finished @mrava87's task in 1m 34s —— View job

---

Code Review: Segment tau/mu parameters (round 2)

• Gather context (PR diff since last review, repo guidelines)
• Review current state of segmentation.py and test_solver.py
• Verify point 1 (mypy narrowing) from previous review
• Verify point 3 (tests) from previous review
• Post updated review feedback

---

Both previously raised points have been addressed, and CI (ruff, mypy --strict, tests, coverage) is all green:

1. mypy narrowing — fixed. pyproximal/optimization/segmentation.py:119-125 now uses the nested per-variable structure I suggested:

if tau is None:
    if mu is None:
        raise ValueError(...)
    tau = 1.0 / (mu * L)
elif mu is None:
    mu = 1.0 / (tau * L)

This narrows correctly under mypy --strict (confirmed passing in CI) and preserves the old default (tau=1.0, mu derived from it) since tau defaults to 1.0, not None.

2. Positional args — not re-litigated, per your note that this doesn't matter pre-v1.

3. Tests — a ValueError test was added (pytests/test_solver.py:190-200, test_Segment_taumu), confirming the raise when both tau and mu are None. One small residual gap: the three par1/par2/par3 fixtures used in @pytest.mark.parametrize don't actually vary anything relevant to this test (y is always built as 2D with np.zeros((n, m)), and dtype/ndims aren't exercised), so it effectively runs the same case three times. Not a blocker, just means there's still no coverage of the success paths (tau only, mu only) or the L = 4 * ndims fix for the 3D case — worth a follow-up if you want that documented behavior actually tested.

Bonus improvement noticed (not previously flagged): the diff also swaps in get_array_module(y) for x0/argmax, so Segment now respects cupy inputs (x0 = ncp.zeros_like(g), cl = ncp.argmax(...)) instead of being hardcoded to numpy. This aligns with the repo convention of dual numpy/cupy support and is a nice fix beyond the original ask.

Nothing else to flag — looks good to merge from my side.
(branch)