Address pr comments

Author: albrja

src/vivarium/framework/randomness/stream.py

@@ -274,15 +274,11 @@ def filter_for_probability(
             return population
 
         # Check for null values in probabilities
-        if isinstance(probability, float):
+        if isinstance(probability, (float, int)):
             if np.isnan(probability):
                 raise ValueError("Probabilities contain null values")
-        elif isinstance(probability, pd.Series):
-            if probability.isna().any():
-                raise ValueError("Probabilities contain null values")
-        elif not isinstance(probability, (int, np.integer)):
-            # Handle lists, tuples, and numpy arrays (but skip int types)
-            if np.any(np.isnan(probability)):
+        else:
+            if np.isnan(probability).any():
                 raise ValueError("Probabilities contain null values")
 
         if isinstance(population, pd.Index):

tests/framework/randomness/test_stream.py

@@ -275,21 +275,26 @@ def test_stream_rate_conversion_config(
     assert sim._randomness._rate_conversion_type == rate_conversion
 
 
-def test_filter_for_probability_error_with_null_values() -> None:
+@pytest.mark.parametrize(
+    "probs",
+    [
+        [0.3, np.nan, 0.5, np.nan, 0.7],
+        np.nan,
+        pd.Series([0.2, 0.3, np.nan, 0.4, 0.5]),
+    ],
+)
+def test_filter_for_probability_error_with_null_values(
+    probs: float | list[float] | pd.Series[float],
+) -> None:
     randomness_stream = RandomnessStream(
         "test", lambda: pd.Timestamp(2020, 1, 1), 1, IndexMap()
     )
     pop = pd.DataFrame({"age": [10, 11, 12, 13, 14], "id": [1, 2, 3, 4, 5]}).set_index("id")
     with pytest.raises(ValueError, match="Probabilities contain null values"):
-        randomness_stream.filter_for_probability(pop, [0.3, np.nan, 0.5, np.nan, 0.7])
+        randomness_stream.filter_for_probability(pop, probs)
 
     with pytest.raises(ValueError, match="Probabilities contain null values"):
-        randomness_stream.filter_for_probability(pop, np.nan)
+        randomness_stream.filter_for_probability(pop, probs)
 
     with pytest.raises(ValueError, match="Probabilities contain null values"):
-        randomness_stream.filter_for_probability(
-            pop, pd.Series([0.2, 0.3, np.nan, 0.4, 0.5], index=pop.index)
-        )
-
-    # Doesn't raise
-    randomness_stream.filter_for_probability(pop, [0.3, 0.4, 0.5, 0.6, 0.7])
+        randomness_stream.filter_for_probability(pop, probs)