From 795b335ddff49021344f72de275ab5a213e155a0 Mon Sep 17 00:00:00 2001
From: Jerome Kelleher <jk@well.ox.ac.uk>
Date: Thu, 14 Jan 2016 10:08:59 +0000
Subject: [PATCH] Changed default for max_memory to unlimited.

---
 msprime/cli.py    |  2 +-
 msprime/trees.py  | 34 ++++++++++++++++++++--------------
 tests/test_cli.py |  2 +-
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/msprime/cli.py b/msprime/cli.py
index 3ac1c1d0c..ffdf68ee2 100644
--- a/msprime/cli.py
+++ b/msprime/cli.py
@@ -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"))
diff --git a/msprime/trees.py b/msprime/trees.py
index 90cfd6104..2ccea8d6d 100644
--- a/msprime/trees.py
+++ b/msprime/trees.py
@@ -24,6 +24,7 @@
 
 import math
 import random
+import sys
 
 try:
     import svgwrite
@@ -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`.
@@ -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`
@@ -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
@@ -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`
@@ -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):
         """
@@ -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):
         """
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 9bcaaa56b..83b97ade9 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -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):