Make algorithms serializable

- a bit more documentation

README.md

@@ -1,40 +1,93 @@
 # torchtree
 
 [![Python package](https://github.com/4ment/torchtree/actions/workflows/python-package.yml/badge.svg)](https://github.com/4ment/torchtree/actions/workflows/python-package.yml)
+[![License: GPL v3](https://img.shields.io/badge/License-GPLv3-blue.svg)](https://www.gnu.org/licenses/gpl-3.0)
+[![docs](https://github.com/4ment/torchtree/actions/workflows/publish_documentation.yml/badge.svg)](https://github.com/4ment/torchtree/actions/workflows/publish_documentation.yml)
+![PyPI](https://img.shields.io/pypi/v/torchtree)
+![PyPI - Python Version](https://img.shields.io/pypi/pyversions/torchtree)
 
-## Installation
 
-### Use an Anaconda environment (Optional)
+torchtree is a program for inferring phylogenetic trees from molecular sequeces. It is implemented in python and uses [PyTorch] to leverage automatic differentiation. Inference algorithms include variational inference, Hamiltonian Monte Carlo, maximum a posteriori and Markov chain Monte Carlo.
+
+- [Getting Started](#getting-started)
+  - [Dependencies](#dependencies)
+  - [Installation](#installation)
+- [Quick start](#quick-start)
+
+## Getting Started
+
+### Dependencies
+ - [DendroPy]
+ - [PyTorch]
+
+ ### Installation
+Use an Anaconda environment (Optional)
 ```bash
 conda env create -f environment.yml
 conda activate torchtree
 ```
 
-### The easy way
-To install the latest stable version, run
+To install the latest stable version you can run
 ```bash
 pip install torchtree
 ```
 
-### Using the source code
+To build torchtree from source you can run
 ```bash
 git clone https://github.com/4ment/torchtree
-cd torchtree
-pip install .
+pip install torchtree/
 ```
 
-## Check install
+Check install
 ```bash
 torchtree --help
 ```
 
 ## Quick start
-torchtree will approximate the posterior distribution of an unrooted tree with a JC69 substitution model using variational inference 
+`torchtree` requires a JSON file contaning models and algorithms. A configutation file can generated using `torchtree-cli`, a command line-based tool. This two-step process allows the user to ajust values in the configuration file such as hyper prior parameters.
+
+### 1 - Generating a configuration file
+Some examples of models using variational inference:
+
+#### Unrooted tree with GTR+W4 model
+*W4* refers to a site model with 4 rates categories coming from a discretized Weibull distribution. This is similar to the more commonly used discretized Gamma distribution site model.
+
+```bash
+torchtree-cli advi -i data/fluA.fa -t data/fluA.tree -m GTR -C 4 > fluA.json
+```
+
+#### Time tree with strict clock and constant coalescent model
 ```bash
-torchtree examples/advi/fluA.json
+torchtree-cli advi -i data/fluA.fa -t data/fluA.tree -m JC69 --clock strict --coalescent constant > fluA.json
 ```
 
-The JSON file can be generated using the torchtree CLI
+### 2 - Running torchtree
+This will generate `sample.csv` and `sample.trees` files containing parameter and tree samples drawn from the variational distribution
 ```bash
-torchtree-cli advi -i data/fluA.fa -t data/fluA.tree > fluA.json
-```
+torchtree fluA.json
+```
+
+## torchtree plug-in
+torchtree can be easily extended without modifying the code base thanks its modular implementation. Some examples of external packages
+- [torchtree-bito]
+- [torchtree-physher]
+- [torchtree-scipy]
+- [torchtree-tensorflow]
+
+## License
+
+Distributed under the GPLv3 License. See [LICENSE](LICENSE) for more information.
+
+## Acknowledgements
+
+torchtree makes use of the following libraries and tools, which are under their own respective licenses:
+
+ - [PyTorch]
+ - [DendroPy]
+
+[DendroPy]: https://github.com/jeetsukumaran/DendroPy
+[PyTorch]: https://pytorch.org
+[torchtree-bito]: https://github.com/4ment/torchtree-bito
+[torchtree-physher]: https://github.com/4ment/torchtree-physher
+[torchtree-scipy]: https://github.com/4ment/torchtree-scipy
+[torchtree-tensorflow]: https://github.com/4ment/torchtree-tensorflow

docs/bibliography/refs.bib

@@ -54,4 +54,58 @@ @article{polson2014bayesian
   year = {2014},
   publisher = {JSTOR},
   doi = {10.1111/rssb.12042}
-}
+}
+
+@inproceedings{morningstar2021automatic,
+  title = { Automatic Differentiation Variational Inference with Mixtures },
+  author = {Morningstar, Warren and Vikram, Sharad and Ham, Cusuh and Gallagher, Andrew and Dillon, Joshua},
+  booktitle = {Proceedings of The 24th International Conference on Artificial Intelligence and Statistics},
+  pages = {3250--3258},
+  year = {2021},
+  editor = {Banerjee, Arindam and Fukumizu, Kenji},
+  volume = {130},
+  series = {Proceedings of Machine Learning Research},
+  month = {13--15 Apr},
+  publisher = {PMLR},
+  pdf = {http://proceedings.mlr.press/v130/morningstar21b/morningstar21b.pdf},
+  url = {https://proceedings.mlr.press/v130/morningstar21b.html}
+}
+
+@book{murphy2012machine,
+  title={Machine learning: a probabilistic perspective},
+  author={Murphy, Kevin P},
+  year={2012},
+  publisher={MIT press}
+}
+
+@inproceedings{prillo2020softsort,
+  title={Softsort: A continuous relaxation for the argsort operator},
+  author={Prillo, Sebastian and Eisenschlos, Julian},
+  booktitle={International Conference on Machine Learning},
+  pages={7793--7802},
+  year={2020},
+  organization={PMLR}
+}
+
+@article{rannala2012tail,
+  title={Tail paradox, partial identifiability, and influential priors in {B}ayesian branch length inference},
+  author={Rannala, Bruce and Zhu, Tianqi and Yang, Ziheng},
+  journal={Molecular biology and evolution},
+  volume={29},
+  number={1},
+  pages={325--335},
+  year={2012},
+  publisher={Oxford University Press}
+}
+
+@inproceedings{li2016renyi,
+ author = {Li, Yingzhen and Turner, Richard E},
+ booktitle = {Advances in Neural Information Processing Systems},
+ editor = {D. Lee and M. Sugiyama and U. Luxburg and I. Guyon and R. Garnett},
+ pages = {},
+ publisher = {Curran Associates, Inc.},
+ title = {R\'{e}nyi Divergence Variational Inference},
+ url = {https://proceedings.neurips.cc/paper_files/paper/2016/file/7750ca3559e5b8e1f44210283368fc16-Paper.pdf},
+ volume = {29},
+ year = {2016}
+}

docs/conf.py

@@ -36,8 +36,12 @@
 extensions = [
     'sphinx.ext.autodoc',
     'sphinx.ext.autosummary',
+    "sphinx.ext.intersphinx",
     'autoapi.extension',
     'sphinxcontrib.bibtex',
+    "sphinx.ext.viewcode",
+    "sphinx.ext.napoleon",
+    "sphinx_tabs.tabs",
 ]
 autosummary_generate = True
 
@@ -80,3 +84,9 @@
 autodoc_typehints = "signature"
 
 bibtex_bibfiles = ['bibliography/refs.bib']
+
+# Example configuration for intersphinx: refer to the Python standard library.
+intersphinx_mapping = {
+    "python": ("https://docs.python.org/3/", None),
+    "torch": ("https://pytorch.org/docs/master/", None),
+}

docs/index.rst

@@ -1,19 +1,52 @@
-Welcome to torchtree's documentation!
+Welcome to torchtree!
 ======================================
 
+.. warning::
+   The documentation corresponds to the current state of the main branch. There may be differences with the latest released version.
 
-Packages
---------
 
-.. autosummary::
-   :toctree: modules
-   :caption: API
-   :recursive:
+torchtree is a program for inferring phylogenetic trees from molecular sequeces.
+It is implemented in Python and uses PyTorch to leverage automatic differentiation.
+Inference algorithms include variational inference, Hamiltonian Monte Carlo, maximum *a posteriori* and Markov chain Monte Carlo.
+
+Installation
+------------
+
+.. tabs::
+
+   .. code-tab:: bash Latest
+
+      git clone https://github.com/4ment/torchtree
+      pip install torchtree/
+
+   .. code-tab:: bash Pypi
 
-   torchtree
+      pip install torchtree
+
+
+Plug-ins
+------------------
+
+torchtree can be easily extended without modifying the code base thanks its modular implementation. Some examples of external packages:
+
+- torchtree-bito_: is a plug-in interfacing the bito_ library for fast gradient calculations with BEAGLE_.
+- torchtree-physher_: is a plug-in interfacing physher_ for fast gradient calculations of tree likelihood and coalescent models.
+- torchtree-scipy_: is a plug-in interfacing the SciPy package.
+- torchtree-tensorflow_: is a plug-in interacting with Tensorflow.
+
+.. _torchtree-bito: http://github.com/4ment/torchtree-bito
+.. _torchtree-physher: http://github.com/4ment/torchtree-physher
+.. _torchtree-scipy: http://github.com/4ment/torchtree-scipy
+.. _torchtree-tensorflow: http://github.com/4ment/torchtree-tensorflow
+.. _physher: http://github.com/4ment/physher
+.. _BEAGLE: https://github.com/beagle-dev/beagle-lib
+.. _bito: https://github.com/phylovi/bito
 
 
 
 .. toctree::
+   :hidden:
+   :caption: API
 
+   API Reference<autoapi/torchtree/index>
    bibliography/bib

docs/requirements.txt

@@ -3,5 +3,6 @@ numpy
 sphinx
 sphinx-autoapi
 sphinx-book-theme
+sphinx-tabs
 sphinxcontrib-bibtex
 torch

test/test_coalescent.py

@@ -12,6 +12,7 @@
     PiecewiseConstantCoalescentGrid,
     PiecewiseConstantCoalescentGridModel,
     PiecewiseConstantCoalescentModel,
+    PiecewiseLinearCoalescentGrid,
     SoftPiecewiseConstantCoalescentGrid,
 )
 from torchtree.evolution.tree_model import ReparameterizedTimeTreeModel
@@ -326,3 +327,19 @@ def test_skygrid_data_json():
 
     skygrid2 = PiecewiseConstantCoalescentGridModel.from_json(example_intervals, {})
     assert -19.594893640219844 == pytest.approx(skygrid2().item(), 0.0001)
+
+
+def test_piecewise_linear(ratios_list):
+    sampling_times = torch.zeros(4, dtype=torch.float64)
+    ratios = torch.tensor(ratios_list, dtype=torch.float64)
+    thetas = torch.tensor(
+        [3.0, 10.0, 4.0, 2.0, 3.0], dtype=torch.float64, requires_grad=True
+    )
+    heights = inverse_transform_homochronous(ratios)
+    heights.requires_grad = True
+    grid = torch.tensor(np.linspace(0, 10.0, num=5)[1:], dtype=torch.float64)
+    print(grid)
+    constant = PiecewiseLinearCoalescentGrid(thetas, grid)
+    log_p = constant.log_prob(torch.cat((sampling_times, heights), -1))
+    assert -11.08185677776700117647 == pytest.approx(log_p.item(), 0.0001)
+

torchtree/_version.py

@@ -2,4 +2,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "1.0.2-dev1"
+__version__ = "1.0.2-dev2"

torchtree/cli/advi.py

@@ -10,6 +10,7 @@
 
 from torchtree import Parameter
 from torchtree.cli import PLUGIN_MANAGER
+from torchtree.cli.argparse_utils import list_or_int
 from torchtree.cli.evolution import (
     create_alignment,
     create_evolution_joint,
@@ -52,13 +53,13 @@ def create_variational_parser(subprasers):
     )
     parser.add_argument(
         '--elbo_samples',
-        type=int,
+        type=list_or_int,
         default=100,
         help="""number of samples for Monte Carlo estimate of ELBO""",
     )
     parser.add_argument(
         '--grad_samples',
-        type=int,
+        type=list_or_int,
         default=1,
         help="""number of samples for Monte Carlo estimate of gradients""",
     )
@@ -124,6 +125,11 @@ def create_variational_parser(subprasers):
         default='ELBO',
         help="""divergence to optimize""",
     )
+    parser.add_argument(
+        "--checkpoint_all",
+        action='store_true',
+        help="""log a new checkpoint file periodically""",
+    )
     parser.set_defaults(func=build_advi)
     return parser
 
@@ -762,6 +768,9 @@ def create_advi(joint, variational, parameters, arg):
         'parameters': parameters,
     }
 
+    if arg.checkpoint_all:
+        advi_dic["checkpoint_all"] = True
+
     if arg.K_elbo_samples > 1:
         elbo_samples = [arg.elbo_samples, arg.K_elbo_samples]
     else: