Fixes a bug in DataSet count retrieval and adds unit tests. (#663) (Develop Merge) (#665)

Note: Since the automatic merge from bugfix into develop didn't work
this is just a manual merge of PR #663 into develop.

Co-authored-by: Erik Nielsen <[email protected]>

Author: coreyostrove

pygsti/data/dataset.py

@@ -575,18 +575,18 @@ def _get_counts(self, timestamp=None, all_outcomes=False):
                         cntDict.setitem_unsafe(ol, cnt)
             else:
                 for ol, i in self.dataset.olIndex.items():
-                    inds = oli_tslc[oli_tslc == i]
+                    inds = _np.where(oli_tslc == i)[0]
                     if len(inds) > 0 or all_outcomes:
                         cntDict.setitem_unsafe(ol, float(sum(self.reps[tslc][inds])))
         else:
             if self.reps is None:
                 for ol_index in oli_tslc:
                     ol = self.dataset.ol[ol_index]
-                    cntDict.setitem_unsafe(ol, 1.0 + cntDict.getitem_unsafe(ol, 0.0))
+                    cntDict.setitem_unsafe(ol, float(1.0 + cntDict.getitem_unsafe(ol, 0.0)))
             else:
                 for ol_index, reps in zip(oli_tslc, self.reps[tslc]):
                     ol = self.dataset.ol[ol_index]
-                    cntDict.setitem_unsafe(ol, reps + cntDict.getitem_unsafe(ol, 0.0))
+                    cntDict.setitem_unsafe(ol, float(reps + cntDict.getitem_unsafe(ol, 0.0)))
 
         return cntDict
 

test/unit/objects/test_dataset.py

@@ -99,6 +99,43 @@ def test_static_constructor_raises_on_missing_counts(self):
         with self.assertRaises(ValueError):
             DataSet(circuits=self.gstrs, outcome_labels=['0', '1'], static=True)
 
+    def test_add_count_dict_order(self):
+        c = "Gxpi2:0@(0)"
+        counts = {'00': 1, '10': 0, '01': 97, '11': 2}
+        ds = DataSet(outcome_labels=["00", "10", "01", "11"], static=False)
+        ds.add_count_dict(c, counts)
+        check = ds[c].counts
+        print(ds)
+
+        self.assertEqual({k[0]: v for k,v in check.items()}, counts)
+
+        # Add more
+        c2 = "Gypi2:0@(0)"
+        counts2 = {'00': 1, '01': 43, '10': 24, '11': 7}
+        ds.add_count_dict(c2, counts2)
+        check = ds[c2].counts
+
+        self.assertEqual({k[0]: v for k,v in check.items()}, counts2)
+
+        ds.done_adding_data()
+        self.assertEqual({k[0]: v for k,v in ds[c].counts.items()}, counts)
+        self.assertEqual({k[0]: v for k,v in ds[c2].counts.items()}, counts2)
+
+    def test_add_count_dict_single_outcome(self):
+        c0 = Circuit("Gi:Q0@Q0")
+        c1 = Circuit("Gi:Q0Gi:Q0@Q0")
+        c0_counts = {("1",): 10.0}
+        c1_counts = {("0",): 4, ("1",): 6}
+
+        ds = DataSet()
+        ds.add_count_dict(c0, c0_counts)
+        ds.add_count_dict(c1, c1_counts)
+
+        self.assertEqual({k: v for k,v in ds[c0].counts.items()}, c0_counts)
+        self.assertEqual({k: v for k,v in ds[c1].counts.items()}, c1_counts)
+
+        
+
 
 class DefaultDataSetInstance(object):
     def setUp(self):