pyGSTi 0.9.12__cp310-cp310-win_amd64.whl → 0.9.13__cp310-cp310-win_amd64.whl

pygsti/layouts/maplayout.py

@@ -11,10 +11,12 @@ Defines the MapCOPALayout class.
 #***************************************************************************************************
 
 import collections as _collections
+import importlib as _importlib
+import numpy as _np
 
 from pygsti.layouts.distlayout import DistributableCOPALayout as _DistributableCOPALayout
 from pygsti.layouts.distlayout import _DistributableAtom
-from pygsti.layouts.prefixtable import PrefixTable as _PrefixTable
+from pygsti.layouts.prefixtable import PrefixTable as _PrefixTable, PrefixTableJacobian as _PrefixTableJacobian
 from pygsti.circuits.circuitlist import CircuitList as _CircuitList
 from pygsti.tools import listtools as _lt
 
@@ -51,19 +53,39 @@ class _MapCOPALayoutAtom(_DistributableAtom):
     """
 
     def __init__(self, unique_complete_circuits, ds_circuits, group, model,
-                 dataset, max_cache_size):
+                 dataset, max_cache_size, 
+                 circuit_param_dependencies=None, param_circuit_dependencies=None, 
+                 expanded_complete_circuit_cache = None):
+        
+        expanded_circuit_info_by_unique = dict()
+        expanded_circuit_set = dict() # only use SeparatePOVMCircuit keys as ordered set
+
+        if expanded_complete_circuit_cache is None:
+            expanded_complete_circuit_cache = dict()
 
-        expanded_circuit_info_by_unique = _collections.OrderedDict()
-        expanded_circuit_set = _collections.OrderedDict()  # only use SeparatePOVMCircuit keys as ordered set
         for i in group:
-            observed_outcomes = None if (dataset is None) else dataset[ds_circuits[i]].outcomes
-            d = unique_complete_circuits[i].expand_instruments_and_separate_povm(model, observed_outcomes)
+            d = expanded_complete_circuit_cache.get(unique_complete_circuits[i], None)
+            if d is None:
+                unique_observed_outcomes = None if (dataset is None) else dataset[ds_circuits[i]].unique_outcomes
+                d = model.expand_instruments_and_separate_povm(unique_complete_circuits[i], unique_observed_outcomes)
             expanded_circuit_info_by_unique[i] = d  # a dict of SeparatePOVMCircuits => tuples of outcome labels
-            expanded_circuit_set.update(d)
+            expanded_circuit_set.update(d)            
 
         expanded_circuits = list(expanded_circuit_set.keys())
+        
         self.table = _PrefixTable(expanded_circuits, max_cache_size)
 
+        #only Build the Jacobian prefix table if we are using the generic evotype.
+        if model.sim.calclib is _importlib.import_module("pygsti.forwardsims.mapforwardsim_calc_generic"):
+            #create a list for storing the model parameter dependencies of expanded circuits
+            expanded_param_circuit_depend = [{} for _ in range(len(param_circuit_dependencies))]
+            for i in group:
+                for param_idx in circuit_param_dependencies[i]:
+                    expanded_param_circuit_depend[param_idx].update(expanded_circuit_info_by_unique[i])
+            expanded_param_circuit_depend = [list(param_circuit_depend_dict.keys()) for  param_circuit_depend_dict in expanded_param_circuit_depend]
+
+            self.jac_table = _PrefixTableJacobian(expanded_circuits, max_cache_size, expanded_param_circuit_depend)
+
         #Create circuit element <=> integer index lookups for speed
         all_rholabels = set()
         all_oplabels = set()
@@ -91,8 +113,7 @@ class _MapCOPALayoutAtom(_DistributableAtom):
         self.outcomes_by_expcircuit = {}
         self.povm_and_elbls_by_expcircuit = {}
 
-        elindex_outcome_tuples = _collections.OrderedDict([
-            (unique_i, list()) for unique_i in range(len(unique_complete_circuits))])
+        elindex_outcome_tuples = {unique_i: list() for unique_i in range(len(unique_complete_circuits))}
 
         #Assign element indices, "global" indices starting at `offset`
         local_offset = 0
@@ -198,42 +219,99 @@ class MapCOPALayout(_DistributableCOPALayout):
 
     resource_alloc : ResourceAllocation, optional
         The resources available for computing circuit outcome probabilities.
+    
+    circuit_partition_cost_functions : tuple of str, optional (default ('size', 'propagations'))
+        A tuple of strings denoting cost function to use in each of the two stages of the algorithm
+        for determining the partitions of the complete circuit set amongst atoms.
+        Allowed options are 'size', which corresponds to balancing the number of circuits, 
+        and 'propagations', which corresponds to balancing the number of state propagations.
 
     verbosity : int or VerbosityPrinter
         Determines how much output to send to stdout.  0 means no output, higher
         integers mean more output.
+
+    layout_creation_circuit_cache : dict, optional (default None)
+        An optional dictionary containing pre-computed circuit structures/modifications which
+        can be used to reduce the overhead of repeated circuit operations during layout creation.
+    
+    load_balancing_parameters : tuple of floats, optional (default (1.2, .1))
+        A tuple of floats used as load balancing parameters when splitting a layout across atoms,
+        as in the multi-processor setting when using MPI. These parameters correspond to the `imbalance_threshold`
+        and `minimum_improvement_threshold` parameters described in the method `find_splitting_new`
+        of the `PrefixTable` class.
     """
 
     def __init__(self, circuits, model, dataset=None, max_cache_size=None,
                  num_sub_tables=None, num_table_processors=1, num_param_dimension_processors=(),
-                 param_dimensions=(), param_dimension_blk_sizes=(), resource_alloc=None, verbosity=0):
+                 param_dimensions=(), param_dimension_blk_sizes=(), resource_alloc=None, 
+                 circuit_partition_cost_functions=('size', 'propagations'), verbosity=0, 
+                 layout_creation_circuit_cache=None, load_balancing_parameters = (1.2, .1)):
 
         unique_circuits, to_unique = self._compute_unique_circuits(circuits)
         aliases = circuits.op_label_aliases if isinstance(circuits, _CircuitList) else None
         ds_circuits = _lt.apply_aliases_to_circuits(unique_circuits, aliases)
