Skip to content

Fix complex-state Itô Milstein correction for real noise - #778

Draft
Afloat16 wants to merge 1 commit into
patrick-kidger:mainfrom
Afloat16:fix/ito-milstein-complex-correction
Draft

Afloat16 wants to merge 1 commit into
patrick-kidger:mainfrom
Afloat16:fix/ito-milstein-complex-correction

Conversation

@Afloat16

@Afloat16 Afloat16 commented Oct 2, 2026 •

Copy link
Copy Markdown

Validation status: the full-test prerequisite in CONTRIBUTING remains unmet. The bounded CPU recovery attempt is closed with incomplete coverage; this draft is not ready for review.

Default collection on this head: 1,004 nodes. The evidence accounts for 704 passed (668 in the new attempt plus 36 preserved checks on unchanged final source), 6 interrupted, and 294 unrun. No completed failing-node outcome was recorded. Interrupted nodes were not counted as passes; coordinator-held files did not start tests.

A single unchanged pristine test_adjoint.py::test_against control on e950aa2 also exited 137 before completing, after 382.093 s, with a native maximum RSS of 3,222,860 KiB. No cancellation signal was sent. The kill cause is unproven; host-wide OOM counters do not identify a per-process cause. This control establishes neither baseline equivalence nor a full-suite pass. The scoped correctness results below remain separate from this unresolved prerequisite.

Problem

For a complex-valued state driven by real Brownian noise, ItoMilstein conjugates the already-computed directional derivative in its final second-order contraction. The correction should be 0.5 * Dg(y)[g(y)] * (dW**2 - dt) in the scalar-noise case. Conjugating Dg(y)[g(y)] changes the step and its parameter gradient.

This can be checked independently by expressing the same complex SDE in real and imaginary coordinates. The current complex step disagrees with that real-coordinate Milstein step. Both array-valued and Lineax-operator diffusion fields are affected.

Change

Remove this final extra conjugation, with regressions for the closed-form linear-SDE step, equivalent real coordinates, real-state controls, and diffusion-parameter gradients. This preserves ControlTerm's existing array conjugation convention from #454 and does not change the construction of the noise increment tensor.

Validation

On upstream e950aa2cf1974c02816c1542f20beabf7cca496b, the six new regression cases produce 4 failed, 2 passed. With the correction, all 6 pass.

Local CPU validation, Python 3.12.14 / JAX 0.11.2:

  • pytest test/test_term.py test/test_milstein.py test/test_solver.py::test_adaptive_dt0_milstein -q: 14 passed
  • pytest test/test_solver.py -q: 26 passed, on both pristine and corrected source
  • pytest test/test_adjoint.py -q -m 'not slow': 10 passed, 1 deselected
  • Project-wide Ruff lint and format checks pass; changed-file Pyright has no errors
  • Project-wide Pyright reports the same pre-existing unused type-ignore error in _misc.py:175 on pristine and corrected source

A separate paired-path CPU experiment used 4,096 samples of dZ = (-0.2+0.3j) Z dt + (0.4+0.8j) Z dW, with Lineax diffusion, the same nested real Wiener increments, and its exact geometric-Brownian solution. For 16/32/64/128/256 steps, successive empirical strong orders changed from 0.48/0.51/0.48/0.53 to 0.97/1.03/0.98/1.03. At 256 steps, RMS error changed from 0.06615 to 0.003326. This is a numerical-correctness experiment, not a training-quality or runtime-performance claim.

An earlier multi-file test process was killed before completion. The full test suite has not completed locally, and accelerator execution has not been tested. These regressions concern complex states with real noise; they do not establish general complex-Brownian support.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant