patrick-kidger / diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/
Apache License 2.0
1.38k stars 124 forks source link

Advice about time-based 'events' #504

Open breisfeld opened 10 hours ago

breisfeld commented 10 hours ago

Hi, I am excited to try difftrax for my application, but am wondering about the recommended approach for the following.

I have an initial value ODE model that describes drug concentrations in various parts of the body over time. At various prescribed intervals, and for various durations, an additional amount of drug is added to a part of the body.

For example: Dose 1: magnitude = 1, time =5, duration= 15 Dose 2: magnitude = 2, time =25, duration= 30 Dose 3: magnitude = 3, time = 100, duration = 0.1

Is there a way to specify certain 'special' time points (in the above example, these would be t = 5, 20, 25, 55, 100, 100.1) such that the integration will be sure to include them?

Alternatively, I suppose that I could break the integration into numerous pieces, each of which would include a separate dosing, or not, regigon.

Your advice would be greatly appreciated on the best way to handle the situation.

Thank you!

patrick-kidger commented 9 hours ago

That is, you want the numerical solver to step exactly to those locations? For that you want either difrrax.StepTo or diffrax.StepSizeController(... step_ts=..., jump_ts=...), depending on desired behavior!

breisfeld commented 8 hours ago

Thank you! I'll try your suggestions.