Initial prototype of plot_lm

[utkarsh-maheshwari]

Description

This PR and #1747 summarizes my work done in Google Summer of Code 2021. I added 2 plots, plot_lm and plot_ts. The first one is covered in this PR and the next one in #1747. The project workflow was like this:

1. Started with a high-level design: Discussed the user API function with the mentors and other community members. Discussion includes input variable names, accepted datatypes, accepted values and other input details. Also showed a sample output visualization. Once it was approved, I moved on to opening this PR. This Github Gist shows the design and some discussions related to it. The design process took an initial week.

2. Submitted a prototype: This PR was earlier a prototype. This step is basically implementing the design decision made in the previous step. In addition to the user API, backend functions are also added. As ArviZ uses 2 backends, artist functions for both of them were added. This step took another week.

3. Improved it according to reviews: In this step, mentors review your code to provide feedback and improvement tips. Learned the best code practices in this step. Ideally, improving never ends. It is important to maintain the added code after it is merged to keep it bug-free. I aim to provide support after GSoC. This step was a bit lengthy and complex and thus took 2 weeks.

4. Tested the functionality: Added unit tests using pytest. Aimed and achieved to cover all of the added functionality under tests. Similar to step 2, added tests for both of the backends. Also, solved some simple to complex bugs that arise while testing. This step took another week.

5. Added examples: Examples are added to the docstring as well as to the examples folder. Please check out the files changed tab to know more about this step. This step was quick, took only half a week.

6. Added documentation: If you want to know how to use plot_lm, checkout out this blog. However, if you want to go on a low level to know the working in detail, I would suggest taking a look over the docstring in ArviZ docs and probably follow the comments sequentially. Another week was consumed in this step.

Output :

Checklist

• Follows official PR format
• Includes a sample plot to visually illustrate the changes (only for plot-related functions)
• New features are properly documented (with an example if appropriate)?
• Includes new or updated tests to cover the new feature
• Code style correct (follows pylint and black guidelines)
• Changes are listed in changelog

[canyon289]

At a high level this looks good, Please start adding tests for functionality so we can more fully shake out bugs

[canyon289]

I would consider splitting this into another function and testing it independently

[canyon289]

Description

Here is my initial prototype for linear regression plot. I am checking if input data is correct in the the user API function. Backend function is quite small, only does the plotting part.

1. y_model : (Necessary) linear relationship in the form of y = m1*x1 + m2*x2 + c
2. idata : (Necessary if any of other args are string) should contain all the necessary variable names in any of the groups.
3. x_data : (Necessary because we cannot identify the x variable name from y_model) seq/str
4. y_data : (Necessary if y_model is Sequence. We can identify the name from y_model) seq/str
5. y_ppc : (Necessary is y_model is sequence. Assumed to be same as y_data if absent) seq/str

The high level design api doesnt seem to match whats in the code but i generally like the way its stated here. It would be nice if the idata was the only "data" argument, and the other arguments are specifiers for sets of data in the idata object

For the model specification instead of parsing it ourselves should we use patsy instead, and perhaps add it an optional dependency? @ahartikainen for a double check here

Overall the plot looks nice. For next steps I think we should

1. Update the API and inputs to be more idata centric
2. Add tests
3. Potentially break up the internal logic into function sized bits, which makes it easier to reuse functionality and test
4. Consider using patsy

[ahartikainen]

I'm not 100% we should parse formulas.

I think everything should be in the InferenceData and we call it.

For cases where we want to calculate the model mean, user given function is probably a safer option.

[OriolAbril]

I'm not 100% we should parse formulas.

I'm 90% sure we don't want to parse formulas at all, 100% sure we don't want to create any kind of parser (not only because it would be a lot of work and could even eat up the whole GSoC or even more time but also because we'd then have to maintain it!).

AFAIK all formula parsing libraries we could depend on use dataframes as input which we don't have so we'd need to convert the data to a correctly formatted dataframe for things to work, and for formula parsing to be really useful we would not only need to parse the formula but also consider link functions for example. I think having a lower level thing that takes samples/functions is preferrable, and it should be easy for libraries like bambi to build on top of that and offer a formula based api to users.

For cases where we want to calculate the model mean, user given function is probably a safer option.

You can take a look at the refitting notebooks that use xarray for examples on a very similar computation.

[utkarsh-maheshwari]

The high level design api doesnt seem to match whats in the code but i generally like the way its stated here. It would be nice if the idata was the only "data" argument, and the other arguments are specifiers for sets of data in the idata object

