Map single-cell transcriptomes to copy number evolutionary trees.

Running SCATrEx twice with the same random seed yields different results. This does not happen when we do not do any ELBO optimization, and just change the parameters in the move proposals and do ELBO evaluations with those. It is irrespective of local in the run_search being True or False. For some reason, some moves lead to ELBO optimization steps with small divergences between runs, and this accumulates leading to different trees in the end.

Commenting out the ancestor nodes' contribution from the node weights seems to fix it. So it might be a numerical instability problem that XLA does not deal with deterministically?

We get memory overflow when we run too many SCATrEx tree inferences in the same Python session. This only happens when using a GPU, and it occurs on the call to do_grad, which is a jitted function.

Full error trace:

RuntimeError                              Traceback (most recent call last)
<ipython-input-216-9b90a41ab615> in <module>
     47                 'random_seed': 1,
     48                 'verbosity': logging.DEBUG}
---> 49 sca.learn_tree(reset=True, batch_key='batch', search_kwargs=search_kwargs)

/cluster/work/bewi/members/pedrof/miniconda3/envs/py38/lib/python3.8/site-packages/scatrex/scatrex.py in learn_tree(self, observed_tree, reset, cell_filter, filter_genes, max_genes, batch_key, search_kwargs)
    289             logger.info("Will continue search from where it left off.")
    290 
--> 291         self.ntssb = self.search.run_search(**search_kwargs)
    292         self.ntssb.create_augmented_tree_dict()
    293 

/cluster/work/bewi/members/pedrof/miniconda3/envs/py38/lib/python3.8/site-packages/scatrex/ntssb/search.py in run_search(self, n_iters, n_iters_elbo, factor_delay, posterior_delay, global_delay, joint_init, thin, local, num_samples, step_size, verbosity, tol, mb_size, max_nodes, debug, callback, alpha, Tmax, anneal, restart_step, move_weights, weighted, merge_n_tries, opt, search_callback, add_rule, add_rule_thres, random_seed, **callback_kwargs)
    290                 "log_baseline_mean"
    291             ] = init_log_baseline
--> 292             self.tree.optimize_elbo(
    293                 root_node=None,
    294                 sticks_only=True,

/cluster/work/bewi/members/pedrof/miniconda3/envs/py38/lib/python3.8/site-packages/scatrex/ntssb/ntssb.py in optimize_elbo(self, root_node, local_node, global_only, sticks_only, unique_node, num_samples, n_iters, thin, step_size, debug, tol, run, max_nodes, init, opt, opt_triplet, mb_size, callback, **callback_kwargs)
   1543                 # data_mask_subset = data_mask
   1544                 # start = time.time()
-> 1545                 opt_state, g, params, elbo = self.update(
   1546                     obs_params,
   1547                     parent_vector,

/cluster/work/bewi/members/pedrof/miniconda3/envs/py38/lib/python3.8/site-packages/scatrex/ntssb/ntssb.py in update(self, obs_params, parent_vector, children_vector, ancestor_nodes_indices, tssb_indices, previous_branches_indices, tssb_weights, dp_alphas, dp_gammas, node_mask, data_mask_subset, indices, do_global, global_only, sticks_only, num_samples, i, opt_state, opt_update, get_params)
   1259         # print("Recompiling update!")
   1260         params = get_params(opt_state)
-> 1261         value, gradient = self.do_grad(
   1262             obs_params,
   1263             parent_vector,

    [... skipping hidden 8 frame]

RuntimeError: INTERNAL: Failed to load in-memory CUBIN: CUDA_ERROR_OUT_OF_MEMORY: out of memory