ticktack 2.0: equinox edition

[benjaminpope]

I have been thinking a little about refactoring ticktack for version 2.0 using equinox to make it Jaxier and allow inference about the CBM.

We could have eqx modules for:

• CBM
  • __init__ initializes the CBM parameters (which might be fixed or dynamical)
  • __call__ just calls run
• Tree
  • __init__ initializes growth season
  • __call__ bins a continuous function to that growth season
• Fitter
  • __init__ loads data and metadata, including choice of model
  • __call__ does the log-likelihood call

This will then integrate easily with diffrax (which is based on equinox) and allow inference about more parameters (including the CBM and tree) by exposing them to gradients. It will also allow it to be more modular - like dLux we can plug and play with stacks of different layers.

[Jordan-Dennis]

I was actually curious about trying this and had already forked the repository with this in mind. As always though it was just finding time and it didn't go anywhere. I still haven't worked out a good routine for the semester yet, but I would like to do this after substituting in diffrax. The fork isn't suitable for merging back in because I condensed all of the branches into one meaning there would be a huge amount of merge conflicts ect. If we are going to implement this here for 2.0 I will delete the fork,

[benjaminpope]

Yeah - I don't think this is urgent and is good for 2.0. I don't know if it is worth forking at this stage before we do stuff like merging diffrax for 1.0. I think if we use equinox we will get a huge improvement in flexibility for inferring carbon box model parameters, which will be great for various Honours theses; but I don't think it will help the science for Paper 1 or our users in Groningen for their paper, so I think it is lower priority than the dLux work.