State Distributions - How to create
===================================

The Emissions matrix is created using State Distribution
objects. A state can be defined as a particular condition the cell is in
(eg. Resistant vs Susceptible to a drug treatment), which is represented by it’s
distribution(s)/parameter(s). Just as a state can be defined in whatever way
one desires, so can the distributions - typically as some type of
physical observation/phenotype.

The following codes present two examples of how to create StateDistribution.
The first example shows that a state of a cell, can be determined by two phenotypes,
that are cell fate and cell lifetime. Cell fate, which is whether a cell dies or gets to divide,
is representated by a Bernoulli distribution which is defined by one parameter, shown as bern_p. 
Cell lifetime, which is the intermitotic time for each cell, is represented by a Gamma distribution
which is defined by two parameters, shown as `gamma_a` and `gamma_scale`.
In this case, they are independent of each other -- but that may not always be the case.

The second example shows that a state of a cell can be determined only using cell size.
We used a Normal distribution to model the cell size, which has two parameters;
mean and standard deviation, shown as `norm_loc` and `norm_scale`, respectively. Please note that
in this case we are assuming all cells should be accounted for, i.e., we have no censorship.

Necessary components to create StateDistribution class compatible with tHMM are listed bellow:

1. The initialization function
------------------------------

This should take in all defining
parameters for all distributions being used. For the first example, Bernoulli and Gamma, it would be
`bern_p`, `gamma_a`, and `gamma_scale`, and for the sencond example, the normal distribution,
this would be `norm_loc`, and `norm_scale`.
These are assigned as instances of the class object. One can assert that the given values
actually make sense for the distribution. For example, a normal
distribution cannot have a negative or zero standard deviation.

.. code:: python

    import numpy as np
    import scipy.stats as sp

    #Bernoulli/Gamma
    def __init__(self, bern_p, gamma_a, gamma_scale):
        """ Initialization function should take in just in the parameters for the observations that comprise the multivariate random variable emission they expect their data to have. """
        self.bern_p = bern_p
        self.gamma_a = gamma_a
        self.gamma_scale = gamma_scale
        self.params = [self.bern_p, self.gamma_a, self.gamma_scale]

.. code:: python

    #Normal 
    def __init__(self, norm_loc, norm_scale):
        """ Initialization function should take in just in the parameters for the observations that comprise the multivariate random variable emission they expect their data to have. """
        self.norm_loc = norm_loc
        assert norm_scale > 0, "A non-valid scale has been given. Please provide a scale > 0"
        self.norm_scale = norm_scale
        self.params = [self.norm_loc, self.norm_scale]


2. Random variable generator function
-------------------------------------

To create synthetic data, we need to generate random variables of the defined distributions as cell observations.
We do that using the rvs (random variables) function, which actually uses the built-in scipy function rvs.
It takes in size, which represents the number of cells in
the lineage, and assigns each one a random variable from the
characteristic distributions of that state. There is one per
distribution. The time_censor variable exists due to the time based
nature of the gamma distribution. Such a variable is unnecessary in the
normal example, but may be needed depending on the type of censoring
done.
While creatig our synthetic data, at first, we assume all cells are observed and there is no censored cells.
The variables `gamma_obs_censor` and `norm_obs_censor` are created and set to 1 for each cell to represent that.
We do the censoring later.
The rvs function returns the observed phenotypes as a tuple of lists.

.. code:: python

    # Bernoulli/Gamma
    def rvs(self, size: int):
        """ User-defined way of calculating a random variable given the parameters of the state stored in their StateType object. """

        bern_obs = sp.bernoulli.rvs(p=self.params[0], size=size)  # bernoulli observations
        gamma_obs = sp.gamma.rvs(a=self.params[1], scale=self.params[2], size=size)  # gamma observations
        gamma_obs_censor = [1] * size  # 1 if observed

        # These tuples of observations will go into the cells in the lineage tree.
        return bern_obs, gamma_obs, gamma_obs_censor

.. code:: python

    # Normal
    def rvs(self, size):  # user has to identify what the multivariate (or univariate if he or she so chooses) random variable looks like
        """ User-defined way of calculating a random variable given the parameters of the state stored in that observation's object. """

        norm_obs = sp.norm.rvs(loc=self.norm_loc, scale=self.norm_scale, size=size)  # normal observations
        norm_obs_censor = [1] * size  # 1 if observed

        # These tuples of observations will go into the cells in the lineage tree.
        return norm_obs, norm_obs_censor


3. PDF
------