-        unique_complete_circuits = [model.complete_circuit(c) for c in unique_circuits]
-        unique_povmless_circuits = [model.split_circuit(c, split_prep=False)[1] for c in unique_complete_circuits]
+
+        #extract subcaches from layout_creation_circuit_cache:
+        if layout_creation_circuit_cache is None:
+            layout_creation_circuit_cache = dict()
+        self.completed_circuit_cache = layout_creation_circuit_cache.get('completed_circuits', None)
+        self.split_circuit_cache = layout_creation_circuit_cache.get('split_circuits', None)
+        self.expanded_and_separated_circuits_cache = layout_creation_circuit_cache.get('expanded_and_separated_circuits', None)
+        
+        if self.completed_circuit_cache is None:
+            unique_complete_circuits = model.complete_circuits(unique_circuits)
+            split_circuits = model.split_circuits(unique_complete_circuits, split_prep=False)
+        else:
+            unique_complete_circuits = []
+            for c in unique_circuits:
+                comp_ckt = self.completed_circuit_cache.get(c, None)
+                if comp_ckt is not None:
+                    unique_complete_circuits.append(comp_ckt)
+                else:
+                    unique_complete_circuits.append(model.complete_circuit(c))
+            split_circuits = []
+            for c, c_complete in zip(unique_circuits,unique_complete_circuits):
+                split_ckt = self.split_circuit_cache.get(c, None)
+                if split_ckt is not None:
+                    split_circuits.append(split_ckt)
+                else:
+                    split_circuits.append(model.split_circuit(c_complete, split_prep=False))
+
+        #construct a map for the parameter dependence for each of the unique_complete_circuits.
+        #returns a dictionary who's keys are the unique completed circuits, and whose
+        #values are lists of model parameters upon which that circuit depends.
+        if model.sim.calclib is _importlib.import_module("pygsti.forwardsims.mapforwardsim_calc_generic") and model.param_interposer is None:
+            circ_param_map, param_circ_map = model.circuit_parameter_dependence(unique_complete_circuits, return_param_circ_map=True)
+            uniq_comp_circs_param_depend = list(circ_param_map.values())
+            uniq_comp_param_circs_depend = param_circ_map
+        else : 
+            circ_param_map = None
+            param_circ_map = None
+            uniq_comp_circs_param_depend = None
+            uniq_comp_param_circs_depend = None
+        #construct list of unique POVM-less circuits.
+        unique_povmless_circuits = [ckt_tup[1] for ckt_tup in split_circuits]
 
         max_sub_table_size = None  # was an argument but never used; remove in future
         if (num_sub_tables is not None and num_sub_tables > 1) or max_sub_table_size is not None:
             circuit_table = _PrefixTable(unique_povmless_circuits, max_cache_size)
