Type-hint and type-overload a few common Model methods#790
Type-hint and type-overload a few common Model methods#790thomasaarholt wants to merge 4 commits intobambinos:mainfrom
Conversation
whitespace fit vi is meanfield fit laplace
|
Thanks for the contribution @thomasaarholt! I'll let you know what I think once I can review. For now, let's run the CI :) |
Codecov ReportAll modified and coverable lines are covered by tests ✅
Additional details and impacted files@@ Coverage Diff @@
## main #790 +/- ##
==========================================
+ Coverage 90.14% 90.17% +0.02%
==========================================
Files 46 46
Lines 3836 3846 +10
==========================================
+ Hits 3458 3468 +10
Misses 378 378 ☔ View full report in Codecov by Sentry. |
|
Great! No rush at all, just had some downtime in my Easter break and thought it would be a nice little contribution. |
| discard_tuned_samples: bool = True, | ||
| omit_offsets: bool = True, | ||
| include_mean: bool = False, | ||
| inference_method: Literal["mcmc", "blackjax_nuts", "numpyro_nuts", "laplace"] = "mcmc", |
There was a problem hiding this comment.
One problem I see here is that after #775, users have access to a huge variety of samplers through bayeux. I'm not sure how this is handled with literals. But, is it possible that what we need is to create a Literal on the fly (at import time) where we check whether users have access to bayeux, and depending on that, the Literal contains different types?
Also, I don't think there's a convenient way to access the list of samplers without first creating a model. I'm going to create an issue now.
| idata: InferenceData, | ||
| kind: Literal["mean", "pps"] = "mean", | ||
| data: Optional[pd.DataFrame] = None, | ||
| inplace: Literal[True] = True, |
There was a problem hiding this comment.
Shouldn't this be bool instead of Literal[True]? I'm not a typing expert so I may be missing some details, but since this can be either True or False, I think it should be bool
| idata: InferenceData, | ||
| kind: Literal["mean", "pps"] = "mean", | ||
| data: Optional[pd.DataFrame] = None, | ||
| inplace: Literal[False] = False, |
Hello!
I have used pymc extensively over the last two years, and have finally had a play with bambi. I really like it - it has some really nice defaults, and the formula syntax is great!
I noticed that some type hinting was missing on the most commonly used Model methods, so I decided to add that. I initially just wanted to add some overload methods to
Model.fit, so that it would clearly be returning anInferenceDataobject when called with the appropriateinference_methods. Then I got a bit carried away, since it wasn't very difficult to add the remaining commonly used methods. I've previously added type hinting to some of the functions in pymc.I am using pyright to check the types. This is what is used in VSCode for type checking, docstrings and autocompletion.
This PR does not change any behaviour or logic, only adds type hinting that make IDEs better at understanding the code.
Here is a screenshot of my editor from before the changes are made:

And here they are after:

Here is the run.py file in a gist.