latentmi.lmi
============

.. py:module:: latentmi.lmi


Classes
-------

.. autoapisummary::

   latentmi.lmi.EarlyStopper


Functions
---------

.. autoapisummary::

   latentmi.lmi.train
   latentmi.lmi.learn_representation
   latentmi.lmi.estimate


Module Contents
---------------

.. py:class:: EarlyStopper(patience=1)

   Early stopping that returns best weights
   trying to replicate the Keras callback


   .. py:method:: early_stop(validation_loss, model)


.. py:function:: train(model, X_train, Y_train, X_test, Y_test, batch_size=512, lr=0.0001, epochs=300, patience=30, quiet=True)

   training loop for LMI models

   :param model: LMI model

   :param X_train: train samples, shape (N_samples, N_features)
   :param Y_train: train samples, shape (N_samples, N_features)
   :param X_test: test samples, shape (N_samples, N_features)
   :param Y_test: test samples, shape (N_samples, N_features)

   :param batch_size: samples per batch, defaults to 512
   :param lr: learning rate for Adam optimizer, defaults to 1e-4
   :param epochs: max number of epochs, defaults to 300
   :param patience: epochs without val. loss decline before early stopping,
                    defaults to 300
   :param quiet: suppress training progress display, defaults to True


.. py:function:: learn_representation(Xs, Ys, train_indices, test_indices, regularizer='models.AECross', alpha=1, lam=1, N_dims=8, batch_size=512, lr=0.0001, epochs=300, validation_split=0.3, patience=30, quiet=True, device='cpu')

   train paired AE model and embed data

   :param Xs: input data, array with shape (N_samples, N_dims). ordering must align with Y.
   :param Ys: input data, array with shape(N_samples, N_dims). ordering must align with X.

   :param train_indices: list, indices of train samples
   :param test_indices:list, indices of test samples

   :param regularizer: type of regularization, defaults to AECross.
                       can be changed to 'models.AEMINE' or
                       'models.AEInfoNCE' but not recommended.
   :param alpha: self-reconstruction loss weight, defaults to 1
   :param lam: cross-reconstruction regularization weight, defaults to 1
   :param N_dims: dimensions in each latent representation, defaults to 8

   :param batch_size: samples per batch, defaults to 512
   :param lr: learning rate for Adam optimizer, defaults to 1e-4
   :param epochs: max number of epochs, defaults to 300
   :param validation_split: fraction train/test split, defaults to 0.5
   :param patience: epochs without val. loss decline before early stopping,
                    defaults to 300
   :param quiet: suppress training progress display, defaults to True


.. py:function:: estimate(Xs, Ys, regularizer='models.AECross', alpha=1, lam=1, N_dims=8, validation_split=0.5, estimate_on_val=True, batch_size=512, lr=0.0001, epochs=300, patience=30, quiet=True, device=None)

   return pMIs (with NaNs for points not included in KSG estimate), embeddings, trained model

   :param Xs: input data, array with shape (N_samples, N_dims). ordering must align with Y.
   :param Ys: input data, array with shape(N_samples, N_dims). ordering must align with X.

   :param regularizer: type of regularization, defaults to AECross.
                       can be changed to 'models.AEMINE' or
                       'models.AEInfoNCE' but not recommended.
   :param alpha: self-reconstruction loss weight, defaults to 1
   :param lam: cross-reconstruction regularization weight, defaults to 1
   :param N_dims: dimensions in each latent representation, defaults to 8

   :param batch_size: samples per batch, defaults to 512
   :param lr: learning rate for Adam optimizer, defaults to 1e-4
   :param epochs: max number of epochs, defaults to 300
   :param validation_split: fraction train/test split, defaults to 0.5
   :param patience: epochs without val. loss decline before early stopping,
                    defaults to 300
   :param quiet: suppress training progress display, defaults to True
   :param device: device for torch to train model, defaults to cuda if available, else cpu.

   :return: array of pointwise mutual information estimates, order aligned with input. NaNs
            values not included in KSG estimate. mean of this array is MI estimate.
   :return: tuple of arrays of coordinates of latent embeddings. first index is X embeddings, second index is Y.
   :return: Pytorch object for trained representation learning model