Merge pull request #269 from probml/docs

Updating documentation

Author: slinderman

README.md

@@ -39,7 +39,7 @@ To install the latest releast of dynamax from PyPi:
 
 ``` {.console}
 pip install dynamax                 # Install dynamax and core dependencies, or
-pip install dynamax[notebooks]      # Install with dep's for demo notebooks
+pip install dynamax[notebooks]      # Install with demo notebook dependencies
 ```
 
 To install the latest development branch:

docs/api.rst

@@ -1,31 +1,147 @@
-
-API
-===
-
 State Space Model (Base class)
 ===============================
 
 .. autoclass:: dynamax.ssm.SSM
   :members:
 
+Parameters
+----------
+
+Parameters and their associated properties are stored as :class:`jax.DeviceArray`
+and :class:`dynamax.parameters.ParameterProperties`, respectively. They are bundled together into a
+:class:`dynamax.parameters.ParameterSet` and a :class:`dynamax.parameters.PropertySet`, which are simply
+aliases for immutable datastructures (in our case,  :class:`NamedTuple`).
+
+.. autoclass:: dynamax.parameters.ParameterSet
+.. autoclass:: dynamax.parameters.PropertySet
+.. autoclass:: dynamax.parameters.ParameterProperties
+
 Hidden Markov Model
 ===================
 
-High-level class
-----------------
+Abstract classes
+------------------
 
 .. autoclass:: dynamax.hidden_markov_model.HMM
+  :show-inheritance:
   :members:
 
+.. autoclass:: dynamax.hidden_markov_model.HMMInitialState
+  :members:
+
+.. autoclass:: dynamax.hidden_markov_model.HMMTransitions
+  :members:
+
+.. autoclass:: dynamax.hidden_markov_model.HMMEmissions
+  :members:
+
+High-level models
+-----------------
+
+The HMM implementations below cover common emission distributions and,
+if the emissions are exponential family distributions, the models implement
+closed form EM updates. For HMMs with emissions outside the non-exponential family,
+these models default to a generic M-step implemented in :class:`HMMEmissions`.
+
+Unless otherwise specified, these models have standard initial distributions and
+transition distributions with conjugate, Bayesian priors on their parameters.
+
+**Initial distribution:**
+
+$$p(z_1 \mid \pi_1) = \mathrm{Cat}(z_1 \mid \pi_1)$$
+$$p(\pi_1) = \mathrm{Dir}(\pi_1 \mid \alpha 1_K)$$
+
+where $\alpha$ is the prior concentration on the initial distribution $\pi_1$.
+
+**Transition distribution:**
+
+$$p(z_t \mid z_{t-1}, \theta) = \mathrm{Cat}(z_t \mid A_{z_{t-1}})$$
+$$p(A) = \prod_{k=1}^K \mathrm{Dir}(A_k \mid \beta 1_K)$$
+
+where $\beta$ is the prior concentration on the rows of the transition matrix $A$.
+
+
+.. autoclass:: dynamax.hidden_markov_model.BernoulliHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.CategoricalHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.GaussianHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.DiagonalGaussianHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.SphericalGaussianHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.SharedCovarianceGaussianHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.LowRankGaussianHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.MultinomialHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.PoissonHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.GaussianMixtureHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.DiagonalGaussianMixtureHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.LinearRegressionHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.LogisticRegressionHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.CategoricalRegressionHMM
+  :show-inheritance:
+  :members: initialize
+
+.. autoclass:: dynamax.hidden_markov_model.LinearAutoregressiveHMM
+  :show-inheritance:
+  :members: initialize, sample, compute_inputs
+
 Low-level inference
 -------------------
 
+.. autoclass:: dynamax.hidden_markov_model.HMMPosterior
+.. autoclass:: dynamax.hidden_markov_model.HMMPosteriorFiltered
+
 .. autofunction:: dynamax.hidden_markov_model.hmm_filter
 .. autofunction:: dynamax.hidden_markov_model.hmm_smoother
 .. autofunction:: dynamax.hidden_markov_model.hmm_two_filter_smoother
 .. autofunction:: dynamax.hidden_markov_model.hmm_fixed_lag_smoother
 .. autofunction:: dynamax.hidden_markov_model.hmm_posterior_mode
 .. autofunction:: dynamax.hidden_markov_model.hmm_posterior_sample
+.. autofunction:: dynamax.hidden_markov_model.parallel_hmm_filter
+.. autofunction:: dynamax.hidden_markov_model.parallel_hmm_smoother
+
+Types
+-----
+
+.. autoclass:: dynamax.hidden_markov_model.HMMParameterSet
+.. autoclass:: dynamax.hidden_markov_model.HMMPropertySet
+
 
 Linear Gaussian SSM
 ====================
