Discussion regarding the BDF solver

[jiweiqi]

The BDF solver is really cool! I have a quick question regarding the AD of BDF solver.

Is it possible to use Auto-Differentiation across the BDF solver, so that we can do efficient adjoint sensitivity analysis?

For adjoint sensitivity analysis, we have the species concentrations Y = F(t), and we define a loss function g(Y) and compute the gradient of g(Y) for optimizing the mechanism.

[skrsna]

Hi @jiweiqi ,

Thanks for asking. I'm currently to trying to write a custom_vjp using an adjoint method similar to this one https://github.com/google/jax/blob/master/jax/experimental/ode.py#L265. Turns out this is not that trivial with BDF solver cause there are lots of constraints in JAX. I'm hoping to get some feedback from JAX's developers on the PR (jax-ml/jax#3781) and finish writing it.

[jiweiqi]

I just read the Jax PR. Regarding the slowness, you might take a look at the clipping and normalization of the mass fraction. Clipping and normalization of the mass fraction could affect the properties of the Jacobian matrix. This has been discussed somewhere else in the Cantera google group. For ReacTorch, we are working on fixing this issue.

[skrsna]

Thanks for the suggestion. I think this is what you are talking about https://github.com/comocheng/jax-reactor/blob/master/jax_reactor/solution.py#L79. I might be wrong but let me know!!

[jiweiqi]

Yes. Besides that, another clipping is in the wdot function

jax-reactor/jax_reactor/solution.py

Line 58 in e13d359

np.matmul(np.log(C + 1e-300), gas_info.reactant_stioch_coeffs))

.

[jiweiqi]

I am not sure if this is the case. An easy way to test it is to print out the number of function evaluations and jacobian evaluations. If you see significantly more computation than Cantera, then that is probably related to this issue.