search

andrrizzi / tfep

A library to perform targeted free energy perturbation with normalizing flows.
https://tfep.readthedocs.io
MIT License
7 stars 2 forks source link

Add user-friendly application layer with Lightning #12

Closed andrrizzi closed 9 months ago

andrrizzi commented 1 year ago

This should include pre-built normalizing flow models as well as LightningModule for training.

We could have for example an app.models.MAFAffineFlow(torch.nn.Module) and app.models.CouplingAffineFlow(torch.nn.Module) that automatically apply the ZMatrix + Cartesian transformation, and a KLTrainedFlow(lightning.LightningModule) that takes a flow as input and trains it using KL divergence (uni- and bi-directional).

To implement a fault-robust training in multimap TFEP, we might also have to implement a custom Lightning Trainer capable of resuming mid-epoch.

andrrizzi commented 9 months ago

The basic framework for this was implemented in #23, #26, and #28.