pymc-devs / pymc-experimental

https://pymc-experimental.readthedocs.io
Other
72 stars 46 forks source link

`TimeSeasonality` component of statespace module lists wrong dimensions #262

Closed kforeman closed 6 months ago

kforeman commented 8 months ago

If you create a structural timeseries model with a TimeSeasonality component, it says the corresponding coefficients should have dim {name}_state, but should actually be {name}_periods.

ws = st.TimeSeasonality(
    season_length=7,  # day of week
    name="weekly",
    innovations=True,
    state_names=["Mon", "Tue", "Wed", "Thur", "Fri", "Sat", "Sun"],
)
sts_mod = (ws).build()

says:

The following parameters should be assigned priors inside a PyMC model block: 
    weekly_coefs -- shape: (6,), constraints: None, dims: (weekly_state, )
    sigma_weekly -- shape: (1,), constraints: Positive, dims: None
    P0 -- shape: (6, 6), constraints: Positive semi-definite, dims: ('state', 'state_aux')

However, if you use weekly_state as the dim, it throws an error:

with pm.Model(coords=sts_mod.coords) as mod:
    weekly_coefs = pm.Normal("weekly-coefs", mu=0, sigma=1, dims=("weekly_state",))

returns

KeyError: "Dimensions {'weekly_state'} are unknown to the model and cannot be used to specify a `shape`."

Changing instead to weekly_periods fixes it:

with pm.Model(coords=sts_mod.coords) as mod:
    weekly_coefs = pm.Normal("weekly-coefs", mu=0, sigma=1, dims=("weekly_periods",))

I thought maybe this was arising from this line but sts_mod.param_dims correctly returns {'weekly_coefs': ('weekly_periods',), 'P0': ('state', 'state_aux')} in this case, so I suppose I'm not quite sure exactly how the recordkeeping for dims works.

jessegrabowski commented 8 months ago

Ah yes there's a typo somewhere in the codebase. I have actually run into this myself, but I will push a fix ASAP.