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

Corrects an indexing bug that would cause DataSet lookups to silently
fail (return incorrect values) when data had been added with
dictionaries that didn't have keys ordered in the same way as the
DataSet's outcome list. This commit fixes this issue and adds several
unit tests used in debugging.

Previously failing minimal example (now included as a unit test) is:
```
import pygsti

c = "Gxpi2:0@(0)"
counts = {'00': 1, '10': 0, '01': 97, '11': 2}
ds = pygsti.data.DataSet(outcome_labels=["00", "10", "01", "11"], static=False)
ds.add_count_dict(c, counts)
check = ds[c].counts
print(ds)

assert {k[0]: v for k,v in check.items()} == counts
```

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):