Source code for demesdraw.size_history

from __future__ import annotations

import demes
import numpy as np
import matplotlib

from demesdraw import utils


[docs]def size_history( graph: demes.Graph, ax: matplotlib.axes.Axes | None = None, colours: utils.ColourOrColourMapping | None = None, log_time: bool = False, log_size: bool = False, title: str | None = None, inf_ratio: float = 0.1, inf_label: bool = False, invert_x: bool = False, annotate_epochs: bool = False, num_points: int = 100, ) -> matplotlib.axes.Axes: """ Plot population size as a function of time for each deme in the graph. :param demes.Graph graph: The demes graph to plot. :param ax: The matplotlib axes onto which the figure will be drawn. If None, an empty axes will be created for the figure. :type ax: Optional[matplotlib.axes.Axes] :param colours: A mapping from deme name to matplotlib colour. Alternately, ``colours`` may be a named colour that will be used for all demes. :type colours: Optional[dict or str] :param log_time: If True, use a log-10 scale for the time axis. If False (*default*), a linear scale will be used. :param bool log_size: If True, use a log-10 scale for the size axis. If False (*default*), a linear scale will be used. :param title: The title of the figure. :param inf_ratio: The proportion of the time axis that will be used for the time interval which stretches towards infinity. :param inf_label: Write "inf" by the arrow that points towards infinity. :param invert_x: If True, the horizontal axis will have infinity on the left and zero on the right, and the vertical axis will be drawn on the right. If False (*default*), the horizontal axis will have zero on the left and infinity on the right, and the vertical axis will be drawn on the left. :param annotate_epochs: If True, annotate the figure with epoch indices over the relevant parts of the lines. This may be useful as a pedagogical tool. If False (*default*), do not annotate the epochs. :return: The matplotlib axes onto which the figure was drawn. """ if ax is None: _, ax = utils.get_fig_axes() if invert_x: arrowhead = "<" else: arrowhead = ">" colours = utils._get_colours(graph, colours) inf_start_time = utils._inf_start_time(graph, inf_ratio, log_time) linestyles = ["solid"] # , "dashed", "dashdot"] linewidths = [2, 4, 8, 1] legend_handles = [] # Top of the z order stacking. z_top = 1 + len(graph.demes) + max(linewidths) for j, deme in enumerate(graph.demes): colour = colours[deme.name] linestyle = linestyles[j % len(linestyles)] linewidth = linewidths[j % len(linewidths)] plot_kwargs = dict( color=colour, linestyle=linestyle, linewidth=linewidth, label=deme.name, alpha=0.7, zorder=z_top - linewidth, capstyle="butt", fill=False, path_effects=[ matplotlib.patheffects.withStroke(linewidth=3, foreground="white") ], ) discontinuity_kwargs = plot_kwargs.copy() discontinuity_kwargs.update(linestyle=":") legend_kwargs = plot_kwargs.copy() legend_kwargs.pop("fill") # Line2D and Patch use different keywords for the capstyle. *Sigh* legend_kwargs.update(solid_capstyle=legend_kwargs.pop("capstyle")) legend_handles.append(matplotlib.lines.Line2D([], [], **legend_kwargs)) # Path for the main line (solid). vertices_main = [] codes_main = [] # Path for the discontinuity lines (dashed). vertices_discontinuity = [] codes_discontinuity = [] for k, epoch in enumerate(deme.epochs): start_time = epoch.start_time if np.isinf(start_time): start_time = inf_start_time end_time = epoch.end_time if log_time: end_time = max(1, end_time) if epoch.size_function == "constant": x = np.array([start_time, end_time]) y = np.array([epoch.start_size, epoch.end_size]) elif epoch.size_function == "exponential": x = np.linspace(start_time, end_time, num=num_points) dt = np.linspace(0, 1, num=num_points) r = np.log(epoch.end_size / epoch.start_size) y = epoch.start_size * np.exp(r * dt) elif epoch.size_function == "linear": x = np.linspace(start_time, end_time, num=num_points) dt = np.linspace(0, 1, num=num_points) y = epoch.start_size + (epoch.end_size - epoch.start_size) * dt else: raise ValueError( f"Don't know how to plot epoch {k} with " f'"{epoch.size_function}" size_function.' ) vertices_main.extend(list(zip(x, y))) if k == 0 or deme.epochs[k - 1].end_size != epoch.start_size: codes_main.append(matplotlib.path.Path.MOVETO) else: codes_main.append(matplotlib.path.Path.LINETO) codes_main.extend([matplotlib.path.Path.LINETO] * (len(x) - 1)) if k > 0 and deme.epochs[k - 1].end_size != epoch.start_size: # Size discontinuity. vertices_discontinuity.extend( [ (deme.epochs[k - 1].end_time, deme.epochs[k - 1].end_size), (epoch.start_time, epoch.start_size), ] ) codes_discontinuity.extend( [matplotlib.path.Path.MOVETO, matplotlib.path.Path.LINETO] ) if annotate_epochs: if log_time: text_x = np.exp((np.log(start_time) + np.log(end_time)) / 2) else: text_x = (start_time + end_time) / 2 if log_size: text_y = np.exp( (np.log(epoch.start_size) + np.log(max(1, epoch.end_size))) / 2 ) else: text_y = (epoch.start_size + epoch.end_size) / 2 ax.annotate( f"epoch {k}", (text_x, text_y), ha="center", va="bottom", xytext=(0, 4 + linewidth / 2), # vertical offset textcoords="offset points", # Give the text some contrast with its background. bbox=dict( boxstyle="round", fc="white", ec="none", alpha=0.6, pad=0 ), # This is only really a useful feature with 1 deme, # but at least try to do something reasonable for more demes. color="black" if len(graph.demes) == 1 else colour, zorder=z_top, ) # Indicate population size discontinuities from ancestor demes. for ancestor_id in deme.ancestors: anc = graph[ancestor_id] anc_N = anc.size_at(deme.start_time) deme_N = deme.epochs[0].start_size if anc_N != deme_N: vertices_discontinuity.extend( [(deme.start_time, anc_N), (deme.start_time, deme_N)] ) codes_discontinuity.extend( [matplotlib.path.Path.MOVETO, matplotlib.path.Path.LINETO] ) size_path_patch = matplotlib.patches.PathPatch( matplotlib.path.Path(vertices_main, codes_main), **plot_kwargs ) ax.add_patch(size_path_patch) if len(vertices_discontinuity) > 0: discontinuity_path_patch = matplotlib.patches.PathPatch( matplotlib.path.Path(vertices_discontinuity, codes_discontinuity), **discontinuity_kwargs, ) ax.add_patch(discontinuity_path_patch) if np.isinf(deme.start_time): # Plot an arrow at the end of the line, to indicate this # line extends towards infinity. ax.plot( inf_start_time, deme.epochs[0].start_size, arrowhead, color=colour, clip_on=False, zorder=z_top, ) if inf_label: ax.annotate( "inf", (inf_start_time, deme.epochs[0].start_size), xytext=(0, -6), # vertical offset textcoords="offset points", clip_on=False, ha="center", va="top", ) # Update the axes view. ax.add_patch() doesn't do this itself. ax.autoscale_view() if len(graph.demes) > 1: leg = ax.legend(handles=legend_handles, ncol=max(1, len(graph.demes) // 2)) leg.set_zorder(z_top) if title is not None: ax.set_title(title) # Arrange the axes spines, ticks and labels. ax.set_xlim(1 if log_time else 0, inf_start_time) if not log_size: ax.set_ylim(bottom=0) for spine in ax.spines.values(): spine.set_zorder(z_top) ax.spines["top"].set_visible(False) if invert_x: ax.spines["left"].set_visible(False) ax.invert_xaxis() ax.yaxis.tick_right() ax.yaxis.set_label_position("right") else: ax.spines["right"].set_visible(False) ax.set_xlabel(f"time ago ({graph.time_units})") # ax.set_ylabel("N", rotation=0, ha="left" if invert_x else "right") ax.set_ylabel("deme\nsize", rotation=0, labelpad=20) if log_time: ax.set_xscale("log", base=10) if log_size: ax.set_yscale("log", base=10) return ax