Do you mean that we should accept x_data, y_data and y_ppc only as string as the names of DataArray present in idata, and not as sequence? Yea, that would be better.

[utkarsh-maheshwari]

I think everything should be in the InferenceData and we call it.

Even the y_model?
By this, what comes in my mind is:

For example y ~ mx + c is the modeling equation. User will first find the posterior values of m and c. Then user will do something like:
idata.posterior_predictive["y_model"] = idata.posterior["m"] * idata.constant_data["x"] + idata.posterior["c"]
And pass it to the function.?

[canyon289]

Just to add my 2 cents, if we dont need to parse formulas I dont think we should either. Itll add an extra dependency to ArviZ, and in the modeling workflow it will be very annoying to restate a complicated model in a formula when you've already done it in a PPL. Im not sure if we can get away without it but glad @OriolAbril and @ahartikainen chimed in saying we probably dont need it!

[utkarsh-maheshwari]

Sure then! Assuming the y_model is present in inferenceData is looking a good option.

[utkarsh-maheshwari]

I think having a lower level thing that takes samples/functions is preferrable, and it should be easy for libraries like bambi to build on top of that and offer a formula based api to users.

Not completely sure what this means. Are you talking about, for now, taking all of the necessary values as references to precomputed DataArrays, stored in idata. And later, extending it to take an equation string and let bambi handle the equation.

[utkarsh-maheshwari]

As suggested by everyone, what I have come up with is very simple than what I expected.
HIgh level :

plot_lm(
    idata : az.InferenceData object
    x : str
    y : str
    y_model : str
    y_ppc : str
    x_group : str
    y_data_group : str
    y_model_group : str
    y_ppc_group : str
    num_ppc_samples : int
)

Function call :

plot_lm(
    idata = idata, 
    x = "mom_iq", 
    y = "kid_score", 
    y_model = "y_model", 
    y_ppc = "kid_score", 
    x_group = "constant_data", 
    y_group = "observed_data", 
    y_model_group = "posterior", 
    y_ppc_group = "posterior_predictive"
)

We can then move forward to reduce input arguments by making certain assumptions(specially for the group arguments) and accepting inputs as sequence too!

I agree on the comments made here and on slack regarding parsing the formula and totally convinced that making a parser is not a good idea.

[canyon289]

Looking good! Youre gonna get so sick of me saying this. Cant wait to see the function implementation and some tests so we can try it out :D

[utkarsh-maheshwari]

I have added a basic test. Not 100% sure if it's the right approach. Actually, there was no pre-built linear regression model so created one in the test function itself. Should it go in the helper?

[codecov]

Codecov Report

Merging #1727 (3a42a6f) into main (6de30a6) will increase coverage by 0.20%.
The diff coverage is 98.00%.

❗ Current head 3a42a6f differs from pull request most recent head bc9c18e. Consider uploading reports for the commit bc9c18e to get more accurate results

@@            Coverage Diff             @@
##             main    #1727      +/-   ##
==========================================
+ Coverage   90.87%   91.07%   +0.20%     
==========================================
  Files         109      112       +3     
  Lines       11844    12094     +250     
==========================================
+ Hits        10763    11015     +252     
+ Misses       1081     1079       -2     

Impacted Files	Coverage Δ
arviz/plots/lmplot.py	96.80% <96.80%> (ø)
arviz/plots/backends/matplotlib/lmplot.py	98.70% <98.70%> (ø)
arviz/plots/backends/bokeh/lmplot.py	98.71% <98.71%> (ø)
arviz/plots/__init__.py	100.00% <100.00%> (ø)
arviz/plots/plot_utils.py	87.11% <0.00%> (+0.51%)	⬆️
arviz/sel_utils.py	89.70% <0.00%> (+1.47%)	⬆️
arviz/plots/backends/bokeh/__init__.py	91.89% <0.00%> (+2.70%)	⬆️
... and 1 more

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 6de30a6...bc9c18e. Read the comment docs.

[canyon289]

I have added a basic test. Not 100% sure if it's the right approach. Actually, there was no pre-built linear regression model so created one in the test function itself. Should it go in the helper?

Looks like a good start to the test. Your instinct is right I would move the "model" to helpers. I also would suggest instead of creating groups directly to use the inferencedata interface so if that changes itll be reflected in this test

https://github.com/arviz-devs/arviz/blob/main/arviz/tests/helpers.py#L22