-            groups = circuit_table.find_splitting(max_sub_table_size, num_sub_tables, verbosity=verbosity)
+            self.complete_circuit_table = circuit_table
+            groups = circuit_table.find_splitting_new(max_sub_table_size, num_sub_tables, verbosity=verbosity,
+                                                      initial_cost_metric=circuit_partition_cost_functions[0],
+                                                      rebalancing_cost_metric=circuit_partition_cost_functions[1],
+                                                      imbalance_threshold = load_balancing_parameters[0],
+                                                      minimum_improvement_threshold = load_balancing_parameters[1])
+            #groups = circuit_table.find_splitting(max_sub_table_size, num_sub_tables, verbosity=verbosity)
         else:
-            groups = [set(range(len(unique_complete_circuits)))]
-
-        #atoms = []
-        #elindex_outcome_tuples = _collections.OrderedDict(
-        #    [(unique_i, list()) for unique_i in range(len(unique_circuits))])
-
-        #offset = 0
-        #for group in groups:
-        #    atoms.append(_MapCOPALayoutAtom(unique_complete_circuits, ds_circuits, to_orig, group,
-        #                                    model, dataset, offset, elindex_outcome_tuples, max_cache_size))
-        #    offset += atoms[-1].num_elements
+            groups = [list(range(len(unique_complete_circuits)))]
+        
 
         def _create_atom(group):
             return _MapCOPALayoutAtom(unique_complete_circuits, ds_circuits, group,
-                                      model, dataset, max_cache_size)
+                                      model, dataset, max_cache_size,
+                                      circuit_param_dependencies= uniq_comp_circs_param_depend,
+                                      param_circuit_dependencies= uniq_comp_param_circs_depend,
+                                      expanded_complete_circuit_cache=self.expanded_and_separated_circuits_cache)
 
         super().__init__(circuits, unique_circuits, to_unique, unique_complete_circuits,
                          _create_atom, groups, num_table_processors,

pygsti/layouts/matrixlayout.py

@@ -61,6 +61,9 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
     model : Model
         The model being used to construct this layout.  Used for expanding instruments
         within the circuits.
+    
+    unique_circuits : list of Circuits
+        A list of the unique :class:`Circuit` objects representing the circuits this layout will include.
 
     dataset : DataSet
         The dataset, used to include only observed circuit outcomes in this atom
@@ -68,7 +71,11 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
     """
 
     def __init__(self, unique_complete_circuits, unique_nospam_circuits, circuits_by_unique_nospam_circuits,
-                 ds_circuits, group, helpful_scratch, model, dataset):
+                 ds_circuits, group, helpful_scratch, model, unique_circuits, dataset=None, expanded_and_separated_circuit_cache=None,
+                 double_expanded_nospam_circuits_cache = None):
+
+        if expanded_and_separated_circuit_cache is None:
+            expanded_and_separated_circuit_cache = dict()
 
         #Note: group gives unique_nospam_circuits indices, which circuits_by_unique_nospam_circuits
         # turns into "unique complete circuit" indices, which the layout via it's to_unique can map
@@ -78,24 +85,24 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
             for i in indices:
                 nospam_c = unique_nospam_circuits[i]
                 for unique_i in circuits_by_unique_nospam_circuits[nospam_c]:  # "unique" circuits: add SPAM to nospam_c
-                    observed_outcomes = None if (dataset is None) else dataset[ds_circuits[unique_i]].unique_outcomes
-                    expc_outcomes = unique_complete_circuits[unique_i].expand_instruments_and_separate_povm(
-                        model, observed_outcomes)
-                    #Note: unique_complete_circuits may have duplicates (they're only unique *pre*-completion)
-
+                    #the cache is indexed into using the (potentially) incomplete circuits
+                    expc_outcomes = expanded_and_separated_circuit_cache.get(unique_circuits[unique_i], None)
+                    if expc_outcomes is None: #fall back on original non-cache behavior.
+                        observed_outcomes = None if (dataset is None) else dataset[ds_circuits[unique_i]].unique_outcomes
+                        expc_outcomes = model.expand_instruments_and_separate_povm(unique_complete_circuits[unique_i], observed_outcomes)
+                        #and add this new value to the cache.
+                        expanded_and_separated_circuit_cache[unique_circuits[unique_i]] = expc_outcomes 
                     for sep_povm_c, outcomes in expc_outcomes.items():  # for each expanded cir from unique_i-th circuit
                         prep_lbl = sep_povm_c.circuit_without_povm[0]
                         exp_nospam_c = sep_povm_c.circuit_without_povm[1:]  # sep_povm_c *always* has prep lbl
                         spam_tuples = [(prep_lbl, elabel) for elabel in sep_povm_c.full_effect_labels]
-                        outcome_by_spamtuple = _collections.OrderedDict([(st, outcome)
-                                                                         for st, outcome in zip(spam_tuples, outcomes)])
+                        outcome_by_spamtuple = {st:outcome for st, outcome in zip(spam_tuples, outcomes)}
 
                         #Now add these outcomes to `expanded_nospam_circuit_outcomes` - note that multiple "unique_i"'s
                         # may exist for the same expanded & without-spam circuit (exp_nospam_c) and so we need to
                         # keep track of a list of unique_i indices for each circut and spam tuple below.
                         if exp_nospam_c not in _expanded_nospam_circuit_outcomes:
-                            _expanded_nospam_circuit_outcomes[exp_nospam_c] = _collections.OrderedDict(
-                                [(st, (outcome, [unique_i])) for st, outcome in zip(spam_tuples, outcomes)])
+                            _expanded_nospam_circuit_outcomes[exp_nospam_c] = {st:(outcome, [unique_i]) for st, outcome in zip(spam_tuples, outcomes)}
                         else:
                             for st, outcome in outcome_by_spamtuple.items():
                                 if st in _expanded_nospam_circuit_outcomes[exp_nospam_c]:
@@ -109,25 +116,28 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
 
         # keys = expanded circuits w/out SPAM layers; values = spamtuple => (outcome, unique_is) dictionary that
         # keeps track of which "unique" circuit indices having each spamtuple / outcome.
-        expanded_nospam_circuit_outcomes = _collections.OrderedDict()
+        expanded_nospam_circuit_outcomes = dict()
         add_expanded_circuits(group, expanded_nospam_circuit_outcomes)
-        expanded_nospam_circuits = _collections.OrderedDict(
-            [(i, cir) for i, cir in enumerate(expanded_nospam_circuit_outcomes.keys())])
+        expanded_nospam_circuits = {i:cir for i, cir in enumerate(expanded_nospam_circuit_outcomes.keys())}
 
         # add suggested scratch to the "final" elements as far as the tree creation is concerned
         # - this allows these scratch element to help balance the tree.
-        expanded_nospam_circuit_outcomes_plus_scratch = expanded_nospam_circuit_outcomes.copy()
-        add_expanded_circuits(helpful_scratch, expanded_nospam_circuit_outcomes_plus_scratch)
-        expanded_nospam_circuits_plus_scratch = _collections.OrderedDict(
-            [(i, cir) for i, cir in enumerate(expanded_nospam_circuit_outcomes_plus_scratch.keys())])
-
-        double_expanded_nospam_circuits_plus_scratch = _collections.OrderedDict()
+        if helpful_scratch:
+            expanded_nospam_circuit_outcomes_plus_scratch = expanded_nospam_circuit_outcomes.copy()
+            add_expanded_circuits(helpful_scratch, expanded_nospam_circuit_outcomes_plus_scratch)
+            expanded_nospam_circuits_plus_scratch = {i:cir for i, cir in enumerate(expanded_nospam_circuit_outcomes_plus_scratch.keys())}
+        else:
+            expanded_nospam_circuits_plus_scratch = expanded_nospam_circuits.copy()
+        
+        if double_expanded_nospam_circuits_cache is None:
+            double_expanded_nospam_circuits_cache = dict()
+        double_expanded_nospam_circuits_plus_scratch = dict()
         for i, cir in expanded_nospam_circuits_plus_scratch.items():
-            cir = cir.copy(editable=True)
-            cir.expand_subcircuits()  # expand sub-circuits for a more efficient tree
-            cir.done_editing()
-            double_expanded_nospam_circuits_plus_scratch[i] = cir
-
+            # expand sub-circuits for a more efficient tree
+            double_expanded_ckt = double_expanded_nospam_circuits_cache.get(cir, None)
+            if double_expanded_ckt is None: #Fall back to standard behavior and do expansion.
+                double_expanded_ckt = cir.expand_subcircuits()
+            double_expanded_nospam_circuits_plus_scratch[i] = double_expanded_ckt
         self.tree = _EvalTree.create(double_expanded_nospam_circuits_plus_scratch)
         #print("Atom tree: %d circuits => tree of size %d" % (len(expanded_nospam_circuits), len(self.tree)))
 
@@ -138,7 +148,7 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
         # quantity plus a spam-tuple. We order the final indices so that all the outcomes corresponding to a
         # given spam-tuple are contiguous.
 
-        tree_indices_by_spamtuple = _collections.OrderedDict()  # "tree" indices index expanded_nospam_circuits
+        tree_indices_by_spamtuple = dict()  # "tree" indices index expanded_nospam_circuits
         for i, c in expanded_nospam_circuits.items():
             for spam_tuple in expanded_nospam_circuit_outcomes[c].keys():
                 if spam_tuple not in tree_indices_by_spamtuple: tree_indices_by_spamtuple[spam_tuple] = []
@@ -147,7 +157,7 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
         #Assign element indices, starting at `offset`
         # now that we know how many of each spamtuple there are, assign final element indices.
         local_offset = 0
-        self.indices_by_spamtuple = _collections.OrderedDict()  # values are (element_indices, tree_indices) tuples.
+        self.indices_by_spamtuple = dict()  # values are (element_indices, tree_indices) tuples.
         for spam_tuple, tree_indices in tree_indices_by_spamtuple.items():
             self.indices_by_spamtuple[spam_tuple] = (slice(local_offset, local_offset + len(tree_indices)),
                                                      _slct.list_to_slice(tree_indices, array_ok=True))
@@ -157,8 +167,7 @@ class _MatrixCOPALayoutAtom(_DistributableAtom):
         element_slice = None  # slice(offset, offset + local_offset)  # *global* (of parent layout) element-index slice
         num_elements = local_offset
 
-        elindex_outcome_tuples = _collections.OrderedDict([
-            (unique_i, list()) for unique_i in range(len(unique_complete_circuits))])
+        elindex_outcome_tuples = {unique_i: list() for unique_i in range(len(unique_complete_circuits))}
 
         for spam_tuple, (element_indices, tree_indices) in self.indices_by_spamtuple.items():
             for elindex, tree_index in zip(_slct.indices(element_indices), _slct.to_array(tree_indices)):
@@ -231,7 +240,7 @@ class MatrixCOPALayout(_DistributableCOPALayout):
     Parameters
     ----------
     circuits : list
-        A list of:class:`Circuit` objects representing the circuits this layout will include.
+        A list of :class:`Circuit` objects representing the circuits this layout will include.
 
     model : Model
         The model that will be used to compute circuit outcome probabilities using this layout.
@@ -271,11 +280,16 @@ class MatrixCOPALayout(_DistributableCOPALayout):
     verbosity : int or VerbosityPrinter
         Determines how much output to send to stdout.  0 means no output, higher
         integers mean more output.
+
+    layout_creation_circuit_cache : dict, optional (default None)
+        A precomputed dictionary serving as a cache for completed
+        circuits. I.e. circuits with prep labels and POVM labels appended.
     """
 
     def __init__(self, circuits, model, dataset=None, num_sub_trees=None, num_tree_processors=1,
                  num_param_dimension_processors=(), param_dimensions=(),
-                 param_dimension_blk_sizes=(), resource_alloc=None, verbosity=0):
+                 param_dimension_blk_sizes=(), resource_alloc=None, verbosity=0, 
+                 layout_creation_circuit_cache = None):
 
         #OUTDATED: TODO - revise this:
         # 1. pre-process => get complete circuits => spam-tuples list for each no-spam circuit (no expanding yet)
