Skip to content

Commit

Permalink
Changed default for max_memory to unlimited.
Browse files Browse the repository at this point in the history
  • Loading branch information
jeromekelleher committed Jan 14, 2016
1 parent 8ac0de7 commit 795b335
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 16 deletions.
2 changes: 1 addition & 1 deletion msprime/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ def add_header_argument(parser):

def add_max_memory_argument(parser):
parser.add_argument(
"--max-memory", "-M", default="1G",
"--max-memory", "-M", default=None,
help=(
"Maximum memory to use. If the simulation exceeds this limit "
"exit with error status. Supports K,M and G suffixes"))
Expand Down
34 changes: 20 additions & 14 deletions msprime/trees.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

import math
import random
import sys

try:
import svgwrite
Expand Down Expand Up @@ -438,7 +439,7 @@ def __ne__(self, other):
def simulate(
sample_size, num_loci=1, scaled_recombination_rate=0.0,
scaled_mutation_rate=None,
population_models=[], random_seed=None, max_memory="1G"):
population_models=[], random_seed=None, max_memory=None):
"""
Simulates the coalescent with recombination under the specified model
parameters and returns the resulting :class:`.TreeSequence`.
Expand All @@ -456,7 +457,8 @@ def simulate(
random seed will be automatically generated.
:param int,str max_memory: The maximum amount of memory used
during the simulation. If this is exceeded, the simulation will
terminate with a :class:`LibraryError` exception.
terminate with a :class:`LibraryError` exception. By default
this is not set, and no memory limits are imposed.
:return: The :class:`.TreeSequence` object representing the results
of the simulation.
:rtype: :class:`.TreeSequence`
Expand All @@ -477,7 +479,7 @@ def simulate(

def simulate_tree(
sample_size, scaled_mutation_rate=None, population_models=[],
random_seed=None, max_memory="1G"):
random_seed=None, max_memory=None):
"""
Simulates the coalescent at a single locus for the specified sample size
under the specified list of population models. Returns a
Expand All @@ -492,7 +494,8 @@ def simulate_tree(
random seed will be automatically generated.
:param int,str max_memory: The maximum amount of memory used
during the simulation. If this is exceeded, the simulation will
terminate with a :class:`LibraryError` exception.
terminate with a :class:`LibraryError` exception. By default
this is not set, and no memory limits are imposed.
:return: The :class:`.SparseTree` object representing the results
of the simulation.
:rtype: :class:`.SparseTree`
Expand Down Expand Up @@ -643,15 +646,18 @@ def set_max_memory(self, max_memory):
to the specified value. This can be suffixed with
K, M or G to specify units of Kibibytes, Mibibytes or Gibibytes.
"""
s = max_memory
d = {"K": 2**10, "M": 2**20, "G": 2**30}
multiplier = 1
value = s
if s.endswith(tuple(d.keys())):
value = s[:-1]
multiplier = d[s[-1]]
n = int(value)
self._max_memory = n * multiplier
if max_memory is None:
self._max_memory = max_memory
else:
s = max_memory
d = {"K": 2**10, "M": 2**20, "G": 2**30}
multiplier = 1
value = s
if s.endswith(tuple(d.keys())):
value = s[:-1]
multiplier = d[s[-1]]
n = int(value)
self._max_memory = n * multiplier

def _set_environment_defaults(self):
"""
Expand Down Expand Up @@ -681,7 +687,7 @@ def _set_environment_defaults(self):
if self._random_seed is None:
self._random_seed = random.randint(0, 2**31 - 1)
if self._max_memory is None:
self._max_memory = 10 * 1024 * 1024 # 10MiB by default
self._max_memory = sys.maxsize # Unlimited

def run(self):
"""
Expand Down
2 changes: 1 addition & 1 deletion tests/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,7 +405,7 @@ def test_simulate_default_values(self):
self.assertEqual(args.mutation_rate, 0.0)
self.assertEqual(args.num_loci, 1)
self.assertEqual(args.random_seed, None)
self.assertEqual(args.max_memory, "1G")
self.assertEqual(args.max_memory, None)
self.assertEqual(args.compress, False)

def test_simulate_short_args(self):
Expand Down

0 comments on commit 795b335

Please sign in to comment.