.. _tablefs:

Calculating frequency spectra from tree sequences
====================================================================

The mutation frequency spectrum is an array-like object representing the
number of mutation events occurring at different frequencies.  Generating
frequency spectra, or ``fs``, from simulations is handled by :func:`fwdpy11.TableCollection.fs`.

We have to define some conventions:

* A sample list is a list of node indexes.  The node indexes
  are stored in a :class:`numpy.ndarray` with :class:`numpy.dtype`
  ``numpy.int32``.
* We support ``fs`` calculated from multiple lists of samples.
* An ``fs`` always contains the zero and fixed bins.  Thus, for
  a sample of ``n`` nodes, there are ``n + 1`` entries in the ``fs``.
  The first and last values are for frequencies zero and ``n``, respectively.
  Singletons start in bin ``1``, etc..
* An ``fs`` from a single sample list is represented by a :class:`numpy.ma.MaskedArray`,
  object with the zero and fixed bins masked.
* An ``fs`` from more than one sample list is stored in a :class:`sparse.COO` sparse
  matrix.

For our examples, we will initialize a :class:`fwdpy11.DiploidPopulation` from
a :class:`tskit.TreeSequence` generated by :func:`msprime.simulate`.

.. note::

    These examples do not show how to get ``fs`` separately
    for neutral and non-neutral mutations.  See
    :func:`fwdpy11.TableCollection.fs` for details.

.. ipython:: python

    import fwdpy11
    import msprime
    import numpy as np

    rng = fwdpy11.GSLrng(4321678)
    config = [msprime.PopulationConfiguration(500),
              msprime.PopulationConfiguration(500)]
    ts = msprime.simulate(population_configurations=config,
                          Ne=500., random_seed=777,
                          migration_matrix=np.array([0,0.1,0.1,0]).reshape(2,2),
                          recombination_rate=0.25)

    pop = fwdpy11.DiploidPopulation.create_from_tskit(ts)
    md = np.array(pop.diploid_metadata, copy=False)
    np.unique(md['deme'], return_counts=True)

    nmuts = fwdpy11.infinite_sites(rng, pop, 0.1)
    nmuts

The following blocks show several methods for obtaining the ``fs`` from lists of nodes.
First, let's get the lists of nodes from the two demes in our population:

.. ipython:: python

    nodes = np.array(pop.tables.nodes, copy=False)
    alive_nodes = pop.alive_nodes
    deme0_nodes = alive_nodes[np.where(nodes['deme'][alive_nodes] == 0)[0]]
    deme1_nodes = alive_nodes[np.where(nodes['deme'][alive_nodes] == 1)[0]]

Get an ``fs`` from nodes found only in deme 0:

.. ipython:: python

    pop.tables.fs([deme0_nodes[:10]])

Get a joint ``fs`` from nodes from each deme:

.. ipython:: python


    fs = pop.tables.fs([deme0_nodes[:10], deme1_nodes[50:55]])
    fs

Obtain the full :class:`numpy.ndarray` for the joint ``fs``:

.. ipython:: python

    fs.todense()

.. warning::

    The joint ``fs`` can take a lot of memory!

We can use standard array operations to get the marginal ``fs`` from our joint ``fs``:

.. ipython:: python

    fs.sum(axis=1).todense()
    fs.sum(axis=0).todense()

.. note::

    Be careful when processing sparse matrix objects!  Naive application of regular
    :mod:`numpy` functions can lead to erroneous results.  Be sure to check the
    :mod:`sparse` documentation.

The marginalization can be tedious for many samples, so you can have it happen automatically,
in which case a :class:`dict` is returned, keyed by sample list index:

.. ipython:: python

    fs = pop.tables.fs([deme0_nodes[:10],deme1_nodes[50:55]],
                       marginalize=True)
    for key, value in fs.items():
        print(key)
        print(value)
        print(value.data)

.. note::

    Marginalizing in this way preserves the convention that the 1-d ``fs``
    objects are instances of :class:`numpy.ma.MaskedArray`.

To see how the :class:`dict` keying works, let's flip the sample lists:

.. ipython:: python

    fs = pop.tables.fs([deme1_nodes[50:55],deme0_nodes[:10]],
                       marginalize=True)
    for key, value in fs.items():
        print(key)
        print(value)
        print(value.data)

If you only want the ``fs`` from particular regions of the genome.  By default,
the ``fs`` is the sum across windows:

.. ipython:: python

    pop.tables.fs([deme0_nodes[:10]],
                  windows=[(0.1, 0.2), (0.8, 0.9)])

You can get the ``fs`` separately by window, too:

.. ipython:: python

    pop.tables.fs([deme0_nodes[:10]], windows=[(0.1, 0.2), (0.8, 0.9)],
                  separate_windows=True)

You can also get a joint ``fs`` marginalized by sample list and separated
by window.  In this case, the return value is a :class:`list` containing
the :class:`dict` for each window:

.. ipython:: python

    pop.tables.fs([deme0_nodes[:10], deme1_nodes[:20]],
                  windows=[(0.1, 0.2), (0.8, 0.9)],
                  marginalize=True,
                  separate_windows=True)

Simplifying to the samples
---------------------------------------------------------

Finally, it is sometimes more efficient to simplify the tree sequences with
respect to the sample nodes.  For example, if there are a vast number of
ancient samples and you are processing each time point separately
(see :func:`fwdpy11.DiploidPopulation.sample_timepoints`), then *not* simplifying
means iterating over trees that are redundant/irrelevant to the history of 
the current time point.  In order to get the ``fs`` from a simplified
tree sequence, pass ``simplify=True`` when calling :func:`fwdpy11.TableCollection.fs`.