@@ -290,57 +304,82 @@ class MatrixCOPALayout(_DistributableCOPALayout):
         unique_circuits, to_unique = self._compute_unique_circuits(circuits)
         aliases = circuits.op_label_aliases if isinstance(circuits, _CircuitList) else None
         ds_circuits = _lt.apply_aliases_to_circuits(unique_circuits, aliases)
-        unique_complete_circuits = [model.complete_circuit(c) for c in unique_circuits]
+
+        #extract subcaches from layout_creation_circuit_cache:
+        if layout_creation_circuit_cache is None:
+            layout_creation_circuit_cache = dict()
+        self.completed_circuit_cache = layout_creation_circuit_cache.get('completed_circuits', None)
+        self.split_circuit_cache = layout_creation_circuit_cache.get('split_circuits', None)
+        self.expanded_and_separated_circuits_cache = layout_creation_circuit_cache.get('expanded_and_separated_circuits', None)
+        self.expanded_subcircuits_no_spam_cache = layout_creation_circuit_cache.get('expanded_subcircuits_no_spam', None)
+        
+        if self.completed_circuit_cache is None:
+            unique_complete_circuits, split_unique_circuits = model.complete_circuits(unique_circuits, return_split=True)
+        else:
+            unique_complete_circuits = []
+            for c in unique_circuits:
+                comp_ckt = self.completed_circuit_cache.get(c, None)
+                if comp_ckt is not None:
+                    unique_complete_circuits.append(comp_ckt)
+                else:
+                    unique_complete_circuits.append(model.complete_circuit(c))
         #Note: "unique" means a unique circuit *before* circuit-completion, so there could be duplicate
         # "unique circuits" after completion, e.g. "rho0Gx" and "Gx" could both complete to "rho0GxMdefault_0".
 
