Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Type-hint and type-overload a few common Model methods #790

Open
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

thomasaarholt
Copy link

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 an InferenceData object when called with the appropriate inference_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:
image

And here they are after:
image

Here is the run.py file in a gist.

@thomasaarholt thomasaarholt changed the title Overload .fit Type-hint and type-overload a few common Model methods Mar 29, 2024
@tomicapretto
Copy link
Collaborator

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-commenter
Copy link

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 90.17%. Comparing base (714ccb7) to head (3afbc00).

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.
📢 Have feedback on the report? Share it here.

@thomasaarholt
Copy link
Author

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",
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same as above

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants