More safety nets for resizing mutable dimensions

[michaelosthege]

• Adds a Model.set_dim method for resizing dimensions that were created by add_coord(..., mutable=True).
• Changes Model.set_data to anticipate that data-induced resizing can target dimensions that were created through add_coord(..., mutable=True) which are not symbolically linked to the data variables.

Closes #5812.

[codecov]

Codecov Report

Merging #5817 (576e306) into main (57654dc) will increase coverage by 0.58%.
The diff coverage is 96.42%.

@@            Coverage Diff             @@
##             main    #5817      +/-   ##
==========================================
+ Coverage   89.40%   89.98%   +0.58%     
==========================================
  Files          74       73       -1     
  Lines       13769    13221     -548     
==========================================
- Hits        12310    11897     -413     
+ Misses       1459     1324     -135     

Impacted Files	Coverage Δ
pymc/model.py	87.26% <96.42%> (+0.73%)	⬆️
pymc/sampling.py	82.52% <0.00%> (-6.19%)	⬇️
pymc/variational/inference.py	85.78% <0.00%> (-1.46%)	⬇️
pymc/distributions/mixture.py	95.72% <0.00%> (-1.04%)	⬇️
pymc/variational/__init__.py	100.00% <0.00%> (ø)
pymc/distributions/logprob.py	97.65% <0.00%> (ø)
pymc/variational/flows.py
pymc/aesaraf.py	91.95% <0.00%> (+0.06%)	⬆️
pymc/smc/smc.py	96.45% <0.00%> (+0.07%)	⬆️
... and 4 more

[lucianopaz]

Thanks, @michaelosthege. I think that we need to use coords values that are supplied to set_data to call set_dim internally, instead of expecting the users to set_dim before they set_data.

[lucianopaz]

I think that the key thing that is missing here is that users can also pass coords when they set_data. Ideally, we should use the supplied coords internally to call set_dim and then update the values of the MutableData instance. This warning should only be raised if the users haven't passed coords along to set_data.

[michaelosthege]

Hmm.. You're welcome to try, but I think this could become more complicated than the current way with the warning.
The set_data method doesn't actually care if the ScalarSharedVariable corresponds to a dimension in Model.dim_lengths or not!

Most important will be to add test cases that showcase the intended call order

[lucianopaz]

What would be the expected flow to change data and make predictions? Before we had to do:

with pm.Model(coords={"A": range(10)}) as m:
    x = pm.MutableData("x", x_values, dims="A")
    y = pm.MutableData("y", y_values, dims="A")
    a = pm.Normal("a", 0, 1)
    b = pm.Normal("b", 0, 1)
    c = pm.HalfNormal("c", 1)
    obs = pm.Normal("obs", mu=a + b * x, sigma=c, observed=y)
    idata = pm.sample()

with m:
    pm.set_data({"x": np.linspace(-2, 3, 100), "y": np.full(100, np.nan)})
    ppc = pm.sample_posterior_predictive(idata)

I would like us to be able to do this:

with pm.Model(coords={"A": range(10)}) as m:
    x = pm.Data("x", x_values, dims="A")
    y = pm.Data("y", y_values, dims="A")
    a = pm.Normal("a", 0, 1)
    b = pm.Normal("b", 0, 1)
    c = pm.HalfNormal("c", 1)
    obs = pm.Normal("obs", mu=a + b * x, sigma=c, observed=y)
    idata = pm.sample()

with m:
    pm.set_data({"x": np.linspace(-2, 3, 100), "y": np.full(100, np.nan)}, coords={"A": range(10, 110)})
    ppc = pm.sample_posterior_predictive(idata)

But due to the choice of default Mutable/Constant coords, we need to:

1. Manually add_coord for "A" and say it is mutable
2. Call set_dim before calling set_data

Or we have to create the coordinate values using MutableData, which I don't really know how to do.

[lucianopaz]

We should add the test where pmodel.set_data("mdata", [1, 2, 3, 4], coords={"mdim": range(4)}).
Ideally, we want to catch the coords fed into set_data and call set_dim with them

[AlexAndorra]

Thanks for the proposed fix @michaelosthege ! I agree with @lucianopaz that it'd be great if the user didn't have to set_dim before they set_data

[michaelosthege]

Thanks for the proposed fix @michaelosthege ! I agree with @lucianopaz that it'd be great if the user didn't have to set_dim before they set_data

Since #5763 (7ec106c to be exact) one can pass coords to pm.Data containers.

So with this, you'd create the "A" dimension as mutable.

with pm.Model() as m:
    x = pm.MutableData("x", x_values, dims="A", coords={"A": range(10)})
    y = pm.MutableData("y", y_values, dims="A")

Then, if you add the coords to set_data, we could also call set_dim internally.

[michaelosthege]

I'll add the test suggested above and look into adding the coords kwarg to set_data.

[michaelosthege]

@lucianopaz please review the new test cases - they should govern which usage patterns are now supported.

The only thing where I'm hestiant is introducing mutable dims through the pm.Model constructor.
For that we could introduce a new kwarg, but I didn't want to make that decision already.
Options:

• pm.Model(coords_mutable=..., coords=...)
• pm.Model(coords_mutable=..., coords_immutable=...)
• more?

Please continue on this branch as you see fit. Let's get it over the finish line today

[ricardoV94]

This is a vague hand waving comment, so apologies for that. I feel the internal code to manage dims is become a bit too convoluted for its own good. Suggestion:

1. Have a coords and mutable_coords kwargs in pm.Model.
2. Remove defining news dims / coords with variables
3. Downsides?

[michaelosthege]

The idea of defining coords through pm.Data came from the export_index_as_coords kwarg.

Thinking this further, with #5796 we could do this:

population = xarray.DataArray(
    "population",
    [10, 200],
    dims="city",
    coords=dict(city=["Atown", "Bvil"]),
)

with pm.Model() as m:
    pop = pm.MutableData("population", population)
    assert "city" in m.coords # True

Note that this would work for multi-dimensional xarray.DataArray objects but not for pd.DataFrame with multi-index.

For wroking with GPs this would be really useful, because one has to carry data and grid coordinates around quite a lot..

---

I agree that the code is a little convoluted - it would probably benefit from a little extraction refactoring.

But the downsides of removing coords from pm.Data would be that we'd have to break the export_index_as_coords feature and we couldn't do what the above example does.

If we add a kwarg to the Model signature, I would prefer coords_mutable because of how things are sorted by code completion tools.

[michaelosthege]

Any decision regarding the default (im)mutable of pm.Model(coords) will build on top of the commits made in this branch.

Since the commits here don't affect the thing we're still in discussion about, can we go ahead and merge this in the interest of closing #5812 and increasing overall safety?

[ricardoV94]

Suggestion

[ricardoV94]

Some suggestions and questions :)

[michaelosthege]

@ricardoV94 @twiecki I won't have time to do any commits here until Monday.
As in "I'm out, please take over the rest here and get this merged".

[twiecki]

What's missing here, can we merge?

[ricardoV94]

Waiting for tests to pass to make sure I didn't break anything accidentally

Stale