-        circuits_by_unique_nospam_circuits = _collections.OrderedDict()
-        for i, c in enumerate(unique_complete_circuits):
-            _, nospam_c, _ = model.split_circuit(c)
-            if nospam_c in circuits_by_unique_nospam_circuits:
-                circuits_by_unique_nospam_circuits[nospam_c].append(i)
-            else:
-                circuits_by_unique_nospam_circuits[nospam_c] = [i]
-        unique_nospam_circuits = list(circuits_by_unique_nospam_circuits.keys())
+        circuits_by_unique_nospam_circuits = dict()
+        if self.completed_circuit_cache is None:
+            for i, (_, nospam_c, _) in enumerate(split_unique_circuits):
+                if nospam_c in circuits_by_unique_nospam_circuits:
+                    circuits_by_unique_nospam_circuits[nospam_c].append(i)
+                else:
+                    circuits_by_unique_nospam_circuits[nospam_c] = [i]
+            #also create the split circuit cache at this point for future use.
+            if self.split_circuit_cache is None:
+                self.split_circuit_cache = {unique_ckt:split_ckt for unique_ckt, split_ckt in zip(unique_circuits, split_unique_circuits)}
 
+        else:
+            if self.split_circuit_cache is None:
+                self.split_circuit_cache = dict()
+            for i, (c_unique_complete, c_unique) in enumerate(zip(unique_complete_circuits, unique_circuits)):
+                split_ckt_tup = self.split_circuit_cache.get(c_unique, None)
+                nospam_c= split_ckt_tup[1] if split_ckt_tup is not None else None
+                if nospam_c is None:
+                    split_ckt_tup = model.split_circuit(c_unique_complete)
+                    nospam_c= split_ckt_tup[1]
+                    #also add this missing circuit to the cache for future use.
+                    self.split_circuit_cache[c_unique] = split_ckt_tup
+                if nospam_c in circuits_by_unique_nospam_circuits:
+                    circuits_by_unique_nospam_circuits[nospam_c].append(i)
+                else:
+                    circuits_by_unique_nospam_circuits[nospam_c] = [i]
+
+        unique_nospam_circuits = list(circuits_by_unique_nospam_circuits.keys())
+        
         # Split circuits into groups that will make good subtrees (all procs do this)
         max_sub_tree_size = None  # removed from being an argument (unused)
         if (num_sub_trees is not None and num_sub_trees > 1) or max_sub_tree_size is not None:
             circuit_tree = _EvalTree.create(unique_nospam_circuits)
             groups, helpful_scratch = circuit_tree.find_splitting(len(unique_nospam_circuits),
                                                                   max_sub_tree_size, num_sub_trees, verbosity - 1)
