Revision becc7b948123f8683c49ed41480ca2682d979a7f authored by Yan Wong on 09 December 2022, 18:50:12 UTC, committed by mergify[bot] on 09 December 2022, 19:02:47 UTC
Co-authored-by: Peter Ralph <petrel.harp@gmail.com>
1 parent 7fbcdce
demography.py
#
# Copyright (C) 2015-2021 University of Oxford
#
# This file is part of msprime.
#
# msprime is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# msprime is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with msprime. If not, see <http://www.gnu.org/licenses/>.
#
"""
Module responsible for defining and debugging demographic models.
"""
from __future__ import annotations
import collections
import copy
import dataclasses
import enum
import inspect
import itertools
import json
import logging
import math
import numbers
import operator
import re
import sys
import textwrap
import warnings
from typing import Any
from typing import ClassVar
from typing import MutableMapping
import demes
import numpy as np
import tskit
from . import ancestry
from . import core
from . import species_trees
logger = logging.getLogger(__name__)
class IncompletePopulationMetadataWarning(UserWarning):
"""
Warning raised when we don't have sufficient information to fill
out population metadata.
"""
class LruCache(collections.OrderedDict):
# LRU example from the OrderedDict documentation
def __init__(self, maxsize=128, *args, **kwds):
self.maxsize = maxsize
super().__init__(*args, **kwds)
def __getitem__(self, key):
value = super().__getitem__(key)
self.move_to_end(key)
return value
def __setitem__(self, key, value):
super().__setitem__(key, value)
if len(self) > self.maxsize:
oldest = next(iter(self))
del self[oldest]
_population_table_cache = LruCache(16)
def _build_population_table(populations):
"""
Return a tskit PopulationTable instance encoding the metadata for the
specified populations. Because encoding metadata is quite expensive
we maintain an LRU cache.
"""
population_metadata = []
for population in populations:
metadata = {
"name": population.name,
"description": population.description,
}
if population.extra_metadata is not None:
intersection = set(population.extra_metadata.keys()) & set(metadata.keys())
if len(intersection) > 0:
printed_list = list(sorted(intersection))
raise ValueError(
f"Cannot set standard metadata key(s) {printed_list} "
"using extra_metadata. Please set using the corresponding "
"property of the Population class."
)
metadata.update(population.extra_metadata)
population_metadata.append(metadata)
# The only thing we store in the Population table is the metadata, so
# we cache based on this.
key = json.dumps(population_metadata, sort_keys=True)
if key not in _population_table_cache:
table = tskit.PopulationTable()
table.metadata_schema = tskit.MetadataSchema(
{
"codec": "json",
"type": "object",
"properties": {
"name": {"type": "string"},
"description": {"type": ["string", "null"]},
},
# The name and description fields are always filled out by
# msprime, so we tell downstream tools this by making them
# "required" by the schema.
"required": ["name", "description"],
"additionalProperties": True,
}
)
for metadata in population_metadata:
table.add_row(metadata=metadata)
_population_table_cache[key] = table
return _population_table_cache[key]
def check_num_populations(num_populations):
"""
Check if an input number of populations is valid.
"""
if num_populations < 1:
raise ValueError("Must have at least one population")
def check_migration_rate(migration_rate):
"""
Check if an input migration rate makes sense.
"""
if migration_rate < 0:
raise ValueError("Migration rates must be non-negative")
@dataclasses.dataclass
class Population:
"""
A single population in a :class:`.Demography`. See the
:ref:`sec_demography_populations` section for more information on
what populations represent, and how they can be used.
.. warning:: This class should not be instantiated directly. Please use
:meth:`.Demography.add_population` method instead.
"""
initial_size: float = 0.0
"""
The absolute size of the population at time zero.
See the :ref:`sec_demography_populations_initial_size` section
for more details and examples.
"""
growth_rate: float = 0.0
"""
The exponential growth rate of the population per generation (forwards in time).
Growth rates can be negative. This is zero for a constant population size,
and positive for a population that has been growing.
See the :ref:`sec_demography_populations_growth_rate` section for more
details and examples.
"""
name: str | None = None
"""
The name of the population. If specified this must be a uniquely
identifying string and must be a valid Python identifier (i.e., could be
used as a variable name in Python code).
See :ref:`sec_demography_populations_identifiers` for more details
and recommendations on best practise.
"""
description: str = ""
"""
A concise description of the population. Defaults to the empty string if not
specified.
"""
extra_metadata: dict = dataclasses.field(default_factory=dict)
"""
A JSON-encodable dictionary of metadata items to be stored in the
associated tskit population object. This dictionary must not contain keys
for any of the pre-defined metadata items.
See the :ref:`sec_demography_populations_metadata` section for more
details and examples.
"""
default_sampling_time: float | None = None
"""
The default time at which samples are drawn from this population. See the
:ref:`sec_demography_populations_default_sampling_time` section for more
details.
"""
initially_active: bool | None = None
"""
If True, this population will always be initially active, regardless
of whether it participates in a :ref:`sec_demography_events_population_split`.
If not set, or None, the initial state of the population will be
set automatically depending on the events declared in the demography.
See the :ref:`sec_demography_populations_life_cycle` section for
more details.
"""
id: int | None = dataclasses.field(default=None) # noqa: A003
"""
The integer ID of this population within the parent :class:`.Demography`.
This attribute is assigned by the Demography class and should not be set
or changed by user code.
"""
def asdict(self):
return dataclasses.asdict(self)
def validate(self):
if self.initial_size < 0:
raise ValueError("Negative population size")
if self.name is None:
raise ValueError("A population name must be set.")
if not self.name.isidentifier():
raise ValueError("A population name must be a valid Python identifier")
@dataclasses.dataclass
class Demography(collections.abc.Mapping):
"""
The definition of a demographic model for an msprime simulation,
consisting of a set of populations, a migration matrix, and a list
of demographic events. See the :ref:`sec_demography` section for
detailed documentation on how to define, debug and simulate with
demography in msprime.
Please see the :ref:`sec_demography_demography_objects` section
for details of how to access and update population information
within a model.
Demography objects implement the Python
:class:`python:collections.abc.Mapping` protocol, in which the
keys are **either** the population ``name`` or
integer ``id`` values (see the :ref:`sec_demography_populations_identifiers`
section for more information) and the values are :class:`.Population`
objects.
In general, population references in methods such as
:meth:`.Demography.add_population_split` can either be string names
or integer IDs, and the two forms can be used interchangeably.
"""
populations: list[Population] = dataclasses.field(default_factory=list)
events: list = dataclasses.field(default_factory=list)
# Until we can use numpy type hints properly, it's not worth adding them
# here. We still have to add in ignores below for indexed assignment errors.
migration_matrix: Any | None = None
def __post_init__(self):
if self.migration_matrix is None:
N = self.num_populations
self.migration_matrix = np.zeros((N, N))
# People might get cryptic errors from passing in copies of the same
# population, so check for it.
if len({id(pop) for pop in self.populations}) != len(self.populations):
raise ValueError("Population objects must be distinct")
# Assign the IDs and default names, if needed.
for j, population in enumerate(self.populations):
if population.id is not None:
raise ValueError(
"Population ID should not be set before using to create "
"a Demography"
)
population.id = j
if population.name is None:
population.name = f"pop_{j}"
self._validate_populations()
def add_population(
self,
*,
initial_size: float,
growth_rate: float | None = None,
name: str | None = None,
description: str | None = None,
extra_metadata: dict | None = None,
default_sampling_time: float | None = None,
initially_active: bool | None = None,
) -> Population:
"""
Adds a new :class:`.Population` to this :class:`.Demography` with the
specified parameters. The new population will have ID equal to the
the number of populations immediately before ``add_population``
is called, such that the first population added has ID 0, the next
ID 1 and so on. If the ``name`` is not specified, this defaults
to ``"pop_{id}"``. An :ref:`sec_demography_populations_initial_size`
value must be specified (but may be zero).
:param float initial_size: The number of individuals of the population
at time zero. See the :ref:`sec_demography_populations_initial_size`
section for more details and examples.
:param float growth_rate: The exponential growth rate of the
population. See the :ref:`sec_demography_populations_growth_rate`
section for more details and examples.
:param str name: The human-readable identifier for this population.
If not specified, defaults to the string ``"pop_{id}"`` where
``id`` is the population's integer ID. See
:ref:`sec_demography_populations_identifiers` for more details
and recommendations on best practise.
:param str description: A concise but informative description of
what this population represents within the wider model. Defaults
to the empty strings.
:param dict extra_metadata: Extra metadata to associate with
this population that will be stored tree sequences output
by :func:`.sim_ancestry`. See the
:ref:`sec_demography_populations_metadata` section for more
details and examples.
:param float default_sampling_time: The time at which samples
will be taken from this population, if a time in not otherwise
specified. By default this is determined by the details
of the model, and whether populations are ancestral in
:ref:`sec_demography_events_population_split` events. See the
:ref:`sec_demography_populations_default_sampling_time` section
for more details.
:param bool initially_active: Whether this population is initially
:ref:`active<sec_demography_populations_life_cycle>`.
By default this is determined by the details
of the model, and whether populations are ancestral in
:ref:`sec_demography_events_population_split` events. See the
:ref:`sec_demography_populations_life_cycle` section
for more details.
:returns: The new :class:`.Population` instance.
:rtype: Population
"""
N = self.num_populations
population = Population(
id=N,
initial_size=initial_size,
growth_rate=0 if growth_rate is None else growth_rate,
name=f"pop_{N}" if name is None else name,
description="" if description is None else description,
extra_metadata={} if extra_metadata is None else extra_metadata,
default_sampling_time=default_sampling_time,
initially_active=initially_active,
)
self.populations.append(population)
# TODO this is inefficient - we should probably store the migration
# matrix in a sparse dictionary form internally.
M = self.migration_matrix
self.migration_matrix = np.zeros((N + 1, N + 1))
self.migration_matrix[:N, :N] = M
self._validate_populations()
return population
def _add_population_from_old_style(
self, pop_config: PopulationConfiguration, name: str | None = None
) -> Population:
population = self.add_population(
name=name,
initial_size=pop_config.initial_size,
growth_rate=pop_config.growth_rate,
)
metadata = pop_config.metadata
if metadata is not None and isinstance(metadata, collections.abc.Mapping):
metadata = metadata.copy()
if "name" in metadata:
population.name = metadata.pop("name")
if name is not None and name != population.name:
# Maybe this should be a warning, or just ignored entirely?
raise ValueError(
"Population name already set in old-style metadata "
f"({name}) and doesn't match supplied name "
f"({population.name})"
)
if "description" in metadata:
population.description = metadata.pop("description")
population.extra_metadata = metadata
return population
def add_event(self, event: DemographicEvent) -> DemographicEvent:
if not isinstance(event, DemographicEvent):
raise TypeError("Events must be instances of DemographicEvent")
event.demography = self
self.events.append(event)
return event
def set_migration_rate(
self, source: str | int, dest: str | int, rate: float
) -> None:
"""
Sets the backwards-time rate of migration from the specified ``source``
population to ``dest`` to the specified value. This has the effect of
setting ``demography.migration_matrix[source, dest] = rate``. It is
the rate at which a lineage currently in ``source`` moves to ``dest``
as one follows the lineage back through time.
.. important:: Note this is the
:ref:`backwards in time<sec_demography_direction_of_time>`;
migration rate and that
``source`` and ``dest`` are from the perspective of lineages in the
coalescent process. See :ref:`sec_demography_migration` for more
details and clarification on this vital point.
The ``source`` and ``dest`` populations can be referred to either by
their integer ``id`` or string ``name`` values.
:param str,int source: The source population from which lineages originate
in the backwards-time process.
:param str,int dest: The destination population where lineages are move
to in the backwards-time process.
:param float rate: The per-generation migration rate.
"""
source = self[source].id
dest = self[dest].id
if source == dest:
raise ValueError("The source and dest populations must be different")
self.migration_matrix[source, dest] = rate # type: ignore
def set_symmetric_migration_rate(
self,
populations: list[str | int],
rate: float,
) -> None:
"""
Sets the symmetric migration rate between all pairs of populations in
the specified list to the specified value. For a given pair of population
IDs ``j`` and ``k``, this sets ``demography.migration_matrix[j, k] = rate``
and ``demography.migration_matrix[k, j] = rate``.
Populations may be specified either by their integer IDs or by
their string names.
:param list populations: An iterable of population identifiers (integer
IDs or string names).
:param float rate: The value to set the migration matrix entries to.
"""
# There's an argument for not checking this so that corner cases on
# single population models can be handled. However, it's nearly always
# going to be a user error where someone forgets the second population
# so it seems better to raise an error to prevent hard-to-detect mistakes.
if len(populations) < 2:
raise ValueError("Must specify at least two populations")
pop_ids = [self[identifier].id for identifier in populations]
for pop_j, pop_k in itertools.combinations(pop_ids, 2):
self.migration_matrix[pop_j, pop_k] = rate # type: ignore
self.migration_matrix[pop_k, pop_j] = rate # type: ignore
# Demographic events.
def _check_population_references(self, populations: list[str | int]):
for pop_ref in populations:
# Slightly unsure whether this call can get optimised out, but
# there doesn't seem to be a better way to do it without duplicating
# the KeyError message
self[pop_ref]
def _add_activate_population_event(
self, time: float, *, population: str | int
) -> ActivatePopulationEvent:
"""
Activates a population at the specified time. The population is expected to be
initially inactive.
:param float time: The time at which this event occurs in generations.
:param str, int population: The population to activate.
"""
self._check_population_references([population])
return self.add_event(ActivatePopulationEvent(time=time, population=population))
def add_population_split(
self, time: float, *, derived: list[str | int], ancestral: str | int
) -> PopulationSplit:
"""
Adds a population split event at the specified time. In a population
split event all lineages from the (more recent) derived populations
move to the (more ancient) ancestral population. Forwards in time,
this corresponds to the ancestral population splitting into the
derived populations.
See the :ref:`sec_demography_events_population_split` section
for more details and examples.
In addition to moving lineages from the derived population(s) into the
ancestral population, a population split has the following additional
effects:
- All derived populations are set to
:ref:`inactive<sec_demography_populations_life_cycle>`.
- All migration rates to and from the derived populations are set to 0.
- Population sizes and growth rates for the derived populations are set
to 0.
- The ``default_sampling_time`` of the ``ancestral`` :class:`.Population`
is set to the time of this event, **if** the ``default_sampling_time``
for the ancestral population has not already been set.
:param float time: The time at which this event occurs in generations.
:param list(str, int) derived: The derived populations.
:param str, int ancestral: The ancestral population.
"""
self._check_population_references(list(derived) + [ancestral])
pop = self[ancestral]
if pop.initially_active is None:
pop.initially_active = False
if pop.default_sampling_time is None:
pop.default_sampling_time = time
return self.add_event(
PopulationSplit(time=time, derived=derived, ancestral=ancestral)
)
def add_admixture(
self,
time: float,
*,
derived: str | int,
ancestral: list[str | int],
proportions: list[float],
) -> Admixture:
"""
Adds an admixture event at the specified time. In an admixture
event all lineages from a (more recent) ``derived`` population
move to a list of (more ancient) ``ancestral`` populations according
to a list of ``proportions``, such that a given lineage has a
probability ``proportions[j]`` of being moved to the population
``ancestral[j]``. This movement of lineages backwards in time
corresponds to the initial state of the admixed derived population
the specified ``time`` being composed of individuals from the
specified ``ancestral`` populations in the specified ``proportions``.
See the :ref:`sec_demography_events_admixture` section
for more details and examples.
In addition to moving lineages from the derived population into the
ancestral population(s), an admixture has the following additional
effects:
- The derived population is set to
:ref:`inactive<sec_demography_populations_life_cycle>`.
- The ancestral populations are set to
:ref:`active<sec_demography_populations_life_cycle>`, if they are
not already active.
- All migration rates to and from the derived population are set to 0.
- Population sizes and growth rates for the derived population are set
to 0, and the population is marked as inactive.
:param float time: The time at which this event occurs in generations.
:param str, int derived: The derived population.
:param list(str, int) ancestral: The ancestral populations.
:param list(float) proportions: The proportion of the derived population
from each of the ancestral populations at the time of the event.
"""
self._check_population_references(list(ancestral) + [derived])
# Useful feature here might be to support taking n - 1 proportion values
# and computing 1 - sum for the last value. Could be tedious for users to
# do this manually.
if not math.isclose(sum(proportions), 1.0):
raise ValueError("Sum of the admixture proportions must be approximately 1")
return self.add_event(
Admixture(
time=time, derived=derived, ancestral=ancestral, proportions=proportions
)
)
def add_mass_migration(
self,
time: float,
*,
source: str | int,
dest: str | int,
proportion: float,
) -> MassMigration:
"""
Adds a mass migration (or "pulse migration") event at the specified
time. In a mass migration event, lineages in the ``source`` population
are moved to the ``dest`` population with probability ``proportion``.
Forwards-in-time, this corresponds to individuals migrating
**from** population ``dest`` **to** population ``source``.
Please see the :ref:`sec_demography_events_mass_migration` section
for more details and examples.
.. warning:: Mass migrations are an advanced feature and should
only be used if the required population dynamics cannot be
modelled by :ref:`sec_demography_events_population_split`
or :ref:`sec_demography_events_admixture` events.
.. important::
Note that ``source`` and ``dest`` are from the perspective of the
coalescent process, i.e.
:ref:`backwards in time<sec_demography_direction_of_time>`;
please see the
:ref:`sec_demography_migration` section for more details.
:param float time: The time at which this event occurs in generations.
:param str, int source: The population **from** which lineages are moved.
:param str, int dest: The population **to** which lineages are moved.
:param float proportion: For each lineage in the ``source`` population,
this is the probability that it moves to the ``dest`` population.
"""
self._check_population_references([source, dest])
return self.add_event(MassMigration(time, source, dest, proportion))
def add_migration_rate_change(
self,
time: float,
*,
rate: float,
source: int | str | None = None,
dest: int | str | None = None,
) -> MigrationRateChange:
"""
Changes the rate of migration from one deme to another to a new value at a
specific time. Migration rates are specified in terms of the rate at which
lineages move from population ``source`` to ``dest`` during the progress of
the simulation.
.. important::
Note that ``source`` and ``dest`` are from the perspective of the
coalescent process, i.e.
:ref:`backwards in time<sec_demography_direction_of_time>`;
please see the
:ref:`sec_demography_migration` section for more details.
By default, ``source=None`` and ``dest=None``, which results in all
non-diagonal elements of the migration matrix being changed to the new
rate. If ``source`` and ``dest`` are specified, they must refer to valid
populations (either integer IDs or string names).
:param float time: The time at which this event occurs in generations.
:param float rate: The new per-generation migration rate.
:param str, int source: The ID of the source population.
:param str, int dest: The ID of the destination population.
:param int source: The source population ID.
"""
if source is None:
source = -1
else:
self._check_population_references([source])
if dest is None:
dest = -1
else:
self._check_population_references([dest])
return self.add_event(
MigrationRateChange(time=time, source=source, dest=dest, rate=rate)
)
def add_symmetric_migration_rate_change(
self, time: float, populations: list[str | int], rate: float
) -> SymmetricMigrationRateChange:
"""
Sets the symmetric migration rate between all pairs of populations in
the specified list to the specified value. For a given pair of population
IDs ``j`` and ``k``, this sets ``migration_matrix[j, k] = rate``
and ``migration_matrix[k, j] = rate``.
Please see the :ref:`sec_demography_migration` section for more details.
Populations may be specified either by their integer IDs or by
their string names.
:param float time: The time at which this event occurs in generations.
:param list populations: An sequence of population identifiers (integer
IDs or string names).
:param float rate: The new migration rate.
"""
self._check_population_references(populations)
return self.add_event(
SymmetricMigrationRateChange(time=time, populations=populations, rate=rate)
)
def add_population_parameters_change(
self,
time: float,
*,
initial_size: float | None = None,
growth_rate: float | None = None,
population: int | None = None,
) -> PopulationParametersChange:
"""
Changes the size parameters of a population (or all populations)
at a given time.
Please see the :ref:`sec_demography_populations` section for more details.
:param float time: The length of time ago at which this event
occurred.
:param float initial_size: The number of individuals in the population
at the beginning of the time slice starting at ``time``. If None,
the initial_size of the population is computed according to
the initial population size and growth rate over the preceding
time slice.
:param float growth_rate: The new per-generation growth rate. If None,
the growth rate is not changed. Defaults to None.
:param str, int population: The ID of the population affected. If
``population`` is None, the changes affect all populations
simultaneously.
"""
if population is not None:
self._check_population_references([population])
event = PopulationParametersChange(
time,
initial_size=initial_size,
growth_rate=growth_rate,
population=population,
)
return self.add_event(event)
def add_simple_bottleneck(
self,
time: float,
population: int | str,
proportion: float | None = None,
) -> SimpleBottleneck:
"""
Adds a population bottleneck at the specified time in which each lineage
has probability equal to ``proportion`` of coalescing into a single
ancestor.
Please see the :ref:`sec_demography_events_simple_bottleneck` section
for more details.
:param float time: The length of time ago at which this event
occurred.
:param str, int population: The ID of the population affected.
:param float proportion: The probability of each lineage coalescing
into a single ancestor. Defaults to 1.0.
"""
proportion = 1.0 if proportion is None else proportion
self._check_population_references([population])
return self.add_event(
SimpleBottleneck(time=time, population=population, proportion=proportion)
)
def add_instantaneous_bottleneck(
self, time: float, *, population: str | int, strength: float
) -> InstantaneousBottleneck:
"""
Adds a bottleneck at the specified time in the specified population
that is equivalent to the coalescent process running for ``strength``
generations.
Please see the :ref:`sec_demography_events_instantaneous_bottleneck`
section for more details.
.. note:: The :ref:`ploidy<sec_ancestry_ploidy_coalescent_time_scales>`
is also use to scale the time scale of the coalescent process
during the bottleneck.
:param float time: The length of time ago at which this event
occurred.
:param str, int population: The ID of the population affected.
:param float strength: The equivalent amount of time in the standard
coalescent.
"""
self._check_population_references([population])
return self.add_event(
InstantaneousBottleneck(time=time, population=population, strength=strength)
)
def add_census(self, time: float) -> CensusEvent:
"""
Adds a "census" event at the specified time. In a census we add a node
to each branch of every tree, thus recording the population that each
lineage is in at the specified time.
This may be used to record all ancestral haplotypes present at that
time, and to extract other information related to these haplotypes: for
instance to trace the local ancestry of a sample back to a set of
contemporaneous ancestors, or to assess whether a subset of samples has
coalesced more recently than the census time. (However, these added
nodes will only represent the portions of ancestral genomes inherited
by the samples, rather than complete ancestral genomes.)
See :ref:`sec_ancestry_census_events` for more details.
.. warning:: When used in the conjunction with the DTWF model
non-integer census times should be used to guarantee that
the census nodes don't coincide with coalescences (and
therefore zero branch length errors).
See :ref:`sec_ancestry_census_events_dtwf` for more details.
:param float time: The time at which the census should occur.
"""
return self.add_event(CensusEvent(time))
def _populations_table(self):
cols = [
("id", ""),
("name", ""),
("description", ""),
("initial_size", ".1f"),
("growth_rate", ".2g"),
("default_sampling_time", ".2g"),
("extra_metadata", ""),
]
data = [
[f"{getattr(pop, attr):{fmt}}" for attr, fmt in cols]
for pop in self.populations
]
return [title for title, _ in cols], data
def _populations_text(self):
col_titles, data = self._populations_table()
alignments = ["^", "<", "<", "<", "^", ">", "<"]
data = [
[
[item.as_text() if isinstance(item, core.TableEntry) else item]
for item in row
]
for row in data
]
return core.text_table(
"Populations", [[title] for title in col_titles], alignments, data
)
def _populations_html(self):
col_titles, data = self._populations_table()
return core.html_table(
f"Populations ({len(self.populations)})", col_titles, data
)
def _migration_rate_info(self, source, dest, rate):
extra = None
if source != dest:
source_name = self.populations[source].name
dest_name = self.populations[dest].name
extra = (
"Backwards in time migration rate from population "
f"{source_name} to {dest_name} = {rate} per generation. "
"Forwards in time, this is the expected number of migrants "
f"moving from {dest_name} to {source_name} "
f"per generation, divided by the size of {source_name}."
)
return core.TableEntry(f"{rate:.4g}", extra)
def _migration_matrix_table(self):
col_titles = [""] + [pop.name for pop in self.populations]
data = []
for j in range(self.num_populations):
row = [self.populations[j].name] + [
self._migration_rate_info(j, k, self.migration_matrix[j, k])
for k in range(self.num_populations)
]
data.append(row)
return col_titles, data
def _migration_matrix_text(self):
col_titles, data = self._migration_matrix_table()
alignments = ">" + "^" * self.num_populations
data = [
[
[item.as_text() if isinstance(item, core.TableEntry) else item]
for item in row
]
for row in data
]
return core.text_table(
"Migration Matrix", [[title] for title in col_titles], alignments, data
)
def _migration_matrix_html(self):
if np.all(self.migration_matrix == 0):
return core.html_table("Migration matrix (all zero)", [], [])
else:
col_titles, data = self._migration_matrix_table()
return core.html_table("Migration matrix", col_titles, data)
def _events_text(self, events, title="Events"):
col_titles = [["time"], ["type"], ["parameters"], ["effect"]]
alignments = "><<<"
data = []
for event in events:
type_text = textwrap.wrap(event._type_str, 15)
description = textwrap.wrap(event._parameters(), 22)
effect = textwrap.wrap(event._effect(), 38)
row = [[f"{event.time:.4g}"], type_text, description, effect]
data.append(row)
return core.text_table(
title, col_titles, alignments, data, internal_hlines=True
)
def _events_html(self, events, title=None):
if title is None:
title = f"Events ({len(events)})"
if len(self.events) == 0:
return core.html_table(title, [], [])
col_titles = ["time", "type", "parameters", "effect"]
data = []
for event in events:
class_name = event.__class__.__name__
camel_case = re.sub(r"(?<!^)(?=[A-Z])", "_", class_name).lower()
add_method = f"msprime.Demography.add_{camel_case}"
# TODO change this to stable when 1.0 is released.
type_html = (
"<a href='https://tskit.dev/msprime/docs/latest/api.html#"
f"{add_method}'>{event._type_str}</a>"
)
row = [f"{event.time:.4g}", type_html, event._parameters(), event._effect()]
data.append(row)
return core.html_table(title, col_titles, data, no_escape=[1])
def _repr_html_(self):
resolved = self.validate()
return (
'<div style="margin-left:20px">'
+ resolved._populations_html()
+ resolved._migration_matrix_html()
+ resolved._events_html(self.events)
+ "</div>"
)
def __str__(self):
resolved = self.validate()
populations = resolved._populations_text()
migration_matrix = resolved._migration_matrix_text()
events = resolved._events_text(self.events)
def indent(table):
lines = table.splitlines()
s = "╟ " + lines[0] + "\n"
for line in lines[1:]:
s += "║ " + line + "\n"
return s
s = (
"Demography\n"
+ indent(populations)
+ indent(migration_matrix)
+ indent(events)
)
return s
def __len__(self):
return len(self.populations)
def __iter__(self):
for pop in self.populations:
yield pop.name
def __getitem__(self, identifier):
"""
Returns the population with the specified ID or name.
"""
if isinstance(identifier, str):
for population in self.populations:
if population.name == identifier:
return population
else:
raise KeyError(f"Population with name '{identifier}' not found")
elif isinstance(identifier, numbers.Integral):
# We don't support negative indexing here because -1 is used as
# way to refer to *all* populations in demographic events, and
# it would be too easy to introduce bugs in old code if we changed
# the meaning of this.
if identifier < 0 or identifier >= self.num_populations:
raise KeyError(f"Population id {identifier} out of bounds")
return self.populations[identifier]
raise KeyError(
"Keys must be either string population names or integer IDs:"
f"identifier '{identifier}' is of type {type(identifier)}"
)
@property
def num_populations(self):
return len(self.populations)
@property
def num_events(self):
return len(self.events)
def _validate_populations(self):
names = set()
for j, population in enumerate(self.populations):
if population.id != j:
raise ValueError(
"Incorrect population ID. ID values should not be updated "
"by users. Please use Demography.add_population to add extra "
"populations after initialisation."
)
population.validate()
if population.name in names:
raise ValueError(f"Duplicate population name: '{population.name}'")
names.add(population.name)
def validate(self):
"""
Checks the demography looks sensible and raises errors/warnings
appropriately, and return a copy in which all default values have
been appropriately resolved.
"""
self._validate_populations()
migration_matrix = np.array(self.migration_matrix)
N = self.num_populations
if migration_matrix.shape != (N, N):
raise ValueError(
"migration matrix must be a N x N square matrix encoded "
"as a list-of-lists or numpy array, where N is the number "
"of populations. The diagonal "
"elements of this matrix must be zero. For example, a "
"valid matrix for a 3 population system is "
"[[0, 1, 1], [1, 0, 1], [1, 1, 0]]"
)
last_event = None
for event in self.events:
if not isinstance(event, DemographicEvent):
raise TypeError(
"Demographic events must be a list of DemographicEvent "
"instances sorted in non-decreasing order of time."
)
if last_event is not None:
if last_event.time > event.time:
raise ValueError(
"Events must be time-sorted. Please use demography.sort_events()"
"if you add events out of order."
)
last_event = event
resolved = copy.deepcopy(self)
for population in resolved.populations:
if population.default_sampling_time is None:
population.default_sampling_time = 0
if population.initially_active is None:
population.initially_active = True
return resolved
def copy(self, populations: list[str] | None = None) -> Demography:
"""
Returns a copy of this model. If the ``populations`` argument is
specified, the populations in the copied model will be in this order.
:param list populations: A list of population identifiers defining the
order of the populations in the new model. If not specified, the
current order is used.
:return: A copy of this Demography.
"""
if populations is None:
populations = range(self.num_populations)
if len(populations) != self.num_populations:
raise ValueError("populations must have Demography.num_populations entries")
copy_demog = Demography()
for identifier in populations:
pop = self[identifier]
copy_demog.add_population(
initial_size=pop.initial_size,
growth_rate=pop.growth_rate,
name=pop.name,
description=pop.description,
extra_metadata=pop.extra_metadata,
default_sampling_time=pop.default_sampling_time,
initially_active=pop.initially_active,
)
for copy_p1, copy_p2 in itertools.combinations(copy_demog.populations, 2):
self_p1 = self[copy_p1.name].id
self_p2 = self[copy_p2.name].id
copy_demog.migration_matrix[copy_p1.id, copy_p2.id] = self.migration_matrix[
self_p1, self_p2
]
copy_demog.migration_matrix[copy_p2.id, copy_p1.id] = self.migration_matrix[
self_p2, self_p1
]
def remap(identifier):
if identifier == -1:
return -1
return self[identifier].name
for event in self.events:
copy_event = copy.deepcopy(event)
if isinstance(event, (MassMigration, MigrationRateChange)):
copy_event.source = remap(event.source)
copy_event.dest = remap(event.dest)
elif isinstance(event, PopulationSplit):
copy_event.derived = [remap(pop) for pop in event.derived]
copy_event.ancestral = remap(event.ancestral)
elif isinstance(event, Admixture):
copy_event.ancestral = [remap(pop) for pop in event.ancestral]
copy_event.derived = remap(event.derived)
elif isinstance(
event,
(PopulationParametersChange, SimpleBottleneck, InstantaneousBottleneck),
):
copy_event.population = remap(event.population)
elif isinstance(event, SymmetricMigrationRateChange):
copy_event.populations = [remap(pop) for pop in event.populations]
elif isinstance(event, CensusEvent):
# Census has no population
pass
else:
raise AssertionError("Event class not implemented in copy")
copy_demog.add_event(copy_event)
return copy_demog
def sort_events(self):
# Sort demographic events by time. Sorting is stable so the relative
# order of events at the same time will be preserved.
self.events.sort(key=lambda de: de.time)
def insert_populations(self, tables):
"""
Insert population definitions for this demography into the specified
set of tables.
:meta private:
"""
assert len(tables.populations) == 0
population_table = _build_population_table(self.populations)
tables.populations.metadata_schema = population_table.metadata_schema
tables.populations.set_columns(
metadata=population_table.metadata,
metadata_offset=population_table.metadata_offset,
)
def insert_extra_populations(self, tables):
"""
Insert additional population definitions for this demography
into the specified set of tables. We assume that the populations
up to len(tables.populations) are identical, and append additional
populations to the tables for any remaining.
:meta private:
"""
# TODO we should be accessing a higher-level API for querying
# the schema here, but there's none available right now.
schema = tables.populations.metadata_schema.schema
if schema is not None:
properties = schema.get("properties", {})
additional_properties = schema.get("additionalProperties", True)
name_in_metadata = "name" in properties or additional_properties
description_in_metadata = (
"description" in properties or additional_properties
)
if not name_in_metadata:
warnings.warn(
"The metadata schema does not have a 'name' property; "
"population names will not be recorded in the output "
"tree sequence",
IncompletePopulationMetadataWarning,
)
if not description_in_metadata:
warnings.warn(
"The metadata schema does not have a 'description' property; "
"population descriptions will not be recorded in the output "
"tree sequence",
IncompletePopulationMetadataWarning,
)
left_out = []
for population in self.populations[len(tables.populations) :]:
md = {}
if name_in_metadata:
md["name"] = population.name
if description_in_metadata:
md["description"] = population.description
for k in population.extra_metadata:
if k in properties or additional_properties:
md[k] = population.extra_metadata[k]
else:
left_out.append(k)
tables.populations.add_row(metadata=md)
if len(left_out) > 0:
warnings.warn(
"The metadata schema does not allow for all properties specified "
"in extra_metadata: these keys have not been recorded in the "
f"output tree sequence: {', '.join(left_out)}",
IncompletePopulationMetadataWarning,
)
else:
warnings.warn(
"No metadata schema present in population table, not recording "
"metadata",
IncompletePopulationMetadataWarning,
)
# No metadata schema, just add bare populations.
for _ in self.populations[len(tables.populations) :]:
tables.populations.add_row()
def asdict(self):
# NOTE: this slightly contradicts the interpretation of the
# Demography object as a mapping storing populations. But
# this is an internal undocumented method, so seems OK.
return {
"populations": [pop.asdict() for pop in self.populations],
"events": [event.asdict() for event in self.events],
"migration_matrix": self.migration_matrix.tolist(),
}
def debug(self):
"""
Returns a :class:`.DemographyDebugger` instance for this demography.
:return: A DemographyDebugger object for this demography.
:rtype: msprime.DemographyDebugger
"""
return DemographyDebugger(demography=self)
def __eq__(self, other):
try:
self.assert_equal(other)
return True
except AssertionError:
return False
def assert_equal(self, other: Demography):
"""
Compares this Demography with specified ``other`` and raises an
AssertionError if they are not exactly equal.
:param Demography other: The other demography to compare against.
"""
# TODO we could potentially do better here with error messages
# by showing a diff of the str() values for objects that differ.
assert isinstance(other, Demography)
assert self.num_populations == other.num_populations
for p1, p2 in zip(self.populations, other.populations):
assert p1 == p2, f"{p1} ≠ {p2}"
assert np.array_equal(
self.migration_matrix, other.migration_matrix
) # type: ignore
assert self.num_events == other.num_events
for e1, e2 in zip(self.events, other.events):
assert e1 == e2, f"{e1} ≠ {e2}"
def is_equivalent(self, other: Demography, rel_tol=None, abs_tol=None):
"""
Compares this demography with the other and return True if they are
equivalent up to the specified numerical tolerances. Two demographies
are equivalent if, they have the same set of epochs defined by demographic
events, and for each epoch:
- The population's ``initial_size``, ``growth_rate`` and ``active``
values are equal in all populations.
- The migration matrices are equal
- The same sequence of lineage movements through population splits, etc.
All numerical comparisons are performed using :func:`python:math.isclose`.
:param Demography other: The other demography to compare against.
:param float rel_tol: The relative tolerance used by math.isclose.
:param float abs_tol: The relative tolerance used by math.isclose.
:return: True if this demography and other are equivalent up to numerical
tolerances.
:rtype bool: bool
"""
try:
self.assert_equivalent(other, rel_tol=rel_tol, abs_tol=abs_tol)
return True
except AssertionError:
return False
def assert_equivalent(
self,
other: Demography,
rel_tol: None | float = None,
abs_tol: None | float = None,
):
# Same defaults as math.isclose
rel_tol = 1e-9 if rel_tol is None else rel_tol
abs_tol = 0 if abs_tol is None else abs_tol
assert isinstance(other, Demography)
self_dbg = self.debug()
other_dbg = other.debug()
if self.num_populations != other.num_populations:
raise AssertionError(
"Number of populations not equal: "
f"{self.num_populations} ≠ {other.num_populations}"
)
# Compare the population attributes.
# NB use the *resolved* versions from the debug objects
for self_pop, other_pop in zip(
self_dbg.demography.populations, other_dbg.demography.populations
):
if self_pop.name != other_pop.name:
raise AssertionError(
f"Population names differ: {self_pop.name} ≠ {other_pop.name}"
)
self_st = (
0
if self_pop.default_sampling_time is None
else self_pop.default_sampling_time
)
other_st = (
0
if other_pop.default_sampling_time is None
else other_pop.default_sampling_time
)
if not math.isclose(self_st, other_st):
raise AssertionError(
f"Sampling times not equal for {self_pop.name}: "
f"{self_pop.default_sampling_time} ≠ "
f"{other_pop.default_sampling_time}"
)
if self_dbg.num_epochs != other_dbg.num_epochs:
raise AssertionError(
"Number of epochs not equal: "
f"{self_dbg.num_epochs} ≠ {other_dbg.num_epochs}"
)
for j, (self_epoch, other_epoch) in enumerate(
zip(self_dbg.epochs, other_dbg.epochs)
):
if not math.isclose(
self_epoch.start_time,
other_epoch.start_time,
rel_tol=rel_tol,
abs_tol=abs_tol,
):
raise AssertionError(
f"Epoch[{j}] at different times: "
f"{self_epoch.start_time} ≠ {other_epoch.start_time}"
)
for self_pop, other_pop in zip(
self_epoch.populations, other_epoch.populations
):
if self_pop.state != other_pop.state:
raise AssertionError(
f"State mismatch in populations in epoch[{j}], {self_pop.name}: "
f"{self_pop.state} ≠ {other_pop.state}"
)
if self_pop.state == PopulationStateMachine.ACTIVE:
if not math.isclose(
self_pop.start_size,
other_pop.start_size,
rel_tol=rel_tol,
abs_tol=abs_tol,
):
raise AssertionError(
"Population start_size not equal to required precision "
f"in epoch[{j}], {self_pop.name}: "
f"{self_pop.start_size} ≠ {other_pop.start_size}"
)
if not math.isclose(
self_pop.growth_rate,
other_pop.growth_rate,
rel_tol=rel_tol,
abs_tol=abs_tol,
):
raise AssertionError(
"Population growth_rate not equal to required precision "
f"in epoch[{j}], {self_pop.name}: "
f"{self_pop.growth_rate} ≠ {other_pop.growth_rate}"
)
m_equal = np.isclose(
self_epoch.migration_matrix,
other_epoch.migration_matrix,
rtol=rel_tol,
atol=abs_tol,
)
if not np.all(m_equal):
differs = "Differences between: "
for source_id, dest_id in zip(*np.where(~m_equal)):
source = self[source_id].name
dest = self[dest_id].name
self_rate = self_epoch.migration_matrix[source_id, dest_id]
other_rate = other_epoch.migration_matrix[source_id, dest_id]
differs += f"({source}, {dest}: {self_rate} ≠ {other_rate}), "
raise AssertionError(
f"Migration matrices in epoch[{j}] not equal: \n"
"self = \n"
f"{self_epoch.migration_matrix}\n"
"other = \n"
f"{other_epoch.migration_matrix}\n"
f"::{differs[:-2]}"
)
self._assert_lineage_movements_equivalent(
self_epoch.events, other_epoch.events, rel_tol=rel_tol, abs_tol=abs_tol
)
self_state_change = [
event
for event in self_epoch.events
if isinstance(event, StateChangeEvent)
]
other_state_change = [
event
for event in other_epoch.events
if isinstance(event, StateChangeEvent)
]
if len(self_state_change) > 0 or len(other_state_change) > 0:
raise ValueError(
"State change events not currently supported in equivalent. "
"Please open an issue on GitHub"
)
def _assert_lineage_movements_equivalent(
self, self_events, other_events, *, rel_tol, abs_tol
):
self_lineage_movements = self._normalise_lineage_movements(self_events)
other_lineage_movements = self._normalise_lineage_movements(other_events)
source_pops = set(self_lineage_movements.keys())
if source_pops != set(other_lineage_movements.keys()):
raise AssertionError(
f"Mismatch in the set of populations affected by lineage movements: "
f"{source_pops} ≠ {set(other_lineage_movements.keys())}"
)
for source in source_pops:
self_out_movements = self_lineage_movements[source]
other_out_movements = other_lineage_movements[source]
if len(self_out_movements) != len(other_out_movements):
raise AssertionError(
"Mismatch in number of normalised lineage movements out of "
f"{source}: {len(self_out_movements)} ≠ {len(other_out_movements)}"
)
for self_lm, other_lm in zip(self_out_movements, other_out_movements):
if self_lm.dest != other_lm.dest:
raise AssertionError(
"Mismatch in lineage movement destination:"
f"{self_lm.dest} ≠ {other_lm.dest}"
)
if not math.isclose(
self_lm.proportion,
other_lm.proportion,
rel_tol=rel_tol,
abs_tol=abs_tol,
):
raise AssertionError(
"Mismatch in normalised lineage movement proportions "
f"from {self_lm.source} to {self_lm.dest}: "
f"{self_lm.proportion} ≠ {other_lm.proportion}"
)
def _normalise_lineage_movements(self, events: list[DemographicEvent]):
"""
Extract the LineageMovementEvent instances from the specified list
and normalise their effects into LineageMovement instances, and
return a dictionary mapping source populations to the list of
sequential lineage movements. For each source population we have a
list of sequential lineage movements, sorted by destination population ID.
"""
assert len({event.time for event in events}) <= 1
ret = collections.defaultdict(list)
for event in events:
if isinstance(event, LineageMovementEvent):
for move in event._as_lineage_movements():
ret[move.source].append(move)
for pop in list(ret.keys()):
# We have a list of *conditional* lineage movements out of a
# population. We canonicalise this by sorting by the
# destination population, so we have to first convert back to
# absolute proportions.
assert all(lm.source == pop for lm in ret[pop])
P = _sequential_to_proportions([pm.proportion for pm in ret[pop]])
id_value_pairs = sorted((lm.dest, p) for lm, p in zip(ret[pop], P))
S = _proportions_to_sequential([p for _, p in id_value_pairs])
ret[pop] = [
LineageMovement(source=pop, dest=id_value_pairs[j][0], proportion=S[j])
for j in range(len(S))
]
return ret
@staticmethod
def from_species_tree(
tree,
initial_size,
*,
time_units="gen",
generation_time=None,
growth_rate=None,
) -> Demography:
"""
Parse a species tree in `Newick
<https://en.wikipedia.org/wiki/Newick_format>`_ format and return the
corresponding :class:`Demography` object. The tree is assumed to be
rooted and ultrametric and branch lengths must be included and
correspond to time, either in units of millions of years, years, or
generations.
The returned :class:`.Demography` object contains a
:class:`.Population` for each node in the species tree. The
population's ``name`` attribute will be either the corresponding node
label from the newick tree, if it exists, or otherwise the name takes
the form "pop_{j}", where j is the position of the given population in
the list. Leaf populations are first in the list, and added in
left-to-right order. Populations corresponding to the internal nodes
are then added in a postorder traversal of the species tree. For each
internal node a :ref:`sec_demography_events_population_split`
event is added so that
lineages move from its child populations at the appropriate time
and rates of continuous migration to and from the child populations is
set to zero. See the :ref:`sec_demography_events_population_split`
section for more details.
The initial sizes and growth rates for the populations in the model are
set via the ``initial_size`` and ``growth_rate`` arguments. These can be
specified in two ways: if a single number is provided, this is used
for all populations. The argument may also be a mapping from population
names to their respective values. For example:
.. code-block:: python
tree = "(A:10.0,B:10.0)C"
initial_size = {"A": 1000, "B": 2000, "C": 100}
demography = msprime.Demography.from_species_tree(tree, initial_size)
Note that it is possible to have default population sizes for unnamed
ancestral populations using a :class:`python:collections.defaultdict`, e.g.,
.. code-block:: python
tree = "(A:10.0,B:10.0)"
initial_size = collections.defaultdict(lambda: 100)
initial_size.update({"A": 1000, "B": 2000})
demography = msprime.Demography.from_species_tree(tree, initial_size)
:param str tree: The tree string in Newick format, with named leaves and branch
lengths.
:param initial_size: Each population's initial_size, in numbers of individuals.
May be a single number or a mapping from population names to their sizes.
:param growth_rate: Each population's growth_rate. May be a single number
or a mapping from population names to their exponential growth rates.
Defaults to zero.
:param str time_units: The units of time in which the species tree's
branch lengths are measured. Allowed branch length units are millions of
years, years, and generations; these should be specified with the strings
``"myr"``, ``"yr"``, or ``"gen"``, respectively. This defaults to
``"gen"``.
:param float generation_time: The number of years per generation. If and only
if the branch lengths are not in units of generations, the generation time
must be specified. This defaults to `None`.
:return: A Demography object representing the specified species tree.
:rtype: msprime.Demography
"""
return species_trees.parse_species_tree(
tree,
initial_size=initial_size,
growth_rate=growth_rate,
time_units=time_units,
generation_time=generation_time,
)
@staticmethod
def from_starbeast(tree, generation_time, time_units="myr") -> Demography:
"""
Parse a species tree produced by the program `TreeAnnotator
<https://www.beast2.org/treeannotator>`_
based on a posterior tree distribution generated with `StarBEAST
<https://academic.oup.com/mbe/article/34/8/2101/3738283>`_ and return
the corresponding Demography object.
Species trees produced by TreeAnnotator are written in `Nexus
<https://en.wikipedia.org/wiki/Nexus_file>`_ format and are rooted,
bifurcating, and ultrametric. Branch lengths usually are in units of
millions of years, but the use of other units is permitted by StarBEAST
(and thus TreeAnnotator). This function allows branch length units of
millions of years or years. Leaves must be named and the tree must
include information on population sizes of leaf and ancestral species
in the form of annotation with the "dmv" tag, which is the case for
trees written by TreeAnnotator based on StarBEAST posterior tree
distributions.
The returned :class:`.Demography` object contains a
:class:`.Population` for each node in the species tree. The
population's ``name`` attribute will be either the corresponding node
label from the newick tree, if it exists, or otherwise the name takes
the form "pop_{j}", where j is the position of the given population in
the list. Leaf populations are first in the list, and added in
left-to-right order. Populations corresponding to the internal nodes
are then added in a postorder traversal of the species tree. For each
internal node a :ref:`sec_demography_events_population_split`
event is added so that
lineages move from its child populations at the appropriate time and
rates of continuous migration to and from the child populations is set
to zero. See the :ref:`sec_demography_events_population_split` section
for more details.
:param str tree: The tree string in Nexus format, with named leaves, branch
lengths, and branch annotation. Typically, this string is the entire content
of a file written by TreeAnnotator.
:param float generation_time: The number of years per generation.
:param str time_units: The units of time in which the species tree's
branch lengths are measured. Allowed branch length units are millions of
years, and years; these should be specified with the strings ``"myr"`` or
``"yr"``, respectively. This defaults to ``"myr"``.
:return: A :class:`.Demography` instance that describing the information in the
specified species tree.
:rtype: msprime.Demography
"""
return species_trees.parse_starbeast(
tree=tree,
generation_time=generation_time,
time_units=time_units,
)
def _from_old_style_map_populations(
population_configurations: list[PopulationConfiguration],
migration_matrix: list[list[float]],
demographic_events: list[DemographicEvent],
population_map: [list[dict[int, str]]],
) -> Demography:
direct_model = Demography._from_old_style_simple(
population_configurations, migration_matrix, demographic_events
)
for id_map in population_map:
if len(set(id_map.values())) != len(id_map):
raise ValueError("Population IDs in old model must be unique")
for old_id in id_map.values():
if old_id < 0 or old_id >= direct_model.num_populations:
raise ValueError(
f"Bad population reference {old_id} in old style model"
)
if len(population_map[0]) != len(population_configurations):
raise ValueError(
"The ID map for the first epoch must have entries for all "
"populations in the old-style model"
)
# Suppress misspecification warnings; we'll give more specific messages
# later.
with warnings.catch_warnings():
warnings.simplefilter("ignore")
dbg = direct_model.debug()
if dbg.num_epochs != len(population_map):
raise ValueError(
"Mismatch in the number of epochs in the old style model "
f"({dbg.num_epochs}) and specified in the population ID map "
f"({len(population_map)})"
)
demography = Demography()
# Set up the initial state of the model
id_map = population_map[0]
for new_name, old_id in id_map.items():
pc = population_configurations[old_id]
demography._add_population_from_old_style(pc, new_name)
for epoch_id_map in population_map[1:]:
for name in epoch_id_map.keys():
if name not in demography:
demography.add_population(name=name, initial_size=0)
# Fill out migration matrix
for pop_1, pop_2 in itertools.combinations(id_map.keys(), 2):
for pop_j, pop_k in [(pop_1, pop_2), (pop_2, pop_1)]:
j = id_map[pop_j]
k = id_map[pop_k]
demography.set_migration_rate(
source=pop_j, dest=pop_k, rate=direct_model.migration_matrix[j, k]
)
# Set the state of populations and migration rates epoch by epoch.
for epoch in dbg.epochs[1:]:
epoch_id_map = population_map[epoch.index]
# Set the population sizes and growth_rates for the extant populations.
for new_id, old_id in epoch_id_map.items():
pop_state = epoch.populations[old_id]
demography.add_population_parameters_change(
epoch.start_time,
population=new_id,
initial_size=pop_state.start_size,
growth_rate=pop_state.growth_rate,
)
# Set the migration rates.
for pop_1, pop_2 in itertools.combinations(epoch_id_map.keys(), 2):
for pop_j, pop_k in [(pop_1, pop_2), (pop_2, pop_1)]:
j = epoch_id_map[pop_j]
k = epoch_id_map[pop_k]
demography.add_migration_rate_change(
time=epoch.start_time,
source=pop_j,
dest=pop_k,
rate=epoch.migration_matrix[j][k],
)
# Add splits/admixtures/pulses according to the mass migrations in the
# original model and ID maps.
last_epoch_id_map = population_map[0]
for epoch in dbg.epochs:
logger.debug(
f"Converting epoch[{epoch.index}] "
f"{epoch.start_time:.2f}-{epoch.end_time:.2f}"
)
# New name -> old ID
epoch_id_map = population_map[epoch.index]
# Old ID -> new name
old_id_map = {value: key for key, value in epoch_id_map.items()}
last_epoch_pops = set(last_epoch_id_map.keys())
epoch_pops = set(epoch_id_map.keys())
derived = set()
events = copy.deepcopy(epoch.events)
old_derived_ids = []
if last_epoch_pops != epoch_pops:
ancestral = epoch_pops - last_epoch_pops
derived = last_epoch_pops - epoch_pops
old_derived_ids = []
if len(ancestral) == 1:
logger.debug(
f"Adding split ancestral={ancestral} derived={derived}"
)
ancestral = ancestral.pop()
demography.add_population_split(
time=epoch.start_time,
derived=list(derived),
ancestral=ancestral,
)
old_derived_ids = [last_epoch_id_map[pop] for pop in derived]
derived_mass_migrations = 0
for event in events:
if isinstance(event, MassMigration):
if event.source in old_derived_ids:
derived_mass_migrations += 1
if event.proportion != 1:
raise ValueError(
"MassMigration associated with population split "
"with proportion != 1"
)
# This check is weak, but it will catch some errors at least.
# We're assuming the old model is well-formed, so this should
# help catch errors where the id map is slightly off.
if derived_mass_migrations < len(old_derived_ids) - 1:
raise ValueError(
"Insufficient MassMigrations found for population split"
)
elif len(derived) == 1:
derived = derived.pop()
ancestral = []
sequential_proportions = []
old_derived_id = last_epoch_id_map[derived]
for event in events:
if isinstance(event, MassMigration):
if event.source == old_derived_id:
ancestral.append(old_id_map[event.dest])
sequential_proportions.append(event.proportion)
proportions = _sequential_to_proportions(sequential_proportions)
if math.isclose(sum(proportions), 1):
if len(ancestral) == 1:
# This is a single split from the ancestral population,
# which continues to exist. We need to override the
# defaults for the behaviour we want here.
logger.debug(
f"Adding single split ancestral={ancestral[0]} "
f"derived={derived}"
)
pop = demography[ancestral[0]]
pop.initially_active = True
pop.default_sampling_time = None
demography.add_population_split(
time=epoch.start_time,
derived=[derived],
ancestral=ancestral[0],
)
else:
logger.debug(
f"Adding admixture ancestral={ancestral} "
f"derived={derived} proportions={proportions}"
)
demography.add_admixture(
time=epoch.start_time,
derived=derived,
ancestral=ancestral,
proportions=proportions,
)
old_derived_ids = [old_derived_id]
else:
raise ValueError(
"Admixture or single population split implied by ID map "
"but absolute population proportions don't sum to 1"
)
for event in events:
if isinstance(event, MassMigration):
if event.source not in old_derived_ids:
demography.add_mass_migration(
time=epoch.start_time,
source=old_id_map[event.source],
dest=old_id_map[event.dest],
proportion=event.proportion,
)
elif not isinstance(
event, (MigrationRateChange, PopulationParametersChange)
):
raise ValueError(
"Only MassMigration, MigrationRateChange and "
"PopulationParametersChange events are supported"
)
# Go through the inactive populations and check the migration
# rates make sense. It's an error to migrate out of an inactive
# population and a warning to migrate out of inactive pops.
pairs = itertools.combinations(range(len(epoch.populations)), 2)
active = set(epoch_id_map.values())
for pop_1, pop_2 in pairs:
for source, dest in [(pop_1, pop_2), (pop_2, pop_1)]:
if epoch.migration_matrix[source, dest] != 0:
if source in active and dest not in active:
raise ValueError(
"Non zero migration from an active population "
f"({source}) to inactive ({dest})"
)
if source not in active:
warnings.warn(
"Migration out of inactive population "
f"({source}) to ({dest}). This may be an "
"error in your model."
)
last_epoch_id_map = epoch_id_map
demography.sort_events()
return demography
@staticmethod
def _from_old_style_simple(
population_configurations=None,
migration_matrix=None,
demographic_events=None,
) -> Demography:
"""
Creates a Demography object from the pre 1.0 style input parameters,
reproducing the old semantics with respect to default values.
"""
demography = Demography()
for pop_config in population_configurations:
name = None
if pop_config.metadata is not None:
# If there's a name defined in the old-style metadata use that
name = pop_config.metadata.get("name", None)
demography._add_population_from_old_style(pop_config, name)
if migration_matrix is not None:
demography.migration_matrix = np.array(migration_matrix)
if demographic_events is not None:
for event in demographic_events:
demography.add_event(copy.deepcopy(event))
return demography
@staticmethod
def from_old_style(
population_configurations=None,
*,
migration_matrix=None,
demographic_events=None,
Ne=1,
ignore_sample_size=False,
population_map: [list[dict[int, str | int]]] | None = None,
) -> Demography:
"""
Creates a Demography object from the pre 1.0 style input parameters,
reproducing the old semantics with respect to default values.
No sample information is stored in the new-style :class:`.Demography`
objects, and therefore if the ``sample_size`` attribute of any
of the input :class:`.PopulationConfiguration` objects is set a
ValueError will be raised by default. However, if the
``ignore_sample_size`` parameter is set to True, this check will
not be performed and the sample sizes specified in the old-style
:class:`.PopulationConfiguration` objects will be ignored.
Each :class:`.PopulationConfiguration` instance in the list of
``population_configurations`` corresponds to the equivalent
:class:`.Population` object in the returned :class:`.Demography`.
If a PopulationConfiguration has ``metadata`` defined and this
dictionary contains a ``name`` field, this will be used as the
:class:`.Population` name. Otherwise, the default population
names will be used.
Please see the :ref:`sec_ancestry_samples` section for details on
how to specify sample locations in :func:`.sim_ancestry`.
.. todo:: Document the remaining parameters.
"""
if population_configurations is None:
pop_configs = [PopulationConfiguration(initial_size=Ne)]
else:
pop_configs = copy.deepcopy(population_configurations)
for pop_config in pop_configs:
if pop_config.initial_size is None:
pop_config.initial_size = Ne
if pop_config.sample_size is not None and not ignore_sample_size:
raise ValueError(
"You have specified a `sample_size` in a "
"PopulationConfiguration object that is to be converted "
"into a new-style Demography object, "
"which does not contain any information about samples. "
"Please use the ``samples`` argument to sim_ancestry "
"instead, which provides flexible options for sampling "
"from different populations"
)
if population_map is None:
return Demography._from_old_style_simple(
pop_configs, migration_matrix, demographic_events
)
else:
return Demography._from_old_style_map_populations(
pop_configs,
migration_matrix,
demographic_events,
population_map,
)
@staticmethod
def from_demes(graph: demes.Graph) -> Demography:
"""
Creates a :class:`.Demography` object from the specified
:ref:`demes graph <demes:sec_introduction>`.
Time values in the demes graph may be specified in any units,
but the returned object has units converted into generations.
See the :ref:`sec_demography_importing_demes` section for details.
.. code::
import demes
graph = demes.load("model.yaml")
demography = msprime.Demography.from_demes(graph)
:param demes.Graph graph: A demes graph.
:return: A :class:`.Demography` instance corresponding to the demes model.
:rtype: msprime.Demography
"""
# Check for Demes features that we don't support.
for deme in graph.demes:
for epoch in deme.epochs:
if epoch.selfing_rate != 0:
raise ValueError(
f"deme {deme.name}: non-zero selfing_rate not supported"
)
if epoch.cloning_rate != 0:
raise ValueError(
f"deme {deme.name}: non-zero cloning_rate not supported"
)
def get_growth_rate(epoch):
ret = 0
if epoch.size_function not in ["constant", "exponential"]:
raise ValueError(
"msprime only supports constant or exponentially changing "
"population sizes"
)
if epoch.end_size != epoch.start_size:
ret = -math.log(epoch.start_size / epoch.end_size) / epoch.time_span
return ret
graph = graph.in_generations()
events = graph.discrete_demographic_events()
demography = Demography()
for deme in graph.demes:
initial_size = 0
growth_rate = 0
last_epoch = deme.epochs[-1]
initially_active = False
to_activate = False
if last_epoch.end_time == 0:
initial_size = last_epoch.end_size
growth_rate = get_growth_rate(last_epoch)
initially_active = True
else:
# add activation if needed
# needed if isn't the ancestor of any split or merger
to_activate = True
for split in events["splits"]:
if deme.name == split.parent:
to_activate = False
for merger in events["mergers"]:
if deme.name in merger.parents:
to_activate = False
demography.add_population(
name=deme.name,
description=deme.description,
growth_rate=growth_rate,
initial_size=initial_size,
default_sampling_time=deme.end_time,
initially_active=initially_active,
)
if to_activate:
demography._add_activate_population_event(
time=deme.end_time, population=deme.name
)
for epoch in reversed(deme.epochs):
new_initial_size = None
if initial_size != epoch.end_size:
new_initial_size = epoch.end_size
new_growth_rate = None
alpha = get_growth_rate(epoch)
if growth_rate != alpha:
new_growth_rate = alpha
if new_growth_rate is not None or new_initial_size is not None:
demography.add_population_parameters_change(
population=deme.name,
time=epoch.end_time,
initial_size=new_initial_size,
growth_rate=new_growth_rate,
)
initial_size = new_initial_size
growth_rate = new_growth_rate
for split in events.pop("splits"):
demography.add_population_split(
time=split.time, ancestral=split.parent, derived=split.children
)
for branch in events.pop("branches"):
demography.add_population_split(
time=branch.time,
ancestral=branch.parent,
derived=[branch.child],
)
# Add pulses in reverse order, so that pulses with the same time
# correspond to the correct backwards-time mass migration ordering.
for pulse in reversed(events.pop("pulses")):
sequential_props = _proportions_to_sequential(pulse.proportions)
for prop, source in zip(sequential_props, pulse.sources):
demography.add_mass_migration(
time=pulse.time,
source=pulse.dest,
dest=source,
proportion=prop,
)
for merger in events.pop("mergers"):
demography.add_admixture(
time=merger.time,
derived=merger.child,
ancestral=merger.parents,
proportions=merger.proportions,
)
for admixture in events.pop("admixtures"):
demography.add_admixture(
time=admixture.time,
derived=admixture.child,
ancestral=admixture.parents,
proportions=admixture.proportions,
)
assert len(events) == 0
# Turn migrations off at the start_time. We schedule all start_time
# events first, and then all end_time events. This ensures that events
# for migrations with an end_time that coincides with the start_time of
# another migration will be scheduled later (backwards in time),
# and thus override the rate=0 setting.
for migration in graph.migrations:
if not math.isinf(migration.start_time):
demography.add_migration_rate_change(
time=migration.start_time,
source=migration.dest,
dest=migration.source,
rate=0,
)
for migration in graph.migrations:
if migration.end_time == 0:
demography.set_migration_rate(
source=migration.dest, dest=migration.source, rate=migration.rate
)
else:
demography.add_migration_rate_change(
time=migration.end_time,
source=migration.dest,
dest=migration.source,
rate=migration.rate,
)
demography.sort_events()
return demography
@staticmethod
def from_tree_sequence(
ts: tskit.TreeSequence, initial_size: float = 0
) -> Demography:
"""
Creates a :class:`.Demography` object based on the information in the
specified :class:`tskit.TreeSequence`. The returned demography will
contain a population for each of the populations in the tree sequence,
in the same order.
The metadata for each population in the tree sequence will be
inspected. If a schema is present and the metadata can be decoded, the
``name`` and ``description`` properties of populations are set if the
corresponding keys are present.
If the metadata cannot be decoded, the default values for ``name``
and ``description`` are used.
The ``initial_size`` of each of the new populations is set to zero
by default, and all other :class:`.Population` attributes
are set to their default values. It is therefore essential to
update the ``initial_size`` and ``growth_rate`` values to reflect
the desired demography.
.. seealso:: See the
:ref:`initial state<sec_ancestry_initial_state_demography>`
section for examples of how this method can be used.
:param tskit.TreeSequence ts: The tree sequence
to extract population information from.
:param float initial_size: The default initial size for the newly
added populations, as a number of individuals (Default=0).
:return: A Demography object representing the populations in the
specified tree sequence.
:rtype: msprime.Demography
"""
demography = Demography()
for population in ts.populations():
name = None
description = None
if isinstance(population.metadata, collections.abc.Mapping):
# We have decoded metadata, extract some useful information
# from it.
name = population.metadata.get("name")
description = population.metadata.get("description")
demography.add_population(
initial_size=initial_size, name=name, description=description
)
return demography
@staticmethod
def isolated_model(initial_size, *, growth_rate=None) -> Demography:
"""
Returns a :class:`.Demography` object representing a collection of
isolated populations with specified initial population sizes and
growth rates. Please see :ref:`sec_demography` for more details on
population sizes and growth rates.
:param list initial_size: the ``initial_size`` value for each
of the :class:`.Population` in the returned model. The length
of the array corresponds to the number of populations.
model.
:param list growth_rate: The exponential growth rate for each
population. Must be either None (the default, resulting a zero
growth rate) or an array with the same length as ``initial_size``.
:return: A Demography object representing this model, suitable as
input to :func:`.sim_ancestry`.
:rtype: msprime.Demography
"""
initial_size = np.array(initial_size, dtype=np.float64)
if len(initial_size.shape) != 1:
raise ValueError(
"The initial_size argument must a 1D array of population size values"
)
if growth_rate is None:
growth_rate = np.zeros_like(initial_size)
else:
growth_rate = np.array(growth_rate, dtype=np.float64)
if initial_size.shape != growth_rate.shape:
raise ValueError(
"If growth_rate is specified it must be a 1D array of the same "
"length as the population_size array"
)
if np.any(initial_size < 0):
raise ValueError("population size values must be nonnegative.")
if not np.all(np.isfinite(initial_size)):
raise ValueError("population size values must be finite.")
if not np.all(np.isfinite(growth_rate)):
raise ValueError("growth_rate values must be finite.")
populations = [
Population(
initial_size=initial_size[j],
growth_rate=growth_rate[j],
)
for j in range(len(initial_size))
]
return Demography(populations=populations)
@staticmethod
def island_model(initial_size, migration_rate, *, growth_rate=None) -> Demography:
"""
Returns a :class:`.Demography` object representing a collection of
populations with specified initial population sizes and growth
rates, with symmetric migration between each pair of populations at the
specified rate. Please see :ref:`sec_demography` for more details on
population sizes and growth rates.
:param list initial_size: the ``initial_size`` value for each
of the :class:`.Population` in the returned model. The length
of the array corresponds to the number of populations.
model.
:param float migration_rate: The migration rate between each pair of
populations.
:param list growth_rate: The exponential growth rate for each
population. Must be either None (the default, resulting a zero
growth rate) or an array with the same length as ``initial_size``.
:return: A Demography object representing this model, suitable as
input to :func:`.sim_ancestry`.
:rtype: msprime.Demography
"""
model = Demography.isolated_model(initial_size, growth_rate=growth_rate)
check_migration_rate(migration_rate)
model.migration_matrix[:] = migration_rate
np.fill_diagonal(model.migration_matrix, 0)
return model
@staticmethod
def stepping_stone_model(
initial_size, migration_rate, *, growth_rate=None, boundaries=False
) -> Demography:
"""
Returns a :class:`.Demography` object representing a collection of
populations with specified initial population sizes and growth
rates, in which adjacent demes exchange migrants at the
specified rate. Please see :ref:`sec_demography` for more details on
population sizes and growth rates.
.. note:: The current implementation only supports a one-dimensional
stepping stone model, but higher dimensions could also be supported.
Please open an issue on GitHub if this feature would be useful to you.
:param list initial_size: the ``initial_size`` value for each
of the :class:`.Population` in the returned model. The length
of the array corresponds to the number of populations.
:param float migration_rate: The migration rate between adjacent pairs
of populations.
:param list growth_rate: The exponential growth rate for each
population. Must be either None (the default, resulting a zero
growth rate) or an array with the same length as ``initial_size``.
:param bool boundaries: If True the stepping stone model has boundary
conditions imposed so that demes at either end of the chain do
not exchange migrants. If False (the default), the set of
populations is "circular" and migration takes place between the
terminal demes.
:return: A Demography object representing this model, suitable as
input to :func:`.sim_ancestry`.
:rtype: msprime.Demography
"""
initial_size = np.array(initial_size, dtype=np.float64)
if len(initial_size.shape) > 1:
raise ValueError(
"Only 1D stepping stone models currently supported. Please open "
"an issue on GitHub if you would like 2D (or more) models"
)
model = Demography.isolated_model(initial_size, growth_rate=growth_rate)
check_migration_rate(migration_rate)
if model.num_populations > 1:
index1 = np.arange(model.num_populations, dtype=int)
index2 = np.mod(index1 + 1, model.num_populations)
model.migration_matrix[index1, index2] = migration_rate
model.migration_matrix[index2, index1] = migration_rate
if boundaries:
model.migration_matrix[0, -1] = 0
model.migration_matrix[-1, 0] = 0
return model
@staticmethod
def _ooa_model():
"""
Returns the Gutenkunst et al three population out-of-Africa model.
This version is included here temporarily as a way to get some
test coverage on the model compared with stdpopsim. Because we
use this model in the documentation, we want make sure that it's
doing what we think. We compare the model defined here then with
the one presented in the docs, to ensure that no errors creep in.
Once the upstream code in stdpopsim is updated to use msprime 1.0
APIs we can remove this model and instead compare directly
to the stdpopsim model with .is_equivalent() or whatever.
"""
# Times are provided in years, so we convert into generations.
generation_time = 25
T_OOA = 21.2e3 / generation_time
T_AMH = 140e3 / generation_time
T_ANC = 220e3 / generation_time
# We need to work out the starting (diploid) population sizes based on
# the growth rates provided for these two populations
r_CEU = 0.004
r_CHB = 0.0055
N_CEU = 1000 / math.exp(-r_CEU * T_OOA)
N_CHB = 510 / math.exp(-r_CHB * T_OOA)
demography = Demography()
demography.add_population(
name="YRI",
description="Yoruba in Ibadan, Nigeria",
initial_size=12300,
)
demography.add_population(
name="CEU",
description=(
"Utah Residents (CEPH) with Northern and Western European Ancestry"
),
initial_size=N_CEU,
growth_rate=r_CEU,
)
demography.add_population(
name="CHB",
description="Han Chinese in Beijing, China",
initial_size=N_CHB,
growth_rate=r_CHB,
)
demography.add_population(
name="OOA",
description="Bottleneck out-of-Africa population",
initial_size=2100,
)
demography.add_population(
name="AMH",
description="Anatomically modern humans",
initial_size=12300,
)
demography.add_population(
name="ANC",
description="Ancestral equilibrium population",
initial_size=7300,
)
# Set the migration rates between extant populations
demography.set_symmetric_migration_rate(["CEU", "CHB"], 9.6e-5)
demography.set_symmetric_migration_rate(["YRI", "CHB"], 1.9e-5)
demography.set_symmetric_migration_rate(["YRI", "CEU"], 3e-5)
demography.add_population_split(
time=T_OOA, derived=["CEU", "CHB"], ancestral="OOA"
)
demography.add_symmetric_migration_rate_change(
time=T_OOA, populations=["YRI", "OOA"], rate=25e-5
)
demography.add_population_split(
time=T_AMH, derived=["YRI", "OOA"], ancestral="AMH"
)
demography.add_population_split(time=T_ANC, derived=["AMH"], ancestral="ANC")
return demography
def _ooa_trunk_model():
"""
Returns the Gutenkunst et al three population out-of-Africa model,
rendered in a more "old-style" way where we merge the various
populations back into Africa.
This version is included here temporarily as a way to get some
test coverage on the model compared with stdpopsim. Because we
use this model in the documentation, we want make sure that it's
doing what we think. We compare the model defined here then with
the one presented in the docs, to ensure that no errors creep in.
Once the upstream code in stdpopsim is updated to use msprime 1.0
APIs we can remove this model and instead compare directly
to the stdpopsim model with .is_equivalent() or whatever.
"""
# Times are provided in years, so we convert into generations.
generation_time = 25
T_OOA = 21.2e3 / generation_time
T_AMH = 140e3 / generation_time
T_ANC = 220e3 / generation_time
# We need to work out the starting (diploid) population sizes based on
# the growth rates provided for these two populations
r_CEU = 0.004
r_CHB = 0.0055
N_CEU = 1000 / math.exp(-r_CEU * T_OOA)
N_CHB = 510 / math.exp(-r_CHB * T_OOA)
demography = Demography()
# This is the "trunk" population that we merge other populations into
demography.add_population(
name="YRI",
description="Africa",
initial_size=12300,
initially_active=True,
)
demography.add_population(
name="CEU",
description="European",
initial_size=N_CEU,
growth_rate=r_CEU,
)
demography.add_population(
name="CHB",
description="East Asian",
initial_size=N_CHB,
growth_rate=r_CHB,
)
demography.add_population(
name="OOA",
description="Bottleneck out-of-Africa population",
initial_size=2100,
)
# Set the migration rates between extant populations
demography.set_symmetric_migration_rate(["CEU", "CHB"], 9.6e-5)
demography.set_symmetric_migration_rate(["YRI", "CHB"], 1.9e-5)
demography.set_symmetric_migration_rate(["YRI", "CEU"], 3e-5)
demography.add_population_split(
time=T_OOA, derived=["CEU", "CHB"], ancestral="OOA"
)
demography.add_symmetric_migration_rate_change(
time=T_OOA, populations=["YRI", "OOA"], rate=25e-5
)
demography.add_population_split(time=T_AMH, derived=["OOA"], ancestral="YRI")
demography.add_population_parameters_change(
time=T_ANC, population="YRI", initial_size=7300
)
return demography
@staticmethod
def _ooa_archaic_model():
"""
See notes for the _ooa model above.
"""
# Implement the OutOfAfricaArchaicAdmixture_5R19 model
# NOTE: this example isn't very well factored and needs more work.
# Times are provided in years, so we convert into generations.
generation_time = 29
T_OOA = 36_000 / generation_time
T_AMH = 60_700 / generation_time
T_ANC = 300_000 / generation_time
T_ArchaicAFR = 499_000 / generation_time
T_Neanderthal = 559_000 / generation_time
T_archaic_migration_start = 18_700 / generation_time
T_archaic_migration_end = 125_000 / generation_time
# We need to work out the starting (diploid) population sizes based on
# the growth rates provided for these two populations
r_CEU = 0.00125
r_CHB = 0.00372
N_CEU = 2300 / math.exp(-r_CEU * T_OOA)
N_CHB = 650 / math.exp(-r_CHB * T_OOA)
demography = Demography()
# This is the "trunk" population that we merge other populations into
demography.add_population(
name="AFR",
description="African population",
initial_size=13900,
initially_active=True,
)
demography.add_population(
name="CEU",
description=(
"Utah Residents (CEPH) with Northern and Western European Ancestry"
),
initial_size=N_CEU,
growth_rate=r_CEU,
)
demography.add_population(
name="CHB",
description="Han Chinese in Beijing, China",
initial_size=N_CHB,
growth_rate=r_CHB,
)
demography.add_population(
name="Neanderthal",
description="Putative Neanderthals",
initial_size=3600,
)
demography.add_population(
name="ArchaicAFR",
description="Putative Archaic Africans",
initial_size=3600,
)
demography.add_population(
name="OOA",
description="Bottleneck out-of-Africa population",
initial_size=880,
)
# Set the migration rates between extant populations
demography.set_symmetric_migration_rate(["CEU", "CHB"], 11.3e-5)
demography.set_symmetric_migration_rate(["AFR", "CEU"], 2.48e-5)
demography.add_symmetric_migration_rate_change(
T_archaic_migration_start, ["CEU", "Neanderthal"], 0.825e-5
)
demography.add_symmetric_migration_rate_change(
T_archaic_migration_start, ["CHB", "Neanderthal"], 0.825e-5
)
demography.add_symmetric_migration_rate_change(
T_archaic_migration_start, ["ArchaicAFR", "AFR"], 1.98e-5
)
demography.add_migration_rate_change(T_archaic_migration_end, rate=0)
demography.add_population_split(
time=T_OOA, derived=["CEU", "CHB"], ancestral="OOA"
)
demography.add_symmetric_migration_rate_change(
time=T_OOA, populations=["AFR", "OOA"], rate=52.2e-5
)
demography.add_symmetric_migration_rate_change(
time=T_OOA, populations=["OOA", "Neanderthal"], rate=0.825e-5
)
demography.add_population_split(time=T_AMH, derived=["OOA"], ancestral="AFR")
demography.add_symmetric_migration_rate_change(
T_AMH, ["ArchaicAFR", "AFR"], 1.98e-5
)
demography.add_population_parameters_change(
time=T_AMH, population="AFR", initial_size=13900
)
demography.add_population_parameters_change(
time=T_ANC, population="AFR", initial_size=3600
)
demography.add_population_split(
time=T_ArchaicAFR, derived=["ArchaicAFR"], ancestral="AFR"
)
demography.add_population_split(
time=T_Neanderthal, derived=["Neanderthal"], ancestral="AFR"
)
demography.sort_events()
return demography
@staticmethod
def _american_admixture_model():
# Implementation of AmericanAdmixture_4B11 model. See notes from the _ooa
# model above as to why this is here.
T_OOA = 920
N_EUR = 34039
r_EUR = 0.0038
N_EAS = 45852
r_EAS = 0.0048
T_ADMIX = 12
N_ADMIX = 54664
r_ADMIX = 0.05
demography = Demography()
demography.add_population(name="AFR", description="African", initial_size=14474)
demography.add_population(
name="EUR",
description="European",
initial_size=N_EUR,
growth_rate=r_EUR,
)
demography.add_population(
name="EAS",
description="East Asian",
initial_size=N_EAS,
growth_rate=r_EAS,
)
demography.add_population(
name="ADMIX",
description="Admixed America",
initial_size=N_ADMIX,
growth_rate=r_ADMIX,
)
demography.add_admixture(
T_ADMIX,
derived="ADMIX",
ancestral=["AFR", "EUR", "EAS"],
proportions=[1 / 6, 2 / 6, 3 / 6],
)
demography.add_population(
name="OOA",
description="Bottleneck out-of-Africa",
initial_size=1861,
)
demography.add_population(
name="AMH",
description="Anatomically modern humans",
initial_size=14474,
)
demography.add_population(
name="ANC",
description="Ancestral equilibrium",
initial_size=7310,
)
demography.set_symmetric_migration_rate(["AFR", "EUR"], 2.5e-5)
demography.set_symmetric_migration_rate(["AFR", "EAS"], 0.78e-5)
demography.set_symmetric_migration_rate(["EUR", "EAS"], 3.11e-5)
demography.add_population_split(T_OOA, derived=["EUR", "EAS"], ancestral="OOA")
demography.add_symmetric_migration_rate_change(
time=T_OOA, populations=["AFR", "OOA"], rate=15e-5
)
demography.add_population_split(2040, derived=["OOA", "AFR"], ancestral="AMH")
demography.add_population_split(5920, derived=["AMH"], ancestral="ANC")
return demography
def to_demes(self) -> demes.Graph:
"""
Creates a :class:`demes.Graph` object from the demography.
See the :ref:`sec_demography_importing_demes` section for details.
.. note::
Demographic models using bottlenecks added via the
:meth:`add_simple_bottleneck` or :meth:`add_instantaneous_bottleneck`
methods are not able to be converted into a demes graph.
.. note::
Demes is stricter than msprime with regard to how a demographic
model is structured, so models that can be simulated with msprime
are not guaranteed to be convertible to a demes graph. In particular,
msprime's legacy API permits setting migrations or other attributes
for a population even after that population has been merged into an
ancestor. Such models are rarely constructed deliberately, so an
error during conversion of a legacy model could indicate model
misspecification.
The returned graph can be saved as a
`Demes-format <https://popsim-consortium.github.io/demes-spec-docs/main/>`_
YAML file using the :ref:`demes <demes:sec_introduction>` API.
.. code::
import demes
demography = msprime.Demography.island_model([1000] * 3, 1e-5)
graph = demography.to_demes()
demes.dump(graph, "island_model.yaml")
Or plotted using the `demesdraw <https://github.com/grahamgower/demesdraw>`_
visualisation package.
.. code::
import demesdraw
demography = msprime.Demography.island_model([1000] * 3, 1e-5)
graph = demography.to_demes()
ax = demesdraw.tubes(graph)
ax.figure.savefig("island_model.pdf")
:return: A :class:`demes.Graph` object corresponding to the demography.
:rtype: demes.Graph
"""
dbg = self.debug()
resolved = dbg.demography
b = demes.Builder()
for pop in resolved.populations:
start_time = max(
epoch.end_time
for epoch in dbg.epochs
if epoch.populations[pop.id].state == PopulationStateMachine.ACTIVE
)
end_time = min(
epoch.start_time
for epoch in dbg.epochs
if epoch.populations[pop.id].state == PopulationStateMachine.ACTIVE
)
assert start_time > end_time
initial_epoch = dict(
end_size=pop.initial_size,
growth_rate=pop.growth_rate,
end_time=end_time,
)
b.add_deme(
pop.name,
description=pop.description if pop.description else None,
start_time=start_time,
epochs=[initial_epoch],
)
deme_map = {pop.name: b.data["demes"][pop.id] for pop in resolved.populations}
# Copied from demes/ms.py
def epoch_resolve(deme, time):
"""
Return the oldest epoch if it has end_time == time. If not, create a
new oldest epoch with end_time=time. Also resolve sizes by dealing
with the growth_rate attribute (if required).
"""
epoch = deme["epochs"][0]
start_time = deme["start_time"]
end_time = epoch["end_time"]
if not (start_time > time >= end_time):
raise ValueError(
f"time {time} outside {deme['name']}'s existence interval "
f"(start_time={start_time}, end_time={end_time}]"
)
if time > end_time:
new_epoch = copy.deepcopy(epoch)
# find size at given time
growth_rate = epoch.pop("growth_rate", 0)
dt = time - epoch["end_time"]
size_at_t = epoch["end_size"] * math.exp(-growth_rate * dt)
epoch["start_size"] = size_at_t
new_epoch["end_size"] = size_at_t
new_epoch["end_time"] = time
deme["epochs"].insert(0, new_epoch)
epoch = new_epoch
return epoch
for time, events_group in itertools.groupby(
resolved.events, operator.attrgetter("time")
):
events_group = list(events_group)
lineage_movements = []
for event in events_group:
if isinstance(event, LineageMovementEvent):
# Collect these so that we can later group them by source.
lineage_movements.extend(event._as_lineage_movements())
elif isinstance(event, PopulationParametersChange):
if event.population == -1:
pids = range(len(resolved.populations))
else:
pids = [event.population]
for pid in pids:
deme = deme_map[resolved[pid].name]
if (
len(deme["epochs"]) == 1
and deme["epochs"][0]["end_size"] == 0
and event.initial_size is not None
and event.initial_size > 0
):
# The population size was 0 at time 0, but is now
# positive, indicating population extinction at the
# current time.
deme["epochs"][0]["end_time"] = time
epoch = epoch_resolve(deme, time)
if event.growth_rate is not None:
epoch["growth_rate"] = event.growth_rate
if event.initial_size is not None:
epoch["end_size"] = event.initial_size
elif isinstance(event, StateChangeEvent):
raise ValueError(f"Cannot convert {event} to demes model.")
# Lineage movement events correspond to either the source deme's
# creation, or otherwise just pulses into the source.
# We distinguish these cases based on the total ancestry
# proportion for the source deme at this time.
lm_by_source = collections.defaultdict(list)
for lm in lineage_movements:
source = resolved[lm.source].name
lm_by_source[source].append(lm)
pulses = set()
for source, lm_group in lm_by_source.items():
ancestors = [resolved[lm.dest].name for lm in lm_group]
proportions = _sequential_to_proportions(
[lm.proportion for lm in lm_group]
)
if math.isclose(sum(proportions), 1):
# Source deme is created from the ancestors.
deme = deme_map[source]
deme.update(
ancestors=ancestors, proportions=proportions, start_time=time
)
for ancestor in ancestors:
if not resolved[ancestor].initially_active:
anc_deme = deme_map[ancestor]
anc_deme["epochs"][-1]["end_time"] = time
else:
# Source deme receives pulses from the ancestors.
pulses.update(
(lm.source, lm.dest, lm.proportion) for lm in lm_group
)
# The order of pulses matters when multiple pulses occur at the
# same time, so we must be careful to add the pulses in the same
# order as the lineage movements were specified. The pulses list
# is later reversed, to correspond with the forwards-time ordering.
for lm in lineage_movements:
if (lm.source, lm.dest, lm.proportion) in pulses:
b.add_pulse(
sources=[resolved[lm.dest].name],
dest=resolved[lm.source].name,
proportions=[lm.proportion],
time=time,
)
# Resolve/remove growth_rate in oldest epochs.
for deme in b.data["demes"]:
start_time = deme.get("start_time", math.inf)
epoch = deme["epochs"][0]
growth_rate = epoch.pop("growth_rate", 0)
if growth_rate != 0:
if math.isinf(start_time):
raise ValueError(
f"{deme['name']}: growth rate for infinite-length "
"epoch is invalid"
)
dt = start_time - epoch["end_time"]
epoch["start_size"] = epoch["end_size"] * math.exp(-dt * growth_rate)
else:
epoch["start_size"] = epoch["end_size"]
# Copied from demes/ms.py
def migrations_from_mm_list(
mm_list: list[Any], end_times: list[float], deme_names: list[str]
) -> list[MutableMapping]:
"""
Convert a list of migration matrices into a list of migration dicts.
"""
assert len(mm_list) == len(end_times)
migrations: list[MutableMapping] = []
current: dict[tuple[int, int], MutableMapping] = dict()
start_time = math.inf
for migration_matrix, end_time in zip(mm_list, end_times):
n = len(migration_matrix)
assert n == len(deme_names)
for j in range(n):
assert n == len(migration_matrix[j])
for k in range(n):
if j == k:
continue
rate = migration_matrix[j][k]
migration_dict = current.get((j, k))
if migration_dict is None:
if rate != 0:
migration_dict = dict(
source=deme_names[j],
dest=deme_names[k],
start_time=start_time,
end_time=end_time,
rate=rate,
)
current[(j, k)] = migration_dict
migrations.append(migration_dict)
else:
if rate == 0:
del current[(j, k)]
elif migration_dict["rate"] == rate:
# extend migration_dict
migration_dict["end_time"] = end_time
else:
migration_dict = dict(
source=deme_names[j],
dest=deme_names[k],
start_time=start_time,
end_time=end_time,
rate=rate,
)
current[(j, k)] = migration_dict
migrations.append(migration_dict)
start_time = end_time
return migrations
mm_list = [epoch.migration_matrix.T for epoch in reversed(dbg.epochs)]
mm_end_times = [epoch.start_time for epoch in reversed(dbg.epochs)]
migrations = migrations_from_mm_list(mm_list, mm_end_times, list(resolved))
b.data["migrations"] = migrations
# Reverse the order of pulses, so that older pulses come first.
# This also fixes the order of pulses that occur simultaneously,
# so that realised ancestry proportions match the msprime model.
if "pulses" in b.data:
b.data["pulses"].reverse()
# Sort demes by their start time (so that ancestors come before descendants).
b.data["demes"] = sorted(
b.data["demes"], key=operator.itemgetter("start_time"), reverse=True
)
return b.resolve()
# This was lifted out of older code as-is. No point in updating it
# to use dataclasses, since all we want to do is maintain compatibility
# with older code.
class PopulationConfiguration:
"""
The initial configuration of a population (or deme) in a simulation.
.. important::
This class is deprecated (but supported indefinitely);
please use the msprime 1.0 :ref:`demography API<sec_demography>`
in new code.
:param int sample_size: The number of initial samples that are drawn
from this population.
:param float initial_size: The absolute size of the population at time
zero. Defaults to the reference population size :math:`N_e`.
:param float growth_rate: The forwards-time exponential growth rate of the
population per generation. Growth rates can be negative. This is zero for a
constant population size, and positive for a population that has been
growing. Defaults to 0.
:param dict metadata: A JSON-encodable dictionary of metadata to associate
with the corresponding Population in the output tree sequence.
If not specified or None, no metadata is stored (i.e., an empty bytes array).
Note that this metadata is ignored when using the ``from_ts`` argument to
:func:`simulate`, as the population definitions in the tree sequence that
is used as the starting point take precedence.
"""
def __init__(
self, sample_size=None, initial_size=None, growth_rate=0.0, metadata=None
):
if initial_size is not None and initial_size < 0:
raise ValueError("Population size must be >= 0")
if sample_size is not None and sample_size < 0:
raise ValueError("Sample size must be >= 0")
self.sample_size = sample_size
self.initial_size = initial_size
self.growth_rate = growth_rate
self.metadata = metadata
def asdict(self):
return dict(
sample_size=self.sample_size,
initial_size=self.initial_size,
growth_rate=self.growth_rate,
metadata=self.metadata,
)
def _list_str(a: list, fmt=None):
"""
Returns the specified items rendered as a string without quotes.
"""
if fmt is None:
joined = ", ".join(str(item) for item in a)
else:
joined = ", ".join(f"{item:{fmt}}" for item in a)
return f"[{joined}]"
@dataclasses.dataclass
class DemographicEvent:
"""
Superclass of demographic events that occur during simulations.
"""
time: float
demography: Demography = dataclasses.field(
init=False, compare=False, default=None, repr=False
)
def _parameters(self):
raise NotImplementedError()
def _effect(self):
raise NotImplementedError()
def asdict(self):
return {
key: getattr(self, key)
for key in inspect.signature(self.__init__).parameters.keys()
if hasattr(self, key)
}
def _convert_id(self, population_ref):
"""
Converts the specified population reference into an integer,
suitable for input into the low-level code. We treat -1 as a special
case because it's used as meaning "all populations" by the events.
"""
if population_ref in [-1, None]:
# Both of these mean "all populations"
return -1
if self.demography is None:
# We need to be able to handle Events that are not associated with
# a Demography to support old code. However, these should only ever
# happen with integer IDs.
if not core.isinteger(population_ref):
raise ValueError(
"Working with demographic events not associated with a "
"Demography object is a legacy-only operation. Population "
"references must be integer IDs"
)
return population_ref
return self.demography[population_ref].id
class ParameterChangeEvent(DemographicEvent):
"""
Superclass of events that change some parameters in the underlying
simulation model but don't actually affect the state in any other
way.
"""
@dataclasses.dataclass
class PopulationParametersChange(ParameterChangeEvent):
"""
Changes the demographic parameters of a population at a given time.
This event generalises the ``-eg``, ``-eG``, ``-en`` and ``-eN``
options from ``ms``. Note that unlike ``ms`` we do not automatically
set growth rates to zero when the population size is changed.
.. important::
This class is deprecated (but supported indefinitely);
please use the :meth:`.Demography.add_population_parameters_change`
method in new code.
:param float time: The length of time ago at which this event
occurred.
:param float initial_size: The absolute diploid size of the population
at the beginning of the time slice starting at ``time``. If None,
this is calculated according to the initial population size and
growth rate over the preceding time slice.
:param float growth_rate: The new per-generation growth rate. If None,
the growth rate is not changed. Defaults to None.
:param int population: The ID of the population affected. If
``population`` is None, the changes affect all populations
simultaneously.
"""
initial_size: float | None = None
growth_rate: float | None = None
# TODO change the default to -1 to match MigrationRateChange.
population: int | None = None
# Deprecated.
# TODO add a formal deprecation notice
population_id: int | None = dataclasses.field(default=None, repr=False)
_type_str: ClassVar[str] = "Population parameter change"
def __post_init__(self):
if self.population_id is not None and self.population is not None:
raise ValueError(
"population_id and population are aliases; cannot supply both."
)
if self.population_id is not None:
self.population = self.population_id
if self.growth_rate is None and self.initial_size is None:
raise ValueError("Must specify one or more of growth_rate and initial_size")
if self.initial_size is not None and self.initial_size < 0:
raise ValueError("Cannot have a population size < 0")
self.population = -1 if self.population is None else self.population
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.2 is out
# https://github.com/tskit-dev/msprime/issues/1037
ret = {
"type": "population_parameters_change",
"time": self.time,
"population": self._convert_id(self.population),
}
if self.growth_rate is not None:
ret["growth_rate"] = self.growth_rate
if self.initial_size is not None:
ret["initial_size"] = self.initial_size
return ret
def _parameters(self):
s = f"population={self.population}, "
if self.initial_size is not None:
s += f"initial_size={self.initial_size}, "
if self.growth_rate is not None:
s += f"growth_rate={self.growth_rate}, "
return s[:-2]
def _effect(self):
s = ""
if self.initial_size is not None:
s += f"initial_size → {self.initial_size:.2g} "
if self.growth_rate is not None:
s += "and "
if self.growth_rate is not None:
s += f"growth_rate → {self.growth_rate:.3g} "
s += "for"
if self.population == -1:
s += " all populations"
else:
s += f" population {self.population}"
return s
@dataclasses.dataclass
class MigrationRateChange(ParameterChangeEvent):
"""
Changes the rate of migration from one deme to another to a new value at a
specific time. Migration rates are specified in terms of the rate at which
lineages move from population ``source`` to ``dest`` during the progress of
the simulation. Note that ``source`` and ``dest`` are from the perspective
of the coalescent process; please see the :ref:`sec_ancestry_models`
section for more details on the interpretation of this migration model.
By default, ``source=-1`` and ``dest=-1``, which results in all
non-diagonal elements of the migration matrix being changed to the new
rate. If ``source`` and ``dest`` are specified, they must refer to valid
population IDs.
.. important::
This class is deprecated (but supported indefinitely);
please use the :meth:`.Demography.add_migration_rate_change`
method in new code.
:param float time: The time at which this event occurs in generations.
:param float rate: The new per-generation migration rate.
:param int source: The ID of the source population.
:param int dest: The ID of the destination population.
"""
rate: float
source: int = -1
dest: int = -1
# Deprecated.
# TODO add a formal deprecation notice
matrix_index: tuple | None = dataclasses.field(default=None, repr=False)
_type_str: ClassVar[str] = "Migration rate change"
def __post_init__(self):
# If the deprecated form is used, it overwrites the values of source
# and dest
if self.matrix_index is not None:
self.source = self.matrix_index[0]
self.dest = self.matrix_index[1]
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "migration_rate_change",
"time": self.time,
# Note: We'd like to change the name here to "rate" but it's best
# to leave this alone until stdpopsim has been moved away from
# using this internal API.
"migration_rate": self.rate,
"source": self._convert_id(self.source),
"dest": self._convert_id(self.dest),
}
def _parameters(self):
return f"source={self.source}, dest={self.dest}, rate={self.rate}"
def _effect(self):
ret = "Backwards-time migration rate "
if self.source == -1 and self.dest == -1:
ret += "for all populations "
else:
ret += f"from {self.source} to {self.dest} "
ret += f"→ {self.rate}"
return ret
@dataclasses.dataclass
class SymmetricMigrationRateChange(ParameterChangeEvent):
"""
Class representing a SymmetricMigrationRateChange. Not part of the
external API as it was added after 1.0.
"""
populations: list[int | str]
rate: float
_type_str: ClassVar[str] = dataclasses.field(
default="Symmetric migration rate change", repr=False
)
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "symmetric_migration_rate_change",
"time": self.time,
"populations": [self._convert_id(pop) for pop in self.populations],
"rate": self.rate,
}
def _parameters(self):
return f"populations={_list_str(self.populations)}, rate={self.rate}"
def _effect(self):
s = "Sets the symmetric migration rate between "
if len(self.populations) == 2:
s += f"{self.populations[0]} and {self.populations[1]} "
else:
s += f"all pairs of populations in {_list_str(self.populations)} "
s += f"to {self.rate} per generation"
return s
@dataclasses.dataclass
class LineageMovement:
"""
A single instantaneous movement of lineages from one population to
another. Note that 'source' and 'dest' are in the backwards-in-time
sense.
"""
source: int
dest: int
proportion: float
class LineageMovementEvent(DemographicEvent):
"""
Superclass of events that move lineages around between populations.
"""
def _as_lineage_movements(self) -> list[LineageMovement]:
"""
Returns the equivalent of this lineage movement event as a
list of lineage movements.
"""
raise NotImplementedError()
@dataclasses.dataclass
class MassMigration(LineageMovementEvent):
"""
A mass migration event in which some fraction of the population in one deme
(the ``source``) simultaneously move to another deme (``dest``) during the
progress of the simulation. Each lineage currently present in the source
population moves to the destination population with probability equal to
``proportion``. Note that ``source`` and ``dest`` are from the perspective
of the coalescent process; please see the :ref:`sec_ancestry_models`
section for more details on the interpretation of this migration model.
This event class generalises the population split (``-ej``) and
admixture (``-es``) events from ``ms``. Note that MassMigrations
do *not* have any side effects on the migration matrix.
.. important::
This class is deprecated (but supported indefinitely);
please use the :meth:`.Demography.add_mass_migration`
method in new code. In addition, please see the new
higher-level :ref:`sec_demography_events_population_split`
and :ref:`sec_demography_events_admixture` events.
:param float time: The time at which this event occurs in generations.
:param int source: The ID of the source population.
:param int dest: The ID of the destination population.
:param float proportion: The probability that any given lineage within
the source population migrates to the destination population.
"""
source: int
# dest only has a default because of the deprecated destination attr.
dest: None | int = None
proportion: float = 1.0
# Deprecated.
# TODO add a formal deprecation notice
destination: int | None = dataclasses.field(default=None, repr=False)
_type_str: ClassVar[str] = dataclasses.field(default="Mass Migration", repr=False)
def __post_init__(self):
if self.dest is not None and self.destination is not None:
raise ValueError("dest and destination are aliases; cannot supply both")
if self.destination is not None:
self.dest = self.destination
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "mass_migration",
"time": self.time,
"source": self._convert_id(self.source),
"dest": self._convert_id(self.dest),
"proportion": self.proportion,
}
def _parameters(self):
return (
f"source={self.source}, dest={self.dest}, proportion={self.proportion:.3g}"
)
def _effect(self):
if self.proportion == 1.0:
ret = (
f"All lineages currently in population {self.source} move "
f"to {self.dest} "
)
else:
ret = (
f"Lineages currently in population {self.source} move to {self.dest} "
f"with probability {self.proportion:.3g} "
)
ret += (
"(equivalent to individuals "
f"migrating from {self.dest} to {self.source} forwards in time)"
)
return ret
def _as_lineage_movements(self):
return [
LineageMovement(
source=self._convert_id(self.source),
dest=self._convert_id(self.dest),
proportion=self.proportion,
)
]
@dataclasses.dataclass
class ActivatePopulationEvent(DemographicEvent):
"""
A population activation event, which changes the state of the population
from inactive to active at the given time.
"""
population: int | str | None = None
"""
The name or ID of the population to activate.
"""
_type_str: ClassVar[str] = dataclasses.field(
default="Activate population event", repr=False
)
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "activate_population_event",
"time": self.time,
"population": self._convert_id(self.population),
}
def _parameters(self):
return f"population={self.population}"
def _effect(self):
s = f"Activates population {self.population}"
return s
@dataclasses.dataclass
class PopulationSplit(LineageMovementEvent):
"""
:param float time: The time at which this event occurs in generations.
:param list(int) derived: The ID(s) of the derived population(s).
:param int ancestral: The ID of the ancestral population.
"""
derived: list[int | str]
ancestral: int | str
_type_str: ClassVar[str] = dataclasses.field(default="Population Split", repr=False)
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "population_split",
"time": self.time,
"derived": [self._convert_id(pop) for pop in self.derived],
"ancestral": self._convert_id(self.ancestral),
}
def _parameters(self):
return f"derived={_list_str(self.derived)}, ancestral={self.ancestral}"
def _effect(self):
s = "Moves all lineages from "
if len(self.derived) == 1:
s += f"the '{self.derived[0]}' derived population "
else:
s += "derived populations "
if len(self.derived) == 2:
s += f"'{self.derived[0]}' and '{self.derived[1]}' "
else:
s += f"{_list_str(self.derived)} "
s += f"to the ancestral '{self.ancestral}' population. "
s += "Also set "
if len(self.derived) == 1:
s += f"'{self.derived[0]}' "
else:
s += "the derived populations "
s += (
"to inactive, and all migration rates to and from "
f"the derived population{'s' if len(self.derived) > 1 else ''} "
"to zero."
)
return s
def _as_lineage_movements(self):
ancestral = self._convert_id(self.ancestral)
return [
LineageMovement(source=self._convert_id(pop), dest=ancestral, proportion=1)
for pop in self.derived
]
@dataclasses.dataclass
class Admixture(LineageMovementEvent):
"""
:param float time: The time at which this event occurs in generations.
"""
derived: int | str
ancestral: list[int | str]
proportions: list[float]
_type_str: ClassVar[str] = dataclasses.field(default="Admixture", repr=False)
@property
def num_ancestral(self):
return len(self.ancestral)
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "admixture",
"time": self.time,
"derived": self._convert_id(self.derived),
"ancestral": [self._convert_id(pop) for pop in self.ancestral],
"proportions": self.proportions,
}
def _parameters(self):
return (
f"derived={self.derived} ancestral={_list_str(self.ancestral)} "
f"proportions={_list_str(self.proportions, '.2f')}"
)
def _effect(self):
move_to = "; ".join(
f"'{pop}' with proba {proba:.3g}"
for pop, proba in zip(self.ancestral, self.proportions)
)
return (
f"Moves all lineages from admixed population '{self.derived}' "
f"to ancestral population{'s' if len(self.ancestral) > 1 else ''}. "
f"Lineages move to {move_to}. Set '{self.derived}' to inactive, "
f"and all migration rates to and from '{self.derived}' to zero."
)
def _as_lineage_movements(self):
derived = self._convert_id(self.derived)
ancestral = [self._convert_id(pop) for pop in self.ancestral]
# Conditioned on having already distributed a fraction q of the
# lineages, the we need a fraction p / (1 - q) of the remaining
# lineages to get an overall proportion of p.
S = _proportions_to_sequential(self.proportions)
return [
LineageMovement(source=derived, dest=ancestral[j], proportion=S[j])
for j in range(self.num_ancestral)
]
class StateChangeEvent(DemographicEvent):
"""
Superclass of events that change the state of the simulation in complex
ways.
"""
# This is an unsupported/undocumented demographic event.
@dataclasses.dataclass
class SimpleBottleneck(StateChangeEvent):
population: int
proportion: float = 1.0
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "simple_bottleneck",
"time": self.time,
"population": self._convert_id(self.population),
"proportion": self.proportion,
}
_type_str: ClassVar[str] = dataclasses.field(
default="Simple Bottleneck", repr=False
)
def _parameters(self):
return f"population={self.population}, proportion={self.proportion}"
def _effect(self):
return (
f"Lineages in population {self.population} coalesce with "
f"probability {self.proportion}"
)
# TODO document
@dataclasses.dataclass
class InstantaneousBottleneck(StateChangeEvent):
population: int
strength: float = 1.0
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "instantaneous_bottleneck",
"time": self.time,
"population": self._convert_id(self.population),
"strength": self.strength,
}
_type_str: ClassVar[str] = dataclasses.field(
default="Instantaneous Bottleneck", repr=False
)
def _parameters(self):
return f"population={self.population}, strength={self.strength}"
def _effect(self):
return f"Equivalent to {self.strength} generations of the coalescent"
@dataclasses.dataclass
class CensusEvent(DemographicEvent):
"""
An event that adds a node to each branch of every tree at a given time
during the simulation. This may be used to record all ancestral haplotypes
present at that time, and to extract other information related to these
haplotypes: for instance to trace the local ancestry of a sample back to a
set of contemporaneous ancestors, or to assess whether a subset of samples
has coalesced more recently than the census time.
See :ref:`sec_ancestry_census_events` for more details.
.. important::
This class is deprecated (but supported indefinitely);
please use the :meth:`.Demography.add_census`
method in new code.
:param float time: The time at which this event occurs in generations.
"""
_type_str: ClassVar[str] = dataclasses.field(default="Census", repr=False)
def get_ll_representation(self, num_populations=None):
# We need to keep the num_populations argument until stdpopsim 0.1 is out
# https://github.com/tskit-dev/msprime/issues/1037
return {
"type": "census_event",
"time": self.time,
}
def _parameters(self):
return ""
def _effect(self):
return "Insert census nodes to record the location of all lineages"
def _sequential_to_proportions(S):
"""
Given a list of sequential lineage proportions out of a population,
return the absolute proportions of the original population this
corresponds to.
"""
P = []
for j in range(len(S)):
P.append(S[j] * (1 - sum(P[:j])))
return P
def _proportions_to_sequential(P):
"""
Given a list of absolute proportions of lineages moving out of
a population, return the sequential conditional movements required
to give them same proportions.
"""
# Conditioned on having already distributed a fraction q of the
# lineages, the we need a fraction p / (1 - q) of the remaining
# lineages to get an overall proportion of p
C = [0 for _ in P]
for j in range(len(P)):
s = sum(P[:j])
if s < 1:
C[j] = P[j] / (1 - s)
return C
def _matrix_exponential(A, n=5):
"""
Returns the matrix exponential of A.
Only works if the offdiagonals of A are nonnegative
and the row sums of A are less than or equal to zero.
Melloy & Bennett (1993), https://doi.org/10.1016/0377-0427(93)90036-B
"""
assert np.max(np.diag(A)) <= 0
assert np.min(np.tril(A, k=-1)) == 0
assert np.min(np.triu(A, k=1)) == 0
assert np.max(np.sum(A, 1)) <= 1e-10 # for floating-point error
dA = (-1) * np.diag(A)
dmax = np.max(dA)
expA = np.eye(A.shape[0])
if dmax > 0:
P = A / dmax
np.fill_diagonal(P, 1 - dA / dmax)
# nscales is the number of scaling-and-squaring steps:
# we compute exp(tA) and then square the result nscales times
nscales = int(max(0, np.ceil(np.log2(dmax) - np.log2(0.2))))
t = dmax / (2**nscales)
# Now, exp(tA) = exp(-t) * (I + tP + (tP)^2 / 2 + ...)
# and the k-th term in this sum is Pk=(tP/k times the previous)
Pk = np.eye(A.shape[0])
for k in range(1, n + 1):
Pk = (t / k) * np.matmul(P, Pk)
expA = expA + Pk
expA *= np.exp(-t)
expA = np.linalg.matrix_power(expA, 2**nscales)
return expA
class PopulationStateMachine(enum.IntEnum):
"""
During a simulation each population has three possible states described
by this state machine. In general, a population follows:
INACTIVE -> ACTIVE -> PREVIOUSLY_ACTIVE
All populations are by default ACTIVE at the start of the simulation,
except if they are they are "ancestral" in a population split event.
In this case populations are initially INACTIVE. A population
then transitions from INACTIVE -> ACTIVE when the corresponding
population split event occurs.
Populations transition from ACTIVE -> PREVIOUSLY_ACTIVE when they
are "derived" in either population split or admixture events.
No other transitions are possible.
"""
INACTIVE = 0
ACTIVE = 1
PREVIOUSLY_ACTIVE = 2
@dataclasses.dataclass
class PopulationState:
"""
Simple class to represent the state of a population in terms of its
demographic parameters. Note: start and end here refer to time flowing
*backwards*!
"""
id: int # noqa: A003
name: str
start_size: float
end_size: float
growth_rate: float
state: int
@property
def active(self):
return self.state == PopulationStateMachine.ACTIVE
@dataclasses.dataclass
class Epoch:
"""
Represents a single epoch in the simulation within which the state
of the demographic parameters are constant.
"""
index: int
start_time: float
end_time: float
populations: list[PopulationState]
migration_matrix: list # TODO numpy array
events: list[DemographicEvent]
def _title_text(self):
return (
f"Epoch[{self.index}]: "
f"[{self.start_time:.3g}, {self.end_time:.3g}) generations"
)
def _population_state_text(self):
return (
f"Populations "
f"(total={len(self.populations)} active={self.num_active_populations})"
)
@property
def demographic_events(self):
# For compatibility with msprime 0.x
return self.events
@property
def active_populations(self):
return [pop for pop in self.populations if pop.active]
@property
def num_active_populations(self):
return len(self.active_populations)
class DemographyDebugger:
"""
Utilities to compute and display information about the state of populations
during the different simulation epochs defined by demographic events.
.. warning:: This class is not intended to be instantiated directly using
the constructor - please use :meth:`.Demography.debug()` to obtain
a DemographyDebugger for a given :class:`.Demography` instead.
"""
def __init__(
self,
# Deprecated pre-1.0 parameters.
Ne=1,
population_configurations=None,
migration_matrix=None,
demographic_events=None,
model=None,
*,
demography=None,
):
if demography is None:
# Support the pre-1.0 syntax
demography = Demography.from_old_style(
population_configurations,
migration_matrix=migration_matrix,
demographic_events=demographic_events,
Ne=Ne,
ignore_sample_size=True,
)
self.demography = demography.validate()
self.num_populations = demography.num_populations
self._make_epochs()
self._check_misspecification()
def _make_epochs(self):
self.epochs = []
simulator = ancestry._parse_sim_ancestry(
demography=self.demography, init_for_debugger=True
)
start_time = 0
end_time = 0
abs_tol = 1e-9
event_index = 0
all_events = self.demography.events
while not math.isinf(end_time):
events = []
while event_index < len(all_events) and math.isclose(
all_events[event_index].time, start_time, abs_tol=abs_tol
):
events.append(all_events[event_index])
event_index += 1
end_time = simulator.debug_demography()
migration_matrix = simulator.migration_matrix
pop_conf = simulator.population_configuration
populations = [
PopulationState(
id=j,
name=self.demography.populations[j].name,
start_size=simulator.compute_population_size(j, start_time),
end_size=simulator.compute_population_size(j, end_time),
growth_rate=pop_conf[j]["growth_rate"],
state=PopulationStateMachine(pop_conf[j]["state"]),
)
for j in range(self.num_populations)
]
epoch_index = len(self.epochs)
self.epochs.append(
Epoch(
epoch_index,
start_time,
end_time,
populations,
migration_matrix,
events,
)
)
start_time = end_time
def _check_misspecification(self):
"""
Check for things that might indicate model misspecification.
"""
merged_pops = set()
for epoch in self.epochs:
for de in epoch.events:
if isinstance(de, MassMigration) and de.proportion == 1:
merged_pops.add(de.source)
mm = epoch.migration_matrix
for pop_k in merged_pops:
k = self.demography[pop_k].id
if any(mm[k, :] != 0) or any(mm[:, k] != 0):
warnings.warn(
"Non-zero migration rates exist after merging "
f"population {k}. This almost certainly indicates "
"demographic misspecification."
)
def _populations_table(self, epoch, as_text=True):
active_populations = epoch.active_populations
column_titles = ["", "start", "end", "growth_rate"]
if len(active_populations) > 1:
column_titles += [pop.name for pop in active_populations]
data = []
for pop in active_populations:
row = [
pop.name,
f"{pop.start_size: .1f}",
f"{pop.end_size: .1f}",
f"{pop.growth_rate: .3g}",
]
if len(active_populations) > 1:
for other_pop in active_populations:
item = self.demography._migration_rate_info(
pop.id,
other_pop.id,
epoch.migration_matrix[pop.id, other_pop.id],
)
if as_text:
row.append(item.as_text())
else:
row.append(item)
data.append(row)
return column_titles, data
def _populations_html(self, epoch):
column_titles, data = self._populations_table(epoch, as_text=False)
return core.html_table(epoch._population_state_text(), column_titles, data)
def _populations_text(self, epoch):
column_titles, data = self._populations_table(epoch)
alignments = ">>><"
if epoch.num_active_populations > 1:
alignments += "^" * epoch.num_active_populations
# Repack the table items as lists
column_titles = [[x] for x in column_titles]
data = [[[x] for x in row] for row in data]
return core.text_table(
epoch._population_state_text(), column_titles, alignments, data
)
def _repr_html_(self):
out = ""
for epoch in self.epochs:
if epoch.index > 0:
assert len(epoch.events) > 0
title = f"Events @ generation {epoch.start_time:.3g}"
out += self.demography._events_html(epoch.events, title)
out += "</div></details>"
else:
assert len(epoch.events) == 0
title = epoch._title_text()
out += f'<details open="true"><summary>{title}</summary>'
# Indent the content div slightly
out += '<div style="margin-left:20px">'
out += self._populations_html(epoch)
out += "</div>"
out += "</details>"
return f"<div>{out}</div>"
def print_history(self, output=sys.stdout):
"""
Prints a summary of the history of the populations.
Deprecated since 1.0: use ``print(debugger)`` instead.
"""
print(self, file=output, end="")
def __str__(self):
def indent(table, header_char="╟", depth=4):
lines = table.splitlines()
s = header_char + (" " * depth) + lines[0] + "\n"
for line in lines[1:]:
s += "║" + (" " * depth) + line + "\n"
return s
def box(title):
N = len(title) + 2
top = "╠" + ("═" * N) + "╗"
bottom = "╠" + ("═" * N) + "╝"
return f"{top}\n║ {title} ║\n{bottom}\n"
out = "DemographyDebugger\n"
for epoch in self.epochs:
if epoch.index > 0:
assert len(epoch.events) > 0
title = f"Events @ generation {epoch.start_time:.3g}"
out += indent(self.demography._events_text(epoch.events, title))
out += box(epoch._title_text())
out += indent(self._populations_text(epoch))
return out
@property
def population_size_history(self):
"""
Returns a (num_pops, num_epochs) numpy array giving the starting population size
for each population in each epoch.
"""
pop_size = np.zeros((self.num_populations, self.num_epochs))
for j, epoch in enumerate(self.epochs):
for k, pop in enumerate(epoch.populations):
pop_size[k, j] = pop.start_size
return pop_size
@property
def epoch_start_time(self):
"""
The array of epoch start_times defined by the demographic model.
"""
return np.array([x.start_time for x in self.epochs])
@property
def num_epochs(self):
"""
Returns the number of epochs defined by the demographic model.
"""
return len(self.epochs)
def population_size_trajectory(self, steps):
"""
Return an array of per-population population sizes,
as defined by the demographic model. These are the `initial_size`
parameters of the model, modified by any population growth rates.
The sizes are computed at the time points given by `steps`.
:param list steps: List of times ago at which the population
size will be computed.
:return: Returns a numpy array of population sizes, with one column per
population, whose [i,j]th entry is the size of population
j at time steps[i] ago.
"""
num_pops = self.num_populations
N_t = np.zeros([len(steps), num_pops])
for j, t in enumerate(steps):
N, _ = self._pop_size_and_migration_at_t(t)
N_t[j] = N
return N_t
def lineage_probabilities(self, steps, sample_time=0):
"""
Returns an array such that P[j, a, b] is the probability that a lineage that
started in population a at time sample_time is in population b at time steps[j]
ago.
This function reports sampling probabilities _before_ mass migration events
(or other events that move lineages) at a step time, if a mass migration
event occurs at one of those times. Migrations will then effect the next
time step.
:param list steps: A list of times to compute probabilities.
:param sample_time: The time of sampling of the lineage. For any times in steps
that are more recent than sample_time, the probability of finding the
lineage in any population is zero.
:return: An array of dimension len(steps) by num pops by num_pops.
"""
num_pops = self.num_populations
# P[i, j] will be the probability that a lineage that started in i is now in j
P = np.eye(num_pops)
# epochs are defined by mass migration events or changes to population sizes
# or migration rates, so we add the epoch interval times to the steps that we
# need to account for
epoch_breaks = [t for t in self.epoch_start_time if t not in steps]
all_steps = np.concatenate([steps, epoch_breaks])
sampling = []
if sample_time not in all_steps:
sampling.append(sample_time)
all_steps = np.concatenate((all_steps, sampling))
ix = np.argsort(all_steps)
all_steps = all_steps[ix]
# keep track of the steps to report in P_out
keep_steps = np.concatenate(
[
np.repeat(True, len(steps)),
np.repeat(False, len(epoch_breaks)),
np.repeat(False, len(sampling)),
]
)[ix]
assert len(np.unique(all_steps)) == len(all_steps)
assert np.all(steps == all_steps[keep_steps])
P_out = np.zeros((len(all_steps), num_pops, num_pops))
first_step = 0
while all_steps[first_step] < sample_time:
first_step += 1
P_out[first_step] = P
# get ordered mass migration events
mass_migration_objects = []
mass_migration_times = []
for demo in self.demography.events:
if isinstance(demo, LineageMovementEvent):
# Convert higher-level lineage movement events like Admixtures
# and PopulationSplits into LineageMovement instances. These are
# equivalent to MassMigrations
for lm in demo._as_lineage_movements():
mass_migration_objects.append(lm)
mass_migration_times.append(demo.time)
for jj in range(first_step, len(all_steps) - 1):
t_j = all_steps[jj]
# apply any mass migration events to P
# so if we sample at this time, we do not account for the instantaneous
# mass migration events that occur at the same time. that will show up
# at the next step
if t_j > sample_time:
for mass_mig_t, mass_mig_e in zip(
mass_migration_times, mass_migration_objects
):
if mass_mig_t == t_j:
S = np.eye(num_pops, num_pops)
S[mass_mig_e.source, mass_mig_e.dest] = mass_mig_e.proportion
S[mass_mig_e.source, mass_mig_e.source] = (
1 - mass_mig_e.proportion
)
P = np.matmul(P, S)
# get continuous migration matrix over next interval
_, M = self._pop_size_and_migration_at_t(t_j)
dt = all_steps[jj + 1] - all_steps[jj]
dM = np.diag([sum(s) for s in M])
# advance to next interval time (dt) taking into account continuous mig
P = P.dot(_matrix_exponential(dt * (M - dM)))
P_out[jj + 1] = P
return P_out[keep_steps]
def possible_lineage_locations(self, samples=None):
"""
Given the sampling configuration, this function determines when lineages are
possibly found within each population over epochs defined by demographic events
and sampling times. If no sampling configuration is given, we assume we sample
lineages from every population at time zero.
The epoch intervals returned are those in which there are *distinct*
configurations of possible lineage locations, and so the number of
returned epochs may be less than the total number of epochs defined
by the demography and will depend on the input sample configuration.
The samples are specified by either a list of population identifiers (
integer IDs or string names) or by a list of :class:`.SampleSet` objects,
allowing sampling times to be specified explicitly. If the ``time`` field
of the :class:`.SampleSet` is not specified (or population IDs are used)
samples are taken at the population's `default_sampling_time`. Only
:class:`.SampleSet` objects with ``num_samples > 0`` are counted as
contributing samples to a particular population.
To support legacy code, :class:`.Sample` objects from the 0.x API
can also provided, although its use is discouraged in new code.
:param list samples: The populations that we sample from. Can be either
a list of population identifiers, :class:`.SampleSet` or
:class:`.Sample` objects.
:return: Returns a dictionary with epoch intervals as keys whose values are a
list with length equal to the number of populations with True and False
indicating which populations could possibly contain lineages over that
epoch. The epoch intervals are given by tuples: (epoch start, epoch end).
The first epoch necessarily starts at time 0, and the final epoch has end
time of infinity.
"""
if samples is None:
samples = [
pop.id
for pop in self.demography.populations
if pop.default_sampling_time == 0
]
# get configuration of sampling times from samples ({time:[pops_sampled_from]})
sampling_times = collections.defaultdict(list)
for sample in samples:
if isinstance(sample, (ancestry.Sample, ancestry.SampleSet)):
pop_id = self.demography[sample.population].id
sample_time = (
self.demography[pop_id].default_sampling_time
if sample.time is None
else sample.time
)
if isinstance(sample, ancestry.SampleSet) and sample.num_samples <= 0:
# If someone specifies 0 samples it should not be counted
continue
else:
# Assume this is a population identifier.
pop = self.demography[sample]
pop_id = pop.id
sample_time = pop.default_sampling_time
sampling_times[sample_time].append(pop_id)
for t in sampling_times.keys():
sampling_times[t] = list(set(sampling_times[t]))
all_steps = sorted(
list(set([t for t in self.epoch_start_time] + list(sampling_times.keys())))
)
epochs = [(x, y) for x, y in zip(all_steps[:-1], all_steps[1:])]
epochs.append((all_steps[-1], np.inf))
# need to go a bit beyond last step and into the final epoch that extends to inf
all_steps.append(all_steps[-1] + 1)
indicators = {e: np.zeros(self.num_populations, dtype=bool) for e in epochs}
for sample_time, pop_ids in sampling_times.items():
P_out = self.lineage_probabilities(all_steps, sample_time=sample_time)
for epoch, P in zip(epochs, P_out[1:]):
if epoch[1] <= sample_time:
# samples shouldn't affect the epoch previous to the sampling time
continue
for pop_id in pop_ids:
indicators[epoch][P[pop_id] > 0] = True
# join epochs if adjacent epochs have same set of possible live populations
combined_indicators = {}
skip = 0
for ii, (epoch, inds) in enumerate(indicators.items()):
if skip > 0:
skip -= 1
continue
this_epoch = epoch
while ii + skip + 1 < len(epochs) and np.all(
indicators[epochs[ii + 1 + skip]] == inds
):
this_epoch = (this_epoch[0], epochs[ii + 1 + skip][1])
skip += 1
combined_indicators[this_epoch] = inds
return combined_indicators
def mean_coalescence_time(
self, lineages, min_pop_size=1, steps=None, rtol=0.005, max_iter=12
):
"""
Compute the mean time until coalescence between pairs of the specified
sample ``lineages``. Sample lineages are specified as a mapping from
populations to the number of **monoploid** sample genomes present in
that population at time zero. See the
:ref:`sec_demography_numerical_coalescence` section for usage examples
and more details.
.. important:: This function assumes a diploid model when computing
coalescence rates (see the
:ref:`sec_ancestry_ploidy_coalescent_time_scales` section for more
information).
The calculation is performed by using
:meth:`~.DemographyDebugger.coalescence_rate_trajectory`
to compute the probability that the lineages have not yet coalesced by time `t`,
and using these to approximate :math:`E[T] = \\int_t^\\infty P(T > t) dt`,
where :math:`T` is the coalescence time. See
:meth:`~.DemographyDebugger.coalescence_rate_trajectory`
for more details.
To compute this, an adequate time discretisation must be arrived at
by iteratively extending or refining the current discretisation.
Debugging information about numerical convergence of this procedure is
logged using the Python :mod:`logging` infrastructure.
The `daiquiri <https://pypi.org/project/daiquiri/>`_ module is a
convenient way to set up logging, and we can use it to make these
messages appear on stderr like this::
import daiquiri
daiquiri.setup(level="DEBUG")
debugger.mean_coalescence_time(1)
Briefly, this outputs iteration number, mean coalescence time, maximum
difference in probability of not having coalesced yet, difference to
last coalescence time, probability of not having coalesced by the final
time point, and whether the last iteration was an extension or
refinement.
:param dict lineages: A mapping of populations (either integer IDs
or string names: see the :ref:`sec_demography_populations_identifiers`
section for more details) to the number of monoploid sample lineages
in that population.
:param int min_pop_size: See
:meth:`~.DemographyDebugger.coalescence_rate_trajectory`.
:param list steps: The time discretisation to start out with (by default,
picks something based on epoch times).
:param float rtol: The relative tolerance to determine mean coalescence time
to (used to decide when to stop subdividing the steps).
:param int max_iter: The maximum number of times to subdivide the steps.
:return: The mean coalescence time (a number).
:rtype: float
"""
def mean_time(steps, P):
# Mean is int_0^infty P(T > t) dt, which we estimate by discrete integration
# assuming that f(t) = P(T > t) is piecewise exponential:
# if f(u) = a exp(bu) then b = log(f(t)/f(s)) / (t-s) for each s < t, so
# \int_s^t f(u) du = (a/b) \int_s^t exp(bu) b du = (a/b)(exp(bt) - exp(bs))
# = (t - s) * (f(t) - f(s)) / log(f(t) / f(s))
# unless b = 0, of course.
assert steps[0] == 0
dt = np.diff(steps)
dP = np.diff(P)
with np.errstate(divide="ignore", invalid="ignore"):
dlogP = np.diff(np.log(P))
nz = np.logical_and(dP < 0, P[1:] * P[:-1] > 0)
const = dP == 0
return np.sum(dt[const] * (P[:-1])[const]) + np.sum(
dt[nz] * dP[nz] / dlogP[nz]
)
if steps is None:
last_N = max(self.population_size_history[:, self.num_epochs - 1])
last_epoch = self.epoch_start_time[-1]
steps = sorted(
list(
set(np.linspace(0, last_epoch + 12 * last_N, 101)).union(
set(self.epoch_start_time)
)
)
)
p_diff = m_diff = np.inf
last_P = np.inf
step_type = "none"
n = 0
logger.debug(
"iter mean P_diff mean_diff last_P adjust_type "
"num_steps last_step"
)
# The factors of 20 here are probably not optimal: clearly, we need to
# compute P accurately, but there's no good reason for this stopping rule.
# If populations have piecewise constant size then we shouldn't need this:
# setting steps equal to the epoch boundaries should suffice; while if
# there is very fast exponential change in some epochs caution is needed.
while n < max_iter and (
last_P > rtol or p_diff > rtol / 20 or m_diff > rtol / 20
):
last_steps = steps
if n == 0:
_, P1 = self.coalescence_rate_trajectory(
steps=last_steps,
lineages=lineages,
min_pop_size=min_pop_size,
double_step_validation=False,
)
m1 = mean_time(last_steps, P1)
if last_P > rtol:
step_type = "extend"
steps = np.concatenate(
[steps, np.linspace(steps[-1], steps[-1] * 1.2, 20)[1:]]
)
else:
step_type = "refine"
inter = steps[:-1] + np.diff(steps) / 2
steps = np.concatenate([steps, inter])
steps.sort()
_, P2 = self.coalescence_rate_trajectory(
steps=steps,
lineages=lineages,
min_pop_size=min_pop_size,
double_step_validation=False,
)
m2 = mean_time(steps, P2)
keep_steps = np.in1d(steps, last_steps)
p_diff = max(np.abs(P1 - P2[keep_steps]))
m_diff = np.abs(m1 - m2) / m2
last_P = P2[-1]
n += 1
# Use the old-style string formatting as this is the logging default
logger.debug(
"%d %g %g %g %g %s %d %d",
n,
m2,
p_diff,
m_diff,
last_P,
step_type,
len(steps),
max(steps),
)
P1 = P2
m1 = m2
if n == max_iter:
raise ValueError(
"Did not converge on an adequate discretisation: "
"Increase max_iter or rtol. Consult the log for "
"debugging information"
)
return m2
def coalescence_rate_trajectory(
self, steps, lineages, min_pop_size=1, double_step_validation=True
):
"""
Calculate the mean coalescence rates and proportions
of uncoalesced lineages between the specified sample lineages,
at each of the times ago listed by steps, in this demographic model.
Sample lineages are specified as a mapping from
populations to the number of **monoploid** sample genomes present in
that population at time zero. See the
:ref:`sec_demography_numerical_trajectories` section for usage examples
and more details.
The coalescence rate at time t in the past is the average rate of coalescence of
as-yet-uncoalesced lineages, computed as follows: let :math:`p(t)` be
the probability that the lineages of a randomly chosen pair of samples
has not yet coalesced by time :math:`t`, let :math:`p(z,t)` be the
probability that the lineages of a randomly chosen pair of samples has
not yet coalesced by time :math:`t` *and* are both in population
:math:`z`, and let :math:`N(z,t)` be the diploid effective population
size of population :math:`z` at time :math:`t`. Then the mean
coalescence rate at time :math:`t` is :math:`r(t) = (\\sum_z p(z,t) /
(2 * N(z,t))) / p(t)`.
The computation is done by approximating population size trajectories
with piecewise constant trajectories between each of the steps. For
this to be accurate, the distance between the steps must be small
enough so that (a) short epochs (e.g., bottlenecks) are not missed, and
(b) populations do not change in size too much over that time, if they
are growing or shrinking. This function optionally provides a simple
check of this approximation by recomputing the coalescence rates on a
grid of steps twice as fine and throwing a warning if the resulting
values do not match to a relative tolerance of 0.001.
:param list steps: The times ago at which coalescence rates will be computed.
:param dict lineages: A mapping of populations (either integer IDs
or string names: see the :ref:`sec_demography_populations_identifiers`
section for more details) to the number of monoploid sample lineages
in that population.
:param int min_pop_size: The smallest allowed population size during
computation of coalescent rates (i.e., coalescence rates are actually
1 / (2 * max(min_pop_size, N(z,t))). Spurious very small population sizes
can occur in models where populations grow exponentially but are unused
before some time in the past, and lead to floating point error.
This should be set to a value smaller than the smallest
desired population size in the model.
:param bool double_step_validation: Whether to perform the check that
step sizes are sufficiently small, as described above. This is highly
recommended, and will take at most four times the computation.
:return: A tuple of arrays whose jth elements, respectively, are the
coalescence rate at the jth time point (denoted r(t[j]) above),
and the probability that a randomly chosen pair of lineages has
not yet coalesced (denoted p(t[j]) above).
:rtype: (numpy.ndarray, numpy.ndarray)
"""
num_samples = np.zeros(self.num_populations, dtype=int)
for population, num_genomes in lineages.items():
pop_id = self.demography[population].id
num_samples[pop_id] += num_genomes
steps = np.array(steps)
if not np.all(np.diff(steps) > 0):
raise ValueError("`steps` must be a sequence of increasing times.")
if np.any(steps < 0):
raise ValueError("`steps` must be non-negative")
r, p_t = self._calculate_coalescence_rate_trajectory(
steps=steps, num_samples=num_samples, min_pop_size=min_pop_size
)
if double_step_validation:
inter = steps[:-1] + np.diff(steps) / 2
double_steps = np.concatenate([steps, inter])
double_steps.sort()
rd, p_td = self._calculate_coalescence_rate_trajectory(
steps=double_steps, num_samples=num_samples, min_pop_size=min_pop_size
)
assert np.all(steps == double_steps[::2])
r_prediction_close = np.allclose(r, rd[::2], rtol=1e-3, equal_nan=True)
p_prediction_close = np.allclose(p_t, p_td[::2], rtol=1e-3, equal_nan=True)
if not (r_prediction_close and p_prediction_close):
warnings.warn(
"Doubling the number of steps has resulted in different "
" predictions, please re-run with smaller step sizes to ensure "
" numerical accuracy."
)
return r, p_t
def _calculate_coalescence_rate_trajectory(self, steps, num_samples, min_pop_size):
num_pops = self.num_populations
P = np.zeros([num_pops**2, num_pops**2])
IA = np.array(range(num_pops**2)).reshape([num_pops, num_pops])
Identity = np.eye(num_pops)
for x in range(num_pops):
for y in range(num_pops):
P[IA[x, y], IA[x, y]] = num_samples[x] * (num_samples[y] - (x == y))
P = P / np.sum(P)
# add epoch breaks if not there already but remember which steps they are
epoch_breaks = list(
set([0.0] + [t for t in self.epoch_start_time if t not in steps])
)
steps_b = np.concatenate([steps, epoch_breaks])
ix = np.argsort(steps_b)
steps_b = steps_b[ix]
keep_steps = np.concatenate(
[np.repeat(True, len(steps)), np.repeat(False, len(epoch_breaks))]
)[ix]
assert np.all(steps == steps_b[keep_steps])
mass_migration_objects = []
mass_migration_times = []
for demo in self.demography.events:
if isinstance(demo, LineageMovementEvent):
# Convert higher-level lineage movement events like Admixtures
# and PopulationSplits into LineageMovement instances. These are
# equivalent to MassMigrations
for lm in demo._as_lineage_movements():
mass_migration_objects.append(lm)
mass_migration_times.append(demo.time)
num_steps = len(steps_b)
# recall that steps_b[0] = 0.0
r = np.zeros(num_steps)
p_t = np.zeros(num_steps)
for j in range(num_steps - 1):
time = steps_b[j]
dt = steps_b[j + 1] - steps_b[j]
N, M = self._pop_size_and_migration_at_t(time)
C = np.zeros([num_pops**2, num_pops**2])
for idx in range(num_pops):
C[IA[idx, idx], IA[idx, idx]] = 1 / (2 * max(min_pop_size, N[idx]))
dM = np.diag([sum(s) for s in M])
for mmt, mmo in zip(mass_migration_times, mass_migration_objects):
if mmt == time:
a = mmo.source
b = mmo.dest
p = mmo.proportion
S = np.eye(num_pops**2, num_pops**2)
for x in range(num_pops):
if x == a:
S[IA[a, a], IA[a, b]] = S[IA[a, a], IA[b, a]] = p * (1 - p)
S[IA[a, a], IA[b, b]] = p**2
S[IA[a, a], IA[a, a]] = (1 - p) ** 2
else:
S[IA[x, a], IA[x, b]] = S[IA[a, x], IA[b, x]] = p
S[IA[x, a], IA[x, a]] = S[IA[a, x], IA[a, x]] = 1 - p
P = np.matmul(P, S)
p_notcoal = np.sum(P)
p_t[j] = p_notcoal
if p_notcoal > 0:
r[j] = np.sum(np.matmul(P, C)) / np.sum(P)
else:
r[j] = np.nan
G = (np.kron(M - dM, Identity) + np.kron(Identity, M - dM)) - C
P = np.matmul(P, _matrix_exponential(dt * G))
p_notcoal = np.sum(P)
p_t[num_steps - 1] = p_notcoal
if p_notcoal > 0:
r[num_steps - 1] = np.sum(np.matmul(P, C)) / p_notcoal
else:
r[num_steps - 1] = np.nan
return r[keep_steps], p_t[keep_steps]
def _pop_size_and_migration_at_t(self, t):
"""
Returns a tuple (N, M) of population sizes (N) and migration rates (M) at
time t ago.
Note: this isn't part of the external API as it is be better to provide
separate methods to access the population size and migration rates, and
needing both together is specialised for internal calculations.
:param float t: The time ago.
:return: A tuple of arrays, of the same form as the population sizes and
migration rate arrays of the demographic model.
"""
j = 0
while self.epochs[j].end_time <= t:
j += 1
N = self.population_size_history[:, j]
for i, pop in enumerate(self.epochs[j].populations):
s = t - self.epochs[j].start_time
g = pop.growth_rate
N[i] *= np.exp(-1 * g * s)
return N, self.epochs[j].migration_matrix
Computing file changes ...
