def test_marginal_counts_no_cregs(self): """Test that marginal_counts without cregs See qiskit-terra/6430.""" raw_counts_1 = { "0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0x12": 8 } data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1)) exp_result_header_1 = QobjExperimentHeader(memory_slots=5) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) result = Result(results=[exp_result_1], **self.base_result_args) _ = marginal_counts(result, indices=[0]) marginal_counts_result = marginal_counts(result, indices=[0]) self.assertEqual(marginal_counts_result.get_counts(), { "0": 27, "1": 27 })
def test_marginal_counts(self): """Test that counts are marginalized correctly.""" raw_counts = { "0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0xE": 8 } data = models.ExperimentResultData(counts=dict(**raw_counts)) exp_result_header = QobjExperimentHeader(creg_sizes=[["c0", 4]], memory_slots=4) exp_result = models.ExperimentResult(shots=54, success=True, data=data, header=exp_result_header) result = Result(results=[exp_result], **self.base_result_args) expected_marginal_counts = {"00": 4, "01": 27, "10": 23} self.assertEqual(marginal_counts(result.get_counts(), [0, 1]), expected_marginal_counts) self.assertEqual(marginal_counts(result.get_counts(), [1, 0]), expected_marginal_counts)
def test_marginal_counts(self): """Test that counts are marginalized correctly.""" raw_counts = { '0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0xE': 8 } data = models.ExperimentResultData(counts=dict(**raw_counts)) exp_result_header = QobjExperimentHeader(creg_sizes=[['c0', 4]], memory_slots=4) exp_result = models.ExperimentResult(shots=54, success=True, data=data, header=exp_result_header) result = Result(results=[exp_result], **self.base_result_args) expected_marginal_counts = {'00': 4, '01': 27, '10': 23} self.assertEqual(marginal_counts(result.get_counts(), [0, 1]), expected_marginal_counts) self.assertEqual(marginal_counts(result.get_counts(), [1, 0]), expected_marginal_counts)
def test_marginal_counts_with_dict(self): """Test the marginal_counts method with dictionary instead of Result object. """ dict_counts_1 = { '0000': 4, '0001': 7, '0010': 10, '0110': 5, '1001': 11, '1101': 9, '1110': 8 } dict_counts_2 = {'10': 5, '11': 8} expected_marginal_counts_1 = {'00': 4, '01': 27, '10': 23} expected_marginal_counts_2 = {'0': 5, '1': 8} self.assertEqual(marginal_counts(dict_counts_1, [0, 1]), expected_marginal_counts_1) self.assertEqual(marginal_counts(dict_counts_2, [0], inplace=True), expected_marginal_counts_2) self.assertNotEqual(dict_counts_2, expected_marginal_counts_2) self.assertRaises( AttributeError, lambda: marginal_counts(dict_counts_1, [0, 1]).get_counts(0))
def test_marginal_counts_result(self): """Test that a Result object containing counts marginalizes correctly.""" raw_counts_1 = { "0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0xE": 8 } data_1 = models.ExperimentResultData(counts=raw_counts_1) exp_result_header_1 = QobjExperimentHeader(creg_sizes=[["c0", 4]], memory_slots=4) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) raw_counts_2 = {"0x2": 5, "0x3": 8} data_2 = models.ExperimentResultData(counts=raw_counts_2) exp_result_header_2 = QobjExperimentHeader(creg_sizes=[["c0", 2]], memory_slots=2) exp_result_2 = models.ExperimentResult(shots=13, success=True, data=data_2, header=exp_result_header_2) result = Result(results=[exp_result_1, exp_result_2], **self.base_result_args) expected_marginal_counts_1 = {"00": 4, "01": 27, "10": 23} expected_marginal_counts_2 = {"0": 5, "1": 8} expected_marginal_counts_none = { "0000": 4, "0001": 7, "0010": 10, "0110": 5, "1001": 11, "1101": 9, "1110": 8, } self.assertEqual( marginal_counts(result, [0, 1]).get_counts(0), expected_marginal_counts_1) self.assertEqual( marginal_counts(result, [0]).get_counts(1), expected_marginal_counts_2) self.assertEqual( marginal_counts(result, None).get_counts(0), expected_marginal_counts_none)
def test_marginal_counts_result_creg_sizes(self): """Test that marginal_counts with Result input properly changes creg_sizes.""" raw_counts = { "0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0xE": 8 } data = models.ExperimentResultData(counts=dict(**raw_counts)) exp_result_header = QobjExperimentHeader(creg_sizes=[["c0", 1], ["c1", 3]], memory_slots=4) exp_result = models.ExperimentResult(shots=54, success=True, data=data, header=exp_result_header) result = Result(results=[exp_result], **self.base_result_args) expected_marginal_counts = {"0 0": 14, "0 1": 18, "1 0": 13, "1 1": 9} expected_creg_sizes = [["c0", 1], ["c1", 1]] expected_memory_slots = 2 marginal_counts_result = marginal_counts(result, [0, 2]) self.assertEqual(marginal_counts_result.results[0].header.creg_sizes, expected_creg_sizes) self.assertEqual(marginal_counts_result.results[0].header.memory_slots, expected_memory_slots) self.assertEqual(marginal_counts_result.get_counts(0), expected_marginal_counts)
def test_marginal_counts_result_format(self): """Test that marginal_counts with format_marginal true properly formats output.""" raw_counts_1 = { "0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0x12": 8 } data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1)) exp_result_header_1 = QobjExperimentHeader(creg_sizes=[["c0", 2], ["c1", 3]], memory_slots=5) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) result = Result(results=[exp_result_1], **self.base_result_args) expected_marginal_counts_1 = { "0_0 _0": 14, "0_0 _1": 18, "0_1 _0": 5, "0_1 _1": 9, "1_0 _0": 8, } marginal_counts_result = marginal_counts(result.get_counts(), [0, 2, 4], format_marginal=True) self.assertEqual(marginal_counts_result, expected_marginal_counts_1)
def _component_data(self, composite_data: List[Dict]) -> List[List[Dict]]: """Return marginalized data for component experiments""" # Marginalize data marginalized_data = {} for datum in composite_data: metadata = datum.get("metadata", {}) # Add marginalized data to sub experiments if "composite_clbits" in metadata: composite_clbits = metadata["composite_clbits"] else: composite_clbits = None for i, index in enumerate(metadata["composite_index"]): if index not in marginalized_data: # Initialize data list for marginalized marginalized_data[index] = [] sub_data = {"metadata": metadata["composite_metadata"][i]} if "counts" in datum: if composite_clbits is not None: sub_data["counts"] = marginal_counts( datum["counts"], composite_clbits[i]) else: sub_data["counts"] = datum["counts"] if "memory" in datum: if composite_clbits is not None: sub_data["memory"] = (np.array( datum["memory"])[composite_clbits[i]]).tolist() else: sub_data["memory"] = datum["memory"] marginalized_data[index].append(sub_data) # Sort by index return [marginalized_data[i] for i in sorted(marginalized_data.keys())]
def format_result(data_index, chunk): """Create new result object from partial result and marginalize.""" new_result = deepcopy(result) new_result.results = [] new_result.results.extend(result.results[data_index:data_index + chunk]) return marginal_counts(new_result, __reserved_registers)
def test_marginal_counts_result_memory_indices_None(self): """Test that a Result object containing memory marginalizes correctly.""" result = self.generate_qiskit_result() memory = "should not be touched" result.results[0].data.memory = memory marginal_result = marginal_counts(result, indices=None) marginal_memory = marginal_result.results[0].data.memory self.assertEqual(marginal_memory, memory)
def test_marginal_counts_result_inplace(self): """Test that a Result object containing memory marginalizes correctly inplace.""" result = self.generate_qiskit_result() marginal_result = marginal_counts(result, indices=[0], inplace=True) self.assertEqual(id(result), id(marginal_result)) marginal_memory = marginal_result.results[0].data.memory self.assertEqual(marginal_memory, [hex(ii % 2) for ii in range(8)])
def test_marginal_counts_result_memory_nonzero_indices(self): """Test that a Result object containing memory marginalizes correctly.""" result = self.generate_qiskit_result() index = 2 marginal_result = marginal_counts(result, indices=[index]) marginal_memory = marginal_result.results[0].data.memory mask = 1 << index expected = [hex((ii & mask) >> index) for ii in range(8)] self.assertEqual(marginal_memory, expected)
def marginalizeCounts(counts, qubits): ''' Input: List of buckets each of which holds multiple dicts denoting the counts Output: All those dicts marginalized...i.e. holding 2 qubit keys instead of 7 qubit ones ''' qubits = [i for i in range(qubits)] for i in range(len(counts)): for j in range(len(counts[i])): counts[i][j] = result.marginal_counts(counts[i][j], qubits) return counts
def test_marginal_counts_result(self): """Test that a Result object containing counts marginalizes correctly.""" raw_counts_1 = { '0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0xE': 8 } data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1)) exp_result_header_1 = QobjExperimentHeader(creg_sizes=[['c0', 4]], memory_slots=4) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) raw_counts_2 = {'0x2': 5, '0x3': 8} data_2 = models.ExperimentResultData(counts=dict(**raw_counts_2)) exp_result_header_2 = QobjExperimentHeader(creg_sizes=[['c0', 2]], memory_slots=2) exp_result_2 = models.ExperimentResult(shots=13, success=True, data=data_2, header=exp_result_header_2) result = Result(results=[exp_result_1, exp_result_2], **self.base_result_args) expected_marginal_counts_1 = {'00': 4, '01': 27, '10': 23} expected_marginal_counts_2 = {'0': 5, '1': 8} self.assertEqual( marginal_counts(result, [0, 1]).get_counts(0), expected_marginal_counts_1) self.assertEqual( marginal_counts(result, [0]).get_counts(1), expected_marginal_counts_2)
def test_marginal_counts_with_dict(self): """Test the marginal_counts method with dictionary instead of Result object.""" dict_counts_1 = { "0000": 4, "0001": 7, "0010": 10, "0110": 5, "1001": 11, "1101": 9, "1110": 8, } dict_counts_2 = {"10": 5, "11": 8} expected_marginal_counts_1 = {"00": 4, "01": 27, "10": 23} expected_marginal_counts_2 = {"0": 5, "1": 8} self.assertEqual(marginal_counts(dict_counts_1, [0, 1]), expected_marginal_counts_1) self.assertEqual(marginal_counts(dict_counts_2, [0], inplace=True), expected_marginal_counts_2) self.assertNotEqual(dict_counts_2, expected_marginal_counts_2) self.assertRaises( AttributeError, lambda: marginal_counts(dict_counts_1, [0, 1]).get_counts(0))
def test_marginal_counts_result_format(self): """Test that marginal_counts with format_marginal true properly formats output.""" raw_counts_1 = {'0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0x12': 8} data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1)) exp_result_header_1 = QobjExperimentHeader(creg_sizes=[['c0', 2], ['c1', 3]], memory_slots=5) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) result = Result(results=[exp_result_1], **self.base_result_args) expected_marginal_counts_1 = {'0_0 _0': 14, '0_0 _1': 18, '0_1 _0': 5, '0_1 _1': 9, '1_0 _0': 8} marginal_counts_result = marginal_counts(result.get_counts(), [0, 2, 4], format_marginal=True) self.assertEqual(marginal_counts_result, expected_marginal_counts_1)
def _fitter_data( data: List[Dict[str, any]] ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, List[np.ndarray]]: """Return list a tuple of basis, frequency, shot data""" outcome_dict = {} meas_size = None prep_size = None for datum in data: # Get basis data metadata = datum["metadata"] meas_element = tuple(metadata["m_idx"]) prep_element = tuple( metadata["p_idx"]) if "p_idx" in metadata else tuple() if meas_size is None: meas_size = len(meas_element) if prep_size is None: prep_size = len(prep_element) # Add outcomes counts = Counts( marginal_counts(datum["counts"], metadata["clbits"])).int_outcomes() basis_key = (meas_element, prep_element) if basis_key in outcome_dict: TomographyAnalysis._append_counts(outcome_dict[basis_key], counts) else: outcome_dict[basis_key] = counts num_basis = len(outcome_dict) measurement_data = np.zeros((num_basis, meas_size), dtype=int) preparation_data = np.zeros((num_basis, prep_size), dtype=int) shot_data = np.zeros(num_basis, dtype=int) outcome_data = [] for i, (basis_key, counts) in enumerate(outcome_dict.items()): measurement_data[i] = basis_key[0] preparation_data[i] = basis_key[1] outcome_arr = np.zeros((len(counts), 2), dtype=int) for j, (outcome, freq) in enumerate(counts.items()): outcome_arr[j] = [outcome, freq] shot_data[i] += freq outcome_data.append(outcome_arr) return outcome_data, shot_data, measurement_data, preparation_data
def _generate_matrices(self, data) -> List[np.array]: num_qubits = len(data[0]["metadata"]["label"]) counts = [None, None] for result in data: for i in range(2): if result["metadata"]["label"] == str(i) * num_qubits: counts[i] = result["counts"] matrices = [] for k in range(num_qubits): matrix = np.zeros([2, 2], dtype=float) marginalized_counts = [] for i in range(2): marginalized_counts.append(marginal_counts(counts[i], [k])) # matrix[i][j] is the probability of counting i for expected j for i in range(2): for j in range(2): matrix[i][j] = marginalized_counts[j][str(i)] / sum( marginalized_counts[j].values()) matrices.append(matrix) return matrices
def test_marginal_counts_result_creg_sizes(self): """Test that marginal_counts with Result input properly changes creg_sizes.""" raw_counts = {'0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0xE': 8} data = models.ExperimentResultData(counts=dict(**raw_counts)) exp_result_header = QobjExperimentHeader(creg_sizes=[['c0', 1], ['c1', 3]], memory_slots=4) exp_result = models.ExperimentResult(shots=54, success=True, data=data, header=exp_result_header) result = Result(results=[exp_result], **self.base_result_args) expected_marginal_counts = {'0 0': 14, '0 1': 18, '1 0': 13, '1 1': 9} expected_creg_sizes = [['c0', 1], ['c1', 1]] expected_memory_slots = 2 marginal_counts_result = marginal_counts(result, [0, 2]) self.assertEqual(marginal_counts_result.results[0].header.creg_sizes, expected_creg_sizes) self.assertEqual(marginal_counts_result.results[0].header.memory_slots, expected_memory_slots) self.assertEqual(marginal_counts_result.get_counts(0), expected_marginal_counts)
def _add_single_data(self, data): """Add data to the experiment""" # TODO: Handle optional marginalizing IQ data metadata = data.get("metadata", {}) if metadata.get("experiment_type") == self._experiment._type: # Add parallel data self._data.append(data) # Add marginalized data to sub experiments if "composite_clbits" in metadata: composite_clbits = metadata["composite_clbits"] else: composite_clbits = None for i, index in enumerate(metadata["composite_index"]): sub_data = {"metadata": metadata["composite_metadata"][i]} if "counts" in data: if composite_clbits is not None: sub_data["counts"] = marginal_counts(data["counts"], composite_clbits[i]) else: sub_data["counts"] = data["counts"] self._composite_expdata[index].add_data(sub_data)
def test_marginal_counts_inplace_false(self): """Test marginal_counts(Result, inplace=False) """ raw_counts_1 = {'0x0': 4, '0x1': 7, '0x2': 10, '0x6': 5, '0x9': 11, '0xD': 9, '0xE': 8} data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1)) exp_result_header_1 = QobjExperimentHeader(creg_sizes=[['c0', 4]], memory_slots=4) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) raw_counts_2 = {'0x2': 5, '0x3': 8} data_2 = models.ExperimentResultData(counts=dict(**raw_counts_2)) exp_result_header_2 = QobjExperimentHeader(creg_sizes=[['c0', 2]], memory_slots=2) exp_result_2 = models.ExperimentResult(shots=13, success=True, data=data_2, header=exp_result_header_2) result = Result(results=[exp_result_1, exp_result_2], **self.base_result_args) expected_marginal_counts = {'0': 27, '1': 27} self.assertEqual(marginal_counts(result, [0], inplace=False).get_counts(0), expected_marginal_counts) self.assertNotEqual(result.get_counts(0), expected_marginal_counts)
def test_marginal_counts_inplace_false(self): """Test marginal_counts(Result, inplace=False)""" raw_counts_1 = { "0x0": 4, "0x1": 7, "0x2": 10, "0x6": 5, "0x9": 11, "0xD": 9, "0xE": 8 } data_1 = models.ExperimentResultData(counts=dict(**raw_counts_1)) exp_result_header_1 = QobjExperimentHeader(creg_sizes=[["c0", 4]], memory_slots=4) exp_result_1 = models.ExperimentResult(shots=54, success=True, data=data_1, header=exp_result_header_1) raw_counts_2 = {"0x2": 5, "0x3": 8} data_2 = models.ExperimentResultData(counts=dict(**raw_counts_2)) exp_result_header_2 = QobjExperimentHeader(creg_sizes=[["c0", 2]], memory_slots=2) exp_result_2 = models.ExperimentResult(shots=13, success=True, data=data_2, header=exp_result_header_2) result = Result(results=[exp_result_1, exp_result_2], **self.base_result_args) expected_marginal_counts = {"0": 27, "1": 27} self.assertEqual( marginal_counts(result, [0], inplace=False).get_counts(0), expected_marginal_counts) self.assertNotEqual(result.get_counts(0), expected_marginal_counts)
def test_marginal_counts_result_marginalize_memory(self): """Test that a Result object containing memory marginalizes correctly inplace.""" result = self.generate_qiskit_result() marginal_result = marginal_counts(result, indices=[0], inplace=True, marginalize_memory=False) self.assertFalse(hasattr(marginal_result.results[0].data, "memory")) result = self.generate_qiskit_result() marginal_result = marginal_counts(result, indices=[0], inplace=True, marginalize_memory=None) self.assertTrue(hasattr(marginal_result.results[0].data, "memory")) result = self.generate_qiskit_result() marginal_result = marginal_counts(result, indices=[0], inplace=True, marginalize_memory=True) self.assertTrue(hasattr(marginal_result.results[0].data, "memory")) result = self.generate_qiskit_result() marginal_result = marginal_counts(result, indices=[0], inplace=False, marginalize_memory=False) self.assertFalse(hasattr(marginal_result.results[0].data, "memory")) marginal_result = marginal_counts(result, indices=[0], inplace=False, marginalize_memory=None) self.assertTrue(hasattr(marginal_result.results[0].data, "memory")) marginal_result = marginal_counts(result, indices=[0], inplace=False, marginalize_memory=True) self.assertTrue(hasattr(marginal_result.results[0].data, "memory"))
def _fitter_data( data: List[Dict[str, any]], measurement_basis: Optional[MeasurementBasis] = None, measurement_qubits: Optional[Tuple[int, ...]] = None, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Return list a tuple of basis, frequency, shot data""" meas_size = None prep_size = None # Construct marginalized tomography count dicts outcome_dict = {} shots_dict = {} for datum in data: # Get basis data metadata = datum["metadata"] meas_element = tuple( metadata["m_idx"]) if "m_idx" in metadata else tuple() prep_element = tuple( metadata["p_idx"]) if "p_idx" in metadata else tuple() if meas_size is None: meas_size = len(meas_element) if prep_size is None: prep_size = len(prep_element) # Add outcomes counts = Counts( marginal_counts(datum["counts"], metadata["clbits"])) shots = datum.get("shots", sum(counts.values())) basis_key = (meas_element, prep_element) if basis_key in outcome_dict: TomographyAnalysis._append_counts(outcome_dict[basis_key], counts) shots_dict[basis_key] += shots else: outcome_dict[basis_key] = counts shots_dict[basis_key] = shots # Construct function for converting count outcome dit-strings into # integers based on the specified number of outcomes of the measurement # bases on each qubit if meas_size == 0: # Trivial case with no measurement num_outcomes = 1 outcome_func = lambda _: 1 elif measurement_basis is None: # If no basis is provided assume N-qubit measurement case num_outcomes = 2**meas_size outcome_func = lambda outcome: int(outcome, 2) else: # General measurement basis case for arbitrary outcome measurements if measurement_qubits is None: measurement_qubits = tuple(range(meas_size)) elif len(measurement_qubits) != meas_size: raise AnalysisError( "Specified number of measurementqubits does not match data." ) outcome_shape = measurement_basis.outcome_shape(measurement_qubits) num_outcomes = np.prod(outcome_shape) outcome_func = _int_outcome_function(outcome_shape) num_basis = len(outcome_dict) measurement_data = np.zeros((num_basis, meas_size), dtype=int) preparation_data = np.zeros((num_basis, prep_size), dtype=int) shot_data = np.zeros(num_basis, dtype=int) outcome_data = np.zeros((num_basis, num_outcomes), dtype=int) for i, (basis_key, counts) in enumerate(outcome_dict.items()): measurement_data[i] = basis_key[0] preparation_data[i] = basis_key[1] shot_data[i] = shots_dict[basis_key] for outcome, freq in counts.items(): outcome_data[i][outcome_func(outcome)] = freq return outcome_data, shot_data, measurement_data, preparation_data
def test_marginal_counts_result_invalid_indices(self): """Test that a Result object containing memory marginalizes correctly inplace.""" result = self.generate_qiskit_result() with self.assertRaises(QiskitError): _ = marginal_counts(result, indices=[0, 1, 100], inplace=True)
def _marginalized_component_data( self, composite_data: List[Dict]) -> List[List[Dict]]: """Return marginalized data for component experiments. Args: composite_data: a list of composite experiment circuit data. Returns: A List of lists of marginalized circuit data for each component experiment in the composite experiment. """ # Marginalize data marginalized_data = {} for datum in composite_data: metadata = datum.get("metadata", {}) # Add marginalized data to sub experiments if "composite_clbits" in metadata: composite_clbits = metadata["composite_clbits"] else: composite_clbits = None # Pre-process the memory if any to avoid redundant calls to format_counts_memory f_memory = self._format_memory(datum, composite_clbits) for i, index in enumerate(metadata["composite_index"]): if index not in marginalized_data: # Initialize data list for marginalized marginalized_data[index] = [] sub_data = {"metadata": metadata["composite_metadata"][i]} if "counts" in datum: if composite_clbits is not None: sub_data["counts"] = marginal_counts( datum["counts"], composite_clbits[i]) else: sub_data["counts"] = datum["counts"] if "memory" in datum: if composite_clbits is not None: # level 2 if f_memory is not None: idx = slice(-1 - composite_clbits[i][-1], -composite_clbits[i][0] or None) sub_data["memory"] = [ shot[idx] for shot in f_memory ] # level 1 else: mem = np.array(datum["memory"]) # Averaged level 1 data if len(mem.shape) == 2: sub_data["memory"] = mem[ composite_clbits[i]].tolist() # Single-shot level 1 data if len(mem.shape) == 3: sub_data["memory"] = mem[:, composite_clbits[ i]].tolist() else: sub_data["memory"] = datum["memory"] marginalized_data[index].append(sub_data) # Sort by index return [marginalized_data[i] for i in sorted(marginalized_data.keys())]