Skip to content

Conversation

@leahaeusel
Copy link
Member

@leahaeusel leahaeusel commented Dec 8, 2025

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

Note: More information on the merge request procedure in QUEENS can be found in the Submit a pull request section in the CONTRIBUTING.md file.

Copy link

Copilot AI left a 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 Array1xN and ArrayNxM in src/queens/utils/type_hinting.py to 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.

@leahaeusel leahaeusel force-pushed the add-type-hinting-to-variational-distributions branch from b16cf83 to b69c2b1 Compare December 9, 2025 08:22
@leahaeusel leahaeusel requested a review from Copilot December 9, 2025 08:23
Copy link

Copilot AI left a 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.

style: add type hinting to variational distributions
@leahaeusel leahaeusel force-pushed the add-type-hinting-to-variational-distributions branch from b69c2b1 to 8dec669 Compare December 9, 2025 08:43
@leahaeusel leahaeusel marked this pull request as ready for review December 9, 2025 08:59
Copy link
Member

@gilrrei gilrrei left a 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 :)

@leahaeusel leahaeusel marked this pull request as draft December 9, 2025 12:09
@leahaeusel leahaeusel force-pushed the add-type-hinting-to-variational-distributions branch 3 times, most recently from d161ffc to 87ad5eb Compare December 10, 2025 12:36
@leahaeusel leahaeusel requested a review from Copilot December 10, 2025 12:37
Copy link

Copilot AI left a 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.

@leahaeusel leahaeusel force-pushed the add-type-hinting-to-variational-distributions branch from 87ad5eb to 79692a8 Compare December 10, 2025 13:20
@leahaeusel leahaeusel marked this pull request as ready for review December 10, 2025 15:24
Comment on lines +24 to +42
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]]
Copy link
Member

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! 🙌

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.

3 participants