search

SciML / DiffEqCallbacks.jl

A library of useful callbacks for hybrid scientific machine learning (SciML) with augmented differential equation solvers
https://docs.sciml.ai/DiffEqCallbacks/stable/
Other
85 stars 44 forks source link

Rework `LinearizingSavingCallback` #176

Closed staticfloat closed 7 months ago

staticfloat commented 8 months ago

The original implementation of LinearizingSavingCallback had two modes; one with a single SavedValues that held all u values, and one with multiple SavedValues structs, one for each u index. Unfortunately, this second version (being the one that most users would want) was inefficient, wasteful in memory, and a pain to deal with after the fact due to each u index having a different time vector.

This PR deletes the first mode and improves the second mode to now have a much more efficient and useful datastructure backing the independently linearized states. In particular, it now preallocates memory in chunks while solving, stores the time vector only once (with a memory-efficient bitmap storing time vector occupancy for each u variable) and provides easy-to-use methods to consume the output both through iteration and explicit sampling at timepoints. In my experiments with a moderate-size ODE with 22 states that takes ~2000 timesteps, adding this callback increases solve time by ~10%, and allocations by ~100%.

ChrisRackauckas commented 7 months ago

Update the documentation as well?

BTW, we should consider having a form of this as a serialization method for DE solutions.

pepijndevos commented 7 months ago

It'd be nice if this provided a way to get derivatives as well. It'd basically be (u₁ - u₀) / (t₁ - t₀) modulo the weird cursor stuff going on in interpolate.

staticfloat commented 7 months ago

I can add derivative calculation fairly easily. I’ll do that when I get back to work.

ChrisRackauckas commented 7 months ago

https://github.com/staticfloat/DiffEqCallbacks.jl/blob/sf/independently_linearized/test/saving_tests.jl#L183-L187

What is broken in finalize? Is there something else we should be fixing downstream? Looks like it's clear that would have nothing chunks so I'm not quite sure what that's trying to show/do.

staticfloat commented 7 months ago

What is broken in finalize?

It should be solved by https://github.com/SciML/OrdinaryDiffEq.jl/pull/2061 now that that is merged.

staticfloat commented 7 months ago

I should have fixed the issue here (all test pass locally, at least), but I am marking this as draft again as I have a feature request from @pepijndevos to add derivative saving as well.

staticfloat commented 7 months ago

Alright, that was easier than expected.