-            #print("%d circuits => tree of size %d" % (len(unique_nospam_circuits), len(circuit_tree)))
         else:
             groups = [set(range(len(unique_nospam_circuits)))]
             helpful_scratch = [set()]
         # (elements of `groups` contain indices into `unique_nospam_circuits`)
 
-        # Divide `groups` into num_tree_processors roughly equal sets (each containing
-        # potentially multiple groups)
-        #my_group_indices, group_owners, grp_subcomm = self._distribute(num_tree_processors, len(groups),
-        #                                                                resource_alloc, verbosity)
-        #my_group_indices = set(my_group_indices)
-
-        #my_atoms = []
-        #elindex_outcome_tuples = _collections.OrderedDict([
-        #    (orig_i, list()) for orig_i in range(len(unique_circuits))])
-        #
-        #offset = 0
-        #for i, (group, helpful_scratch_group) in enumerate(zip(groups, helpful_scratch)):
-        #    if i not in my_group_indices: continue
-        #    my_atoms.append(_MatrixCOPALayoutAtom(unique_complete_circuits, unique_nospam_circuits,
-        #                                       circuits_by_unique_nospam_circuits, ds_circuits,
-        #                                       group, helpful_scratch_group, model, dataset, offset,
-        #                                       elindex_outcome_tuples))
-        #    offset += my_atoms[-1].num_elements
-
         def _create_atom(args):
             group, helpful_scratch_group = args
             return _MatrixCOPALayoutAtom(unique_complete_circuits, unique_nospam_circuits,
                                          circuits_by_unique_nospam_circuits, ds_circuits,
-                                         group, helpful_scratch_group, model, dataset)
+                                         group, helpful_scratch_group, model, 
+                                         unique_circuits, dataset,
+                                         self.expanded_and_separated_circuits_cache,
+                                         self.expanded_subcircuits_no_spam_cache)
 
         super().__init__(circuits, unique_circuits, to_unique, unique_complete_circuits,
                          _create_atom, list(zip(groups, helpful_scratch)), num_tree_processors,
                          num_param_dimension_processors, param_dimensions,
                          param_dimension_blk_sizes, resource_alloc, verbosity)
+
+
+
+