I wont have time for a full review just this moment but will get to it in next 24 to 72 hours.

[OriolAbril]

I would try to get something like:

plot_lm(x, y, y_model, y_ppc, num_pp_samples, plot_dim):
    """Posterior predictive plots for regression-like data.

    idata : az.InferenceData object
    x : str or array_like, optional
        If str, it's a variable name from constant_data group so we get a dataarray. 
        Therefore it should not be much of a change to also take directly dataarrays or arrays. 
        It is optional because if the x values are the coordinate values from y I think the best is to 
        pass y only and no x to indicate that x should be the coordinates in y
    y : str or array_like
        If str, variable name from observed data, could also be 1d array or dataarray.
    y_model : str or array_like, optional
        If str, variable name from posterior group
    y_ppc : str or array_like, optional
        If str, variable name from posterior_predictive group
    num_pp_samples : int
        Note it's pp_samples! not ppc_samples, it should be the same as the arg in plot_ppc.
    plot_dim/x_dim : str, optional
        The dimension name that contains the x/y values. i.e. if y_ppc has dims (chain, draw, group, time) 
        and we set plot_dim=time, then we'd generate n_group subplots, one per value of group coordinate. 
        Note that we might have multiple covariates (x could be a list of str) 
        maybe even multiple y instead of multidimensional y? so things could bet a bit tricky. 
        This argument is required unless x/y are 1d and y_model/y_ppc are 3d 
        (or 2d if dataarrays with `__sample__, plot_dim` dimensions? not sure if we should support this case, 
        I think it could be useful for cases where generating pp samples is hard so users can use extract function and 
        generate pp samples only for the small random subset).
    """
    # code function

facetting/iterating to generate the multiple subplots is going to be challenging and therefore I think we should take that into account (or a specific subset we want to support) from the beginning, so we don't define somehting for 1 subplot case and then try to force facetting on it. Regarding the group arguments, I strongly feel we should not have them, for several reasons: variables should have a group where they belong and having that seems to indicate they don't; we don't have that in any of the similar plots (i.e. the group argument in plot_ppc is posterior and prior yet we never get samples from either group there, prior means use data from prior_predictive and observed_data, same for posterior, we could have this here too, but not a group per variable); it's not difficult to move variables from one group to another (we might need to improve docs though) and it should not duplicate memory.

I am also unsure about the hdi for the mean. Most people I have seen use that were doing so inadvertedly thinking it was the posterior predictive hdi. What information is that part supposed to transmit? I understand and know how to interpret spagetti plots of either pp samples directly or regression lines but I don't really know how to interpret these mean hdi ones so everything might come from my own ignorance.

[utkarsh-maheshwari]

x : str or array_like, optional
If str, it's a variable name from constant_data group so we get a dataarray.
Therefore it should not be much of a change to also take directly dataarrays or arrays.
It is optional because if the x values are the coordinate values from y I think the best is to
pass y only and no x to indicate that x should be the coordinates in y

I really like this approach. Thanks for putting this.

y_model : str or array_like, optional

If optional, would we be plotting just x and y?

Regarding the group arguments, I strongly feel we should not have them

So, then we would assume variables to be present in a certain group. If any of the variable is not found in that particular group, we should not search in other groups and return immediately with an error.? or should we search in other groups too?
(for example should we search for x only in canstant_data? )

[OriolAbril]

If optional, would we be plotting just x and y?