The third required function is the probability distribution function (pdf).
For univariate and independent multivariate distributions, it is fairly
simple and can just use the already implemented pdf functions in scipy.
For more complex multivariate distributions, the pdf might be more
complicated and require a custom function. It is to calculate the likelihood of the observations.

In the Bernoulli/Gamma example, we assume the two phenotypes are independent
and we add their Log-lilelihood to find the total log-likelihood,
which is equivalent to multiplying the likelihoods.

In this function, we consider the censorship of the observations, 
based on the integer value we assigned to them to show whether they are censored or not.
The fully observed cells are fed to `logpdf` to calculate the likelihood,
and those cells that have missing information are fed to `logsf`.
Those cells that died are then removed in the first example that cell's fate matters.

.. code:: python

    # Bernoulli/Gamma
    def logpdf(self, x: np.ndarray):
        """ User-defined way of calculating the likelihood of the observation stored in a cell.
        """
        ll = np.zeros(x.shape[0])

        # Update uncensored Gamma
        ll[x[:, 2] == 1] += sp.gamma.logpdf(x[x[:, 2] == 1, 1], a=self.params[1], scale=self.params[2])

        # Update censored Gamma
        ll[x[:, 2] == 0] += sp.gamma.logsf(x[x[:, 2] == 0, 1], a=self.params[1], scale=self.params[2])

        # Remove dead cells
        ll[x[:, 0] == 0] = 0.0

        # Update for observed Bernoulli
        ll[np.isfinite(x[:, 0])] += sp.bernoulli.logpmf(x[np.isfinite(x[:, 0]), 0], self.params[0])

        return ll

.. code:: python

    # Normal
    def logpdf(self, tuple_of_obs):  # user has to define how to calculate the likelihood
        """ User-defined way of calculating the likelihood of the observation stored in a cell. """
        
        ll = np.zeros(x.shape[0])

        ll += sp.norm.logpdf(tuple_of_obs[0], self.norm_loc, self.norm_scale)

        return ll


4. The estimator
----------------

The `estimator` method provides estimation of distribution parameters given the observations.
The user must define their own estimator function. 
Typically this would just be the maximum likelihood estimate,
but in our fisr example, due the the incorporation of the gammas term 
and that we have censorship, it is slightly more complicated.
One must find the MLE by taking the product of the pdf over all
obervations (the likelihood), then taking the log, then the derivative
and setting equal to zero to find the optimal value. For the Bernoulli,
for example, the likelihood is the product from i=1 to n (where there
are n observations) of p^x_i \* (1-p)^1-x_i.

