reduce default `SolverParams.max_steps` and add more descriptive errors.

Currently our max_steps parameter is set to 5,000,000. While we should never really be hitting this ceiling, it is set too high to be useful for identifying a mis specified simulation solve.

to get an idea of how many steps a relatively complex model takes, I fit a currently COVID experiment model (6 strains, 1100 day fitting window, seasonality, external introductions, vaccinations, discontinuities) using SVI and logged the statistics of each of the solves for the first 100 or so iterations. Here is a section of output:

stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(882, dtype=int64), 'num_rejected_steps': Array(176, dtype=int64), 'num_steps': Array(1058, dtype=int64)}
 12%|█████████████▌                                                                                                 | 98/800 [04:35<34:14,  2.93s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(899, dtype=int64), 'num_rejected_steps': Array(176, dtype=int64), 'num_steps': Array(1075, dtype=int64)}
 12%|█████████████▋                                                                                                 | 99/800 [04:38<33:23,  2.86s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(892, dtype=int64), 'num_rejected_steps': Array(167, dtype=int64), 'num_steps': Array(1059, dtype=int64)}
 12%|█████████████▊                                                                                                | 100/800 [04:41<32:31,  2.79s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(884, dtype=int64), 'num_rejected_steps': Array(181, dtype=int64), 'num_steps': Array(1065, dtype=int64)}
 13%|█████████████▉                                                                                                | 101/800 [04:43<32:09,  2.76s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(886, dtype=int64), 'num_rejected_steps': Array(173, dtype=int64), 'num_steps': Array(1059, dtype=int64)}
 13%|██████████████                                                                                                | 102/800 [04:46<31:44,  2.73s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(877, dtype=int64), 'num_rejected_steps': Array(168, dtype=int64), 'num_steps': Array(1045, dtype=int64)}
 13%|██████████████▏                                                                                               | 103/800 [04:48<30:55,  2.66s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(900, dtype=int64), 'num_rejected_steps': Array(172, dtype=int64), 'num_steps': Array(1072, dtype=int64)}
 13%|██████████████▎                                                                                               | 104/800 [04:51<30:11,  2.60s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(889, dtype=int64), 'num_rejected_steps': Array(178, dtype=int64), 'num_steps': Array(1067, dtype=int64)}
 13%|██████████████▍                                                                                               | 105/800 [04:53<29:34,  2.55s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(879, dtype=int64), 'num_rejected_steps': Array(167, dtype=int64), 'num_steps': Array(1046, dtype=int64)}
 13%|██████████████▌                                                                                               | 106/800 [04:56<29:09,  2.52s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(899, dtype=int64), 'num_rejected_steps': Array(174, dtype=int64), 'num_steps': Array(1073, dtype=int64)}
 13%|██████████████▋                                                                                               | 107/800 [04:58<29:49,  2.58s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(883, dtype=int64), 'num_rejected_steps': Array(165, dtype=int64), 'num_steps': Array(1048, dtype=int64)}
 14%|██████████████▊                                                                                               | 108/800 [05:01<30:57,  2.68s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(887, dtype=int64), 'num_rejected_steps': Array(175, dtype=int64), 'num_steps': Array(1062, dtype=int64)}
 14%|██████████████▉                                                                                               | 109/800 [05:04<31:08,  2.70s/it, init loss: 1363089.9402, avg. loss [41-80]: 27029.4839]stats : {'max_steps': Array(5000000, dtype=int64), 'num_accepted_steps': Array(893, dtype=int64), 'num_rejected_steps': Array(169, dtype=int64), 'num_steps': Array(1062, dtype=int64)}
 14%|██████████████▉      

scroll right to see stats ^^^

Now it appears we are rejecting around 170 steps, and accepting around 880-890 each time, meaning a success rate of approximately 83%, which isnt horrible in my mind maybe @SamuelBrand1 would have some thoughts on that rate.

Now the number of steps taken as well as the proportion of accepted to rejected is important when identifying if a component of the model is mis specified in some way. For example if we have a discontinuity in the model that is not properly marked by a jump_ts to diffrax.diffeqsolve.

I think it is important to set the ceiling much lower, and then add more descriptive errors to the user so that they may be led in the right direction as to WHY they hit the ceiling.

Some useful resources for solving this problem:

I have taken a stab at this problem in the uva-descriptive-simulation-errors branch. What I did was pass a throw=False flag to diffrax.diffeqsolve in order to delay its raising of an error when the max_steps ceiling was hit. Then I tried to use diffrax.is_okay(solution.result) to determine if the solution is valid, and if not, log the accepted/rejected proportion. This works fine when simulating individual timeseries, but when wrapping calls to dynode.simulate() inside of an dynode.InferenceProcess() you get the following error:

jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function body_fn at /home/uva5/repos/DynODE-Experiments/covid_6strain_2202_2501/.venv/lib/python3.10/site-packages/numpyro/infer/svi.py:383 for jit. This concrete value was not available in Python because it depends on the values of the arguments svi_state.optim_state[1][0][0][0], svi_state.optim_state[1][0][1][0], and svi_state.rng_key.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError

in short, diffrax.Solution.result becomes a jax tracer, due to the Just-in-time compilation scheme of Jax, and so you cant easily check if the solution was valid.

In order for this process to work we need to find a jax compatible way to check for a valid solution.

Option 2: look into the store_rejected_steps flag within the diffrax clip stepsize controller in order to return to the user which simulation days are causing the most rejected steps, which may still be useful to them to identify a bug.

All in all this problem is not the highest priority, but its possible that setting a lower ceiling will help us identify inefficiencies in our ODEs, and thus improve performance by having less rejected steps.

Start with dynode.simulate() and dynode.InferenceProcess(), then trace how they call diffrax.diffeqsolve. Review the attempted uva-descriptive-simulation-errors approach, including throw=False and diffrax.is_okay(solution.result), and compare it with ClipStepSizeController's store_rejected_steps option. Done means selecting a JAX-compatible diagnostic approach, lowering the ceiling, and producing descriptive errors.