MemoryError: Unable to allocate 4.18 TiB...

[Erythroxylum]

Hello, I am trying to estimate tdiv and tmigration parameters and receiving errors as above. The crux is that even when I have reduced the number of leaves and time parameters, I am not observing a decrease in the necessary memory allocation, it sometimes increases. I have altered set_data(mem_chunk_size=###), but this seems to have no effect on my problem.

Why am I not seeing the allocation and the array shape also reduce under a smaller model? Are all of these models just way too large? Is there another way to limit the memory allocation?

The largest model has seven leaves with seven size parameters and eleven time parameters and requests 1.7 TB of memory.

A smaller model has six leaves and seven time parameters, requesting 4.18 TB of memory:

I am getting the following error:

---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-6-2e658921fbd7> in <module>
      1 #fixSizeAdmix_red.optimize(options=dict(maxiter=500))
----> 2 fixSizeAdmix_red.optimize(options=dict(maxiter=500), method="L-BFGS-B")
      3 #fixSizeAdmix_red.stochastic_optimize(options=dict(maxiter=500), num_iters=10, snps_per_minibatch=2000, svrg_epoch=5)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/demo_model.py in optimize(self, method, jac, hess, hessp, printfreq, **kwargs)
    950             jac=jac, hess=hess, hessp=hessp,
    951             bounds=bounds, callback=callback,
--> 952             **kwargs)
    953 
    954         self._set_x(res.x)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in find_mle(self, x0, method, jac, hess, hessp, bounds, callback, **kwargs)
    267         return _find_minimum(fun, x0, scipy.optimize.minimize,
    268                              bounds=bounds, callback=callback,
--> 269                              opt_kwargs=opt_kwargs, gradmakers=gradmakers, replacefun=replacefun)
    270 
    271 

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/optimizers.py in _find_minimum(f, start_params, optimizer, bounds, callback, opt_kwargs, **kwargs)
     87 
     88     ret = _find_minimum_helper(
---> 89         f, start_params, optimizer, opt_kwargs, **kwargs)
     90     if fixed_params:
     91         ret.x = get_x(ret.x)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/optimizers.py in _find_minimum_helper(f, start_params, optimizer, opt_kwargs, gradmakers, replacefun)
    101     if replacefun:
    102         f = replacefun(f)
