Cast all component parameters to lists in Component.__init__

Author: jessegrabowski

pymc_extras/statespace/models/structural/core.py

@@ -507,12 +507,14 @@ def __init__(
         self.k_posdef = k_posdef
         self.measurement_error = measurement_error
 
-        self.state_names = state_names if state_names is not None else []
-        self.observed_state_names = observed_state_names if observed_state_names is not None else []
-        self.data_names = data_names if data_names is not None else []
-        self.shock_names = shock_names if shock_names is not None else []
-        self.param_names = param_names if param_names is not None else []
-        self.exog_names = exog_names if exog_names is not None else []
+        self.state_names = list(state_names) if state_names is not None else []
+        self.observed_state_names = (
+            list(observed_state_names) if observed_state_names is not None else []
+        )
+        self.data_names = list(data_names) if data_names is not None else []
+        self.shock_names = list(shock_names) if shock_names is not None else []
+        self.param_names = list(param_names) if param_names is not None else []
+        self.exog_names = list(exog_names) if exog_names is not None else []
 
         self.needs_exog_data = len(self.exog_names) > 0
         self.coords = {}
@@ -741,13 +743,15 @@ def make_slice(name, x, o_x):
 
     def _combine_property(self, other, name, allow_duplicates=True):
         self_prop = getattr(self, name)
+        other_prop = getattr(other, name)
+
         if isinstance(self_prop, list) and allow_duplicates:
-            return self_prop + getattr(other, name)
+            return self_prop + other_prop
         elif isinstance(self_prop, list) and not allow_duplicates:
-            return self_prop + [x for x in getattr(other, name) if x not in self_prop]
+            return self_prop + [x for x in other_prop if x not in self_prop]
         elif isinstance(self_prop, dict):
             new_prop = self_prop.copy()
-            new_prop.update(getattr(other, name))
+            new_prop.update(other_prop)
             return new_prop
 
     def _combine_component_info(self, other):

tests/statespace/models/structural/test_core.py

@@ -1,3 +1,5 @@
+from functools import partial
+
 import numpy as np
 import pandas as pd
 import pymc as pm
@@ -162,3 +164,27 @@ def test_extract_multiple_observed(rng):
 
     missing = set(comp_states) - set(expected_states)
     assert len(missing) == 0, missing
+
+
+@pytest.mark.parametrize(
+    "arg_type", [tuple, list, set, np.array], ids=["tuple", "list", "set", "array"]
+)
+def test_sequence_type_component_arguments(arg_type):
+    state_names = list("ABCDEFG")
+    components = [
+        st.LevelTrendComponent,
+        partial(st.CycleComponent, cycle_length=12),
+        st.AutoregressiveComponent,
+        partial(st.FrequencySeasonality, season_length=12),
+        partial(st.TimeSeasonality, season_length=12),
+        st.MeasurementError,
+    ]
+
+    components = [
+        components[i](observed_state_names=arg_type(state_names))
+        for i in np.random.choice(len(components), size=3, replace=False)
+    ]
+    ss_mod = sum(components[1:], start=components[0]).build(verbose=False)
+
+    assert ss_mod.k_endog == len(state_names)
+    assert sorted(ss_mod.observed_states) == sorted(list(state_names))