Merge pull request #277 from ihmeuw/collijk/bugfix/index-map-bug-for-…

…non-unique-initial-values

Fix bug for non-unique initial IndexMap.update

src/vivarium/framework/randomness/index_map.py

@@ -39,8 +39,10 @@ def update(self, new_keys: pd.Index) -> None:
         """
         if new_keys.empty or not self._use_crn:
             return  # Nothing to do
-        elif not self._map.index.intersection(new_keys).empty:
-            raise KeyError("Non-unique keys in index")
+
+        new_index = self._map.index.append(new_keys)
+        if len(new_index) != len(new_index.unique()):
+            raise RandomnessError("Non-unique keys in index")
 
         mapping_update = self._hash(new_keys)
         if self._map.empty:

src/vivarium/framework/randomness/manager.py

@@ -16,7 +16,7 @@ class RandomnessManager:
     configuration_defaults = {
         "randomness": {
             "map_size": 1_000_000,
-            "key_columns": ["entrance_time"],
+            "key_columns": [],
             "random_seed": 0,
             "additional_seed": None,
         }

tests/conftest.py

@@ -21,7 +21,8 @@ def base_config():
                 },
                 "end": {"year": 2010},
                 "step_size": 30.5,
-            }
+            },
+            "randomness": {"key_columns": ["entrance_time", "age"]},
         },
         **metadata(__file__),
     )

tests/framework/randomness/test_index_map.py

@@ -164,31 +164,49 @@ def test_hash_uniformity(map_size_and_hashed_values):
     assert p > 0.05, "Data not uniform"
 
 
-def test_update(mocker):
+@pytest.fixture(scope="function")
+def index_map(mocker):
     m = IndexMap()
-    keys = generate_keys(10000)
 
     def hash_mock(k, salt=0):
         seed = 123456
         rs = np.random.RandomState(seed=seed + salt)
         return pd.Series(rs.randint(0, len(k) * 10, size=len(k)), index=k)
 
     mocker.patch.object(m, "_hash", side_effect=hash_mock)
-    m.update(keys)
-    assert len(m._map) == len(keys), "All keys not in mapping"
-    assert m._map.index.difference(keys).empty, "All keys not in mapping"
-    assert len(m._map.unique()) == len(keys), "Duplicate values in mapping"
-
-    # Can't have duplicate keys.
-    with pytest.raises(KeyError):
-        m.update(keys)
-
-    new_unique_keys = generate_keys(1000).difference(keys)
-    m.update(new_unique_keys)
-    assert len(m._map) == len(keys) + len(new_unique_keys), "All keys not in mapping"
-    assert m._map.index.difference(
-        keys.union(new_unique_keys)
-    ).empty, "All keys not in mapping"
-    assert len(m._map.unique()) == len(keys) + len(
-        new_unique_keys
-    ), "Duplicate values in mapping"
+
+    return m
+
+
+def test_update_empty_bad_keys(index_map):
+    keys = pd.Index(["a"] * 10)
+    with pytest.raises(RandomnessError):
+        index_map.update(keys)
+
+
+def test_update_nonempty_bad_keys(index_map):
+    keys = generate_keys(1000)
+
+    index_map.update(keys)
+    with pytest.raises(RandomnessError):
+        index_map.update(keys)
+
+
+def test_update_empty_good_keys(index_map):
+    keys = generate_keys(1000)
+    index_map.update(keys)
+    assert len(index_map._map) == len(keys), "All keys not in mapping"
+    assert index_map._map.index.difference(keys).empty, "All keys not in mapping"
+    assert len(index_map._map.unique()) == len(keys), "Duplicate values in mapping"
+
+
+def test_update_nonempty_good_keys(index_map):
+    keys = generate_keys(2000)
+    keys1, keys2 = keys[:1000], keys[1000:]
+
+    index_map.update(keys1)
+    index_map.update(keys2)
+
+    assert len(index_map._map) == len(keys), "All keys not in mapping"
+    assert index_map._map.index.difference(keys).empty, "All keys not in mapping"
+    assert len(index_map._map.unique()) == len(keys), "Duplicate values in mapping"