--> 103     return optimizer(f, start_params, **opt_kwargs)
    104 
    105 stochastic_opts = {}

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/_minimize.py in minimize(fun, x0, args, method, jac, hess, hessp, bounds, constraints, tol, callback, options)
    616     elif meth == 'l-bfgs-b':
    617         return _minimize_lbfgsb(fun, x0, args, jac, bounds,
--> 618                                 callback=callback, **options)
    619     elif meth == 'tnc':
    620         return _minimize_tnc(fun, x0, args, jac, bounds, callback=callback,

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/lbfgsb.py in _minimize_lbfgsb(fun, x0, args, jac, bounds, disp, maxcor, ftol, gtol, eps, maxfun, maxiter, iprint, callback, maxls, finite_diff_rel_step, **unknown_options)
    306     sf = _prepare_scalar_function(fun, x0, jac=jac, args=args, epsilon=eps,
    307                                   bounds=new_bounds,
--> 308                                   finite_diff_rel_step=finite_diff_rel_step)
    309 
    310     func_and_grad = sf.fun_and_grad

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/optimize.py in _prepare_scalar_function(fun, x0, jac, args, bounds, epsilon, finite_diff_rel_step, hess)
    260     # calculation reduces overall function evaluations.
    261     sf = ScalarFunction(fun, x0, args, grad, hess,
--> 262                         finite_diff_rel_step, bounds, epsilon=epsilon)
    263 
    264     return sf

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py in __init__(self, fun, x0, args, grad, hess, finite_diff_rel_step, finite_diff_bounds, epsilon)
     74 
     75         self._update_fun_impl = update_fun
---> 76         self._update_fun()
     77 
     78         # Gradient evaluation

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py in _update_fun(self)
    164     def _update_fun(self):
    165         if not self.f_updated:
--> 166             self._update_fun_impl()
    167             self.f_updated = True
    168 

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py in update_fun()
     71 
     72         def update_fun():
---> 73             self.f = fun_wrapped(self.x)
     74 
     75         self._update_fun_impl = update_fun

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/_differentiable_functions.py in fun_wrapped(x)
     68         def fun_wrapped(x):
     69             self.nfev += 1
---> 70             return fun(x, *args)
     71 
     72         def update_fun():

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/optimize.py in __call__(self, x, *args)
     72     def __call__(self, x, *args):
     73         """ returns the the function value """
---> 74         self._compute_if_needed(x, *args)
     75         return self._value
     76 

~/anaconda3/envs/momi/lib/python3.7/site-packages/scipy/optimize/optimize.py in _compute_if_needed(self, x, *args)
     66         if not np.all(x == self.x) or self._value is None or self.jac is None:
     67             self.x = np.asarray(x).copy()
---> 68             fg = self.fun(x, *args)
     69             self.jac = fg[1]
     70             self._value = fg[0]

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/wrap_util.py in nary_f(*args, **kwargs)
     18             else:
     19                 x = tuple(args[i] for i in argnum)
---> 20             return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
     21         return nary_f
     22     return nary_operator

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/differential_operators.py in value_and_grad(fun, x)
    133     """Returns a function that returns both value and gradient. Suitable for use
    134     in scipy.optimize"""
--> 135     vjp, ans = _make_vjp(fun, x)
    136     if not vspace(ans).size == 1:
    137         raise TypeError("value_and_grad only applies to real scalar-output "

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/core.py in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in trace(start_node, fun, x)
      8     with trace_stack.new_trace() as t:
      9         start_box = new_box(x, t, start_node)
---> 10         end_box = fun(start_box)
     11         if isbox(end_box) and end_box._trace == start_box._trace:
     12             return end_box._value, end_box._node

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/wrap_util.py in unary_f(x)
     13                 else:
     14                     subargs = subvals(args, zip(argnum, x))
---> 15                 return fun(*subargs, **kwargs)
     16             if isinstance(argnum, int):
     17                 x = args[argnum]

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in fun(x)
    261         @functools.wraps(self.kl_div)
    262         def fun(x):
--> 263             ret = self.kl_div(x)
    264             hist.recent_vals += [(x, ret)]
    265             return ret

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in kl_div(self, x)
    173         Returns KL-Divergence(Empirical || Theoretical(x)).
    174         """
--> 175         log_lik = self.log_lik(x)
    176         #ret = -log_lik + self.sfs.n_snps() * self.sfs._entropy + _entropy_mut_term(self.mut_rate, self.sfs, self.p_missing, self.use_pairwise_diffs)
    177         ret = -log_lik + self.sfs.n_snps() * self.sfs._entropy

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in log_lik(self, x, vector)
     98         Returns the composite log-likelihood of the data at the point x.
     99         """
--> 100         ret = self._log_lik(x, vector=vector)
    101         logger.debug("log-likelihood = {0}".format(ret))
    102         return ret

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in _log_lik(self, x, vector)
    121     def _log_lik(self, x, vector):
    122         demo = self._get_multipop_moran(x)
--> 123         ret = self._get_multinom_loglik(demo, vector=vector) + self._mut_factor(demo, vector=vector)
    124         if vector:
    125             ret = ret + self._log_prior(x) / len(ret)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in _get_multinom_loglik(self, demo, vector)
    146                     cache, G, batch,
    147                     self.truncate_probs, self.folded,
--> 148                     self.error_matrices, vector)
    149         else:
    150             ret = _composite_log_likelihood(

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in _raw_log_lik(cache, G, data, truncate_probs, folded, error_matrices, vector)
    514         ## checkpointing for hessian,
    515         ## but doesn't work for vectorized output
--> 516         return rearrange_dict_grad(wrapped_fun)(cache)
    517 
    518 

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in wrapped_fun(xdict)
    501     @functools.wraps(fun)
    502     def wrapped_fun(xdict):
--> 503         return wrapped_fun_helper(ag.dict(xdict), lambda:None)
    504     return wrapped_fun
    505 

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     42             parents = tuple(box._node for _     , box in boxed_args)
     43             argnums = tuple(argnum    for argnum, _   in boxed_args)
---> 44             ans = f_wrapped(*argvals, **kwargs)
     45             node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
     46             return new_box(ans, trace, node)

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     46             return new_box(ans, trace, node)
     47         else:
---> 48             return f_raw(*args, **kwargs)
     49     f_wrapped.fun = f_raw
     50     f_wrapped._is_autograd_primitive = True

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in wrapped_fun_helper(xdict, dummy)
    487         ## ag.value_and_grad() to avoid second forward pass
    488         ## ag.checkpoint() ensures hessian gets properly checkpointed
--> 489         val, grad = ag.checkpoint(ag.value_and_grad(fun))(xdict)
    490         assert len(val.shape) == 0
    491         dummy.cache = grad

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     46             return new_box(ans, trace, node)
     47         else:
---> 48             return f_raw(*args, **kwargs)
     49     f_wrapped.fun = f_raw
     50     f_wrapped._is_autograd_primitive = True

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/wrap_util.py in nary_f(*args, **kwargs)
     18             else:
     19                 x = tuple(args[i] for i in argnum)
---> 20             return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
     21         return nary_f
     22     return nary_operator

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/differential_operators.py in value_and_grad(fun, x)
    133     """Returns a function that returns both value and gradient. Suitable for use
    134     in scipy.optimize"""
--> 135     vjp, ans = _make_vjp(fun, x)
    136     if not vspace(ans).size == 1:
    137         raise TypeError("value_and_grad only applies to real scalar-output "

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/core.py in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in trace(start_node, fun, x)
      8     with trace_stack.new_trace() as t:
      9         start_box = new_box(x, t, start_node)
---> 10         end_box = fun(start_box)
     11         if isbox(end_box) and end_box._trace == start_box._trace:
     12             return end_box._value, end_box._node

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/wrap_util.py in unary_f(x)
     13                 else:
     14                     subargs = subvals(args, zip(argnum, x))
---> 15                 return fun(*subargs, **kwargs)
     16             if isinstance(argnum, int):
     17                 x = args[argnum]

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in wrapped_fun(cache)
    507     def wrapped_fun(cache):
    508         demo = Demography(G, cache=cache)
--> 509         return _composite_log_likelihood(data, demo, truncate_probs=truncate_probs, folded=folded, error_matrices=error_matrices, vector=vector)
    510     if vector:
    511         return ag.checkpoint(wrapped_fun)(cache)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/likelihood.py in _composite_log_likelihood(data, demo, mut_rate, truncate_probs, vector, p_missing, use_pairwise_diffs, **kwargs)
    418         sfs = data
    419 
--> 420     sfs_probs = np.maximum(expected_sfs(demo, sfs.configs, normalized=True, **kwargs),
    421                            truncate_probs)
    422     log_lik = sfs._integrate_sfs(np.log(sfs_probs), vector=vector)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in expected_sfs(demography, configs, mut_rate, normalized, folded, error_matrices)
     56     expected_sfs_tensor_prod : compute summary statistics of SFS
     57     """
---> 58     sfs, denom = _expected_sfs(demography, configs, folded, error_matrices)
     59     if normalized:
     60         sfs = sfs / denom

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in _expected_sfs(demography, configs, folded, error_matrices)
     74         vecs = _apply_error_matrices(vecs, error_matrices)
     75 
---> 76     vals = expected_sfs_tensor_prod(vecs, demography)
     77 
     78     sfs = vals[idxs['idx_2_row']]

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in expected_sfs_tensor_prod(vecs, demography, mut_rate, sampled_pops)
    246             for v, n in zip(vecs, demography.sampled_n)]
    247 
--> 248     res = _expected_sfs_tensor_prod(vecs, demography, mut_rate=mut_rate)
    249 
    250     # subtract out mass for all ancestral/derived state

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in _expected_sfs_tensor_prod(vecs, demography, mut_rate)
    262 
    263     res = LikelihoodTensorList.compute_sfs(
--> 264         leaf_states, demography)
    265 
    266     return res * mut_rate

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in compute_sfs(cls, leaf_states, demo)
    273         for event in nx.dfs_postorder_nodes(
    274                 demo._event_tree):
--> 275             liklist._process_event(event)
    276         assert len(liklist.likelihood_list) == 1
    277         lik, = liklist.likelihood_list

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in _process_event(self, event)
    297             self._process_merge_subpops_likelihood(event)
    298         elif e_type == 'merge_clusters':
--> 299             self._process_merge_clusters_likelihood(event)
    300         elif e_type == 'pulse':
    301             self._process_pulse_likelihood(event)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in _process_merge_clusters_likelihood(self, event)
    364 
    365         assert not self._in_same_lik(*child_pops)
--> 366         self._merge_pops(parent_pop, child_pops)
    367 
    368     def _process_merge_subpops_likelihood(self, event):

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in _merge_pops(self, newpopname, child_pops, n)
    347        else:
    348            self.likelihood_list.remove(lik2)
--> 349            lik1.convolve_trailing_axes(lik2)
    350 
    351        lik1.mul_trailing_binoms(divide=True)

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/compute_sfs.py in convolve_trailing_axes(self, other)
    470 
    471         assert reshaped0.lik.shape[0] == reshaped1.lik.shape[0]
--> 472         convolved = convolve_trailing_axes(reshaped0.lik, reshaped1.lik)
    473         self.liks = np.reshape(
    474             convolved,

~/anaconda3/envs/momi/lib/python3.7/site-packages/momi/math_functions.py in convolve_trailing_axes(A, B)
     15     A = np.reshape(A, list(A.shape) + [1])
     16     B = np.reshape(B, list(B.shape) + [1])
---> 17     return convolve_sum_axes(A, B)
     18 
     19 convolve_sum_axes = primitive(convolve_sum_axes)

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     42             parents = tuple(box._node for _     , box in boxed_args)
     43             argnums = tuple(argnum    for argnum, _   in boxed_args)
---> 44             ans = f_wrapped(*argvals, **kwargs)
     45             node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
     46             return new_box(ans, trace, node)

~/anaconda3/envs/momi/lib/python3.7/site-packages/autograd/tracer.py in f_wrapped(*args, **kwargs)
     46             return new_box(ans, trace, node)
     47         else:
---> 48             return f_raw(*args, **kwargs)
     49     f_wrapped.fun = f_raw
     50     f_wrapped._is_autograd_primitive = True

momi/convolution.pyx in momi.convolution.convolve_sum_axes()

MemoryError: Unable to allocate 4.18 TiB for an array with shape (2005, 103, 15375, 181) and data type float64

Unix 5.4.0-52-generic Python 3.7.9 (default, Aug 31 2020, 12:42:55) [GCC 7.3.0] numpy 1.19.5

[jackkamm]

Yes, I think the model may be too large. It looks like you have 6 populations linked by migration, and if each population has n samples, this will require creating an array of size n^7 at some point.

mem_chunk_size affects the number of SNPs simultaneously evaluated, but at a certain point even evaluating a single SNP at a time is too memory intensive.

One thing you can try is to make n smaller, say n=2 per population. You can use the SnpAlleleCounts.down_sample() method to do this: https://momi2.readthedocs.io/en/latest/api.html#momi.SnpAlleleCounts.down_sample

[Erythroxylum]

Well that is a helpful comment since n=20 for each pop. I will try it out. Thanks!