-
Notifications
You must be signed in to change notification settings - Fork 30
style: add type hinting to variational distributions #261
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
base: main
Are you sure you want to change the base?
style: add type hinting to variational distributions #261
Conversation
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
This PR adds comprehensive type hinting to the variational distributions module in QUEENS, introducing two new type aliases (Array1xN and ArrayNxM) for better shape indication in numpy arrays. The PR also refactors FullRankNormal.reconstruct_distribution_parameters() to satisfy mypy's requirement for consistent return types by creating a separate method for cases where the Cholesky decomposition is needed, and enables mypy checking for the variational distributions module by removing it from exclusion lists.
Key Changes
- Introduced new type aliases
Array1xNandArrayNxMinsrc/queens/utils/type_hinting.pyto indicate array shapes - Added type hints throughout the variational distributions hierarchy including the base class and all concrete implementations
- Split
FullRankNormal.reconstruct_distribution_parameters()into two methods to resolve mypy's varying return count issue
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 27 comments.
Show a summary per file
| File | Description |
|---|---|
src/queens/utils/type_hinting.py |
New file introducing Array1xN and ArrayNxM type aliases for shaped numpy arrays |
src/queens/variational_distributions/_variational_distribution.py |
Added type hints to base class including Array1xN for variational parameters and proper return types |
src/queens/variational_distributions/particle.py |
Added type hints to all methods, updated docstrings, and modified initialization to pass n_parameters |
src/queens/variational_distributions/mixture_model.py |
Added type hints, improved variable naming (logpdf_lst, parameters), and initialized FIM properly |
src/queens/variational_distributions/mean_field_normal.py |
Added type hints to all methods and updated docstrings for consistency |
src/queens/variational_distributions/full_rank_normal.py |
Added type hints, split reconstruct_distribution_parameters() into two methods for mypy compatibility |
src/queens/variational_distributions/joint.py |
Added type hints, fixed logpdf initialization to use proper numpy array instead of scalar |
src/queens/variational_distributions/__init__.py |
Added type hints to __getattr__ function and imported base class for type checking |
pyproject.toml |
Removed variational_distributions from mypy exclusion list to enable type checking |
.pre-commit-config.yaml |
Removed variational_distributions from mypy exclusion pattern in pre-commit hooks |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
b16cf83 to
b69c2b1
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
Copilot reviewed 10 out of 10 changed files in this pull request and generated 5 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
src/queens/variational_distributions/_variational_distribution.py
Outdated
Show resolved
Hide resolved
style: add type hinting to variational distributions
b69c2b1 to
8dec669
Compare
gilrrei
left a comment
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think the return statements could be more precise :)
src/queens/variational_distributions/_variational_distribution.py
Outdated
Show resolved
Hide resolved
src/queens/variational_distributions/_variational_distribution.py
Outdated
Show resolved
Hide resolved
src/queens/variational_distributions/_variational_distribution.py
Outdated
Show resolved
Hide resolved
d161ffc to
87ad5eb
Compare
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Pull request overview
Copilot reviewed 9 out of 9 changed files in this pull request and generated 23 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
src/queens/variational_distributions/_variational_distribution.py
Outdated
Show resolved
Hide resolved
87ad5eb to
79692a8
Compare
| NDims = TypeVar("NDims", bound=int) | ||
| NSamples = TypeVar("NSamples", bound=int) | ||
| NParams = TypeVar("NParams", bound=int) | ||
| NParamsComponent = TypeVar("NParamsComponent", bound=int) | ||
|
|
||
| # Vectors | ||
| ArrayNDims: TypeAlias = np.ndarray[tuple[NDims], np.dtype[np.floating]] | ||
| ArrayNParams: TypeAlias = np.ndarray[tuple[NParams], np.dtype[np.floating]] | ||
| ArrayNParamsComponent: TypeAlias = np.ndarray[tuple[NParamsComponent], np.dtype[np.floating]] | ||
| ArrayNSamples: TypeAlias = np.ndarray[tuple[NSamples], np.dtype[np.floating]] | ||
|
|
||
| # Matrices | ||
| Array1XNParams: TypeAlias = np.ndarray[tuple[Literal[1], NParams], np.dtype[np.floating]] | ||
| ArrayNDimsX1: TypeAlias = np.ndarray[tuple[NDims, Literal[1]], np.dtype[np.floating]] | ||
| ArrayNDimsXNDims: TypeAlias = np.ndarray[tuple[NDims, NDims], np.dtype[np.floating]] | ||
| ArrayNParamsXNParams: TypeAlias = np.ndarray[tuple[NParams, NParams], np.dtype[np.floating]] | ||
| ArrayNParamsXNSamples: TypeAlias = np.ndarray[tuple[NParams, NSamples], np.dtype[np.floating]] | ||
| ArrayNSamplesXNDims: TypeAlias = np.ndarray[tuple[NSamples, NDims], np.dtype[np.floating]] | ||
| ArrayNSamplesXNParams: TypeAlias = np.ndarray[tuple[NSamples, NParams], np.dtype[np.floating]] |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I love this!
I was always hoping we would be able to do something like this. It will help us a lot! 🙌
Description and Context:
What and Why?
This PR adds type hinting to our variational distributions.
I have introduced new type aliases to indicate the shapes of numpy arrays, since the docstrings provided this information for some of the variables. In some cases, the shapes were outdated, though, as I realized when I checked the shapes with our tests. In my opinion, it would be really helpful to eventually replace all or most of the
np.ndarrays in our code base with these shape-indicating aliases.FYI: I had to adapt
FullRankNormal.reconstruct_distribution_parameters()because mypy doesn't like functions with a varying number of returned variables.Related Issues and Pull Requests
Interested Parties
@gilrrei