Skip to content

feat: closed-loop ESN Jacobian API (#169) - #540

Open
Saswatsusmoy wants to merge 2 commits into
SciML:masterfrom
Saswatsusmoy:Saswatsusmoy/169-jacobian-trajectory
Open

Saswatsusmoy wants to merge 2 commits into
SciML:masterfrom
Saswatsusmoy:Saswatsusmoy/169-jacobian-trajectory

Conversation

@Saswatsusmoy

Copy link
Copy Markdown
Contributor

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the contributor guidelines, in particular the SciML Style Guide and COLPRAC
  • Any new documentation only uses public API

Additional context

Adds jacobian / jacobian! / jacobians for the closed-loop reservoir map of a discrete ESN (in_dims == out_dims), so Jacobians can be taken along generative trajectories for Lyapunov analysis (Pathak2017).

  • Analytical default for standard activations and modifiers ((), NLAT1/2/3, Pad, PartialSquare, ExtendedSquare)
  • Optional backend=:forwarddiff via RCForwardDiffExt (ForwardDiff weakdep)
  • jacobians uses the same autoregressive feedback as predict
  • Deep / Hybrid / Continuous / Extend are out of scope for this PR
J, st = jacobian(esn, state, ps, st)
Js, outputs, st = jacobians(esn, steps, ps, st; initialdata)

Closes #169

Add jacobian/jacobian!/jacobians for the autonomous reservoir map of a
trained ESN (analytical default; ForwardDiff weakdep fallback), with
Models tests and API docs.

@chatgpt-codex-connector chatgpt-codex-connector Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 9f1575f7f7

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Comment thread src/jacobian.jl

T = eltype(x)
leak = __format_leak(T, cell.leak_coefficient)
x_new = __one_minus_leak(T, leak) .* x .+ leak .* cell.activation.(preactivation)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep heterogeneous-leak closed-loop states vector-shaped

When leak_coefficient is a vector, __format_leak reshapes it to n×1, so broadcasting it with the vector x and vector preactivation produces an n×n x_new. Consequently backend=:forwarddiff differentiates a length- map and cannot assign its result into the required n×n Jacobian; the newly added vector-leak finite-difference test also reaches this malformed map. Flatten the leak for this vector-only closed-loop path (or otherwise keep the operands one-dimensional).

Useful? React with 👍 / 👎.

Comment thread src/jacobian.jl
Comment on lines +389 to +390
@inbounds for i in 1:threshold
M[i, i] = 2 * x[i]

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Bound the PartialSquare derivative to the state length

For PartialSquare(eta) with eta > 1, the existing forward modifier simply squares every state component, but this loop uses floor(Int, eta * n) without capping it and writes past the n×n matrix. Such a model can run normally in prediction yet jacobian(...; backend=:analytical) throws a BoundsError; cap the derivative loop at n (or validate eta consistently in the modifier constructor).

Useful? React with 👍 / 👎.

Collapse row-scale/leak update into one finalize step, reuse LinearReadout
for the pure closed-loop map, and drive jacobians through a single rollout
loop shared with predict semantics.
@MartinuzziFrancesco

Copy link
Copy Markdown
Collaborator

This is a nice direction, but I don't think it should be in main library. Perhaps it is time to think about a ReservoirComputingUtilities library

However, I do think it would make for a fantastic example in the documentation!

@Saswatsusmoy

Copy link
Copy Markdown
Contributor Author

Agreed — I'll convert this PR to a documentation example instead of adding Jacobian APIs to the main library.

Plan:

  • drop the core jacobian / jacobian! / jacobians API and the ForwardDiff extension
  • add a docs example (closed-loop ESN trajectory → Jacobian / Lyapunov-style analysis with existing Lux + predict APIs)
  • keep #169 in mind as docs-only for now; a utilities package can wait until that direction is clearer

I'll push the reworked branch shortly.

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.

Feature Request: Jacobian along Trajectory

2 participants