x, y and y_ppc (which we should probably call y_hat instead to keep the same naming convention from plot_ppc and plot_loo_pit

So, then we would assume variables to be present in a certain group

Yes, I think it will make the code much simpler and easier to maintain with little to no drawback. You have the x in observed_data instead of constant_data? use plot_lm(x=idata.observed_data["x_var_name"], ...)

[utkarsh-maheshwari]

What information is that part supposed to transmit?

According to me, It depicts the uncertainty of mean line just like spagetti plots do, but instead by plotting credible interval. We could do it like following as well.

About y_ppc we can plot posterior_predictive samples in different colour than observed values.

[canyon289]

What information is that part supposed to transmit?

According to me, It depicts the uncertainty of mean line just like spagetti plots do, but instead by plotting credible interval. We could do it like following as well.

About y_ppc we can plot posterior_predictive samples in different colour than observed values.

It is possible to plot the uncertainty three ways using an arguments? Two of the ways you've shown, either an hdi band, or samples from the posterior predictive. The third idea is a band but where the opacity changes as a function of the HDI. this last one dont implement in this PR just throwing the idea out there.

I can see different situations where a user would want the first two, so being able to switch, or plot both together, with an argument would be nice. The third Im less sure about but what to suggest it now while were in the early PR phase, although I want to stress again don't implement the third one. If method is made flexible enough to switch between points and band uncertainty representation we can test the third idea later

[ahartikainen]

Added couple of comments

[ahartikainen]

At some point these could be used with kwargs (with some setdefault magic)

[utkarsh-maheshwari]

Yes, this remains to be done.

[utkarsh-maheshwari]

Different kwargs for observed data, PPC and model?

[ahartikainen]

I think this also needs some thinking when we want to do this and when not

[utkarsh-maheshwari]

Yes , I have just taken it from from your function. I need to think about it because I have less idea where is it useful.

[ahartikainen]

This should also be some kwargs

[utkarsh-maheshwari]

plot_dim/x_dim : str, optional
The dimension name that contains the x/y values. i.e. if y_ppc has dims (chain, draw, group, time)
and we set plot_dim=time, then we'd generate n_group subplots, one per value of group coordinate.
Note that we might have multiple covariates (x could be a list of str)
maybe even multiple y instead of multidimensional y? so things could bet a bit tricky.
This argument is required unless x/y are 1d and y_model/y_ppc are 3d
(or 2d if dataarrays with __sample__, plot_dim dimensions? not sure if we should support this case,
I think it could be useful for cases where generating pp samples is hard so users can use extract function and
generate pp samples only for the small random subset).

@OriolAbril I have written some code for it. Can you share a model to test this argument?

[utkarsh-maheshwari]

I have added some tests. Function is now working fine for multidimensional y as well as x. Now some labelling and touch-ups are required so that graph looks nice. It gets complicated at some places in user API, which I think need some review because there may be better ways to process the inputs. Thank you for being so helpful so far. 😄

[canyon289]

• The visual style doesnt seem to match the plots in the ArviZ gallery
• The "parallel plot style, uncertainty seems a bit hard to read. Im not totally opposed to it just curious in the motivation and thinking behind it
• The posterior predictive samples in the bottom plot are really hard to see by default

1. Really have very less idea how to achieve it. When I plot other plots' tests results, even they don't  follow the visual style it. But plots in examples gallery follows. grey bg and all. I guess calling `az.style.use("arviz-darkgrid")` before calling `plot_lm` will do the job?

2. As I have mentioned earlier, it is because of random data. Should I add a well-defined model in helpers to tackle this?

3. Solved

1. I think it looks fine now. Thanks for fixing it.
2. Sorry for missing that. I agree Lets defer this for now
3. Solved

Overall this code looks great and I like the plots. This is fantastic add to the ArviZ toolset. Youre nearly most of the way there just some comments below.

Need to fix or address comments

Returned object should always be a plot

This examples from the docs returns a numpy array. It should return a plot

data = az.from_dict(observed_data = { "y": np.random.normal(size=(5, 7)) }, 
                     posterior_predictive = {"y": np.random.randn(4, 1000, 5, 7) / 2},
                     dims={"y": ["dim1", "dim2"]}, coords={"dim1": range(5), "dim2": range(7)} )

Extra code documentation

There are a couple code comments that are very helpful. Add a couple more like the one noted above. nothing fancy, just one liners on major sections of logic. If youd like i can note where to help out

Nitpicky Plot comment

This doesnt need to be fixed in this pr but the legend will need a box around it. Its really hard to read otherwise. I know we dont need this across all plots but I think we really should as some of our other plots suffer as well.
https://arviz-devs.github.io/arviz/examples/plot_parallel_rank.html

Future PRs: Lots of Documentation

First off thanks for adding both a complete api documentation, code examples, and tests. It is a fantastic start. Some of the arguments are not particularly intuitive, which is not you fault just the nature of the best, so documentation will help.

In a followup PR can you add docs that show the following? Perhaps start with a notebook for prototyping and we can see how complicated it is and decide the best final home for the code

Things that will help are

• Showing more complicated y_models that include multiple covariates
• Non linear models
• Are counterfactuals supported?

[canyon289]

Also this branch might need a rebase on main to fix the failing tests. Once we have passing tests we can assess the test coverage as well

[utkarsh-maheshwari]

This examples from the docs returns a numpy array

This one is interesting. This example is plotting multiple plots. So, it's returning an array of plots. Checked it for other plots like plot_trace, it too returns NumPy array.

But it is even returning an array of shape (1,1) in case there is only one plot. It should return a plot, not an array of 1 plot. This needs some thinking, why is it happening

[canyon289]

This examples from the docs returns a numpy array

This one is interesting. This example is plotting multiple plots. So, it's returning an array of plots. Checked it for other plots like plot_trace, it too returns NumPy array.

But it is even returning an array of shape (1,1) in case there is only one plot. It should return a plot, not an array of 1 plot. This needs some thinking, why is it happening

Oh youre right ignore my comment, no need to return a 1d array. I apologize my mistake there

[utkarsh-maheshwari]

arviz/data/io_pyro.py:111:20: E1102: pyro.poutine.trace is not callable (not-callable)

Is it an error from pyro's side?

[canyon289]

arviz/data/io_pyro.py:111:20: E1102: pyro.poutine.trace is not callable (not-callable)

Is it an error from pyro's side?

Yea this is a linting error from pyro. Did you rebase on master? I think all the issues are fixed on master? It was passing as of 13 days ago, not sure if this is a new issue
https://dev.azure.com/ArviZ/ArviZ/_build/results?buildId=4558&view=results

[utkarsh-maheshwari]

Did you rebase on master

Yes, I did. I think this is some new error. Other PR's are also failing.

[canyon289]

Master is fixed. Please rebase or merge so we can see if CI passes

[utkarsh-maheshwari]

All green fInally!! Good to see passed tests after a long time xD. Thanks to all who worked to fix CI!

[OriolAbril]

The comments don't necessarily need to be addressed in this PR, but they need to be addressed at some point.

This doesnt need to be fixed in this pr but the legend will need a box around it. Its really hard to read otherwise. I know we dont need this across all plots but I think we really should as some of our other plots suffer as well.
https://arviz-devs.github.io/arviz/examples/plot_parallel_rank.html

This may actually an issue with the arviz-darkgrid theme: https://github.com/arviz-devs/arviz/blob/main/arviz/plots/styles/arviz-darkgrid.mplstyle#L11. Do we only want this plot to have a box? Or is it something more general that should be handled by the theme?

This may also be related to the grid parameter, I have focused on docstring only and don't know how it works, but it may be there only to override the value set by the theme.

This doesn't mean that we can't override theme set values, if we think this plot should always/never have a grid/legend box independently of the theme for example, just wanted to comment that so we are all aware.

[OriolAbril]

sequence is not a valid type according to y docstring, so I'm not sure what this means.

[utkarsh-maheshwari]

If any of y, x, y_hat, y_model is str, then idata is necessary. Yea, This line is unclear.

[OriolAbril]

This explanation should be the one on the docstring! It's different from the current one too which seeing this looks wrong

[OriolAbril]

can you make plot and circle hyperlinks using intersphinx? same for all other kwargs.

see https://arviz-devs.github.io/arviz/contributing/developer_guide.html#hyperlinks and #1188 for guidance and examples.

[utkarsh-maheshwari]

Done!

[OriolAbril]

these need an explanation, even if a single sentence.

[OriolAbril]

according to sphinxinv, the objects in bokeh docs containing circle are these:

:py:class:`bokeh.models.glyphs.Circle`
:py:attribute:`bokeh.models.glyphs.Circle.x`
:py:attribute:`bokeh.models.glyphs.Circle.y`
:py:function:`bokeh.models.markers.Circle`
:py:function:`bokeh.models.markers.CircleX`
:py:function:`bokeh.models.markers.CircleY`
:py:method:`bokeh.plotting.Figure.circle`
:py:method:`bokeh.plotting.Figure.circle_x`
:py:method:`bokeh.plotting.Figure.circle_y`
:py:method:`bokeh.plotting.GMap.circle`
:py:method:`bokeh.plotting.GMap.circle_x`
:std:label:`bokeh.models.glyphs.circle`
:std:label:`bokeh.models.markers.circle`
:std:label:`bokeh.models.markers.circledot`
:std:label:`bokeh.models.markers.circlex`
:std:label:`bokeh.models.markers.circley`

the one you used is not an option, maybe it was in an old version?

[OriolAbril]

this also needs some explanation as well as the intersphinx references.

[OriolAbril]

Suggested change

	Passed to az.plot_hdi()
	Passed to {func}`~arviz.plot_hdi`

[utkarsh-maheshwari]

I don't know how I missed this one!

[OriolAbril]

the links to the methods where these are passed are still missing, you can take them from #1728 to avoid having to track them down in the source code.