UT-Covid / episimlab

Framework for development of epidemiological models
https://ut-covid.github.io/episimlab/
BSD 3-Clause "New" or "Revised" License
3 stars 1 forks source link

Output transition matrix #58

Closed kellypierce closed 2 years ago

kellypierce commented 2 years ago

The current version of episimlab doesn't save the transition matrix, which contains the information on incident infections, hospitalizations, deaths, etc.

Feature request is for addition of the tm attribute of ComptModel to the xarray output. https://github.com/UT-Covid/episimlab/blob/3ead2e8f0934094a829004d61a4314ed88e5c217/episimlab/compt_model.py#L26

@ethho I can take a stab at this if you can point me to the best way to add this to the xarray without breaking the application of edge weights.

ethho commented 2 years ago

Hi, @kellypierce, hope you're doing well!

If I understand your feature request correctly, this is already supported. You just need to specify tm as an output variable in the model:

https://github.com/UT-Covid/episimlab/blob/dbef34b8c4c92b3df898ee62e0f01b16f57b6bf1/examples/example_sir.py#L84-L86

and then you can access the tm at every timestep (this is the matrix just before it's added to the current state):

https://github.com/UT-Covid/episimlab/blob/dbef34b8c4c92b3df898ee62e0f01b16f57b6bf1/examples/example_sir.py#L110-L111

If you want to take snapshots of tm less frequently than every step, you could define another clock: https://xarray-simlab.readthedocs.io/en/latest/run_model.html?highlight=clock#create-a-new-setup-from-scratch

kellypierce commented 2 years ago

Perfect -- thanks @ethho !