@@ -50,8 +166,9 @@ Types
 .. autoclass:: dynamax.linear_gaussian_ssm.ParamsLGSSMInitial
 .. autoclass:: dynamax.linear_gaussian_ssm.ParamsLGSSMDynamics
 .. autoclass:: dynamax.linear_gaussian_ssm.ParamsLGSSMEmissions
-.. autoclass:: dynamax.linear_gaussian_ssm.PosteriorLGSSMFiltered
-.. autoclass:: dynamax.linear_gaussian_ssm.PosteriorLGSSMSmoothed
+
+.. autoclass:: dynamax.linear_gaussian_ssm.PosteriorGSSMFiltered
+.. autoclass:: dynamax.linear_gaussian_ssm.PosteriorGSSMSmoothed
 
 Nonlinear Gaussian GSSM
 ========================
@@ -74,6 +191,12 @@ Low-level inference
 .. autofunction:: dynamax.nonlinear_gaussian_ssm.unscented_kalman_filter
 .. autofunction:: dynamax.nonlinear_gaussian_ssm.unscented_kalman_smoother
 
+Types
+-----
+
+.. autoclass:: dynamax.nonlinear_gaussian_ssm.ParamsNLGSSM
+
+
 Generalized Gaussian GSSM
 ==========================
 
@@ -89,4 +212,9 @@ Low-level inference
 .. autofunction:: dynamax.generalized_gaussian_ssm.conditional_moments_gaussian_filter
 .. autofunction:: dynamax.generalized_gaussian_ssm.iterated_conditional_moments_gaussian_filter
 .. autofunction:: dynamax.generalized_gaussian_ssm.conditional_moments_gaussian_smoother
-.. autofunction:: dynamax.generalized_gaussian_ssm.iterated_conditional_moments_gaussian_smoother
+.. autofunction:: dynamax.generalized_gaussian_ssm.iterated_conditional_moments_gaussian_smoother
+
+Types
+-----
+
+.. autoclass:: dynamax.generalized_gaussian_ssm.ParamsGGSSM

docs/conf.py

@@ -99,4 +99,5 @@
 
 autosummary_generate = True
 autodoc_typehints = "description"
-add_module_names = False
+add_module_names = False
+autodoc_member_order = "bysource"

docs/index.rst

@@ -31,7 +31,7 @@ To install the latest releast of dynamax from PyPi:
 .. code-block:: console
 
    pip install dynamax                 # Install dynamax and core dependencies, or
-   pip install dynamax[notebooks]      # Install with dep's for demo notebooks
+   pip install dynamax[notebooks]      # Install with demo notebook dependencies
 
 
 To install the latest development branch:
@@ -83,23 +83,23 @@ The corresponding joint distribution has the following form
 
 .. math::
 
-      p(y_{1:T}, z_{1:T} | u_{1:T}) = p(z_1 | u_1) p(y_1 | z_1, u_1) \prod_{t=1}^T p(z_t | z_{t-1}, u_t) p(y_t | z_t, u_t)
+      p(y_{1:T}, z_{1:T} \mid u_{1:T}) = p(z_1 \mid u_1) \prod_{t=2}^T p(z_t \mid z_{t-1}, u_t) \prod_{t=1}^T p(y_t \mid z_t, u_t)
 
 
-Here :math:`p(z_t | z_{t-1}, u_t)` is called the transition or dynamics model,
-and :math:`p(y_t | z_{t}, u_t)` is called the observation or emission model.
+Here :math:`p(z_t \mid z_{t-1}, u_t)` is called the transition or dynamics model,
+and :math:`p(y_t \mid z_{t}, u_t)` is called the observation or emission model.
 (In both cases, the inputs :math:`u_t` are optional;
 furthermore, the observation model may have auto-regressive dependencies,
-in which case we write  :math:`p(y_t | z_{t}, u_t, y_{1:t-1})`.)
+in which case we write  :math:`p(y_t \mid z_{t}, u_t, y_{1:t-1})`.)
 
 We assume that we see the observations :math:`y_{1:T}`,
 and want to infer the hidden states, either
-using online filtering (i.e., computing  :math:`p(z_t|y_{1:t})`)
-or offline smoothing (i.e., computing  :math:`p(z_t|y_{1:T})`).
+using online filtering (i.e., computing  :math:`p(z_t \mid y_{1:t})`)
+or offline smoothing (i.e., computing  :math:`p(z_t \mid y_{1:T})`).
 We may also be interested in predicting future states,
-:math:`p(z_{t+h}|y_{1:t})`,
+:math:`p(z_{t+h} \mid y_{1:t})`,
 or future observations,
-:math:`p(y_{t+h}|y_{1:t})`,
+:math:`p(y_{t+h} \mid y_{1:t})`,
 where h is the forecast horizon.
 (Note that by using a hidden state  to represent the past observations,
 the model can have "infinite" memory, unlike a standard auto-regressive model.)
@@ -204,7 +204,7 @@ API documentation
 ==================
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 3
    :caption: API Documentation
 
    api