For the Bernoulli it becomes the product from i=1 to n of (p^x_i \*
(1-p)\ :sup:`1-x_i)`\ z_i. The MLE is then calculated as normal. Once
this modified MLE is found for each parameter, the estimator function
must simply calculate and return it given the observations and gammas.
The Gamma estimator function takes in the observations, and uses the 
minimize function of the scipy.optimize to find the parameters and
the function is located in the lineage/states/stateCommon.py


.. code:: python

    # Bernoulli/Gamma
    def estimator(self, x: list, gammas: np.array):
        """ User-defined way of estimating the parameters given a list of the tuples of observations from a group of cells. """

        # getting the observations as individual lists
        x = np.array(x)
        bern_obs = x[:, 0].astype('bool')
        γ_obs = x[:, 1]
        gamma_obs_censor = x[:, 2]

        b_mask = np.isfinite(bern_obs)
        # Both unoberved and dead cells should be removed from gamma
        g_mask = np.logical_and(np.isfinite(γ_obs), bern_obs)
        assert np.sum(g_mask) > 0, f"All the cells are eliminated from the Gamma estimator."

        # Handle an empty state
        if np.sum(gammas[b_mask]) == 0.0:
            self.params[0] = np.average(bern_obs[b_mask])
        else:
            self.params[0] = np.average(bern_obs[b_mask], weights=gammas[b_mask])

        # Don't allow Bernoulli to hit extremes
        self.params[0] = np.clip(self.params[0], 0.00001, 0.99999) # bernoulli parameter

        self.params[1], self.params[2] = gamma_estimator(γ_obs[g_mask], gamma_obs_censor[g_mask], gammas[g_mask], self.params[1:3]) # gamma shape and scale

.. code:: python

    # Normal
    def estimator(self, x: list, gammas: np.array):
        """ User-defined way of estimating the parameters given a list of the tuples of observations from a group of cells. """

        # getting the observations as individual lists
        x = np.array(x)
        norm_obs = x[:, 0]
        norm_obs_censor = x[:, 2]

        # mask for shape
        s_mask = np.isfinite(norm_obs)
        assert np.sum(s_mask) > 0, f"All the cells are eliminated from the Gamma estimator."
        self.params[0] = np.average(norm_obs[s_mask], weights=gammas[s_mask]) # mean
        self.params[1] = ((np.sum(gammas[s_mask]*(norm_obs-self.params[0])**2) + 1e-10)/ (np.sum(gammas[s_mask])+ 1e-10))**.5 # std


Example
-------

Now that we have a functioning Gaussian StateDistribution, let's try it
with the overall model. As in the overview, we first define our initial probability vector and
the state transition probability matrices.

.. code:: python

    from lineage.LineageTree import LineageTree

    pi = np.array([0.6, 0.4], dtype="float")
    
    T = np.array([[0.75, 0.25],
                  [0.25, 0.75]], dtype="float")

Creating the Emissions for two states:

.. code:: python

    # E: states are defined as StateDistribution objects
    
    # Normal distribution state 0 parameters "Basal"
    norm_loc0 = 14
    norm_scale0 = 2
    
    # Normal distribution state 1 parameters "Luminal"
    norm_loc1 = 19
    norm_scale1 = 3
    
    state_obj0 = StateDistribution(norm_loc0, norm_scale0)
    state_obj1 = StateDistribution(norm_loc1, norm_scale1)
    
    E = [state_obj0, state_obj1]

Creating the lineage tree:

.. code:: python

    lineage1 = LineageTree.rand_init(pi, T, E, desired_num_cells=2**5 - 1)
    # These are the minimal arguments required to instantiate lineages
    print(lineage1)
    print("\n")

Below is the analysis for a single lineage. Note that the state objects
are merely switched. However, the model fairly accurately predicts the
transition matrix and state parameters.

.. code:: python

    from lineage.Analyze import Analyze_list

    X = [lineage1] # population just contains one lineage
    tHMMobj_list, LL, gammas = Analyze_list([X], 2, write_states=True) # find two states
    tHMMobj = tHMMobj_list[0]

.. code:: python

    print(tHMMobj.estimate.pi)

.. code:: python

    print(tHMMobj.estimate.T)

.. code:: python

    for state in range(tHMMobj.num_states):
        print("State {}:".format(state))
        print("                    estimated state:", tHMMobj.estimate.E[state])
        print("original parameters given for state:", E[state])
        print("\n")

Competing risks: tying lifetime to fate
---------------------------------------

The Bernoulli/Gamma example above treats a cell's fate and its lifetime as two
independent observations, and it discards the lifetime of any cell that dies --
that is what ``ll[x[:, 0] == 0] = 0.0`` does. It also scores a time-censored cell
with ``logsf`` of the division clock alone, which states that the cell has not yet
divided but says nothing about it not having died. Those two choices together leave
the emission unnormalized: summed over the outcomes of a cell watched to a finite
horizon, its total probability exceeds one whenever there is censoring.

``lineage.states.StateDistributionCR`` shows the alternative. Each phase carries two
latent clocks -- a division clock ``T_D`` and a death clock ``T_X`` -- and what we
observe is ``min(T_D, T_X)`` together with which one fired. The likelihood has the
three standard competing-risks cases:

.. code:: python

    # transition seen at t:   f_D(t) * S_X(t)
    # death seen at t:        f_X(t) * S_D(t)
    # censored at t:          S_D(t) * S_X(t)

    divided, died, censored = event_masks(x)
    timed = divided | died | censored

    ll = np.zeros(x.shape[0])
    # every timed cell survived both clocks up to t ...
    ll[timed] += div.logsf(t[timed]) + death.logsf(t[timed])
    # ... and whichever clock fired swaps its survival term for a density
    ll[divided] += div.logpdf(t[divided]) - div.logsf(t[divided])
    ll[died] += death.logpdf(t[died]) - death.logsf(t[died])

Two things follow. Death times now carry information rather than being thrown away,
and the division probability stops being a free parameter: it is derived as
``P(divide) = int f_D(t) S_X(t) dt``, so the fraction of cells that die and the times
at which they die are forced to agree with each other.

The estimator is no harder than before. In the complete-data likelihood the two
clocks separate, so each is an independently weighted right-censored fit over exactly
the same cells -- once with "transition observed" as the event indicator, once with
"death observed". Both reuse ``gamma_estimator`` from ``stateCommon``.

Choosing the death clock is an empirical question, and the answer differs by phase in
the Heiser lab data (see ``lineage/figures/figureS18.py``). G1 death times are
memoryless, so that clock is a one-parameter exponential and the phase costs no more
degrees of freedom than the Bernoulli/Gamma form did. G2 death times have a strongly
increasing hazard, so that clock keeps a free shape.