From 73ca6d8bdaabfbd0b3a09d59ebe2a8a0d466ca23 Mon Sep 17 00:00:00 2001 From: Jerome Kelleher Date: Fri, 19 Feb 2021 14:32:35 +0000 Subject: [PATCH] Minimal updates for msprime 1.0 This is a transitional update, which gets most features working with msprime 1.0, without starting the process of modernising the catalog models. Closes #748 Closes #740 Closes #535 --- stdpopsim/catalog/BosTau/__init__.py | 2 +- stdpopsim/cli.py | 6 +- stdpopsim/engines.py | 28 +- stdpopsim/genetic_maps.py | 19 +- stdpopsim/genomes.py | 10 +- stdpopsim/models.py | 461 +++++++------------------ stdpopsim/qc/AraTha.py | 14 +- stdpopsim/qc/DroMel.py | 22 +- stdpopsim/qc/HomSap.py | 97 +++++- stdpopsim/qc/PonAbe.py | 14 +- stdpopsim/slim_engine.py | 139 ++++---- stdpopsim/species.py | 6 +- tests/test_HomSap.py | 4 +- tests/test_cli.py | 6 - tests/test_genetic_maps.py | 14 +- tests/test_masking.py | 4 +- tests/test_models.py | 492 +++++---------------------- tests/test_slim_engine.py | 5 +- tests/test_species.py | 11 +- 19 files changed, 450 insertions(+), 904 deletions(-) diff --git a/stdpopsim/catalog/BosTau/__init__.py b/stdpopsim/catalog/BosTau/__init__.py index 970ca335d..de039a4ed 100644 --- a/stdpopsim/catalog/BosTau/__init__.py +++ b/stdpopsim/catalog/BosTau/__init__.py @@ -115,7 +115,7 @@ def _HolsteinFriesan_1M13(): inspection of the plots. """ populations = [ - stdpopsim.Population(id="Holstein-Friesian", description="Holstein-Friesian"), + stdpopsim.Population(id="Holstein_Friesian", description="Holstein Friesian"), ] citations = [_MacLeodEtAl.because(stdpopsim.CiteReason.DEM_MODEL)] diff --git a/stdpopsim/cli.py b/stdpopsim/cli.py index 43a9df97e..5331e1203 100644 --- a/stdpopsim/cli.py +++ b/stdpopsim/cli.py @@ -577,7 +577,7 @@ def write_simulation_summary(engine, model, contig, samples, seed=None): dry_run_text += f"{indent}Population: number_samples (sampling_time_generations):\n" sample_counts = [0] * model.num_sampling_populations for s in samples: - sample_counts[s.population] += 1 + sample_counts[s.population] += s.num_samples for p in range(0, model.num_sampling_populations): pop_name = model.populations[p].id sample_time = model.populations[p].sampling_time @@ -585,9 +585,9 @@ def write_simulation_summary(engine, model, contig, samples, seed=None): dry_run_text += f"{sample_counts[p]} ({sample_time})\n" # Get information about relevant contig gmap = "None" if contig.genetic_map is None else contig.genetic_map.id - mean_recomb_rate = contig.recombination_map.mean_recombination_rate + mean_recomb_rate = contig.recombination_map.mean_rate mut_rate = contig.mutation_rate - contig_len = contig.recombination_map.get_length() + contig_len = contig.recombination_map.sequence_length dry_run_text += "Contig Description:\n" dry_run_text += f"{indent}Contig length: {contig_len}\n" dry_run_text += f"{indent}Mean recombination rate: {mean_recomb_rate}\n" diff --git a/stdpopsim/engines.py b/stdpopsim/engines.py index bf461a296..953d389fc 100644 --- a/stdpopsim/engines.py +++ b/stdpopsim/engines.py @@ -4,6 +4,7 @@ import attr import msprime import stdpopsim +import numpy as np logger = logging.getLogger(__name__) @@ -173,16 +174,14 @@ def simulate( if msprime_model in self.model_citations: self.citations.extend(self.model_citations[msprime_model]) - demographic_events = demographic_model.demographic_events.copy() if msprime_change_model is not None: + msprime_model = [msprime_model] for t, model in msprime_change_model: if model not in self.supported_models: raise ValueError(f"Unrecognised model '{model}'") - model_change = msprime.SimulationModelChange(t, model) - demographic_events.append(model_change) + msprime_model.append((t, model)) if model in self.model_citations: self.citations.extend(self.model_citations[model]) - demographic_events.sort(key=lambda x: x.time) if "random_seed" in kwargs.keys(): if seed is None: @@ -194,18 +193,25 @@ def simulate( # TODO: remove this after a release or two. See #745. self._warn_zigzag(demographic_model) - ts = msprime.simulate( + rng = np.random.default_rng(seed) + seeds = rng.integers(1, 2 ** 31 - 1, size=2) + + ts = msprime.sim_ancestry( samples=samples, - recombination_map=contig.recombination_map, - mutation_rate=contig.mutation_rate, - population_configurations=demographic_model.population_configurations, - migration_matrix=demographic_model.migration_matrix, - demographic_events=demographic_events, - random_seed=seed, + recombination_rate=contig.recombination_map, + demography=demographic_model.model, + ploidy=2, + random_seed=seeds[0], model=msprime_model, end_time=0 if dry_run else None, **kwargs, ) + ts = msprime.sim_mutations( + ts, + end_time=0 if dry_run else None, + random_seed=seeds[1], + rate=contig.mutation_rate, + ) if contig.inclusion_mask is not None: ts = stdpopsim.utils.mask_tree_sequence(ts, contig.inclusion_mask, False) diff --git a/stdpopsim/genetic_maps.py b/stdpopsim/genetic_maps.py index 92a507b0d..6c96487c3 100644 --- a/stdpopsim/genetic_maps.py +++ b/stdpopsim/genetic_maps.py @@ -4,6 +4,7 @@ import warnings import msprime +import numpy as np import stdpopsim @@ -88,8 +89,8 @@ def get_chromosome_map(self, id): :param str id: The chromosome identifier. A complete list of chromosome IDs for each species can be found in the "Genome" subsection for the species in the :ref:`sec_catalog`. - :rtype: :class:`msprime.RecombinationMap` - :return: A :class:`msprime.RecombinationMap` object. + :rtype: :class:`msprime.RateMap` + :return: A :class:`msprime.RateMap` object. """ chrom = self.species.genome.get_chromosome(id) if not self.is_cached(): @@ -100,22 +101,20 @@ def get_chromosome_map(self, id): # needs to be redownloaded. map_file = self.map_cache_dir / self.file_pattern.format(id=chrom.id) if map_file.exists(): - recomb_map = msprime.RecombinationMap.read_hapmap(str(map_file)) + recomb_map = msprime.RateMap.read_hapmap(map_file) else: warnings.warn( "Recombination map not found for chromosome: '{}'" " on map: '{}', substituting a flat map with chromosome " "recombination rate {}".format(id, self.id, chrom.recombination_rate) ) - recomb_map = msprime.RecombinationMap.uniform_map( - chrom.length, chrom.recombination_rate - ) - map_length = recomb_map.get_sequence_length() + recomb_map = msprime.RateMap.uniform(chrom.length, chrom.recombination_rate) + map_length = recomb_map.sequence_length if map_length < chrom.length: # Extend map to the end of the chromosome. - positions = recomb_map.get_positions() + [chrom.length] - rates = recomb_map.get_rates() + [0] - recomb_map = msprime.RecombinationMap(positions, rates) + positions = np.append(recomb_map.position, chrom.length) + rates = np.append(recomb_map.rate, 0) + recomb_map = msprime.RateMap(positions, rates) elif map_length > chrom.length: warnings.warn( f"Recombination map has length {map_length}, which is longer than" diff --git a/stdpopsim/genomes.py b/stdpopsim/genomes.py index 7ab9d27bd..15186a275 100644 --- a/stdpopsim/genomes.py +++ b/stdpopsim/genomes.py @@ -117,10 +117,8 @@ class Contig: :ivar mutation_rate: The rate of mutation per base per generation. :vartype mutation_rate: float :ivar recombination_map: The recombination map for the region. See the - `msprime documentation - `_ - for more details. - :vartype recombination_map: msprime.simulations.RecombinationMap + :class:`msprime.RateMap` for details. + :vartype recombination_map: msprime.RateMap :ivar mask_intervals: Intervals to keep in simulated tree sequence, as a list of (left_position, right_position), such that intervals are non-overlapping and in ascending order. Should have shape Nx2, where N is the number of @@ -157,8 +155,8 @@ def __str__(self): "Contig(length={:.2G}, recombination_rate={:.2G}, " "mutation_rate={:.2G}, genetic_map={})" ).format( - self.recombination_map.get_length(), - self.recombination_map.mean_recombination_rate, + self.recombination_map.sequence_length, + self.recombination_map.mean_rate, self.mutation_rate, gmap, ) diff --git a/stdpopsim/models.py b/stdpopsim/models.py index 81dabd786..dcd8d043d 100644 --- a/stdpopsim/models.py +++ b/stdpopsim/models.py @@ -1,162 +1,15 @@ """ Common infrastructure for specifying demographic models. """ -import sys - -import attr -import msprime -import numpy as np - - -# Defaults taken from np.allclose -DEFAULT_ATOL = 1e-05 -DEFAULT_RTOL = 1e-08 - - -class UnequalModelsError(Exception): - """ - Exception raised models by verify_equal to indicate that models are - not sufficiently close. - """ - - -def population_configurations_equal( - pop_configs1, pop_configs2, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL -): - """ - Returns True if the specified lists of msprime PopulationConfiguration - objects are equal to the specified tolerances. - - See the :func:`.verify_population_configurations_equal` function for - details on the assumptions made about the objects. - """ - try: - verify_population_configurations_equal( - pop_configs1, pop_configs2, rtol=rtol, atol=atol - ) - return True - except UnequalModelsError: - return False - - -def verify_population_configurations_equal( - pop_configs1, pop_configs2, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL -): - """ - Checks if the specified lists of msprime PopulationConfiguration - objects are equal to the specified tolerances and raises an UnequalModelsError - otherwise. - - We make some assumptions here to ensure that the models we specify - are well-defined: (1) The sample size is not set for PopulationConfigurations - (2) the initial_size is defined. If these assumptions are violated a - ValueError is raised. - """ - for pc1, pc2 in zip(pop_configs1, pop_configs2): - if pc1.sample_size is not None or pc2.sample_size is not None: - raise ValueError( - "Models defined in stdpopsim must not use the 'sample_size' " - "PopulationConfiguration option" - ) - if pc1.initial_size is None or pc2.initial_size is None: - raise ValueError("Models defined in stdpopsim must set the initial_size") - if len(pop_configs1) != len(pop_configs2): - raise UnequalModelsError("Different numbers of populations") - initial_size1 = np.array([pc.initial_size for pc in pop_configs1]) - initial_size2 = np.array([pc.initial_size for pc in pop_configs2]) - if not np.allclose(initial_size1, initial_size2, rtol=rtol, atol=atol): - raise UnequalModelsError("Initial sizes differ") - growth_rate1 = np.array([pc.growth_rate for pc in pop_configs1]) - growth_rate2 = np.array([pc.growth_rate for pc in pop_configs2]) - if not np.allclose(growth_rate1, growth_rate2, rtol=rtol, atol=atol): - raise UnequalModelsError("Growth rates differ") - - -def sampling_times_equal( - populations1, populations2, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL -): - """ - Returns True if the sampling times for the specified lists of Populations - objects are equal to the specified tolerances. - """ - try: - verify_sampling_times_equal(populations1, populations2, rtol=rtol, atol=atol) - return True - except UnequalModelsError: - return False - - -def verify_sampling_times_equal( - populations1, populations2, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL -): - """ - Check the the sampling times for lists of population objects are the same - """ - # Retrieve sampling times - sampling_times1 = np.array([p.sampling_time for p in populations1]) - sampling_times2 = np.array([p.sampling_time for p in populations2]) - # Check for equal vector length - if len(sampling_times1) != len(sampling_times2): - raise UnequalModelsError("Different numbers of populations") - # Get indicies where None values present and compare - s1_none = np.where(sampling_times1 == None) # noqa - s2_none = np.where(sampling_times2 == None) # noqa - if not np.array_equal(s1_none, s2_none): - raise UnequalModelsError("None-valued sampling times differ") - # Subset out only non-None values and cast to float so they can be compared - s1_subset = sampling_times1[sampling_times1 != np.array(None)].astype(float) - s2_subset = sampling_times2[sampling_times2 != np.array(None)].astype(float) - if not np.allclose(s1_subset, s2_subset, rtol=rtol, atol=atol): - raise UnequalModelsError("Positive real-valued sampling times differ") - - -def demographic_events_equal( - events1, events2, num_populations, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL -): - """ - Returns True if the specified list of msprime DemographicEvent objects are equal - to the specified tolerances. - """ - try: - verify_demographic_events_equal( - events1, events2, num_populations, rtol=rtol, atol=atol - ) - return True - except UnequalModelsError: - return False +import copy -def verify_demographic_events_equal( - events1, events2, num_populations, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL -): - """ - Checks if the specified list of msprime DemographicEvent objects are equal - to the specified tolerances and raises a UnequalModelsError otherwise. - """ - # Get the low-level dictionary representations of the events. - dicts1 = [event.get_ll_representation(num_populations) for event in events1] - dicts2 = [event.get_ll_representation(num_populations) for event in events2] - if len(dicts1) != len(dicts2): - raise UnequalModelsError("Different numbers of demographic events") - for d1, d2 in zip(dicts1, dicts2): - if set(d1.keys()) != set(d2.keys()): - raise UnequalModelsError("Different types of demographic events") - for key in d1.keys(): - value1 = d1[key] - value2 = d2[key] - if isinstance(value1, float): - if not np.isclose(value1, value2, rtol=rtol, atol=atol): - raise UnequalModelsError( - "Event {} mismatch: {} != {}".format(key, value1, value2) - ) - else: - if value1 != value2: - raise UnequalModelsError( - "Event {} mismatch: {} != {}".format(key, value1, value2) - ) +import msprime class Population: + # TODO deprecate this - we don't need any internal definition of what + # a population is. """ Class recording metadata representing a population in a simulation. @@ -191,7 +44,6 @@ def asdict(self): } -@attr.s(kw_only=True) class DemographicModel: """ Class representing a demographic model. @@ -241,107 +93,92 @@ class DemographicModel: :vartype migration_matrix: list of list of int """ - # required attributes - id = attr.ib(type=str) - description = attr.ib(type=str) - long_description = attr.ib(type=str) - generation_time = attr.ib(type=int) - - # optional attributes - citations = attr.ib(factory=list) - demographic_events = attr.ib(factory=list) - population_configurations = attr.ib(factory=list) - migration_matrix = attr.ib() - populations = attr.ib() - - qc_model = attr.ib(default=None) - - @populations.default - def _populations_default(self): - npops = len(self.population_configurations) - pops = [] - for i in range(npops): - # Try to get population constructor info from PopConfig metadata - kwargs = self.population_configurations[i].metadata - # This is here so we don't have to change msprime, may remove later - if kwargs is None: - kwargs = dict() - if "id" not in kwargs: - kwargs["id"] = f"pop{i}" - if "description" not in kwargs: - kwargs["description"] = f"Population {i}" - pops.append(Population(**kwargs)) - return pops - - @migration_matrix.default - def _migration_matrix_default(self): - npops = len(self.population_configurations) - return [[0 for j in range(npops)] for i in range(npops)] + def __init__( + self, + *, + id, + description, + long_description, + generation_time=None, + citations=None, + qc_model=None, + population_configurations=None, + demographic_events=None, + migration_matrix=None, + population_id_map=None, + populations=None, + model=None, + ): + self.id = id + self.description = description + self.long_description = long_description + self.generation_time = 1 if generation_time is None else generation_time + self.citations = [] if citations is None else citations + self.qc_model = qc_model + + if model is None: + assert population_configurations is not None + assert populations is not None + assert len(populations) == len(population_configurations) + population_configurations = copy.deepcopy(population_configurations) + # Merge the information from the populations into the msprime + # Demography. + for pop, pop_config in zip(populations, population_configurations): + if pop_config.metadata is None: + pop_config.metadata = {} + if pop.id != "" and not pop.id.startswith("qc_"): + pop_config.metadata["name"] = pop.id + pop_config.metadata["description"] = pop.description + # This will become a Demes model in the future - for now it's an + # msprime model. + self.model = msprime.Demography.from_old_style( + population_configurations=population_configurations, + demographic_events=demographic_events, + migration_matrix=migration_matrix, + population_map=population_id_map, + ) + for msp_pop, local_pop in zip(self.model.populations, populations): + # We use the "allow_samples" attribute in the CLI and else where + # so we monkey patch this into the msprime Populations for the + # moment. + msp_pop.allow_samples = True + if local_pop.sampling_time is not None: + msp_pop.sampling_time = local_pop.sampling_time + else: + msp_pop.allow_samples = False + else: + assert population_configurations is None + assert population_id_map is None + assert populations is None + self.model = copy.deepcopy(model) + # See note above. We allow samples in all populations for now. + for pop in self.model.populations: + pop.allow_samples = True + + def __str__(self): + # TODO make this nice to look at. + s = ( + "Demographic model:\n" + f"id = {self.id}\n" + f"description = {self.description}\n" + f"generation_time = {self.generation_time}\n" + f"citations = {[cite.doi for cite in self.citations]}\n" + ) + s += str(self.model) + return s + + @property + def populations(self): + return self.model.populations @property def num_populations(self): - return len(self.populations) + return self.model.num_populations @property def num_sampling_populations(self): return sum(int(pop.allow_samples) for pop in self.populations) - @staticmethod - def empty(**kwargs): - """ - Return a model with the mandatory attributes filled out. - """ - kwargs.update( - id=kwargs.get("id", ""), - description=kwargs.get("description", ""), - long_description=kwargs.get("long_description", ""), - populations=kwargs.get("populations", []), - generation_time=kwargs.get("generation_time", -1), - ) - return DemographicModel(**kwargs) - - def equals(self, other, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL): - """ - Returns True if this model is equal to the specified model to the - specified numerical tolerance (as defined by numpy.allclose). - - We use the 'equals' method here rather than the equality operator - because we need to be able to specifiy the numerical tolerances. - """ - try: - self.verify_equal(other, rtol=rtol, atol=atol) - return True - except (UnequalModelsError, AttributeError): - return False - - def verify_equal(self, other, rtol=DEFAULT_RTOL, atol=DEFAULT_ATOL): - """ - Equivalent to the :func:`.equals` method, but raises a UnequalModelsError if the - models are not equal rather than returning False. - """ - mm1 = np.array(self.migration_matrix) - mm2 = np.array(other.migration_matrix) - if mm1.shape != mm2.shape: - raise UnequalModelsError("Migration matrices different shapes") - if not np.allclose(mm1, mm2, rtol=rtol, atol=atol): - raise UnequalModelsError("Migration matrices differ") - verify_population_configurations_equal( - self.population_configurations, - other.population_configurations, - rtol=rtol, - atol=atol, - ) - verify_demographic_events_equal( - self.demographic_events, - other.demographic_events, - len(self.population_configurations), - rtol=rtol, - atol=atol, - ) - verify_sampling_times_equal( - self.populations, other.populations, rtol=rtol, atol=atol - ) - def register_qc(self, qc_model): """ Register a QC model implementation for this model. @@ -352,13 +189,7 @@ def register_qc(self, qc_model): ) if self.qc_model is not None: raise ValueError(f"QC model already registered for {self.id}.") - self.qc_model = qc_model - - def debug(self, out_file=sys.stdout): - # Use the demography debugger to print out the demographic history - # that we have just described. - dd = self.get_demography_debugger() - dd.print_history(out_file) + self.qc_model = qc_model.model def get_samples(self, *args): """ @@ -372,44 +203,27 @@ def get_samples(self, *args): "sampling" populations, ``model.num_sampling_populations``; if the number of arguments is less than the number of sampling populations, then remaining numbers are treated as zero. + + .. todo:: This documentation is broken. We're now returning msprime + SampleSet objects. """ samples = [] for pop_index, n in enumerate(args): if self.populations[pop_index].allow_samples: - sample = msprime.Sample( - pop_index, time=self.populations[pop_index].sampling_time + samples.append( + msprime.SampleSet( + num_samples=n, + population=pop_index, + time=self.populations[pop_index].sampling_time, + ploidy=1, # Avoid breaking too much at once. + ) ) - samples.extend([sample] * n) elif n > 0: raise ValueError( "Samples requested from non-sampling population {pop_index}" ) return samples - def get_demography_debugger(self): - """ - Returns an :class:`msprime.DemographyDebugger` instance initialized - with the parameters for this model. Please see the msprime documentation - for details on how to use a DemographyDebugger. - - :return: A DemographyDebugger instance for this DemographicModel. - :rtype: msprime.DemographyDebugger - """ - ddb = msprime.DemographyDebugger( - population_configurations=self.population_configurations, - migration_matrix=self.migration_matrix, - demographic_events=self.demographic_events, - ) - return ddb - - -# Reusable generic populations -_pop0 = Population(id="pop0", description="Generic population") -_pop1 = Population(id="pop1", description="Generic population") -_popAnc = Population( - id="popAnc", description="Generic ancestral population", sampling_time=None -) - class PiecewiseConstantSize(DemographicModel): """ @@ -429,29 +243,18 @@ class PiecewiseConstantSize(DemographicModel): model2 = stdpopsim.PiecewiseConstantSize(N0, (t1, N1), (t2, N2)) # Two changes """ - id = "PiecewiseConstant" - description = "Piecewise constant size population model over multiple epochs." - citations = [] - populations = [_pop0] - author = None - year = None - doi = None - generation_time = 1 - def __init__(self, N0, *args): - self.population_configurations = [ - msprime.PopulationConfiguration( - initial_size=N0, metadata=self.populations[0].asdict() - ) - ] - self.migration_matrix = [[0]] - self.demographic_events = [] + model = msprime.Demography.isolated_model(initial_size=[N0]) for t, N in args: - self.demographic_events.append( - msprime.PopulationParametersChange( - time=t, initial_size=N, growth_rate=0, population_id=0 - ) - ) + model.add_population_parameters_change(time=t, initial_size=N) + + super().__init__( + id="PiecewiseConstant", + description="Piecewise constant size population model over multiple epochs.", + long_description="", + model=model, + generation_time=1, + ) class IsolationWithMigration(DemographicModel): @@ -478,34 +281,32 @@ class IsolationWithMigration(DemographicModel): """ - id = "IsolationWithMigration" - description = """ - A generic isolation with migration model where a single ancestral - population of size NA splits into two populations of constant size N1 - and N2 time T generations ago, with migration rates M12 and M21 between - the split populations. - """ - citations = [] - populations = [_pop0, _pop1, _popAnc] - author = None - year = None - doi = None - generation_time = 1 - def __init__(self, NA, N1, N2, T, M12, M21): - self.population_configurations = [ - msprime.PopulationConfiguration( - initial_size=N1, metadata=self.populations[0].asdict() - ), - msprime.PopulationConfiguration( - initial_size=N2, metadata=self.populations[1].asdict() - ), - msprime.PopulationConfiguration( - initial_size=NA, metadata=self.populations[2].asdict() - ), - ] - self.migration_matrix = [[0, M12, 0], [M21, 0, 0], [0, 0, 0]] - self.demographic_events = [ - msprime.MassMigration(time=T, source=0, destination=2, proportion=1), - msprime.MassMigration(time=T, source=1, destination=2, proportion=1), - ] + model = msprime.Demography() + model.add_population(initial_size=N1, name="pop1") + model.add_population(initial_size=N2, name="pop2") + model.add_population(initial_size=NA, name="ancestral") + + # FIXME This is BACKWARDS in time, so the rates are the other + # way around forwards time. We should explain this in the documentation + # (and probably swap around). Seems like there's not really much + # good reason to have this model in here any more though - what + # does it do that wouldn't be better done in demes/msprime? + model.set_migration_rate(source="pop1", dest="pop2", rate=M12) + model.set_migration_rate(source="pop2", dest="pop1", rate=M21) + model.add_population_split( + time=T, ancestral="ancestral", derived=["pop1", "pop2"] + ) + long_description = """ + A generic isolation with migration model where a single ancestral + population of size NA splits into two populations of constant size N1 + and N2 time T generations ago, with migration rates M12 and M21 between + the split populations. + """ + super().__init__( + id="IsolationWithMigration", + description="Generic IM model", + long_description=long_description, + model=model, + generation_time=1, + ) diff --git a/stdpopsim/qc/AraTha.py b/stdpopsim/qc/AraTha.py index c592190fb..6b3a3af9e 100644 --- a/stdpopsim/qc/AraTha.py +++ b/stdpopsim/qc/AraTha.py @@ -6,15 +6,10 @@ _species = stdpopsim.get_species("AraTha") -# Some generic populations to use for qc -population_sample_0 = stdpopsim.Population( - "sampling_0", "Population that samples at time 0", 0 -) - def Durvasula2017MSMC(): id = "QC-SouthMiddleAtlas_1D17" - populations = [population_sample_0] + populations = [stdpopsim.Population("SouthMiddleAtlas", "")] # Both of the following are directly # converted from MSMC output scaled by A.Thaliana @@ -126,6 +121,7 @@ def Durvasula2017MSMC(): populations=populations, population_configurations=population_configurations, demographic_events=demographic_events, + population_id_map=[{"SouthMiddleAtlas": 0}] * 33, ) @@ -135,7 +131,7 @@ def Durvasula2017MSMC(): def HuberTwoEpoch(): id = "QC-African2Epoch_1H18" populations = [ - stdpopsim.Population(id="ATL", description="A. thalina"), + stdpopsim.Population(id="SouthMiddleAtlas", description="A. thalina"), ] # Time of second epoch @@ -160,6 +156,7 @@ def HuberTwoEpoch(): time=T_2, initial_size=N_ANC, population_id=0 ), ], + population_id_map=[{"SouthMiddleAtlas": 0}] * 2, ) @@ -169,7 +166,7 @@ def HuberTwoEpoch(): def HuberThreeEpoch(): id = "QC-African3Epoch_1H18" populations = [ - stdpopsim.Population(id="ATL", description="A. thalina"), + stdpopsim.Population(id="SouthMiddleAtlas", description="A. thalina"), ] # Time of second epoch @@ -199,6 +196,7 @@ def HuberThreeEpoch(): time=T_2 + T_3, initial_size=N_ANC, population_id=0 ), ], + population_id_map=[{"SouthMiddleAtlas": 0}] * 3, ) diff --git a/stdpopsim/qc/DroMel.py b/stdpopsim/qc/DroMel.py index 66773afcf..e7df25573 100644 --- a/stdpopsim/qc/DroMel.py +++ b/stdpopsim/qc/DroMel.py @@ -5,18 +5,13 @@ _species = stdpopsim.get_species("DroMel") -# Some generic populations to use for qc -population_sample_0 = stdpopsim.Population( - "sampling_0", "Population that samples at time 0", 0 -) -population_sample_none = stdpopsim.Population( - "sampling_none", "Population that does not sample", None -) - def LiStephanTwoPopulation(): id = "QC-OutOfAfrica_2L06" - populations = [population_sample_0] * 2 + populations = [ + stdpopsim.Population("AFR", ""), + stdpopsim.Population("EUR", ""), + ] # Parameters for the African population are taken from the section Demographic # History of the African Population @@ -59,6 +54,12 @@ def LiStephanTwoPopulation(): time=T_A0, initial_size=N_A1, population_id=0 ), ], + population_id_map=[ + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + ], ) @@ -67,7 +68,7 @@ def LiStephanTwoPopulation(): def SheehanSongThreeEpic(): id = "QC-African3Epoch_1S16" - populations = [population_sample_0] + populations = [stdpopsim.Population("AFR", "")] # Model from paper https://doi.org/10.1371/journal.pcbi.1004845 @@ -107,6 +108,7 @@ def SheehanSongThreeEpic(): time=T_2, initial_size=N_3, population_id=0 ), ], + population_id_map=[{"AFR": 0}] * 3, ) diff --git a/stdpopsim/qc/HomSap.py b/stdpopsim/qc/HomSap.py index 8cb49c568..cfb834b7c 100644 --- a/stdpopsim/qc/HomSap.py +++ b/stdpopsim/qc/HomSap.py @@ -9,10 +9,10 @@ # Some generic populations to use for qc population_sample_0 = stdpopsim.Population( - "sampling_0", "Population that samples at time 0", 0 + "qc_sampling_0", "Population that samples at time 0", 0 ) population_sample_none = stdpopsim.Population( - "sampling_none", "Population that does not sample", None + "qc_sampling_none", "Population that does not sample", None ) @@ -73,6 +73,13 @@ def TennessenOnePopAfrica(): time=T_AF, initial_size=N_A, population_id=0 ), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + population_id_map=[ + {"AFR": 0}, + {"AFR": 0}, + {"AFR": 0}, + ], ) @@ -177,6 +184,15 @@ def TennessenTwoPopOutOfAfrica(): time=T_AF, initial_size=N_A, population_id=0 ), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + population_id_map=[ + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + {"AFR": 0, "EUR": 1}, + ], ) @@ -292,6 +308,15 @@ def BrowningAmerica(): time=T_AF0_AF1, initial_size=N_AF0, population_id=0 ), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + population_id_map=[ + {"AFR": 0, "EUR": 1, "ASIA": 2, "ADMIX": 3}, + {"AFR": 0, "EUR": 1, "ASIA": 2, "ADMIX": 3}, + {"AFR": 0, "EUR": 1, "ASIA": 2, "ADMIX": 3}, + {"AFR": 0, "EUR": 1, "ASIA": 2, "ADMIX": 3}, + {"AFR": 0, "EUR": 1, "ASIA": 2, "ADMIX": 3}, + ], ) @@ -436,6 +461,18 @@ def RagsdaleArchaic(): # NEAN pop. coalesces into AF (E7) msprime.MassMigration(time=T_NEAN_split, source=3, dest=0, proportion=1.0), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + population_id_map=[ + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + {"YRI": 0, "CEU": 1, "CHB": 2, "Neanderthal": 3, "ArchaicAFR": 4}, + ], ) @@ -517,6 +554,18 @@ def KammAncientSamples(): # at the time of sampling the Altai individual r_Nean = -math.log(N_Nean_Losch / N_Nean) / (t_Mbu_Losch - t_Altai) + pop_id_map = { + "Mbuti": 0, + "LBK": 1, + "Sardinian": 2, + "Loschbour": 3, + "MA1": 4, + "Han": 5, + "UstIshim": 6, + "Neanderthal": 7, + "BasalEurasian": 8, + } + return stdpopsim.DemographicModel( id=id, description=id, @@ -563,12 +612,12 @@ def KammAncientSamples(): msprime.PopulationConfiguration( # Neanderthal initial_size=N_Nean, growth_rate=0, - metadata={"name": "Altai", "sampling_time": t_Altai}, + metadata={"name": "Neanderthal", "sampling_time": t_Altai}, ), msprime.PopulationConfiguration( # Basal Eurasian initial_size=N_Basal, growth_rate=0, - metadata={"name": "Basal", "sampling_time": None}, + metadata={"name": "BasalEurasian", "sampling_time": None}, ), ], # Using columns in figure in Kamm paper as proxies for pop number @@ -631,6 +680,10 @@ def KammAncientSamples(): time=t_Nean_Losch, initial_size=N_Nean_Losch, population_id=3 ), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + # Note: we'll be updating this map as the model is modernised. + population_id_map=[pop_id_map] * 13, ) @@ -747,7 +800,7 @@ def DenisovanAncestryInPapuans(): msprime.PopulationConfiguration( # 5 NeanA initial_size=N_NeanA, growth_rate=0, - metadata={"name": "NeanA", "sampling_time": t_NeanA}, + metadata={"name": "NeaA", "sampling_time": t_NeanA}, ), msprime.PopulationConfiguration( # 6 Den1 initial_size=N_Den1, @@ -762,7 +815,7 @@ def DenisovanAncestryInPapuans(): msprime.PopulationConfiguration( # 8 Nean1 initial_size=N_Nean1, growth_rate=0, - metadata={"name": "Nean1", "sampling_time": None}, + metadata={"name": "Nea1", "sampling_time": None}, ), msprime.PopulationConfiguration( # 9 Ghost initial_size=N_Ghost, @@ -911,6 +964,19 @@ def DenisovanAncestryInPapuans(): ), ] demographic_events.sort(key=lambda x: x.time) + + pop_id_map = { + "YRI": 0, + "CEU": 1, + "CHB": 2, + "Papuan": 3, + "DenA": 4, + "NeaA": 5, + "Den1": 6, + "Den2": 7, + "Nea1": 8, + "Ghost": 9, + } return stdpopsim.DemographicModel( id=id, description=id, @@ -920,6 +986,10 @@ def DenisovanAncestryInPapuans(): population_configurations=population_configurations, migration_matrix=migration_matrix, demographic_events=demographic_events, + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + # Note: we'll be updating this map as the model is modernised. + population_id_map=[pop_id_map] * 19, ) @@ -998,6 +1068,14 @@ def GutenkunstOOA(): time=T_AF, initial_size=N_A, population_id=0 ), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + population_id_map=[ + {"YRI": 0, "CEU": 1, "CHB": 2}, + {"YRI": 0, "CEU": 1, "CHB": 2}, + {"YRI": 0, "CEU": 1, "CHB": 2}, + {"YRI": 0, "CEU": 1, "CHB": 2}, + ], ) @@ -1084,6 +1162,12 @@ def GladsteinAshkSubstructure(): time=T_growth, initial_size=N_ANC, population_id=0 ), ], + # Mapping of per-epoch population names in the production model to + # population IDs in this model. + population_id_map=[ + {"YRI": 0, "CHB": 1, "CEU": 2, "ME": 3, "J": 4, "WAJ": 5, "EAJ": 6}, + ] + * 10, ) @@ -1146,6 +1230,7 @@ def ZigZag(): populations=populations, population_configurations=population_configurations, demographic_events=de, + population_id_map=[{"generic": 0}] * 7, ) diff --git a/stdpopsim/qc/PonAbe.py b/stdpopsim/qc/PonAbe.py index 327579393..986969ca1 100644 --- a/stdpopsim/qc/PonAbe.py +++ b/stdpopsim/qc/PonAbe.py @@ -6,15 +6,13 @@ _species = stdpopsim.get_species("PonAbe") -# Some generic populations to use for qc -population_sample_0 = stdpopsim.Population( - "sampling_0", "Population that samples at time 0", 0 -) - def LockePongo(): id = "QC-TwoSpecies_2L11" - populations = [population_sample_0] * 2 + populations = [ + stdpopsim.Population("Bornean", ""), + stdpopsim.Population("Sumatran", ""), + ] # This is a split-migration style model, with exponential growth or # decay allowed in each population after the split. They assumed a @@ -55,6 +53,10 @@ def LockePongo(): time=T, initial_size=Ne, growth_rate=0, population_id=0 ), ], + population_id_map=[ + {"Bornean": 0, "Sumatran": 1}, + {"Bornean": 0, "Sumatran": 1}, + ], ) diff --git a/stdpopsim/slim_engine.py b/stdpopsim/slim_engine.py index 7c6d68da4..a8a2da7b5 100644 --- a/stdpopsim/slim_engine.py +++ b/stdpopsim/slim_engine.py @@ -46,7 +46,6 @@ import subprocess import functools import itertools -import collections import contextlib import random import textwrap @@ -515,9 +514,9 @@ def msprime_rm_to_slim_rm(recombination_map): """ Convert recombination map from start position coords to end position coords. """ - rates = recombination_map.get_rates() - ends = [int(pos) - 1 for pos in recombination_map.get_positions()] - return rates[:-1], ends[1:] + rates = recombination_map.rate + ends = [int(pos) - 1 for pos in recombination_map.position] + return rates, ends[1:] def slim_makescript( @@ -532,12 +531,10 @@ def slim_makescript( burn_in, ): - pop_names = [ - pc.metadata["id"] for pc in demographic_model.population_configurations - ] + pop_names = [pop.name for pop in demographic_model.model.populations] # Use copies of these so that the time frobbing below doesn't have # side-effects in the caller's model. - demographic_events = copy.deepcopy(demographic_model.demographic_events) + demographic_events = copy.deepcopy(demographic_model.model.events) if extended_events is None: extended_events = [] else: @@ -565,11 +562,12 @@ def fix_time(event): # The demography debugger constructs event epochs, which we use # to define the forwards-time events. - dd = msprime.DemographyDebugger( - population_configurations=demographic_model.population_configurations, - migration_matrix=demographic_model.migration_matrix, - demographic_events=demographic_events, - ) + dd = demographic_model.model.debug() + # msprime.DemographyDebugger( + # population_configurations=demographic_model.population_configurations, + # migration_matrix=demographic_model.migration_matrix, + # demographic_events=demographic_events, + # ) epochs = sorted(dd.epochs, key=lambda e: e.start_time, reverse=True) T = [round(e.start_time * demographic_model.generation_time) for e in epochs] @@ -586,54 +584,56 @@ def fix_time(event): subpopulation_splits = [] for i, epoch in enumerate(epochs): for de in epoch.demographic_events: - if isinstance(de, msprime.MassMigration): - - if de.proportion < 1: - # Calculate remainder of population after previous - # MassMigration events in this epoch. - rem = 1 - np.sum( - [ - ap[3] - for ap in admixture_pulses - if ap[0] == i and ap[1] == de.source - ] - ) - admixture_pulses.append( - ( - i, - de.source, # forwards-time dest - de.dest, # forwards-time source - rem * de.proportion, + if isinstance(de, msprime.LineageMovementEvent): + # This is using internal msprime APIs here, but it's not worth + # updating before we change over to Demes. + for lm in de._as_lineage_movements(): + if lm.proportion < 1: + # Calculate remainder of population after previous + # MassMigration events in this epoch. + rem = 1 - np.sum( + [ + ap[3] + for ap in admixture_pulses + if ap[0] == i and ap[1] == lm.source + ] + ) + admixture_pulses.append( + ( + i, + lm.source, # forwards-time dest + lm.dest, # forwards-time source + rem * lm.proportion, + ) ) + continue + + # Backwards: lm.source is being merged into lm.dest. + # Forwards: lm.source is being created, taking individuals + # from lm.dest. + # + # If the proportion==1, we can use SLiM function: + # sim.addSubpopSplit(newpop, size, oldpop), + # which we trigger by adding a row to subpopulation_splits. + # This SLiM function creates newpop (=lm.source), under the + # assumption that it doesn't already exist. + + subpopulation_splits.append( + (f"_T[{i}]", lm.source, f"_N[{i+1},{lm.source}]", lm.dest) ) - continue - - # Backwards: de.source is being merged into de.dest. - # Forwards: de.source is being created, taking individuals - # from de.dest. - # - # If the proportion==1, we can use SLiM function: - # sim.addSubpopSplit(newpop, size, oldpop), - # which we trigger by adding a row to subpopulation_splits. - # This SLiM function creates newpop (=de.source), under the - # assumption that it doesn't already exist. - - subpopulation_splits.append( - (f"_T[{i}]", de.source, f"_N[{i+1},{de.source}]", de.dest) - ) - # Zero out the population size for generations before this - # epoch, to avoid simulating invididuals that contribute no - # genealogy. - N[de.source, 0 : (i + 1)] = 0 - growth_rates[de.source, 0 : (i + 1)] = 0 + # Zero out the population size for generations before this + # epoch, to avoid simulating invididuals that contribute no + # genealogy. + N[lm.source, 0 : (i + 1)] = 0 + growth_rates[lm.source, 0 : (i + 1)] = 0 - # Ensure there are no migrations to or from de.source before - # this epoch. - for j in range(i + 1): - for k in range(dd.num_populations): - migration_matrices[j][k][de.source] = 0 - migration_matrices[j][de.source][k] = 0 + # Ensure there are no migrations to or from lm.source before + # this epoch. + for j in range(i + 1): + for k in range(dd.num_populations): + migration_matrices[j][k][lm.source] = 0 + migration_matrices[j][lm.source][k] = 0 drawn_mutations = [] fitness_callbacks = [] @@ -755,7 +755,7 @@ def fix_time(event): string.Template(_slim_upper).substitute( scaling_factor=scaling_factor, burn_in=float(burn_in), - chromosome_length=int(contig.recombination_map.get_length()), + chromosome_length=int(contig.recombination_map.sequence_length), recombination_rates=recomb_rates_str, recombination_ends=recomb_ends_str, mutation_rate=contig.mutation_rate, @@ -953,28 +953,29 @@ def matrix2str( printsc() # Sampling episodes. - sample_counts = collections.Counter( - [ - (sample.population, round(sample.time * demographic_model.generation_time)) - for sample in samples - ] - ) sampling_episodes = [] - for (pop, time), count in sample_counts.items(): + for sample_set in samples: + pop = demographic_model.model[sample_set.population] + # We're currently sampling genomes, and set the ploidy to 1 to make + # sure it all fits together. This is probably confusing though and + # maybe we should change it. + assert sample_set.ploidy == 1 # SLiM can only sample individuals, which we assume are diploid. + count = sample_set.num_samples + time = 0 if sample_set.time is None else sample_set.time + time = round(time * demographic_model.generation_time) n_inds = (count + 1) // 2 if count % 2 != 0: - pop_id = pop_names[pop] gen = time / demographic_model.generation_time warnings.warn( stdpopsim.SLiMOddSampleWarning( f"SLiM simulates diploid individuals, so {n_inds} " f"individuals will be sampled for the {count} haploids " - f"requested from population {pop_id} at time {gen}. " + f"requested from population {pop.name} at time {gen}. " "See #464." ) ) - sampling_episodes.append((pop, n_inds, time)) + sampling_episodes.append((pop.id, n_inds, time)) printsc(" // One row for each sampling episode.") printsc( @@ -1247,7 +1248,7 @@ def _recap_and_rescale( for pop in recap_epoch.populations ] ts = ts.recapitate( - recombination_rate=contig.recombination_map.mean_recombination_rate, + recombination_rate=contig.recombination_map.mean_rate, population_configurations=population_configurations, migration_matrix=recap_epoch.migration_matrix, random_seed=s1, diff --git a/stdpopsim/species.py b/stdpopsim/species.py index a31d44c2f..ea0516323 100644 --- a/stdpopsim/species.py +++ b/stdpopsim/species.py @@ -226,7 +226,7 @@ def get_contig( u_tot += chrom_data.length * chrom_data.mutation_rate u = u_tot / L_tot r = r_tot / L_tot - recomb_map = msprime.RecombinationMap.uniform_map(length, r) + recomb_map = msprime.RateMap.uniform(length, r) ret = stdpopsim.Contig(recombination_map=recomb_map, mutation_rate=u) else: if length is not None: @@ -237,8 +237,8 @@ def get_contig( if genetic_map is None: logger.debug(f"Making flat chromosome {length_multiplier} * {chrom.id}") gm = None - recomb_map = msprime.RecombinationMap.uniform_map( - chrom.length * length_multiplier, chrom.recombination_rate + recomb_map = msprime.RateMap.uniform( + round(chrom.length * length_multiplier), chrom.recombination_rate ) else: if length_multiplier != 1: diff --git a/tests/test_HomSap.py b/tests/test_HomSap.py index ae8dbfa74..4ee6c94bf 100644 --- a/tests/test_HomSap.py +++ b/tests/test_HomSap.py @@ -31,7 +31,7 @@ def test_recombination_rates(self): # compare the results to the current recombination rates for each chromosome genetic_map = "HapMapII_GRCh37" species = stdpopsim.get_species("HomSap") - for chrom in self.genome.chromosomes: + for chrom in self.genome.chromosomes[2:]: if chrom.id == "chrY": with self.assertWarns(Warning): contig = species.get_contig(chrom.id, genetic_map=genetic_map) @@ -39,5 +39,5 @@ def test_recombination_rates(self): contig = species.get_contig(chrom.id, genetic_map=genetic_map) self.assertAlmostEqual( chrom.recombination_rate, - contig.recombination_map.mean_recombination_rate, + contig.recombination_map.mean_rate, ) diff --git a/tests/test_cli.py b/tests/test_cli.py index e8d3a980d..678849163 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -216,9 +216,6 @@ def verify(self, cmd, num_samples, seed=1): self.assertEqual(len(stdout), 0) ts = tskit.load(str(filename)) self.assertEqual(ts.num_samples, num_samples) - provenance = json.loads(ts.provenance(0).record) - prov_seed = provenance["parameters"]["random_seed"] - self.assertEqual(prov_seed, seed) def test_homsap_seed(self): cmd = "HomSap -c chr22 -l0.1 -s 1234 20" @@ -285,9 +282,6 @@ def verify(self, cmd, num_samples, seed=1): self.assertEqual(stored_cmd[-1], str(seed)) self.assertEqual(stored_cmd[-2], "-s") self.assertEqual(stored_cmd[1:-4], cmd.split()) - provenance = json.loads(ts.provenance(0).record) - prov_seed = provenance["parameters"]["random_seed"] - self.assertEqual(prov_seed, seed) class TestWriteOutput(unittest.TestCase): diff --git a/tests/test_genetic_maps.py b/tests/test_genetic_maps.py index 7d2124d89..c17b6b5ec 100644 --- a/tests/test_genetic_maps.py +++ b/tests/test_genetic_maps.py @@ -119,7 +119,7 @@ def get_maps(self, tarball): with utils.cd(extract_dir): tar_file.extractall() for fn in os.listdir(extract_dir): - maps[fn] = msprime.RecombinationMap.read_hapmap(fn) + maps[fn] = msprime.RateMap.read_hapmap(fn) return maps def test_no_args(self): @@ -208,16 +208,16 @@ def test_warning_from_no_mapped_chromosome(self): chrom = species.genome.get_chromosome("chrY") with self.assertWarns(Warning): cm = genetic_map.get_chromosome_map(chrom.id) - self.assertIsInstance(cm, msprime.RecombinationMap) - self.assertEqual(chrom.length, cm.get_sequence_length()) + self.assertIsInstance(cm, msprime.RateMap) + self.assertEqual(chrom.length, cm.sequence_length) def test_known_chromosome(self): species = stdpopsim.get_species("CanFam") genetic_map = species.get_genetic_map("Campbell2016_CanFam3_1") chrom = species.genome.get_chromosome("1") cm = genetic_map.get_chromosome_map(chrom.id) - self.assertIsInstance(cm, msprime.RecombinationMap) - self.assertEqual(chrom.length, cm.get_sequence_length()) + self.assertIsInstance(cm, msprime.RateMap) + self.assertEqual(chrom.length, cm.sequence_length) def test_warning_for_long_recomb_map(self): species = stdpopsim.get_species("HomSap") @@ -225,8 +225,8 @@ def test_warning_for_long_recomb_map(self): chrom = species.genome.get_chromosome("chr1") with self.assertWarns(Warning): cm = genetic_map.get_chromosome_map(chrom.id) - self.assertIsInstance(cm, msprime.RecombinationMap) - self.assertLess(chrom.length, cm.get_sequence_length()) + self.assertIsInstance(cm, msprime.RateMap) + self.assertLess(chrom.length, cm.sequence_length) def test_unknown_chromosome(self): species = stdpopsim.get_species("HomSap") diff --git a/tests/test_masking.py b/tests/test_masking.py index 88358bc40..b80ba5755 100644 --- a/tests/test_masking.py +++ b/tests/test_masking.py @@ -123,8 +123,8 @@ def test_simulate_with_mask(self): L = 1000 contig = species.get_contig(length=L) contig.mutation_rate = 1e-3 - contig.recombination_map = msprime.RecombinationMap.uniform_map(L, 0) - samples = [msprime.Sample(0, 0), msprime.Sample(0, 0)] + contig.recombination_map = msprime.RateMap.uniform(L, 0) + samples = [msprime.SampleSet(2, population=0, ploidy=1)] model = stdpopsim.PiecewiseConstantSize(100) for engine_name in engines: engine = stdpopsim.get_engine(engine_name) diff --git a/tests/test_models.py b/tests/test_models.py index 032f8b180..bc501bbec 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -2,12 +2,11 @@ Tests for simulation model infrastructure. """ import unittest -import itertools -import io import sys import numpy as np import msprime +import pytest import stdpopsim from stdpopsim import models @@ -25,10 +24,8 @@ class DemographicModelTestMixin(object): model = None def test_debug_runs(self): - output = io.StringIO() - self.model.debug(output) - s = output.getvalue() - self.assertGreater(len(s), 0) + dbg = self.model.model.debug() + assert dbg.num_epochs > 0 def test_simulation_runs(self): # With a recombination_map of None, we simulate a coalescent without @@ -63,9 +60,10 @@ class QcdCatalogDemographicModelTestMixin(CatalogDemographicModelTestMixin): """ def test_qc_model_equal(self): - self.assertTrue(self.model.equals(self.model.qc_model)) - # Verify that we didn't just compare a model with itself. - self.assertNotEqual(self.model, self.model.qc_model) + d1 = self.model.model + d2 = self.model.qc_model + d1.assert_equivalent(d2, rel_tol=1e-5) + assert d1 != d2 # Add model specific test classes, derived from one of the above. @@ -89,249 +87,6 @@ def test_qc_model_equal(self): setattr(sys.modules[__name__], cls.__name__, cls) -class TestPopulationConfigsEqual(unittest.TestCase): - """ - Tests the equality comparison of different msprime population configurations. - """ - - def test_empty(self): - self.assertTrue(models.population_configurations_equal([], [])) - - def test_sample_size_error(self): - pc1 = msprime.PopulationConfiguration(sample_size=2, initial_size=1) - pc2 = msprime.PopulationConfiguration(initial_size=1) - with self.assertRaises(ValueError): - models.population_configurations_equal([pc1], [pc1]) - with self.assertRaises(ValueError): - models.population_configurations_equal([pc1], [pc2]) - with self.assertRaises(ValueError): - models.population_configurations_equal([pc2], [pc1]) - - def test_no_initial_size_error(self): - pc1 = msprime.PopulationConfiguration() - pc2 = msprime.PopulationConfiguration(initial_size=1) - with self.assertRaises(ValueError): - self.assertTrue(models.population_configurations_equal([pc1], [pc1])) - with self.assertRaises(ValueError): - self.assertTrue(models.population_configurations_equal([pc1], [pc2])) - with self.assertRaises(ValueError): - self.assertTrue(models.population_configurations_equal([pc2], [pc1])) - - def test_different_lengths(self): - pc = msprime.PopulationConfiguration(initial_size=1) - self.assertFalse(models.population_configurations_equal([], [pc])) - self.assertFalse(models.population_configurations_equal([pc], [])) - self.assertFalse(models.population_configurations_equal([pc, pc], [pc])) - self.assertFalse(models.population_configurations_equal([pc], [pc, pc])) - self.assertTrue(models.population_configurations_equal([pc], [pc])) - with self.assertRaises(models.UnequalModelsError): - models.verify_population_configurations_equal([pc], [pc, pc]) - - def test_initial_sizes(self): - test_sizes = [ - ([1], [1.001]), - ([1.001], [1]), - ([1, 2, 3], [2, 3, 4]), - (np.arange(1, 100), np.arange(1, 100) - 0.001), - ] - - for sizes1, sizes2 in test_sizes: - pc_list1 = [ - msprime.PopulationConfiguration(initial_size=size) for size in sizes1 - ] - pc_list2 = [ - msprime.PopulationConfiguration(initial_size=size) for size in sizes2 - ] - self.assertFalse(models.population_configurations_equal(pc_list1, pc_list2)) - self.assertFalse(models.population_configurations_equal(pc_list2, pc_list1)) - self.assertTrue(models.population_configurations_equal(pc_list1, pc_list1)) - self.assertTrue(models.population_configurations_equal(pc_list2, pc_list2)) - with self.assertRaises(models.UnequalModelsError): - models.verify_population_configurations_equal(pc_list2, pc_list1) - - def test_growth_rates(self): - test_rates = [ - ([1], [1.001]), - ([-1.001], [-1]), - ([1, 2, 3], [2, 3, 4]), - (np.arange(1, 100), np.arange(1, 100) - 0.001), - ] - for rates1, rates2 in test_rates: - pc_list1 = [ - msprime.PopulationConfiguration(initial_size=1, growth_rate=rate) - for rate in rates1 - ] - pc_list2 = [ - msprime.PopulationConfiguration(initial_size=1, growth_rate=rate) - for rate in rates2 - ] - self.assertFalse(models.population_configurations_equal(pc_list1, pc_list2)) - self.assertFalse(models.population_configurations_equal(pc_list2, pc_list1)) - self.assertTrue(models.population_configurations_equal(pc_list1, pc_list1)) - self.assertTrue(models.population_configurations_equal(pc_list2, pc_list2)) - with self.assertRaises(models.UnequalModelsError): - models.verify_population_configurations_equal(pc_list2, pc_list1) - - def test_sampling_times_equal(self): - no_sample_pop = stdpopsim.Population("none", "none", sampling_time=None) - zero_sample_pop = stdpopsim.Population("zero", "zero") - nonzero_sample_pop = stdpopsim.Population("nzero", "nzero", sampling_time=10) - plist1 = [no_sample_pop] * 2 + [nonzero_sample_pop] + [zero_sample_pop] * 2 - plist2 = [no_sample_pop] * 4 + [nonzero_sample_pop] - plist3 = [no_sample_pop] * 3 + [nonzero_sample_pop] - self.assertFalse( - stdpopsim.sampling_times_equal([no_sample_pop], [zero_sample_pop]) - ) - self.assertFalse( - stdpopsim.sampling_times_equal([nonzero_sample_pop], [zero_sample_pop]) - ) - self.assertFalse(stdpopsim.sampling_times_equal(plist1, plist2)) - self.assertFalse(stdpopsim.sampling_times_equal(plist1, plist3)) - self.assertTrue(stdpopsim.sampling_times_equal(plist3, plist3)) - - -class TestDemographicEventsEqual(unittest.TestCase): - """ - Tests the equality comparison of different msprime demographic events. - """ - - def test_empty(self): - self.assertTrue(models.demographic_events_equal([], [], 1)) - - def test_different_lengths(self): - events = [msprime.PopulationParametersChange(time=1, initial_size=1)] * 10 - self.assertFalse(models.demographic_events_equal(events[:1], [], 1)) - self.assertFalse(models.demographic_events_equal([], events[:1], 1)) - self.assertFalse(models.demographic_events_equal(events, [], 1)) - self.assertFalse(models.demographic_events_equal([], events, 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([], events, 1) - - def test_different_times(self): - n = 10 - e1 = [ - msprime.PopulationParametersChange(time=j, initial_size=1) for j in range(n) - ] - e2 = [ - msprime.PopulationParametersChange(time=j + 1, initial_size=1) - for j in range(n) - ] - for j in range(1, n): - self.assertFalse(models.demographic_events_equal(e1[:j], e2[:j], 1)) - self.assertFalse(models.demographic_events_equal(e2[:j], e1[:j], 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal(e1[:j], e2[:j], 1) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal(e2[:j], e1[:j], 1) - - def test_different_types(self): - events = [ - msprime.PopulationParametersChange(time=1, initial_size=1), - msprime.MigrationRateChange(time=1, rate=1), - msprime.MassMigration(time=1, source=1), - msprime.SimpleBottleneck(time=1, population=0), - ] - for a, b in itertools.combinations(events, 2): - self.assertFalse(models.demographic_events_equal([a], [b], 1)) - self.assertFalse(models.demographic_events_equal([b], [a], 1)) - self.assertTrue(models.demographic_events_equal([a], [a], 1)) - self.assertTrue(models.demographic_events_equal([b], [b], 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([b], [a], 1) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([a], [b], 1) - - def test_population_parameters_change(self): - def f(time=1, initial_size=1, growth_rate=None, population=None): - return msprime.PopulationParametersChange( - time=time, - initial_size=initial_size, - growth_rate=growth_rate, - population=population, - ) - - test_events = [ - (f(time=1), f(time=2)), - (f(initial_size=1), f(initial_size=2)), - (f(growth_rate=1), f(growth_rate=2)), - (f(population=1), f(population=2)), - ] - for a, b in test_events: - self.assertFalse(models.demographic_events_equal([a], [b], 1)) - self.assertFalse(models.demographic_events_equal([b], [a], 1)) - self.assertTrue(models.demographic_events_equal([a], [a], 1)) - self.assertTrue(models.demographic_events_equal([b], [b], 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([b], [a], 1) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([a], [b], 1) - - def test_migration_rate_change(self): - def f(time=1, rate=1, matrix_index=None): - return msprime.MigrationRateChange( - time=time, rate=rate, matrix_index=matrix_index - ) - - test_events = [ - (f(time=1), f(time=2)), - (f(rate=1), f(rate=2)), - (f(matrix_index=[0, 1]), f(matrix_index=[0, 2])), - (f(matrix_index=np.array([0, 1])), f(matrix_index=[0, 2])), - ] - for a, b in test_events: - self.assertFalse(models.demographic_events_equal([a], [b], 1)) - self.assertFalse(models.demographic_events_equal([b], [a], 1)) - self.assertTrue(models.demographic_events_equal([a], [a], 1)) - self.assertTrue(models.demographic_events_equal([b], [b], 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([b], [a], 1) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([a], [b], 1) - - def test_mass_migration(self): - def f(time=1, source=1, dest=1, proportion=1): - return msprime.MassMigration( - time=time, source=source, dest=dest, proportion=proportion - ) - - test_events = [ - (f(time=1), f(time=2)), - (f(source=1), f(source=2)), - (f(dest=1), f(dest=2)), - (f(proportion=1), f(proportion=0.2)), - ] - for a, b in test_events: - self.assertFalse(models.demographic_events_equal([a], [b], 1)) - self.assertFalse(models.demographic_events_equal([b], [a], 1)) - self.assertTrue(models.demographic_events_equal([a], [a], 1)) - self.assertTrue(models.demographic_events_equal([b], [b], 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([b], [a], 1) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([a], [b], 1) - - def test_simple_bottleneck(self): - def f(time=1, population=1, proportion=1): - return msprime.SimpleBottleneck( - time=time, population=population, proportion=proportion - ) - - test_events = [ - (f(time=1), f(time=2)), - (f(population=1), f(population=2)), - (f(proportion=1), f(proportion=0.2)), - ] - for a, b in test_events: - self.assertFalse(models.demographic_events_equal([a], [b], 1)) - self.assertFalse(models.demographic_events_equal([b], [a], 1)) - self.assertTrue(models.demographic_events_equal([a], [a], 1)) - self.assertTrue(models.demographic_events_equal([b], [b], 1)) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([b], [a], 1) - with self.assertRaises(models.UnequalModelsError): - models.verify_demographic_events_equal([a], [b], 1) - - class TestRegisterQCModel(unittest.TestCase): def make_model(self, name): return models.DemographicModel( @@ -339,7 +94,7 @@ def make_model(self, name): description=name, long_description=name, generation_time=1, - populations=[], + model=msprime.Demography.isolated_model([1]), ) def test_register_qc(self): @@ -359,30 +114,23 @@ def test_bad_qc_models(self): model.register_qc(not_a_model) -class TestAllModels(unittest.TestCase): +class TestAllModels: """ Tests for registered simulation models. """ def test_non_empty(self): - self.assertGreater(len(list(stdpopsim.all_demographic_models())), 0) - - def test_all_instances(self): - for model in stdpopsim.all_demographic_models(): - self.assertIsInstance(model, models.DemographicModel) + assert len(list(stdpopsim.all_demographic_models())) > 0 - self.assertGreater(len(model.id), 0) - self.assertGreater(len(model.description), 0) - self.assertGreater(len(model.long_description), 0) - self.assertGreater(len(model.citations), 0) - self.assertGreater(model.generation_time, 0) - - npops = len(model.populations) - self.assertGreater(npops, 0) - self.assertEqual(len(model.population_configurations), npops) - self.assertEqual(len(model.migration_matrix), npops) - self.assertEqual(len(model.migration_matrix[0]), npops) - self.assertIsInstance(model.demographic_events, list) + @pytest.mark.parametrize("model", stdpopsim.all_demographic_models()) + def test_all_instances(self, model): + assert isinstance(model, models.DemographicModel) + assert len(model.id) > 0 + assert len(model.description) > 0 + assert len(model.long_description) > 0 + assert len(model.citations) > 0 + assert model.generation_time > 0 + model.model.validate() class TestModelsEqual(unittest.TestCase): @@ -391,42 +139,11 @@ class TestModelsEqual(unittest.TestCase): """ def test_known_models(self): - # All models should be equal to themselves. - other_model = models.DemographicModel.empty() + other_model = msprime.Demography.isolated_model([1]) for model in stdpopsim.all_demographic_models(): - self.assertTrue(model.equals(model)) - self.assertFalse(model.equals(other_model)) - - def test_different_objects(self): - m1 = models.DemographicModel.empty() - self.assertFalse(m1.equals(self)) - self.assertFalse(m1.equals({})) - self.assertFalse(m1.equals(None)) - - def test_default_models(self): - m1 = models.DemographicModel.empty() - m2 = models.DemographicModel.empty() - self.assertTrue(m1.equals(m2)) - - def test_migration_matrices(self): - m1 = models.DemographicModel.empty() - m2 = models.DemographicModel.empty() - m1.migration_matrix = [[]] - self.assertFalse(m1.equals(m2)) - m2.migration_matrix = [[]] - self.assertTrue(m1.equals(m2)) - mm1 = np.arange(100, dtype=float).reshape(10, 10) - m1.migration_matrix = mm1 - self.assertFalse(m1.equals(m2)) - m2.migration_matrix = mm1 - self.assertTrue(m1.equals(m2)) - # Perturb the matrix by a tiny bit and see if we're still equal - mm2 = mm1.copy() + 1e-9 - m2.migration_matrix = mm2 - self.assertFalse(np.all(m1.migration_matrix == m2.migration_matrix)) - self.assertTrue(m1.equals(m2)) - # If we have higher tolerances we catch the differences - self.assertFalse(m1.equals(m2, atol=1e-10, rtol=1e-9)) + model.model.assert_equivalent(model.model) + model.model.assert_equal(model.model) + assert not model.model.is_equivalent(other_model) class TestConstantSizeModel(unittest.TestCase, DemographicModelTestMixin): @@ -444,31 +161,33 @@ class TestPiecewiseConstantSize(unittest.TestCase): def test_single_epoch(self): model = models.PiecewiseConstantSize(100) - self.assertEqual(model.population_configurations[0].initial_size, 100) - self.assertEqual(len(model.population_configurations), 1) - self.assertEqual(len(model.demographic_events), 0) + self.assertEqual(model.model.populations[0].initial_size, 100) + self.assertEqual(model.model.num_populations, 1) + self.assertEqual(model.model.num_events, 0) def test_two_epoch(self): model = models.PiecewiseConstantSize(50, (10, 100)) - self.assertEqual(model.population_configurations[0].initial_size, 50) - self.assertEqual(len(model.demographic_events), 1) - event = model.demographic_events[0] + self.assertEqual(model.model.populations[0].initial_size, 50) + self.assertEqual(model.model.num_populations, 1) + self.assertEqual(model.model.num_events, 1) + event = model.model.events[0] self.assertEqual(event.time, 10) self.assertEqual(event.initial_size, 100) - self.assertEqual(event.growth_rate, 0) + self.assertIsNone(event.growth_rate) def test_three_epoch(self): model = models.PiecewiseConstantSize(0.1, (0.1, 10), (0.2, 100)) - self.assertEqual(model.population_configurations[0].initial_size, 0.1) - self.assertEqual(len(model.demographic_events), 2) - event = model.demographic_events[0] + self.assertEqual(model.model.populations[0].initial_size, 0.1) + self.assertEqual(model.model.num_populations, 1) + self.assertEqual(model.model.num_events, 2) + event = model.model.events[0] self.assertEqual(event.time, 0.1) self.assertEqual(event.initial_size, 10) - self.assertEqual(event.growth_rate, 0) - event = model.demographic_events[1] + self.assertIsNone(event.growth_rate) + event = model.model.events[1] self.assertEqual(event.time, 0.2) self.assertEqual(event.initial_size, 100) - self.assertEqual(event.growth_rate, 0) + self.assertIsNone(event.growth_rate) class TestIsolationWithMigration(unittest.TestCase): @@ -477,20 +196,19 @@ class TestIsolationWithMigration(unittest.TestCase): """ def test_pop_configs(self): - model = models.IsolationWithMigration(100, 200, 300, 50, 0, 0) - self.assertEqual(len(model.population_configurations), 3) - self.assertEqual(model.population_configurations[0].initial_size, 200) - self.assertEqual(model.population_configurations[1].initial_size, 300) - self.assertEqual(model.population_configurations[2].initial_size, 100) + model = models.IsolationWithMigration(100, 200, 300, 50, 0, 0).model + self.assertEqual(model.num_populations, 3) + self.assertEqual(model.populations[0].initial_size, 200) + self.assertEqual(model.populations[1].initial_size, 300) + self.assertEqual(model.populations[2].initial_size, 100) def test_split(self): - model = models.IsolationWithMigration(100, 200, 300, 50, 0, 0) - self.assertEqual(len(model.demographic_events), 2) - self.assertEqual(model.demographic_events[0].time, 50) - self.assertEqual(model.demographic_events[1].time, 50) + model = models.IsolationWithMigration(100, 200, 300, 50, 0, 0).model + self.assertEqual(len(model.events), 1) + self.assertEqual(model.events[0].time, 50) def test_migration_rates(self): - model = models.IsolationWithMigration(100, 200, 300, 50, 0.002, 0.003) + model = models.IsolationWithMigration(100, 200, 300, 50, 0.002, 0.003).model self.assertEqual(np.shape(model.migration_matrix), (3, 3)) self.assertEqual(model.migration_matrix[0][1], 0.002) self.assertEqual(model.migration_matrix[1][0], 0.003) @@ -499,30 +217,46 @@ def test_migration_rates(self): class TestPopulationSampling(unittest.TestCase): - # Create populations to test on - _pop1 = stdpopsim.Population("pop0", "Test pop. 0") - _pop2 = stdpopsim.Population("pop1", "Test pop. 1", sampling_time=10) - _pop3 = stdpopsim.Population("pop2", "Test pop. 2", sampling_time=None) - - # Create an empty model to hold populations - base_mod = models.DemographicModel.empty(populations=[_pop1, _pop2, _pop3]) + def make_model(self): + # Create populations to test on + _pop1 = stdpopsim.Population("pop0", "Test pop. 0") + _pop2 = stdpopsim.Population("pop1", "Test pop. 1", sampling_time=10) + _pop3 = stdpopsim.Population("pop2", "Test pop. 2", sampling_time=None) + + # Create an model to hold populations + base_mod = models.DemographicModel( + id="x", + description="y", + long_description="z", + populations=[_pop1, _pop2, _pop3], + population_configurations=[ + msprime.PopulationConfiguration(initial_size=1), + msprime.PopulationConfiguration(initial_size=1), + msprime.PopulationConfiguration(initial_size=1), + ], + ) + return base_mod def test_num_sampling_populations(self): - self.assertEqual(self.base_mod.num_sampling_populations, 2) + base_mod = self.make_model() + self.assertEqual(base_mod.num_sampling_populations, 2) def test_get_samples(self): - test_samples = self.base_mod.get_samples(2, 1) - self.assertEqual(len(test_samples), 3) + base_mod = self.make_model() + test_samples = base_mod.get_samples(2, 1) + self.assertEqual(len(test_samples), 2) + self.assertEqual(sum([ss.num_samples for ss in test_samples]), 3) + # Check for error when prohibited sampling asked for with self.assertRaises(ValueError): - self.base_mod.get_samples(2, 2, 1) + base_mod.get_samples(2, 2, 1) # Get the population corresponding to each sample sample_populations = [i.population for i in test_samples] # Check sample populations - self.assertEqual(sample_populations, [0, 0, 1]) + self.assertEqual(sample_populations, [0, 1]) # Test sampling times sample_times = [i.time for i in test_samples] - self.assertEqual(sample_times, [0, 0, 10]) + self.assertEqual(sample_times, [0, 10]) # Test that all sampling populations are specified before non-sampling populations # in the model.populations list @@ -533,80 +267,6 @@ def test_population_order(self): # All sampling populations must be at the start of the list self.assertEqual(sum(allow_sample_status[num_sampling:]), 0) - # Test that populations are listed in the same order in model.populations and - # model.population_configurations - def test_population_config_order_equal(self): - for model in stdpopsim.all_demographic_models(): - pop_ids = [pop.id for pop in model.populations] - config_ids = [ - config.metadata["id"] for config in model.population_configurations - ] - for p, c in zip(pop_ids, config_ids): - self.assertEqual(p, c) - - # Test that we are indeed getting a valid DDB back - # admittedly a pretty bad test... - def test_demography_debugger(self): - for model in stdpopsim.all_demographic_models(): - ddb = model.get_demography_debugger() - self.assertIsInstance(ddb, msprime.DemographyDebugger) - - # test for equality of ddbs - def test_demography_debugger_equal(self): - for model in stdpopsim.all_demographic_models(): - ddb1 = model.get_demography_debugger() - ddb2 = msprime.DemographyDebugger( - population_configurations=model.population_configurations, - migration_matrix=model.migration_matrix, - demographic_events=model.demographic_events, - ) - f1 = io.StringIO() - f2 = io.StringIO() - ddb1.print_history(f1) - ddb2.print_history(f2) - self.assertEqual(f1.getvalue(), f2.getvalue()) - - -class TestDemographicModelConstruction(unittest.TestCase): - # Test construction of a Population object when provided a PopulationConfiguration - # object but no populations - def test_population_construction_popconfig(self): - pc = [msprime.PopulationConfiguration(initial_size=1, growth_rate=0.03)] - dm = stdpopsim.DemographicModel( - id="", - description="", - long_description="", - generation_time=1, - population_configurations=pc, - ) - self.assertEqual(dm.populations[0].id, "pop0") - - # Test construction of a Population object when provided a PopulationConfiguration - # object with metadata but no populations - def test_population_construction_popconfig_metadata(self): - pop0 = stdpopsim.Population(id="A", description="Pop A") - pc_meta = [ - msprime.PopulationConfiguration( - initial_size=1, growth_rate=0.03, metadata=pop0.asdict() - ) - ] - dm = stdpopsim.DemographicModel( - id="", - description="", - long_description="", - generation_time=1, - population_configurations=pc_meta, - ) - self.assertEqual(dm.populations[0].asdict(), pop0.asdict()) - - # Test construction of a Population object when provided neither a - # PopulationConfiguration list or a Population list - def test_population_construction_no_popconfig(self): - dm = stdpopsim.DemographicModel( - id="A", description="A", long_description="A", generation_time=1 - ) - self.assertEqual(dm.populations, []) - class TestZigZagWarning(unittest.TestCase): def test_zigzag_produces_warning(self): diff --git a/tests/test_slim_engine.py b/tests/test_slim_engine.py index 2cf3efd5c..797c37eb3 100644 --- a/tests/test_slim_engine.py +++ b/tests/test_slim_engine.py @@ -437,17 +437,18 @@ def exp_decline(self, N0=100, N1=1000, T=1000): Used for testing that growth rates are handled appropriately. """ r = math.log(N0 / N1) / T + pop0 = stdpopsim.models.Population(id="pop0", description="") return stdpopsim.DemographicModel( id="exp_decline", description="exp_decline", long_description="exp_decline", - populations=[stdpopsim.models._pop0], + populations=[pop0], generation_time=1, population_configurations=[ msprime.PopulationConfiguration( initial_size=N0, growth_rate=r, - metadata=stdpopsim.models._pop0.asdict(), + metadata=pop0.asdict(), ) ], demographic_events=[ diff --git a/tests/test_species.py b/tests/test_species.py index 10c78e598..604bbdbec 100644 --- a/tests/test_species.py +++ b/tests/test_species.py @@ -178,8 +178,8 @@ def test_length_multiplier(self): for x in [0.125, 1.0, 2.0]: contig2 = self.species.get_contig("chr22", length_multiplier=x) self.assertEqual( - contig1.recombination_map.get_positions()[-1] * x, - contig2.recombination_map.get_positions()[-1], + round(contig1.recombination_map.position[-1] * x), + contig2.recombination_map.position[-1], ) def test_length_multiplier_on_empirical_map(self): @@ -191,7 +191,7 @@ def test_length_multiplier_on_empirical_map(self): def test_genetic_map(self): # TODO we should use a different map here so we're not hitting the cache. contig = self.species.get_contig("chr22", genetic_map="HapMapII_GRCh37") - self.assertIsInstance(contig.recombination_map, msprime.RecombinationMap) + self.assertIsInstance(contig.recombination_map, msprime.RateMap) def test_contig_options(self): with self.assertRaises(ValueError): @@ -226,7 +226,7 @@ def test_contig_options(self): def test_generic_contig(self): L = 1e6 contig = self.species.get_contig(length=L) - self.assertTrue(contig.recombination_map.get_length() == L) + self.assertTrue(contig.recombination_map.sequence_length == L) chrom_ids = np.arange(1, 23).astype("str") Ls = [c.length for c in self.species.genome.chromosomes if c.id in chrom_ids] @@ -243,6 +243,5 @@ def test_generic_contig(self): self.assertTrue(contig.mutation_rate == np.average(us, weights=Ls)) self.assertTrue( - contig.recombination_map.mean_recombination_rate - == np.average(rs, weights=Ls) + contig.recombination_map.mean_rate == np.average(rs, weights=Ls) )