From 8dec6696da74843214284e2e9c7820f834b60df9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lea=20H=C3=A4usel?= Date: Mon, 1 Dec 2025 16:55:21 +0100 Subject: [PATCH 1/2] style: add type hinting to variational distributions style: add type hinting to variational distributions --- .pre-commit-config.yaml | 2 +- pyproject.toml | 2 +- src/queens/utils/type_hinting.py | 28 +++ .../variational_distributions/__init__.py | 5 +- .../_variational_distribution.py | 114 +++++++--- .../full_rank_normal.py | 209 ++++++++++-------- src/queens/variational_distributions/joint.py | 133 +++++------ .../mean_field_normal.py | 163 +++++++------- .../mixture_model.py | 138 ++++++------ .../variational_distributions/particle.py | 96 ++++---- 10 files changed, 510 insertions(+), 380 deletions(-) create mode 100644 src/queens/utils/type_hinting.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 44d4d0ff7..a081578cc 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -100,7 +100,7 @@ repos: ^(.gitlab|.github|config|doc|tests|test_utils)/| ^src/(example_simulator_functions|queens_interfaces)/| ^src/queens/(data_processors|drivers|iterators|models)/| - ^src/queens/(schedulers|stochastic_optimizers|variational_distributions|visualization)/| + ^src/queens/(schedulers|stochastic_optimizers|visualization)/| ^src/queens/(main.py|global_settings.py) ).*$ - repo: https://github.com/kynan/nbstripout diff --git a/pyproject.toml b/pyproject.toml index 73b2ecdcc..29b9cbae8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -171,7 +171,7 @@ exclude = '''(?x)( ^(.gitlab|.github|config|doc|tests|test_utils)/| ^src/(example_simulator_functions|queens_interfaces)/| ^src/queens/(data_processors|drivers|iterators|models)/| - ^src/queens/(schedulers|stochastic_optimizers|variational_distributions|visualization)/| + ^src/queens/(schedulers|stochastic_optimizers|visualization)/| ^src/queens/(main.py|global_settings.py) ).*$''' [[tool.mypy.overrides]] diff --git a/src/queens/utils/type_hinting.py b/src/queens/utils/type_hinting.py new file mode 100644 index 000000000..7dc54d66a --- /dev/null +++ b/src/queens/utils/type_hinting.py @@ -0,0 +1,28 @@ +# +# SPDX-License-Identifier: LGPL-3.0-or-later +# Copyright (c) 2024-2025, QUEENS contributors. +# +# This file is part of QUEENS. +# +# QUEENS is free software: you can redistribute it and/or modify it under the terms of the GNU +# Lesser General Public License as published by the Free Software Foundation, either version 3 of +# the License, or (at your option) any later version. QUEENS is distributed in the hope that it will +# be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or +# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details. You +# should have received a copy of the GNU Lesser General Public License along with QUEENS. If not, +# see . +# +"""Utilities for type hinting.""" + +from typing import Literal, TypeAlias + +import numpy as np + +# pylint: disable=invalid-name + +ArrayN: TypeAlias = np.ndarray[tuple[int], np.dtype[np.floating]] +Array1xN: TypeAlias = np.ndarray[tuple[Literal[1], int], np.dtype[np.floating]] +ArrayNx1: TypeAlias = np.ndarray[tuple[int, Literal[1]], np.dtype[np.floating]] +ArrayNxM: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.floating]] + +Array1D: TypeAlias = ArrayN | Array1xN | ArrayNx1 diff --git a/src/queens/variational_distributions/__init__.py b/src/queens/variational_distributions/__init__.py index 233d836a6..44f4e4fb0 100644 --- a/src/queens/variational_distributions/__init__.py +++ b/src/queens/variational_distributions/__init__.py @@ -16,11 +16,14 @@ Modules containing probability distributions for variational inference. """ +from __future__ import annotations + from typing import TYPE_CHECKING from queens.utils.imports import extract_type_checking_imports, import_class_from_class_module_map if TYPE_CHECKING: + from queens.variational_distributions._variational_distribution import Variational from queens.variational_distributions.full_rank_normal import FullRankNormal from queens.variational_distributions.joint import Joint from queens.variational_distributions.mean_field_normal import MeanFieldNormal @@ -30,5 +33,5 @@ class_module_map = extract_type_checking_imports(__file__) -def __getattr__(name): +def __getattr__(name: str) -> Variational: return import_class_from_class_module_map(name, class_module_map, __name__) diff --git a/src/queens/variational_distributions/_variational_distribution.py b/src/queens/variational_distributions/_variational_distribution.py index f5b5f1edd..353b7bde3 100644 --- a/src/queens/variational_distributions/_variational_distribution.py +++ b/src/queens/variational_distributions/_variational_distribution.py @@ -15,95 +15,151 @@ """Variational Distribution.""" import abc +from typing import Any + +import numpy as np + +from queens.utils.type_hinting import ArrayN, ArrayNxM class Variational: """Base class for probability distributions for variational inference. Attributes: - dimension (int): dimension of the distribution + dimension: Dimension of the distribution + n_parameters: Number of variational parameters """ - def __init__(self, dimension): - """Initialize variational distribution.""" + def __init__(self, dimension: int, n_parameters: int) -> None: + """Initialize variational distribution. + + Args: + dimension: Dimension of the variational distribution + n_parameters: Number of variational parameters + """ self.dimension = dimension + self.n_parameters = n_parameters + + @abc.abstractmethod + def construct_variational_parameters(self, *args: Any) -> np.ndarray: + """Construct variational parameters from distribution parameters. + + Args: + args: Distribution parameters + + Returns: + Variational parameters + """ @abc.abstractmethod - def reconstruct_distribution_parameters(self, variational_parameters): + def reconstruct_distribution_parameters(self, variational_parameters: ArrayN) -> Any: """Reconstruct distribution parameters from variational parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters + + Returns: + Distribution parameters """ @abc.abstractmethod - def draw(self, variational_parameters, n_draws=1): + def draw( + self, + variational_parameters: ArrayN, + n_draws: int = 1, + ) -> ArrayNxM: """Draw *n_draws* samples from distribution. Args: - variational_parameters (np.ndarray): variational parameters (1 x n_params) - n_draws (int): Number of samples + variational_parameters: Variational parameters of shape (n_params,) + n_draws: Number of samples + + Returns: + Drawn samples of shape (n_draws, n_dim) """ @abc.abstractmethod - def logpdf(self, variational_parameters, x): - """Evaluate the natural logarithm of the logpdf at sample. + def logpdf( + self, + variational_parameters: ArrayN, + x: ArrayNxM, + ) -> np.ndarray: + """Evaluate the natural logarithm of the PDF. Args: - variational_parameters (np.ndarray): variational parameters (1 x n_params) - x (np.ndarray): Locations to evaluate (n_samples x n_dim) + variational_parameters: Variational parameters of shape (n_params,) + x: Locations to evaluate of shape (n_samples, n_dim) + + Returns: + Row vector of the Log-PDF values """ @abc.abstractmethod - def pdf(self, variational_parameters, x): - """Evaluate the probability density function (pdf) at sample. + def pdf( + self, + variational_parameters: ArrayN, + x: ArrayNxM, + ) -> np.ndarray: + """Evaluate the probability density function (PDF). Args: - variational_parameters (np.ndarray): variational parameters (1 x n_params) - x (np.ndarray): Locations to evaluate (n_samples x n_dim) + variational_parameters: Variational parameters of shape (n_params,) + x: Locations to evaluate of shape (n_samples, n_dim) + + Returns: + Row vector of the PDF values """ @abc.abstractmethod - def grad_params_logpdf(self, variational_parameters, x): - """Logpdf gradient w.r.t. the variational parameters. + def grad_params_logpdf( + self, + variational_parameters: ArrayN, + x: ArrayNxM, + ) -> np.ndarray: + """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. Args: - variational_parameters (np.ndarray): variational parameters (1 x n_params) - x (np.ndarray): Locations to evaluate (n_samples x n_dim) + variational_parameters: Variational parameters of shape (n_params,) + x: Locations to evaluate of shape (n_samples, n_dim) + + Returns: + Gradient of the log-PDF w.r.t. the variational parameters """ @abc.abstractmethod - def fisher_information_matrix(self, variational_parameters): + def fisher_information_matrix(self, variational_parameters: ArrayN) -> ArrayNxM: """Compute the fisher information matrix. Depends on the variational distribution for the given parameterization. Args: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + variational_parameters: Variational parameters of shape (n_params,) + + Returns: + Fisher information matrix of shape (n_params, n_params) """ @abc.abstractmethod - def initialize_variational_parameters(self, random=False): + def initialize_variational_parameters(self, random: bool = False) -> ArrayN: """Initialize variational parameters. Args: - random (bool, optional): If True, a random initialization is used. Otherwise the - default is selected + random: If True, a random initialization is used. Otherwise the default is selected. Returns: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + Variational parameters of shape (n_params,) """ @abc.abstractmethod - def export_dict(self, variational_parameters): + def export_dict(self, variational_parameters: np.ndarray) -> dict: """Create a dict of the distribution based on the given parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - export_dict (dictionary): Dict containing distribution information + Dictionary containing distribution information """ diff --git a/src/queens/variational_distributions/full_rank_normal.py b/src/queens/variational_distributions/full_rank_normal.py index e0006df10..3274632d0 100644 --- a/src/queens/variational_distributions/full_rank_normal.py +++ b/src/queens/variational_distributions/full_rank_normal.py @@ -19,6 +19,7 @@ from numba import njit from queens.utils.logger_settings import log_init_args +from queens.utils.type_hinting import ArrayN, ArrayNx1, ArrayNxM from queens.variational_distributions._variational_distribution import Variational @@ -37,20 +38,20 @@ class FullRankNormal(Variational): The Journal of Machine Learning Research 18.1 (2017): 430-474. Attributes: - n_parameters (int): Number of parameters used in the parameterization. + n_parameters: Number of parameters used in the parameterization. """ @log_init_args - def __init__(self, dimension): + def __init__(self, dimension: int) -> None: """Initialize variational distribution. Args: - dimension (int): dimension of the RV + dimension: Dimension of the RV """ - super().__init__(dimension) - self.n_parameters = (dimension * (dimension + 1)) // 2 + dimension + n_parameters = (dimension * (dimension + 1)) // 2 + dimension + super().__init__(dimension, n_parameters) - def initialize_variational_parameters(self, random=False): + def initialize_variational_parameters(self, random: bool = False) -> ArrayN: r"""Initialize variational parameters. Default initialization: @@ -60,11 +61,10 @@ def initialize_variational_parameters(self, random=False): :math:`\mu=Uniform(-0.1,0.1)` :math:`L=diag(Uniform(0.9,1.1))` where :math:`\Sigma=LL^T` Args: - random (bool, optional): If True, a random initialization is used. Otherwise the - default is selected + random: If True, a random initialization is used. Otherwise the default is selected. Returns: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + Variational parameters of shape (n_params,) """ if random: cholesky_covariance = np.eye(self.dimension) + 0.1 * ( @@ -84,16 +84,17 @@ def initialize_variational_parameters(self, random=False): return variational_parameters - @staticmethod - def construct_variational_parameters(mean, covariance): + def construct_variational_parameters( # pylint: disable=arguments-differ + self, mean: ArrayNx1 | ArrayN, covariance: ArrayNxM + ) -> ArrayN: """Construct the variational parameters from mean and covariance. Args: - mean (np.ndarray): Mean values of the distribution (n_dim x 1) - covariance (np.ndarray): Covariance matrix of the distribution (n_dim x n_dim) + mean: Mean values of the distribution of shape (n_dim, 1) or (n_dim,) + covariance: Covariance matrix of the distribution of shape (n_dim, n_dim) Returns: - variational_parameters (np.ndarray): Variational parameters + Variational parameters """ if len(mean) == len(covariance): cholesky_covariance = np.linalg.cholesky(covariance) @@ -109,16 +110,33 @@ def construct_variational_parameters(mean, covariance): ) return variational_parameters - def reconstruct_distribution_parameters(self, variational_parameters, return_cholesky=False): + def reconstruct_distribution_parameters( + self, variational_parameters: ArrayN + ) -> tuple[ArrayNx1, ArrayNxM]: """Reconstruct mean value, covariance and its Cholesky decomposition. Args: - variational_parameters (np.ndarray): Variational parameters - return_cholesky (bool, optional): Return the L if desired + variational_parameters: Variational parameters Returns: - mean (np.ndarray): Mean value of the distribution (n_dim x 1) - cov (np.ndarray): Covariance of the distribution (n_dim x n_dim) - L (np.ndarray): Cholesky decomposition of the covariance matrix (n_dim x n_dim) + Mean value of the distribution of shape (n_dim, 1) + Covariance of the distribution of shape (n_dim, n_dim) + """ + mean, cov, _ = self.reconstruct_distribution_parameters_with_cholesky( + variational_parameters + ) + return mean, cov + + def reconstruct_distribution_parameters_with_cholesky( + self, variational_parameters: ArrayN + ) -> tuple[ArrayNx1, ArrayNxM, ArrayNxM]: + """Reconstruct mean value, covariance and its Cholesky decomposition. + + Args: + variational_parameters: Variational parameters + Returns: + Mean value of the distribution of shape (n_dim, 1) + Covariance of the distribution of shape (n_dim, n_dim) + Cholesky decomposition of the covariance matrix of shape (n_dim, n_dim) """ mean = variational_parameters[: self.dimension].reshape(-1, 1) cholesky_covariance_array = variational_parameters[self.dimension :] @@ -127,54 +145,50 @@ def reconstruct_distribution_parameters(self, variational_parameters, return_cho cholesky_covariance[idx] = cholesky_covariance_array cov = np.matmul(cholesky_covariance, cholesky_covariance.T) - if return_cholesky: - return mean, cov, cholesky_covariance + return mean, cov, cholesky_covariance - return mean, cov - - def _grad_reconstruct_distribution_parameters(self): + def _grad_reconstruct_distribution_parameters(self) -> np.ndarray: """Gradient of the parameter reconstruction. Returns: - grad_reconstruct_params (np.ndarray): Gradient vector of the reconstruction - w.r.t. the variational parameters + Gradient vector of the reconstruction w.r.t. the variational parameters """ grad_reconstruct_params = np.ones((1, self.n_parameters)) return grad_reconstruct_params - def draw(self, variational_parameters, n_draws=1): + def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: """Draw *n_draw* samples from the variational distribution. Args: - variational_parameters (np.ndarray): Variational parameters - n_draws (int): Number of samples to draw + variational_parameters: Variational parameters + n_draws: Number of samples to draw Returns: - samples (np.ndarray): Row-wise samples of the variational distribution + Samples of shape (n_draws, n_dim) """ - mean, _, cholesky = self.reconstruct_distribution_parameters( - variational_parameters, return_cholesky=True + mean, _, cholesky = self.reconstruct_distribution_parameters_with_cholesky( + variational_parameters ) sample = np.dot(cholesky, np.random.randn(self.dimension, n_draws)).T + mean.reshape(1, -1) return sample - def logpdf(self, variational_parameters, x): - """Logpdf evaluated using the at samples *x*. + def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF evaluated at the samples *x*. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - logpdf (np.ndarray): Row vector of the logpdfs + Row vector of the Log-PDF values """ - mean, cov, cholesky = self.reconstruct_distribution_parameters( - variational_parameters, return_cholesky=True + mean, cov, cholesky = self.reconstruct_distribution_parameters_with_cholesky( + variational_parameters ) x = np.atleast_2d(x) u = np.linalg.solve(cov, (x.T - mean)) - def col_dot_prod(x, y): + def col_dot_prod(x: np.ndarray, y: np.ndarray) -> np.ndarray: return np.sum(x * y, axis=0) logpdf = ( @@ -184,35 +198,35 @@ def col_dot_prod(x, y): ) return logpdf.flatten() - def pdf(self, variational_parameters, x): - """Pdf of evaluated at given samples *x*. + def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """PDF of evaluated at given samples *x*. - First computes the logpdf, which is numerically more stable for exponential distributions. + First computes the log-PDF, which is numerically more stable for exponential distributions. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - pdf (np.ndarray): Row vector of the pdfs + Row vector of the PDF values """ pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters, x): - """Logpdf gradient w.r.t. to the variational parameters. + def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF gradient w.r.t. to the variational parameters. Evaluated at samples *x*. Also known as the score function. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - score (np.ndarray): Column-wise scores + Column-wise scores """ - mean, cov, cholesky = self.reconstruct_distribution_parameters( - variational_parameters, return_cholesky=True + mean, cov, cholesky = self.reconstruct_distribution_parameters_with_cholesky( + variational_parameters ) x = np.atleast_2d(x) # Helper variable @@ -236,17 +250,19 @@ def grad_params_logpdf(self, variational_parameters, x): score = np.vstack((dlogpdf_dmu, dlogpdf_dsigma)) return score - def total_grad_params_logpdf(self, variational_parameters, standard_normal_sample_batch): - """Total logpdf reparameterization gradient. + def total_grad_params_logpdf( + self, variational_parameters: ArrayN, standard_normal_sample_batch: np.ndarray + ) -> np.ndarray: + """Total log-PDF reparameterization gradient. - Total logpdf reparameterization gradient w.r.t. the variational parameters. + Total log-PDF reparameterization gradient w.r.t. the variational parameters. Args: - variational_parameters (np.ndarray): Variational parameters - standard_normal_sample_batch (np.ndarray): Standard normal distributed sample batch + variational_parameters: Variational parameters + standard_normal_sample_batch: Standard normal distributed sample batch Returns: - total_grad (np.ndarray): Total Logpdf reparameterization gradient + Total log-PDF reparameterization gradient """ idx = np.tril_indices(self.dimension, k=0, m=self.dimension) cholesky_diagonal_idx = np.where(np.equal(*idx))[0] + self.dimension @@ -254,22 +270,21 @@ def total_grad_params_logpdf(self, variational_parameters, standard_normal_sampl total_grad[:, cholesky_diagonal_idx] = -1 / variational_parameters[cholesky_diagonal_idx] return total_grad - def grad_sample_logpdf(self, variational_parameters, sample_batch): - """Computes the gradient of the logpdf w.r.t. to the *x*. + def grad_sample_logpdf( + self, variational_parameters: ArrayN, sample_batch: np.ndarray + ) -> np.ndarray: + """Computes the gradient of the log-PDF w.r.t. to the sample *x*. Args: - variational_parameters (np.ndarray): Variational parameters - sample_batch (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + sample_batch: Row-wise samples Returns: - gradients_batch (np.ndarray): Gradients of the log-pdf w.r.t. the - sample *x*. The first dimension of the - array corresponds to the different samples. - The second dimension to different dimensions - within one sample. (Third dimension is empty - and just added to keep slices two-dimensional.) + Gradients of the log-pdf w.r.t. the sample *x*. The first dimension of the array + corresponds to the different samples. The second dimension to different dimensions + within one sample. (Third dimension is empty and just added to keep slices + two-dimensional.) """ - # pylint: disable-next=unbalanced-tuple-unpacking mean, cov = self.reconstruct_distribution_parameters(variational_parameters) gradient_lst = [] for sample in sample_batch: @@ -280,20 +295,20 @@ def grad_sample_logpdf(self, variational_parameters, sample_batch): gradients_batch = np.array(gradient_lst) return gradients_batch.reshape(sample_batch.shape) - def fisher_information_matrix(self, variational_parameters): + def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarray: """Compute the Fisher information matrix analytically. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - FIM (np.ndarray): Matrix (num parameters x num parameters) + Fisher information matrix (num parameters x num parameters) """ - _, cov, cholesky = self.reconstruct_distribution_parameters( - variational_parameters, return_cholesky=True + _, cov, cholesky = self.reconstruct_distribution_parameters_with_cholesky( + variational_parameters ) - def fim_blocks(dimension): + def fim_blocks(dimension: int) -> tuple[np.ndarray, np.ndarray]: """Compute the blocks of the FIM.""" mu_block = np.linalg.inv(cov + 1e-8 * np.eye(len(cov))) n_params_chol = (dimension * (dimension + 1)) // 2 @@ -324,16 +339,14 @@ def fim_blocks(dimension): return scipy.linalg.block_diag(mu_block, sigma_block) - def export_dict(self, variational_parameters): + def export_dict(self, variational_parameters: ArrayN) -> dict: """Create a dict of the distribution based on the given parameters. Args: - variational_parameters (np.ndarray): Variational parameters - + variational_parameters: Variational parameters Returns: - export_dict (dictionary): Dict containing distribution information + Dictionary containing distribution information """ - # pylint: disable-next=unbalanced-tuple-unpacking mean, cov = self.reconstruct_distribution_parameters(variational_parameters) export_dict = { "type": "fullrank_Normal", @@ -343,38 +356,42 @@ def export_dict(self, variational_parameters): } return export_dict - def conduct_reparameterization(self, variational_parameters, n_samples): + def conduct_reparameterization( + self, variational_parameters: np.ndarray, n_samples: int + ) -> tuple[np.ndarray, np.ndarray]: """Conduct a reparameterization. Args: - variational_parameters (np.ndarray): Array with variational parameters - n_samples (int): Number of samples for current batch + variational_parameters: Array with variational parameters + n_samples: Number of samples for current batch Returns: - samples_mat (np.ndarray): Array of actual samples from the variational - distribution + Actual samples from the variational distribution + Standard normal distributed samples used for the reparameterization """ standard_normal_sample_batch = np.random.normal(0, 1, size=(n_samples, self.dimension)) - mean, _, cholesky = self.reconstruct_distribution_parameters( - variational_parameters, return_cholesky=True + mean, _, cholesky = self.reconstruct_distribution_parameters_with_cholesky( + variational_parameters ) samples_mat = mean + np.dot(cholesky, standard_normal_sample_batch.T) return samples_mat.T, standard_normal_sample_batch def grad_params_reparameterization( - self, variational_parameters, standard_normal_sample_batch, upstream_gradient - ): + self, + variational_parameters: ArrayN, + standard_normal_sample_batch: np.ndarray, + upstream_gradient: np.ndarray, + ) -> np.ndarray: r"""Calculate the gradient of the reparameterization. Args: - variational_parameters (np.ndarray): Variational parameters - standard_normal_sample_batch (np.ndarray): Standard normal distributed sample batch - upstream_gradient (np.array): Upstream gradient + variational_parameters: Variational parameters + standard_normal_sample_batch: Standard normal distributed sample batch + upstream_gradient: Upstream gradient Returns: - gradient (np.ndarray): Gradient of the upstream function w.r.t. the variational - parameters. + Gradient of the upstream function w.r.t. the variational parameters. **Note:** We assume that *grad_reconstruct_params* is a row-vector containing the partial diff --git a/src/queens/variational_distributions/joint.py b/src/queens/variational_distributions/joint.py index 7ec2caa51..5777f68ce 100644 --- a/src/queens/variational_distributions/joint.py +++ b/src/queens/variational_distributions/joint.py @@ -17,6 +17,7 @@ import numpy as np import scipy +from queens.utils.type_hinting import ArrayN, ArrayNxM from queens.variational_distributions._variational_distribution import Variational @@ -30,32 +31,29 @@ class Joint(Variational): is a generalization of the mean field distribution Attributes: - distributions (list): List of variational distribution objects for the different - independent distributions. - n_parameters (int): Total number of parameters used in the parameterization. - distributions_n_parameters (np.ndarray): Number of parameters per distribution - distributions_dimension (np.ndarray): Number of dimension per distribution + distributions: Variational distribution objects for the different independent distributions. + n_parameters: Total number of parameters used in the parameterization. + distributions_n_parameters: Number of parameters per distribution + distributions_dimension: Number of dimension per distribution """ - def __init__(self, distributions, dimension): + def __init__(self, distributions: list, dimension: int) -> None: """Initialize joint distribution. Args: - dimension (int): Dimension of the random variable - distributions (list): List of variational distribution objects for the different - independent distributions. + distributions: Variational distribution objects for the different independent + distributions. + dimension: Dimension of the random variable """ - super().__init__(dimension) self.distributions = distributions - self.distributions_n_parameters = np.array( - [distribution.n_parameters for distribution in distributions] + [distribution.n_parameters for distribution in self.distributions] ).astype(int) - self.n_parameters = int(np.sum(self.distributions_n_parameters)) + super().__init__(dimension, n_parameters=int(np.sum(self.distributions_n_parameters))) self.distributions_dimension = np.array( - [distribution.dimension for distribution in distributions] + [distribution.dimension for distribution in self.distributions] ).astype(int) if dimension != np.sum(self.distributions_dimension): @@ -64,17 +62,16 @@ def __init__(self, distributions, dimension): f"dimensions of the subdistributions {np.sum(self.distributions_dimension)}" ) - def initialize_variational_parameters(self, random=False): + def initialize_variational_parameters(self, random: bool = False) -> ArrayN: r"""Initialize variational parameters. The distribution initialization is handle by the component itself. Args: - random (bool, optional): If True, a random initialization is used. Otherwise the - default is selected + random: If True, a random initialization is used. Otherwise the default is selected Returns: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + Variational parameters of shape (n_params,) """ variational_parameters = np.concatenate( [ @@ -85,25 +82,29 @@ def initialize_variational_parameters(self, random=False): return variational_parameters - def construct_variational_parameters(self, distributions_parameters_list): + def construct_variational_parameters( # pylint: disable=arguments-differ + self, distributions_parameters: list + ) -> ArrayN: """Construct the variational parameters from the distribution list. Args: - distributions_parameters_list (list): List of the parameters of the distributions + distributions_parameters: Parameters of the distributions Returns: - variational_parameters (np.ndarray): Variational parameters + Variational parameters """ variational_parameters = [] for parameters, distribution in zip( - distributions_parameters_list, self.distributions, strict=True + distributions_parameters, self.distributions, strict=True ): variational_parameters.append( distribution.construct_variational_parameters(*parameters) ) return np.concatenate(variational_parameters) - def _construct_distributions_variational_parameters(self, variational_parameters): + def _construct_distributions_variational_parameters( + self, variational_parameters: ArrayN + ) -> list: """Reconstruct the parameters of the distributions. Creates a list containing the variational parameters of the different components. @@ -111,27 +112,26 @@ def _construct_distributions_variational_parameters(self, variational_parameters The list is nested, each entry correspond to the parameters of a distribution. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - variational_parameters_list (list): List of the variational parameters of the components + Variational parameters of the components """ variational_parameters_list = split_array_by_chunk_sizes( variational_parameters, self.distributions_n_parameters ) return variational_parameters_list - def reconstruct_distribution_parameters(self, variational_parameters): + def reconstruct_distribution_parameters(self, variational_parameters: ArrayN) -> list[list]: """Reconstruct the parameters of distributions. The list is nested, each entry correspond to the parameters of a distribution. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - distribution_parameters_list (list): List of the distribution parameters of the - components + Distribution parameters of the components """ distribution_parameters_list = [] @@ -141,19 +141,18 @@ def reconstruct_distribution_parameters(self, variational_parameters): distribution_parameters_list.append( distribution.reconstruct_distribution_parameters(parameters) ) - return [distribution_parameters_list] - def _zip_variational_parameters_distributions(self, variational_parameters): + def _zip_variational_parameters_distributions(self, variational_parameters: ArrayN) -> zip: """Zip parameters and distributions. This helper function creates a generator for variational parameters and subdistribution. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - zip: of variational parameters and distributions + Zip of variational parameters and distributions """ return zip( split_array_by_chunk_sizes(variational_parameters, self.distributions_n_parameters), @@ -161,18 +160,20 @@ def _zip_variational_parameters_distributions(self, variational_parameters): strict=True, ) - def _zip_variational_parameters_distributions_samples(self, variational_parameters, samples): + def _zip_variational_parameters_distributions_samples( + self, variational_parameters: ArrayN, samples: np.ndarray + ) -> zip: """Zip parameters, samples and distributions. This helper function creates a generator for variational parameters, samples and subdistribution. Args: - variational_parameters (np.ndarray): Variational parameters - samples (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + samples: Row-wise samples Returns: - zip: of variational parameters, samples and distributions + Zip of variational parameters, samples and distributions """ return zip( split_array_by_chunk_sizes(variational_parameters, self.distributions_n_parameters), @@ -181,15 +182,15 @@ def _zip_variational_parameters_distributions_samples(self, variational_paramete strict=True, ) - def draw(self, variational_parameters, n_draws=1): + def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: """Draw *n_draw* samples from the variational distribution. Args: - variational_parameters (np.ndarray): Variational parameters - n_draws (int): Number of samples to draw + variational_parameters: Variational parameters + n_draws: Number of samples to draw Returns: - samples (np.ndarray): Row wise samples of the variational distribution + Samples of shape (n_draws, n_dim) """ sample_array = [] for parameters, distribution in self._zip_variational_parameters_distributions( @@ -198,17 +199,17 @@ def draw(self, variational_parameters, n_draws=1): sample_array.append(distribution.draw(parameters, n_draws)) return np.column_stack(sample_array) - def logpdf(self, variational_parameters, x): - """Logpdf evaluated using the variational parameters at samples *x*. + def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF evaluated using the variational parameters at samples *x*. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - logpdf (np.ndarray): Row vector of the logpdfs + Row vector of the Log-PDF values """ - logpdf = 0 + logpdf = np.zeros_like(x[:, 0], dtype=float) for ( parameters, samples, @@ -217,32 +218,32 @@ def logpdf(self, variational_parameters, x): logpdf += distribution.logpdf(parameters, samples) return logpdf - def pdf(self, variational_parameters, x): + def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: """Pdf evaluated using the variational parameters at given samples `x`. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - pdf (np.ndarray): Row vector of the pdfs + Row vector of the PDF values """ pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters, x): - """Logpdf gradient w.r.t. the variational parameters. + def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. Is a general implementation using the score functions of the components. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - score (np.ndarray): Column-wise scores + Column-wise scores """ score = [] for ( @@ -254,14 +255,14 @@ def grad_params_logpdf(self, variational_parameters, x): return np.row_stack(score) - def fisher_information_matrix(self, variational_parameters): + def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarray: """Approximate the Fisher information matrix using Monte Carlo. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - FIM (np.ndarray): Matrix (num parameters x num parameters) + Fisher information matrix (num parameters x num parameters) """ fim = [] for parameters, distribution in self._zip_variational_parameters_distributions( @@ -271,14 +272,14 @@ def fisher_information_matrix(self, variational_parameters): return scipy.linalg.block_diag(*fim) - def export_dict(self, variational_parameters): + def export_dict(self, variational_parameters: ArrayN) -> dict: """Create a dict of the distribution based on the given parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - export_dict (dictionary): Dict containing distribution information + Dictionary containing distribution information """ export_dict = { "type": "joint", @@ -294,15 +295,15 @@ def export_dict(self, variational_parameters): return export_dict -def split_array_by_chunk_sizes(array, chunk_sizes): +def split_array_by_chunk_sizes(array: np.ndarray, chunk_sizes: np.ndarray) -> list: """Split up array by a list of chunk sizes. Args: - array (np.ndarray): Array to be split - chunk_sizes (np.ndarray): List of chunk sizes + array: Array to be split + chunk_sizes: Chunk sizes Returns: - list: with the chunks + Chunks of the array """ if array.ndim > 2: raise ValueError( diff --git a/src/queens/variational_distributions/mean_field_normal.py b/src/queens/variational_distributions/mean_field_normal.py index db0e0626a..aafcf32fc 100644 --- a/src/queens/variational_distributions/mean_field_normal.py +++ b/src/queens/variational_distributions/mean_field_normal.py @@ -17,6 +17,7 @@ import numpy as np from queens.utils.logger_settings import log_init_args +from queens.utils.type_hinting import ArrayN, ArrayNx1, ArrayNxM from queens.variational_distributions._variational_distribution import Variational @@ -32,20 +33,19 @@ class MeanFieldNormal(Variational): The Journal of Machine Learning Research 18.1 (2017): 430-474. Attributes: - n_parameters (int): Number of parameters used in the parameterization. + n_parameters: Number of parameters used in the parameterization. """ @log_init_args - def __init__(self, dimension): + def __init__(self, dimension: int) -> None: """Initialize variational distribution. Args: - dimension (int): Dimension of RV. + dimension: Dimension of random variable. """ - super().__init__(dimension) - self.n_parameters = 2 * dimension + super().__init__(dimension, n_parameters=2 * dimension) - def initialize_variational_parameters(self, random=False): + def initialize_variational_parameters(self, random: bool = False) -> ArrayN: r"""Initialize variational parameters. Default initialization: @@ -55,11 +55,10 @@ def initialize_variational_parameters(self, random=False): :math:`\mu=Uniform(-0.1,0.1)` and :math:`\sigma^2=Uniform(0.9,1.1)` Args: - random (bool, optional): If True, a random initialization is used. Otherwise the - default is selected + random: If True, a random initialization is used. Otherwise the default is selected Returns: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + Variational parameters of shape (n_params,) """ if random: variational_parameters = np.hstack( @@ -73,16 +72,17 @@ def initialize_variational_parameters(self, random=False): return variational_parameters - @staticmethod - def construct_variational_parameters(mean, covariance): + def construct_variational_parameters( # pylint: disable=arguments-differ + self, mean: ArrayNx1 | ArrayN, covariance: ArrayNxM + ) -> ArrayN: """Construct the variational parameters from mean and covariance. Args: - mean (np.ndarray): Mean values of the distribution (n_dim x 1) - covariance (np.ndarray): Covariance matrix of the distribution (n_dim x n_dim) + mean: Mean values of the distribution of shape (n_dim, 1) or (n_dim,) + covariance: Covariance matrix of the distribution of shape (n_dim, n_dim) Returns: - variational_parameters (np.ndarray): Variational parameters + Variational parameters """ if len(mean) == len(covariance): variational_parameters = np.hstack((mean.flatten(), 0.5 * np.log(np.diag(covariance)))) @@ -93,15 +93,16 @@ def construct_variational_parameters(mean, covariance): ) return variational_parameters - def reconstruct_distribution_parameters(self, variational_parameters): + def reconstruct_distribution_parameters( + self, variational_parameters: ArrayN + ) -> tuple[ArrayNx1, ArrayNxM]: """Reconstruct mean and covariance from the variational parameters. Args: - variational_parameters (np.ndarray): Variational parameters - + variational_parameters: Variational parameters Returns: - mean (np.ndarray): Mean value of the distribution (n_dim x 1) - cov (np.ndarray): Covariance matrix of the distribution (n_dim x n_dim) + Mean value of the distribution of shape (n_dim, 1) + Covariance matrix of the distribution of shape (n_dim, n_dim) """ mean, cov = ( variational_parameters[: self.dimension], @@ -109,30 +110,31 @@ def reconstruct_distribution_parameters(self, variational_parameters): ) return mean.reshape(-1, 1), np.diag(cov) - def _grad_reconstruct_distribution_parameters(self, variational_parameters): + def _grad_reconstruct_distribution_parameters( + self, variational_parameters: ArrayN + ) -> np.ndarray: """Gradient of the parameter reconstruction. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - grad_reconstruct_params (np.ndarray): Gradient vector of the reconstruction - w.r.t. the variational parameters + Gradient vector of the reconstruction w.r.t. the variational parameters """ grad_mean = np.ones((1, self.dimension)) grad_std = (np.exp(variational_parameters[self.dimension :])).reshape(1, -1) grad_reconstruct_params = np.hstack((grad_mean, grad_std)) return grad_reconstruct_params - def draw(self, variational_parameters, n_draws=1): + def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: """Draw *n_draw* samples from the variational distribution. Args: - variational_parameters (np.ndarray): Variational parameters - n_draws (int): Number of samples to draw + variational_parameters: Variational parameters + n_draws: Number of samples to draw Returns: - samples (np.ndarray): Row-wise samples of the variational distribution + Samples of shape (n_draws, n_dim) """ mean, cov = self.reconstruct_distribution_parameters(variational_parameters) samples = np.random.randn(n_draws, self.dimension) * np.sqrt(np.diag(cov)).reshape( @@ -140,15 +142,15 @@ def draw(self, variational_parameters, n_draws=1): ) + mean.reshape(1, -1) return samples - def logpdf(self, variational_parameters, x): - """Logpdf evaluated using the variational parameters at samples `x`. + def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF evaluated using the variational parameters at samples `x`. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - logpdf (np.ndarray): Row vector of the logpdfs + Row vector of the Log-PDF values """ mean, cov = self.reconstruct_distribution_parameters(variational_parameters) mean = mean.flatten() @@ -161,32 +163,32 @@ def logpdf(self, variational_parameters, x): ) return logpdf.flatten() - def pdf(self, variational_parameters, x): - """Pdf of the variational distribution evaluated at samples *x*. + def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """PDF of the variational distribution evaluated at samples *x*. - First computes the logpdf, which is numerically more stable for exponential distributions. + First computes the log-PDF, which is numerically more stable for exponential distributions. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - pdf (np.ndarray): Row vector of the pdfs + Row vector of the PDF values """ pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters, x): - """Logpdf gradient w.r.t. the variational parameters. + def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - score (np.ndarray): Column-wise scores + Column-wise scores """ mean, cov = self.reconstruct_distribution_parameters(variational_parameters) mean = mean.flatten() @@ -201,35 +203,38 @@ def grad_params_logpdf(self, variational_parameters, x): ) return score - def total_grad_params_logpdf(self, variational_parameters, standard_normal_sample_batch): - """Total logpdf reparameterization gradient. + def total_grad_params_logpdf( + self, variational_parameters: ArrayN, standard_normal_sample_batch: np.ndarray + ) -> np.ndarray: + """Total log-PDF reparameterization gradient. - Total logpdf reparameterization gradient w.r.t. the variational parameters. + Total log-PDF reparameterization gradient w.r.t. the variational parameters. Args: - variational_parameters (np.ndarray): Variational parameters - standard_normal_sample_batch (np.ndarray): Standard normal distributed sample batch + variational_parameters: Variational parameters + standard_normal_sample_batch: Standard normal distributed sample batch Returns: - total_grad (np.ndarray): Total Logpdf reparameterization gradient + Total log-PDF reparameterization gradient """ total_grad = np.zeros((standard_normal_sample_batch.shape[0], variational_parameters.size)) total_grad[:, self.dimension :] = -1.0 return total_grad - def grad_sample_logpdf(self, variational_parameters, sample_batch): - """Computes the gradient of the logpdf w.r.t. to the *x*. + def grad_sample_logpdf( + self, variational_parameters: ArrayN, sample_batch: np.ndarray + ) -> np.ndarray: + """Computes the gradient of the log-PDF w.r.t. to the sample *x*. Args: - variational_parameters (np.ndarray): Variational parameters - sample_batch (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + sample_batch: Row-wise samples Returns: - gradients_batch (np.ndarray): Gradients of the log-pdf w.r.t. the - sample *x*. The first dimension of the array corresponds to - the different samples. The second dimension to different dimensions - within one sample. (Third dimension is empty and just added to - keep slices two dimensional.) + Gradients of the log-PDF w.r.t. the sample *x*. The first dimension of the array + corresponds to the different samples. The second dimension to different dimensions + within one sample. (Third dimension is empty and just added to keep slices two + dimensional.) """ mean, cov = self.reconstruct_distribution_parameters(variational_parameters) gradients_batch = -(sample_batch - mean.reshape(1, self.dimension)) / np.diag(cov).reshape( @@ -237,27 +242,27 @@ def grad_sample_logpdf(self, variational_parameters, sample_batch): ) return gradients_batch - def fisher_information_matrix(self, variational_parameters): + def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarray: r"""Compute the Fisher information matrix analytically. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - FIM (np.ndarray): Matrix (n_parameters x n_parameters) + Fisher information matrix of shape (n_parameters, n_parameters) """ fisher_diag = np.exp(-2 * variational_parameters[self.dimension :]) fisher_diag = np.hstack((fisher_diag, 2 * np.ones(self.dimension))) return np.diag(fisher_diag) - def export_dict(self, variational_parameters): + def export_dict(self, variational_parameters: ArrayN) -> dict: """Create a dict of the distribution based on the given parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - export_dict (dictionary): Dict containing distribution information + Dictionary containing distribution information """ mean, cov = self.reconstruct_distribution_parameters(variational_parameters) sd = cov**0.5 @@ -270,18 +275,18 @@ def export_dict(self, variational_parameters): } return export_dict - def conduct_reparameterization(self, variational_parameters, n_samples): + def conduct_reparameterization( + self, variational_parameters: ArrayN, n_samples: int + ) -> tuple[np.ndarray, np.ndarray]: """Conduct a reparameterization. Args: - variational_parameters (np.ndarray): Array with variational parameters - n_samples (int): Number of samples for current batch + variational_parameters: Array with variational parameters + n_samples: Number of samples for current batch Returns: - * samples_mat (np.ndarray): Array of actual samples from the - variational distribution - * standard_normal_sample_batch (np.ndarray): Standard normal - distributed sample batch + Actual samples from the variational distribution + Standard normal distributed sample batch """ standard_normal_sample_batch = np.random.normal(0, 1, size=(n_samples, self.dimension)) mean, cov = self.reconstruct_distribution_parameters(variational_parameters) @@ -290,18 +295,20 @@ def conduct_reparameterization(self, variational_parameters, n_samples): return samples_mat, standard_normal_sample_batch def grad_params_reparameterization( - self, variational_parameters, standard_normal_sample_batch, upstream_gradient - ): + self, + variational_parameters: ArrayN, + standard_normal_sample_batch: np.ndarray, + upstream_gradient: np.ndarray, + ) -> np.ndarray: r"""Calculate the gradient of the reparameterization. Args: - variational_parameters (np.ndarray): Variational parameters - standard_normal_sample_batch (np.ndarray): Standard normal distributed sample batch - upstream_gradient (np.array): Upstream gradient + variational_parameters: Variational parameters + standard_normal_sample_batch: Standard normal distributed sample batch + upstream_gradient: Upstream gradient Returns: - gradient (np.ndarray): Gradient of the upstream function w.r.t. the variational - parameters. + Gradient of the upstream function w.r.t. the variational parameters. Note: We assume that *grad_reconstruct_params* is a row-vector containing the partial diff --git a/src/queens/variational_distributions/mixture_model.py b/src/queens/variational_distributions/mixture_model.py index 75a3174f4..de81b329f 100644 --- a/src/queens/variational_distributions/mixture_model.py +++ b/src/queens/variational_distributions/mixture_model.py @@ -14,8 +14,11 @@ # """Mixture Model Variational Distribution.""" +from typing import Iterable + import numpy as np +from queens.utils.type_hinting import Array1D, ArrayN, ArrayNxM from queens.variational_distributions._variational_distribution import Variational @@ -32,25 +35,24 @@ class MixtureModel(Variational): This allows the weight parameters :math:`\lambda_{weights}` to be unconstrained. Attributes: - n_components (int): Number of mixture components. + n_components: Number of mixture components. base_distribution: Variational distribution object for the components. - n_parameters (int): Number of parameters used in the parameterization. + n_parameters: Number of parameters used in the parameterization. """ - def __init__(self, base_distribution, dimension, n_components): + def __init__(self, base_distribution: Variational, dimension: int, n_components: int) -> None: """Initialize mixture model. Args: - dimension (int): Dimension of the random variable - n_components (int): Number of mixture components base_distribution: Variational distribution object for the components + dimension: Dimension of the random variable + n_components: Number of mixture components """ - super().__init__(dimension) + super().__init__(dimension, n_parameters=n_components * base_distribution.n_parameters) self.n_components = n_components self.base_distribution = base_distribution - self.n_parameters = n_components * base_distribution.n_parameters - def initialize_variational_parameters(self, random=False): + def initialize_variational_parameters(self, random: bool = False) -> ArrayN: r"""Initialize variational parameters. Default weights initialization: @@ -63,11 +65,10 @@ def initialize_variational_parameters(self, random=False): The component initialization is handle by the component itself. Args: - random (bool, optional): If True, a random initialization is used. Otherwise the - default is selected + random: If True, a random initialization is used. Otherwise the default is selected Returns: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + Variational parameters of shape (n_params,) """ variational_parameters_components = ( self.base_distribution.initialize_variational_parameters(random) @@ -87,25 +88,29 @@ def initialize_variational_parameters(self, random=False): return np.concatenate([variational_parameters_components, variational_parameters_weights]) - def construct_variational_parameters(self, component_parameters_list, weights): + def construct_variational_parameters( # pylint: disable=arguments-differ + self, parameters_per_component: list[Iterable[np.ndarray]], weights: Array1D + ) -> ArrayN: """Construct the variational parameters from the probabilities. Args: - component_parameters_list (list): List of the component parameters of the components - weights (np.ndarray): Probabilities of the distribution + parameters_per_component: Distribution parameters per component + weights: Probabilities of the distribution Returns: - variational_parameters (np.ndarray): Variational parameters + Variational parameters """ variational_parameters = [] - for component_parameters in component_parameters_list: + for parameters in parameters_per_component: variational_parameters.append( - self.base_distribution.construct_variational_parameters(*component_parameters) + self.base_distribution.construct_variational_parameters(*parameters) ) variational_parameters.append(np.log(weights).flatten()) return np.concatenate(variational_parameters) - def _construct_component_variational_parameters(self, variational_parameters): + def _construct_component_variational_parameters( + self, variational_parameters: ArrayN + ) -> tuple[list, np.ndarray]: """Reconstruct the weights and parameters of the mixture components. Creates a list containing the variational parameters of the different components. @@ -113,11 +118,11 @@ def _construct_component_variational_parameters(self, variational_parameters): The list is nested, each entry correspond to the parameters of a component. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - variational_parameters_list (list): List of the variational parameters of the components - weights (np.ndarray): Weights of the mixture + Variational parameters of the components + Weights of the mixture """ n_parameters_comp = self.base_distribution.n_parameters variational_parameters_list = [] @@ -131,18 +136,19 @@ def _construct_component_variational_parameters(self, variational_parameters): weights = weights / np.sum(weights) return variational_parameters_list, weights - def reconstruct_distribution_parameters(self, variational_parameters): + def reconstruct_distribution_parameters( + self, variational_parameters: ArrayN + ) -> tuple[list, ArrayN]: """Reconstruct the weights and parameters of the mixture components. The list is nested, each entry correspond to the parameters of a component. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - distribution_parameters_list (list): List of the distribution parameters of the - components - weights (np.ndarray): Weights of the mixture + Distribution parameters of the components + Weights of the mixture """ n_parameters_comp = self.base_distribution.n_parameters distribution_parameters_list = [] @@ -159,7 +165,7 @@ def reconstruct_distribution_parameters(self, variational_parameters): weights = weights / np.sum(weights) return distribution_parameters_list, weights - def draw(self, variational_parameters, n_draws=1): + def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: """Draw *n_draw* samples from the variational distribution. Uses a two-step process: @@ -167,13 +173,13 @@ def draw(self, variational_parameters, n_draws=1): 2. Sample from the selected component Args: - variational_parameters (np.ndarray): Variational parameters - n_draws (int): Number of samples to draw + variational_parameters: Variational parameters + n_draws: Number of samples to draw Returns: - samples (np.ndarray): Row wise samples of the variational distribution + Samples of shape (n_draws, n_dim) """ - parameters_list, weights = self._construct_component_variational_parameters( + parameters, weights = self._construct_component_variational_parameters( variational_parameters ) samples = [] @@ -181,15 +187,15 @@ def draw(self, variational_parameters, n_draws=1): # Select component to draw from component = np.argmax(np.random.multinomial(1, weights)) # Draw a sample of this component - sample = self.base_distribution.draw(parameters_list[component], 1) + sample = self.base_distribution.draw(parameters[component], 1) samples.append(sample) samples = np.concatenate(samples, axis=0) return samples - def logpdf(self, variational_parameters, x): - """Logpdf evaluated using the variational parameters at samples *x*. + def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF evaluated using the variational parameters at samples *x*. - Is a general implementation using the logpdf function of the components. Uses the + Is a general implementation using the log-PDF function of the components. Uses the log-sum-exp trick [1] in order to reduce floating point issues. References: @@ -197,55 +203,55 @@ def logpdf(self, variational_parameters, x): Review for Statisticians, Journal of the American Statistical Association, 112:518 Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - logpdf (np.ndarray): Row vector of the logpdfs + Row vector of the Log-PDF values """ - parameters_list, weights = self._construct_component_variational_parameters( + parameters, weights = self._construct_component_variational_parameters( variational_parameters ) - logpdf = [] + logpdf_lst = [] x = np.atleast_2d(x) # Parameter for the log-sum-exp trick max_logpdf = -np.inf * np.ones(len(x)) for j in range(self.n_components): - logpdf.append(np.log(weights[j]) + self.base_distribution.logpdf(parameters_list[j], x)) - max_logpdf = np.maximum(max_logpdf, logpdf[-1]) - logpdf = np.array(logpdf) - np.tile(max_logpdf, (self.n_components, 1)) + logpdf_lst.append(np.log(weights[j]) + self.base_distribution.logpdf(parameters[j], x)) + max_logpdf = np.maximum(max_logpdf, logpdf_lst[-1]) + logpdf = np.array(logpdf_lst) - np.tile(max_logpdf, (self.n_components, 1)) logpdf = np.sum(np.exp(logpdf), axis=0) logpdf = np.log(logpdf) + max_logpdf return logpdf - def pdf(self, variational_parameters, x): + def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: """Pdf evaluated using the variational parameters at given samples `x`. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - pdf (np.ndarray): Row vector of the pdfs + Row vector of the PDF values """ pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters, x): - """Logpdf gradient w.r.t. the variational parameters. + def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. Is a general implementation using the score functions of the components. Args: - variational_parameters (np.ndarray): Variational parameters - x (np.ndarray): Row-wise samples + variational_parameters: Variational parameters + x: Row-wise samples Returns: - score (np.ndarray): Column-wise scores + Column-wise scores """ - parameters_list, weights = self._construct_component_variational_parameters( + parameters, weights = self._construct_component_variational_parameters( variational_parameters ) x = np.atleast_2d(x) @@ -258,9 +264,9 @@ def grad_params_logpdf(self, variational_parameters, x): logpdf = self.logpdf(variational_parameters, x) for j in range(self.n_components): # coefficient for the score term of every component - precoeff = np.exp(self.base_distribution.logpdf(parameters_list[j], x) - logpdf) + precoeff = np.exp(self.base_distribution.logpdf(parameters[j], x) - logpdf) # Score function of the jth component - score_comp = self.base_distribution.grad_params_logpdf(parameters_list[j], x) + score_comp = self.base_distribution.grad_params_logpdf(parameters[j], x) component_block.append( weights[j] * np.tile(precoeff, (len(score_comp), 1)) * score_comp ) @@ -270,34 +276,36 @@ def grad_params_logpdf(self, variational_parameters, x): score = np.vstack((np.concatenate(component_block, axis=0), weights_block)) return score - def fisher_information_matrix(self, variational_parameters, n_samples=10000): + def fisher_information_matrix( + self, variational_parameters: ArrayN, n_samples: int = 10000 + ) -> np.ndarray: """Approximate the Fisher information matrix using Monte Carlo. Args: - variational_parameters (np.ndarray): Variational parameters - n_samples (int, optional): number of samples for a MC FIM estimation + variational_parameters: Variational parameters + n_samples: Number of samples for a MC FIM estimation Returns: - FIM (np.ndarray): Matrix (num parameters x num parameters) + Fisher information matrix (num parameters x num parameters) """ samples = self.draw(variational_parameters, n_samples) scores = self.grad_params_logpdf(variational_parameters, samples) - fim = 0 + fim = np.zeros((scores.shape[0], scores.shape[0])) for j in range(n_samples): fim = fim + np.outer(scores[:, j], scores[:, j]) fim = fim / n_samples return fim - def export_dict(self, variational_parameters): + def export_dict(self, variational_parameters: ArrayN) -> dict: """Create a dict of the distribution based on the given parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - export_dict (dictionary): Dict containing distribution information + Dictionary containing distribution information """ - parameters_list, weights = self._construct_component_variational_parameters( + parameters, weights = self._construct_component_variational_parameters( variational_parameters ) export_dict = { @@ -309,7 +317,7 @@ def export_dict(self, variational_parameters): } # Loop over the components for j in range(self.n_components): - component_dict = self.base_distribution.export_dict(parameters_list[j]) + component_dict = self.base_distribution.export_dict(parameters[j]) component_key = "component_" + str(j) export_dict.update({component_key: component_dict}) return export_dict diff --git a/src/queens/variational_distributions/particle.py b/src/queens/variational_distributions/particle.py index cb2b89559..8cb22db34 100644 --- a/src/queens/variational_distributions/particle.py +++ b/src/queens/variational_distributions/particle.py @@ -14,9 +14,12 @@ # """Particle Variational Distribution.""" +from typing import Sequence, Sized + import numpy as np from queens.distributions.particle import Particle as ParticleDistribution +from queens.utils.type_hinting import ArrayN, ArrayNxM from queens.variational_distributions._variational_distribution import Variational @@ -27,31 +30,36 @@ class Particle(Variational): :math:`p_i=p(\lambda_i)=\frac{\exp(\lambda_i)}{\sum_k exp(\lambda_k)}` Attributes: - particles_obj (Particle): Particle distribution object - dimension (int): Number of random variables + particles_obj: Particle distribution object + dimension: Number of random variables """ - def __init__(self, sample_space): - """Initialize variational distribution.""" + def __init__(self, sample_space: np.ndarray | Sequence[Sized]) -> None: + """Initialize variational distribution. + + Args: + sample_space: Sample space of the variational distribution + """ self.particles_obj = ParticleDistribution(np.ones(len(sample_space)), sample_space) - super().__init__(self.particles_obj.dimension) - self.n_parameters = len(sample_space) + super().__init__(self.particles_obj.dimension, n_parameters=len(sample_space)) - def construct_variational_parameters(self, probabilities, sample_space): + def construct_variational_parameters( # pylint: disable=arguments-differ + self, probabilities: ArrayN, sample_space: np.ndarray + ) -> ArrayN: """Construct the variational parameters from the probabilities. Args: - probabilities (np.ndarray): Probabilities of the distribution - sample_space (np.ndarray): Sample space of the distribution + probabilities: Probabilities of the distribution + sample_space: Sample space of the distribution Returns: - variational_parameters (np.ndarray): Variational parameters + Variational parameters """ self.particles_obj = ParticleDistribution(probabilities, sample_space) variational_parameters = np.log(probabilities).flatten() return variational_parameters - def initialize_variational_parameters(self, random=False): + def initialize_variational_parameters(self, random: bool = False) -> ArrayN: r"""Initialize variational parameters. Default initialization: @@ -62,11 +70,10 @@ def initialize_variational_parameters(self, random=False): distribution with :math:`N_\text{experiments}` Args: - random (bool, optional): If True, a random initialization is used. Otherwise the - default is selected + random: If True, a random initialization is used. Otherwise the default is selected Returns: - variational_parameters (np.ndarray): variational parameters (1 x n_params) + Variational parameters of shape (n_params,) """ if random: variational_parameters = ( @@ -78,61 +85,64 @@ def initialize_variational_parameters(self, random=False): return variational_parameters - def reconstruct_distribution_parameters(self, variational_parameters): + def reconstruct_distribution_parameters( + self, variational_parameters: ArrayN + ) -> tuple[ArrayN, np.ndarray]: """Reconstruct probabilities from the variational parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - probabilities (np.ndarray): Probabilities of the distribution + Probabilities of the distribution + Sample space of the distribution """ probabilities = np.exp(variational_parameters) probabilities /= np.sum(probabilities) self.particles_obj = ParticleDistribution(probabilities, self.particles_obj.sample_space) return probabilities, self.particles_obj.sample_space - def draw(self, variational_parameters, n_draws=1): + def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: """Draw *n_draws* samples from distribution. Args: - variational_parameters (np.ndarray): Variational parameters of the distribution - n_draws (int): Number of samples + variational_parameters: Variational parameters of the distribution + n_draws: Number of samples Returns: - samples (np.ndarray): samples (n_draws x n_dim) + Samples of shape (n_draws, n_dim) """ self.reconstruct_distribution_parameters(variational_parameters) return self.particles_obj.draw(n_draws) - def logpdf(self, variational_parameters, x): - """Evaluate the natural logarithm of the logpdf at sample. + def logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: + """Evaluate the natural logarithm of the PDF. Args: - variational_parameters (np.ndarray): Variational parameters of the distribution - x (np.ndarray): Locations at which to evaluate the distribution (n_samples x n_dim) + variational_parameters: Variational parameters of the distribution + x: Locations at which to evaluate the distribution of shape (n_samples, n_dim) Returns: - logpdf (np.ndarray): Logpdfs at the locations x + Log-PDF values at the locations x """ self.reconstruct_distribution_parameters(variational_parameters) return self.particles_obj.logpdf(x) - def pdf(self, variational_parameters, x): - """Evaluate the probability density function (pdf) at sample. + def pdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: + """Evaluate the probability density function (PDF). Args: - variational_parameters (np.ndarray): Variational parameters of the distribution - x (np.ndarray): Locations at which to evaluate the distribution (n_samples x n_dim) + variational_parameters: Variational parameters of the distribution + x: Locations at which to evaluate the distribution of shape (n_samples, n_dim) Returns: - logpdf (np.ndarray): Pdfs at the locations x + Row vector of the PDF values """ self.reconstruct_distribution_parameters(variational_parameters) return self.particles_obj.pdf(x) - def grad_params_logpdf(self, variational_parameters, x): - r"""Logpdf gradient w.r.t. the variational parameters. + def grad_params_logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: + r"""Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. @@ -140,11 +150,11 @@ def grad_params_logpdf(self, variational_parameters, x): :math:`\nabla_{\lambda_i}\ln p(\theta_j | \lambda)=\delta_{ij}-p_i` Args: - variational_parameters (np.ndarray): Variational parameters of the distribution - x (np.ndarray): Locations at which to evaluate the distribution (n_samples x n_dim) + variational_parameters: Variational parameters of the distribution + x: Locations at which to evaluate the distribution of shape (n_samples, n_dim) Returns: - score_function (np.ndarray): Score functions at the locations x + Score functions at the locations x """ self.reconstruct_distribution_parameters(variational_parameters) index = np.array( @@ -162,30 +172,30 @@ def grad_params_logpdf(self, variational_parameters, x): # Get the samples return sample_scores[index].T - def fisher_information_matrix(self, variational_parameters): - r"""Compute the fisher information matrix. + def fisher_information_matrix(self, variational_parameters: ArrayN) -> ArrayNxM: + r"""Compute the Fisher information matrix. For the given parameterization, the Fisher information yields: :math:`\text{FIM}_{ij}=\delta_{ij} p_i -p_i p_j` Args: - variational_parameters (np.ndarray): Variational parameters of the distribution + variational_parameters: Variational parameters of the distribution Returns: - fim (np.ndarray): Fisher information matrix (n_params x n_params) + Fisher information matrix of shape (n_parameters, n_parameters) """ probabilities, _ = self.reconstruct_distribution_parameters(variational_parameters) fim = np.diag(probabilities) - np.outer(probabilities, probabilities) return fim - def export_dict(self, variational_parameters): + def export_dict(self, variational_parameters: ArrayN) -> dict: """Create a dict of the distribution based on the given parameters. Args: - variational_parameters (np.ndarray): Variational parameters + variational_parameters: Variational parameters Returns: - export_dict (dictionary): Dict containing distribution information + Dictionary containing distribution information """ self.reconstruct_distribution_parameters(variational_parameters) export_dict = { From 79692a83c8e75c0169ac7f53c519a6557c8b523c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Lea=20H=C3=A4usel?= Date: Tue, 9 Dec 2025 15:04:25 +0100 Subject: [PATCH 2/2] style: Switch to TypeVars for numpy arrays --- src/queens/distributions/__init__.py | 2 +- .../parameters/random_fields/__init__.py | 2 +- src/queens/utils/type_hinting.py | 28 ------ .../variational_distributions/__init__.py | 2 +- .../_variational_distribution.py | 90 ++++++++++------- .../full_rank_normal.py | 97 +++++++++++-------- src/queens/variational_distributions/joint.py | 78 ++++++++++----- .../mean_field_normal.py | 85 +++++++++------- .../mixture_model.py | 68 ++++++++----- .../variational_distributions/particle.py | 49 ++++++---- 10 files changed, 292 insertions(+), 209 deletions(-) delete mode 100644 src/queens/utils/type_hinting.py diff --git a/src/queens/distributions/__init__.py b/src/queens/distributions/__init__.py index 24075fdab..4d9bd4542 100644 --- a/src/queens/distributions/__init__.py +++ b/src/queens/distributions/__init__.py @@ -41,5 +41,5 @@ class_module_map = extract_type_checking_imports(__file__) -def __getattr__(name: str) -> Distribution: +def __getattr__(name: str) -> type[Distribution]: return import_class_from_class_module_map(name, class_module_map, __name__) diff --git a/src/queens/parameters/random_fields/__init__.py b/src/queens/parameters/random_fields/__init__.py index db2576c62..62af70f77 100644 --- a/src/queens/parameters/random_fields/__init__.py +++ b/src/queens/parameters/random_fields/__init__.py @@ -33,5 +33,5 @@ class_module_map = extract_type_checking_imports(__file__) -def __getattr__(name: str) -> RandomField: +def __getattr__(name: str) -> type[RandomField]: return import_class_from_class_module_map(name, class_module_map, __name__) diff --git a/src/queens/utils/type_hinting.py b/src/queens/utils/type_hinting.py deleted file mode 100644 index 7dc54d66a..000000000 --- a/src/queens/utils/type_hinting.py +++ /dev/null @@ -1,28 +0,0 @@ -# -# SPDX-License-Identifier: LGPL-3.0-or-later -# Copyright (c) 2024-2025, QUEENS contributors. -# -# This file is part of QUEENS. -# -# QUEENS is free software: you can redistribute it and/or modify it under the terms of the GNU -# Lesser General Public License as published by the Free Software Foundation, either version 3 of -# the License, or (at your option) any later version. QUEENS is distributed in the hope that it will -# be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or -# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details. You -# should have received a copy of the GNU Lesser General Public License along with QUEENS. If not, -# see . -# -"""Utilities for type hinting.""" - -from typing import Literal, TypeAlias - -import numpy as np - -# pylint: disable=invalid-name - -ArrayN: TypeAlias = np.ndarray[tuple[int], np.dtype[np.floating]] -Array1xN: TypeAlias = np.ndarray[tuple[Literal[1], int], np.dtype[np.floating]] -ArrayNx1: TypeAlias = np.ndarray[tuple[int, Literal[1]], np.dtype[np.floating]] -ArrayNxM: TypeAlias = np.ndarray[tuple[int, int], np.dtype[np.floating]] - -Array1D: TypeAlias = ArrayN | Array1xN | ArrayNx1 diff --git a/src/queens/variational_distributions/__init__.py b/src/queens/variational_distributions/__init__.py index 44f4e4fb0..763ad4a32 100644 --- a/src/queens/variational_distributions/__init__.py +++ b/src/queens/variational_distributions/__init__.py @@ -33,5 +33,5 @@ class_module_map = extract_type_checking_imports(__file__) -def __getattr__(name: str) -> Variational: +def __getattr__(name: str) -> type[Variational]: return import_class_from_class_module_map(name, class_module_map, __name__) diff --git a/src/queens/variational_distributions/_variational_distribution.py b/src/queens/variational_distributions/_variational_distribution.py index 353b7bde3..4a04eb9b9 100644 --- a/src/queens/variational_distributions/_variational_distribution.py +++ b/src/queens/variational_distributions/_variational_distribution.py @@ -15,11 +15,31 @@ """Variational Distribution.""" import abc -from typing import Any +from typing import Any, Literal, TypeAlias, TypeVar import numpy as np -from queens.utils.type_hinting import ArrayN, ArrayNxM +# pylint: disable=invalid-name + +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]] class Variational: @@ -30,7 +50,7 @@ class Variational: n_parameters: Number of variational parameters """ - def __init__(self, dimension: int, n_parameters: int) -> None: + def __init__(self, dimension: NDims, n_parameters: NParams) -> None: """Initialize variational distribution. Args: @@ -41,7 +61,7 @@ def __init__(self, dimension: int, n_parameters: int) -> None: self.n_parameters = n_parameters @abc.abstractmethod - def construct_variational_parameters(self, *args: Any) -> np.ndarray: + def construct_variational_parameters(self, *args: Any) -> ArrayNParams: """Construct variational parameters from distribution parameters. Args: @@ -52,7 +72,7 @@ def construct_variational_parameters(self, *args: Any) -> np.ndarray: """ @abc.abstractmethod - def reconstruct_distribution_parameters(self, variational_parameters: ArrayN) -> Any: + def reconstruct_distribution_parameters(self, variational_parameters: ArrayNParams) -> Any: """Reconstruct distribution parameters from variational parameters. Args: @@ -63,98 +83,96 @@ def reconstruct_distribution_parameters(self, variational_parameters: ArrayN) -> """ @abc.abstractmethod - def draw( - self, - variational_parameters: ArrayN, - n_draws: int = 1, - ) -> ArrayNxM: + def draw(self, variational_parameters: ArrayNParams, n_draws: NSamples) -> ArrayNSamplesXNDims: """Draw *n_draws* samples from distribution. Args: - variational_parameters: Variational parameters of shape (n_params,) + variational_parameters: Variational parameters n_draws: Number of samples Returns: - Drawn samples of shape (n_draws, n_dim) + Drawn samples """ @abc.abstractmethod def logpdf( self, - variational_parameters: ArrayN, - x: ArrayNxM, - ) -> np.ndarray: + variational_parameters: ArrayNParams, + x: ArrayNSamplesXNDims, + ) -> ArrayNSamples: """Evaluate the natural logarithm of the PDF. Args: - variational_parameters: Variational parameters of shape (n_params,) - x: Locations to evaluate of shape (n_samples, n_dim) + variational_parameters: Variational parameters + x: Locations to evaluate Returns: - Row vector of the Log-PDF values + Log-PDF values """ @abc.abstractmethod def pdf( self, - variational_parameters: ArrayN, - x: ArrayNxM, - ) -> np.ndarray: + variational_parameters: ArrayNParams, + x: ArrayNSamplesXNDims, + ) -> ArrayNSamples: """Evaluate the probability density function (PDF). Args: - variational_parameters: Variational parameters of shape (n_params,) - x: Locations to evaluate of shape (n_samples, n_dim) + variational_parameters: Variational parameters + x: Locations to evaluate Returns: - Row vector of the PDF values + PDF values """ @abc.abstractmethod def grad_params_logpdf( self, - variational_parameters: ArrayN, - x: ArrayNxM, - ) -> np.ndarray: + variational_parameters: ArrayNParams, + x: ArrayNSamplesXNDims, + ) -> ArrayNParamsXNSamples: """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. Args: - variational_parameters: Variational parameters of shape (n_params,) - x: Locations to evaluate of shape (n_samples, n_dim) + variational_parameters: Variational parameters + x: Locations to evaluate Returns: Gradient of the log-PDF w.r.t. the variational parameters """ @abc.abstractmethod - def fisher_information_matrix(self, variational_parameters: ArrayN) -> ArrayNxM: - """Compute the fisher information matrix. + def fisher_information_matrix( + self, variational_parameters: ArrayNParams + ) -> ArrayNParamsXNParams: + """Compute the Fisher information matrix. Depends on the variational distribution for the given parameterization. Args: - variational_parameters: Variational parameters of shape (n_params,) + variational_parameters: Variational parameters Returns: - Fisher information matrix of shape (n_params, n_params) + Fisher information matrix """ @abc.abstractmethod - def initialize_variational_parameters(self, random: bool = False) -> ArrayN: + def initialize_variational_parameters(self, random: bool = False) -> ArrayNParams: """Initialize variational parameters. Args: random: If True, a random initialization is used. Otherwise the default is selected. Returns: - Variational parameters of shape (n_params,) + Variational parameters """ @abc.abstractmethod - def export_dict(self, variational_parameters: np.ndarray) -> dict: + def export_dict(self, variational_parameters: ArrayNParams) -> dict: """Create a dict of the distribution based on the given parameters. Args: diff --git a/src/queens/variational_distributions/full_rank_normal.py b/src/queens/variational_distributions/full_rank_normal.py index 3274632d0..38e298369 100644 --- a/src/queens/variational_distributions/full_rank_normal.py +++ b/src/queens/variational_distributions/full_rank_normal.py @@ -19,8 +19,21 @@ from numba import njit from queens.utils.logger_settings import log_init_args -from queens.utils.type_hinting import ArrayN, ArrayNx1, ArrayNxM -from queens.variational_distributions._variational_distribution import Variational +from queens.variational_distributions._variational_distribution import ( + Array1XNParams, + ArrayNDims, + ArrayNDimsX1, + ArrayNDimsXNDims, + ArrayNParams, + ArrayNParamsXNParams, + ArrayNParamsXNSamples, + ArrayNSamples, + ArrayNSamplesXNDims, + ArrayNSamplesXNParams, + NDims, + NSamples, + Variational, +) class FullRankNormal(Variational): @@ -42,7 +55,7 @@ class FullRankNormal(Variational): """ @log_init_args - def __init__(self, dimension: int) -> None: + def __init__(self, dimension: NDims) -> None: """Initialize variational distribution. Args: @@ -51,7 +64,7 @@ def __init__(self, dimension: int) -> None: n_parameters = (dimension * (dimension + 1)) // 2 + dimension super().__init__(dimension, n_parameters) - def initialize_variational_parameters(self, random: bool = False) -> ArrayN: + def initialize_variational_parameters(self, random: bool = False) -> ArrayNParams: r"""Initialize variational parameters. Default initialization: @@ -64,7 +77,7 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: random: If True, a random initialization is used. Otherwise the default is selected. Returns: - Variational parameters of shape (n_params,) + Variational parameters """ if random: cholesky_covariance = np.eye(self.dimension) + 0.1 * ( @@ -85,13 +98,13 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: return variational_parameters def construct_variational_parameters( # pylint: disable=arguments-differ - self, mean: ArrayNx1 | ArrayN, covariance: ArrayNxM - ) -> ArrayN: + self, mean: ArrayNDimsX1 | ArrayNDims, covariance: ArrayNDimsXNDims + ) -> ArrayNParams: """Construct the variational parameters from mean and covariance. Args: - mean: Mean values of the distribution of shape (n_dim, 1) or (n_dim,) - covariance: Covariance matrix of the distribution of shape (n_dim, n_dim) + mean: Mean values of the distribution + covariance: Covariance matrix of the distribution Returns: Variational parameters @@ -111,15 +124,15 @@ def construct_variational_parameters( # pylint: disable=arguments-differ return variational_parameters def reconstruct_distribution_parameters( - self, variational_parameters: ArrayN - ) -> tuple[ArrayNx1, ArrayNxM]: - """Reconstruct mean value, covariance and its Cholesky decomposition. + self, variational_parameters: ArrayNParams + ) -> tuple[ArrayNDimsX1, ArrayNDimsXNDims]: + """Reconstruct mean value and covariance. Args: variational_parameters: Variational parameters Returns: - Mean value of the distribution of shape (n_dim, 1) - Covariance of the distribution of shape (n_dim, n_dim) + Mean value of the distribution + Covariance of the distribution """ mean, cov, _ = self.reconstruct_distribution_parameters_with_cholesky( variational_parameters @@ -127,16 +140,16 @@ def reconstruct_distribution_parameters( return mean, cov def reconstruct_distribution_parameters_with_cholesky( - self, variational_parameters: ArrayN - ) -> tuple[ArrayNx1, ArrayNxM, ArrayNxM]: + self, variational_parameters: ArrayNParams + ) -> tuple[ArrayNDimsX1, ArrayNDimsXNDims, ArrayNDimsXNDims]: """Reconstruct mean value, covariance and its Cholesky decomposition. Args: variational_parameters: Variational parameters Returns: - Mean value of the distribution of shape (n_dim, 1) - Covariance of the distribution of shape (n_dim, n_dim) - Cholesky decomposition of the covariance matrix of shape (n_dim, n_dim) + Mean value of the distribution + Covariance of the distribution + Cholesky decomposition of the covariance matrix """ mean = variational_parameters[: self.dimension].reshape(-1, 1) cholesky_covariance_array = variational_parameters[self.dimension :] @@ -147,7 +160,7 @@ def reconstruct_distribution_parameters_with_cholesky( return mean, cov, cholesky_covariance - def _grad_reconstruct_distribution_parameters(self) -> np.ndarray: + def _grad_reconstruct_distribution_parameters(self) -> Array1XNParams: """Gradient of the parameter reconstruction. Returns: @@ -156,7 +169,7 @@ def _grad_reconstruct_distribution_parameters(self) -> np.ndarray: grad_reconstruct_params = np.ones((1, self.n_parameters)) return grad_reconstruct_params - def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: + def draw(self, variational_parameters: ArrayNParams, n_draws: NSamples) -> ArrayNSamplesXNDims: """Draw *n_draw* samples from the variational distribution. Args: @@ -164,7 +177,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: n_draws: Number of samples to draw Returns: - Samples of shape (n_draws, n_dim) + Samples """ mean, _, cholesky = self.reconstruct_distribution_parameters_with_cholesky( variational_parameters @@ -172,7 +185,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: sample = np.dot(cholesky, np.random.randn(self.dimension, n_draws)).T + mean.reshape(1, -1) return sample - def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def logpdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Log-PDF evaluated at the samples *x*. Args: @@ -180,7 +193,7 @@ def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: x: Row-wise samples Returns: - Row vector of the Log-PDF values + Log-PDF values """ mean, cov, cholesky = self.reconstruct_distribution_parameters_with_cholesky( variational_parameters @@ -198,7 +211,7 @@ def col_dot_prod(x: np.ndarray, y: np.ndarray) -> np.ndarray: ) return logpdf.flatten() - def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def pdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """PDF of evaluated at given samples *x*. First computes the log-PDF, which is numerically more stable for exponential distributions. @@ -213,7 +226,9 @@ def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def grad_params_logpdf( + self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims + ) -> ArrayNParamsXNSamples: """Log-PDF gradient w.r.t. to the variational parameters. Evaluated at samples *x*. Also known as the score function. @@ -251,8 +266,10 @@ def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> n return score def total_grad_params_logpdf( - self, variational_parameters: ArrayN, standard_normal_sample_batch: np.ndarray - ) -> np.ndarray: + self, + variational_parameters: ArrayNParams, + standard_normal_sample_batch: ArrayNSamplesXNDims, + ) -> ArrayNSamplesXNParams: """Total log-PDF reparameterization gradient. Total log-PDF reparameterization gradient w.r.t. the variational parameters. @@ -271,8 +288,8 @@ def total_grad_params_logpdf( return total_grad def grad_sample_logpdf( - self, variational_parameters: ArrayN, sample_batch: np.ndarray - ) -> np.ndarray: + self, variational_parameters: ArrayNParams, sample_batch: ArrayNSamplesXNDims + ) -> ArrayNSamplesXNDims: """Computes the gradient of the log-PDF w.r.t. to the sample *x*. Args: @@ -295,14 +312,16 @@ def grad_sample_logpdf( gradients_batch = np.array(gradient_lst) return gradients_batch.reshape(sample_batch.shape) - def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarray: + def fisher_information_matrix( + self, variational_parameters: ArrayNParams + ) -> ArrayNParamsXNParams: """Compute the Fisher information matrix analytically. Args: variational_parameters: Variational parameters Returns: - Fisher information matrix (num parameters x num parameters) + Fisher information matrix """ _, cov, cholesky = self.reconstruct_distribution_parameters_with_cholesky( variational_parameters @@ -339,7 +358,7 @@ def fim_blocks(dimension: int) -> tuple[np.ndarray, np.ndarray]: return scipy.linalg.block_diag(mu_block, sigma_block) - def export_dict(self, variational_parameters: ArrayN) -> dict: + def export_dict(self, variational_parameters: ArrayNParams) -> dict: """Create a dict of the distribution based on the given parameters. Args: @@ -357,8 +376,8 @@ def export_dict(self, variational_parameters: ArrayN) -> dict: return export_dict def conduct_reparameterization( - self, variational_parameters: np.ndarray, n_samples: int - ) -> tuple[np.ndarray, np.ndarray]: + self, variational_parameters: ArrayNParams, n_samples: NSamples + ) -> tuple[ArrayNSamplesXNDims, ArrayNSamplesXNDims]: """Conduct a reparameterization. Args: @@ -379,10 +398,10 @@ def conduct_reparameterization( def grad_params_reparameterization( self, - variational_parameters: ArrayN, - standard_normal_sample_batch: np.ndarray, - upstream_gradient: np.ndarray, - ) -> np.ndarray: + variational_parameters: ArrayNParams, + standard_normal_sample_batch: ArrayNSamplesXNDims, + upstream_gradient: ArrayNSamplesXNDims, + ) -> ArrayNSamplesXNParams: r"""Calculate the gradient of the reparameterization. Args: diff --git a/src/queens/variational_distributions/joint.py b/src/queens/variational_distributions/joint.py index 5777f68ce..735cffcdb 100644 --- a/src/queens/variational_distributions/joint.py +++ b/src/queens/variational_distributions/joint.py @@ -14,14 +14,32 @@ # """Joint Variational Distribution.""" +from typing import Generic, Iterator, TypeAlias, TypeVar + import numpy as np import scipy - -from queens.utils.type_hinting import ArrayN, ArrayNxM -from queens.variational_distributions._variational_distribution import Variational - - -class Joint(Variational): +from numpy.typing import ArrayLike + +from queens.variational_distributions._variational_distribution import ( + ArrayNParams, + ArrayNParamsComponent, + ArrayNParamsXNParams, + ArrayNParamsXNSamples, + ArrayNSamples, + ArrayNSamplesXNDims, + NDims, + NSamples, + Variational, +) + +NDimsComponent = TypeVar("NDimsComponent", bound=int) +V = TypeVar("V", bound=Variational) +ArrayNSamplesXNDimsComponent: TypeAlias = np.ndarray[ # pylint: disable=invalid-name + tuple[NSamples, NDimsComponent], np.dtype[np.floating] +] + + +class Joint(Variational, Generic[V]): r"""Joint variational distribution class. This distribution allows to join distributions in an independent fashion: @@ -37,7 +55,7 @@ class Joint(Variational): distributions_dimension: Number of dimension per distribution """ - def __init__(self, distributions: list, dimension: int) -> None: + def __init__(self, distributions: list[V], dimension: NDims) -> None: """Initialize joint distribution. Args: @@ -62,7 +80,7 @@ def __init__(self, distributions: list, dimension: int) -> None: f"dimensions of the subdistributions {np.sum(self.distributions_dimension)}" ) - def initialize_variational_parameters(self, random: bool = False) -> ArrayN: + def initialize_variational_parameters(self, random: bool = False) -> ArrayNParams: r"""Initialize variational parameters. The distribution initialization is handle by the component itself. @@ -71,7 +89,7 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: random: If True, a random initialization is used. Otherwise the default is selected Returns: - Variational parameters of shape (n_params,) + Variational parameters """ variational_parameters = np.concatenate( [ @@ -84,7 +102,7 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: def construct_variational_parameters( # pylint: disable=arguments-differ self, distributions_parameters: list - ) -> ArrayN: + ) -> ArrayNParams: """Construct the variational parameters from the distribution list. Args: @@ -103,8 +121,8 @@ def construct_variational_parameters( # pylint: disable=arguments-differ return np.concatenate(variational_parameters) def _construct_distributions_variational_parameters( - self, variational_parameters: ArrayN - ) -> list: + self, variational_parameters: ArrayNParams + ) -> list[ArrayNParamsComponent]: """Reconstruct the parameters of the distributions. Creates a list containing the variational parameters of the different components. @@ -122,7 +140,9 @@ def _construct_distributions_variational_parameters( ) return variational_parameters_list - def reconstruct_distribution_parameters(self, variational_parameters: ArrayN) -> list[list]: + def reconstruct_distribution_parameters( + self, variational_parameters: ArrayNParams + ) -> list[list[tuple[list | np.ndarray]]]: """Reconstruct the parameters of distributions. The list is nested, each entry correspond to the parameters of a distribution. @@ -143,7 +163,9 @@ def reconstruct_distribution_parameters(self, variational_parameters: ArrayN) -> ) return [distribution_parameters_list] - def _zip_variational_parameters_distributions(self, variational_parameters: ArrayN) -> zip: + def _zip_variational_parameters_distributions( + self, variational_parameters: ArrayNParams + ) -> Iterator[tuple[ArrayNParamsComponent, V]]: """Zip parameters and distributions. This helper function creates a generator for variational parameters and subdistribution. @@ -161,8 +183,8 @@ def _zip_variational_parameters_distributions(self, variational_parameters: Arra ) def _zip_variational_parameters_distributions_samples( - self, variational_parameters: ArrayN, samples: np.ndarray - ) -> zip: + self, variational_parameters: ArrayNParams, samples: ArrayNSamplesXNDims + ) -> Iterator[tuple[ArrayNParamsComponent, ArrayNSamplesXNDimsComponent, V]]: """Zip parameters, samples and distributions. This helper function creates a generator for variational parameters, samples and @@ -182,7 +204,7 @@ def _zip_variational_parameters_distributions_samples( strict=True, ) - def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: + def draw(self, variational_parameters: ArrayNParams, n_draws: NSamples) -> ArrayNSamplesXNDims: """Draw *n_draw* samples from the variational distribution. Args: @@ -190,7 +212,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: n_draws: Number of samples to draw Returns: - Samples of shape (n_draws, n_dim) + Samples """ sample_array = [] for parameters, distribution in self._zip_variational_parameters_distributions( @@ -199,7 +221,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: sample_array.append(distribution.draw(parameters, n_draws)) return np.column_stack(sample_array) - def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def logpdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Log-PDF evaluated using the variational parameters at samples *x*. Args: @@ -209,7 +231,7 @@ def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: Returns: Row vector of the Log-PDF values """ - logpdf = np.zeros_like(x[:, 0], dtype=float) + logpdf = np.zeros(x.shape[0], dtype=float) for ( parameters, samples, @@ -218,7 +240,7 @@ def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: logpdf += distribution.logpdf(parameters, samples) return logpdf - def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def pdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Pdf evaluated using the variational parameters at given samples `x`. Args: @@ -231,7 +253,9 @@ def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def grad_params_logpdf( + self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims + ) -> ArrayNParamsXNSamples: """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. @@ -255,14 +279,16 @@ def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> n return np.row_stack(score) - def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarray: + def fisher_information_matrix( + self, variational_parameters: ArrayNParams + ) -> ArrayNParamsXNParams: """Approximate the Fisher information matrix using Monte Carlo. Args: variational_parameters: Variational parameters Returns: - Fisher information matrix (num parameters x num parameters) + Fisher information matrix """ fim = [] for parameters, distribution in self._zip_variational_parameters_distributions( @@ -272,7 +298,7 @@ def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarra return scipy.linalg.block_diag(*fim) - def export_dict(self, variational_parameters: ArrayN) -> dict: + def export_dict(self, variational_parameters: ArrayNParams) -> dict: """Create a dict of the distribution based on the given parameters. Args: @@ -295,7 +321,7 @@ def export_dict(self, variational_parameters: ArrayN) -> dict: return export_dict -def split_array_by_chunk_sizes(array: np.ndarray, chunk_sizes: np.ndarray) -> list: +def split_array_by_chunk_sizes(array: np.ndarray, chunk_sizes: ArrayLike) -> list: """Split up array by a list of chunk sizes. Args: diff --git a/src/queens/variational_distributions/mean_field_normal.py b/src/queens/variational_distributions/mean_field_normal.py index aafcf32fc..7f94d7c9f 100644 --- a/src/queens/variational_distributions/mean_field_normal.py +++ b/src/queens/variational_distributions/mean_field_normal.py @@ -17,8 +17,21 @@ import numpy as np from queens.utils.logger_settings import log_init_args -from queens.utils.type_hinting import ArrayN, ArrayNx1, ArrayNxM -from queens.variational_distributions._variational_distribution import Variational +from queens.variational_distributions._variational_distribution import ( + Array1XNParams, + ArrayNDims, + ArrayNDimsX1, + ArrayNDimsXNDims, + ArrayNParams, + ArrayNParamsXNParams, + ArrayNParamsXNSamples, + ArrayNSamples, + ArrayNSamplesXNDims, + ArrayNSamplesXNParams, + NDims, + NSamples, + Variational, +) class MeanFieldNormal(Variational): @@ -37,7 +50,7 @@ class MeanFieldNormal(Variational): """ @log_init_args - def __init__(self, dimension: int) -> None: + def __init__(self, dimension: NDims) -> None: """Initialize variational distribution. Args: @@ -45,7 +58,7 @@ def __init__(self, dimension: int) -> None: """ super().__init__(dimension, n_parameters=2 * dimension) - def initialize_variational_parameters(self, random: bool = False) -> ArrayN: + def initialize_variational_parameters(self, random: bool = False) -> ArrayNParams: r"""Initialize variational parameters. Default initialization: @@ -58,7 +71,7 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: random: If True, a random initialization is used. Otherwise the default is selected Returns: - Variational parameters of shape (n_params,) + Variational parameters """ if random: variational_parameters = np.hstack( @@ -73,13 +86,13 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: return variational_parameters def construct_variational_parameters( # pylint: disable=arguments-differ - self, mean: ArrayNx1 | ArrayN, covariance: ArrayNxM - ) -> ArrayN: + self, mean: ArrayNDimsX1 | ArrayNDims, covariance: ArrayNDimsXNDims + ) -> ArrayNParams: """Construct the variational parameters from mean and covariance. Args: - mean: Mean values of the distribution of shape (n_dim, 1) or (n_dim,) - covariance: Covariance matrix of the distribution of shape (n_dim, n_dim) + mean: Mean values of the distribution + covariance: Covariance matrix of the distribution Returns: Variational parameters @@ -94,15 +107,15 @@ def construct_variational_parameters( # pylint: disable=arguments-differ return variational_parameters def reconstruct_distribution_parameters( - self, variational_parameters: ArrayN - ) -> tuple[ArrayNx1, ArrayNxM]: + self, variational_parameters: ArrayNParams + ) -> tuple[ArrayNDimsX1, ArrayNDimsXNDims]: """Reconstruct mean and covariance from the variational parameters. Args: variational_parameters: Variational parameters Returns: - Mean value of the distribution of shape (n_dim, 1) - Covariance matrix of the distribution of shape (n_dim, n_dim) + Mean value of the distribution + Covariance matrix of the distribution """ mean, cov = ( variational_parameters[: self.dimension], @@ -111,8 +124,8 @@ def reconstruct_distribution_parameters( return mean.reshape(-1, 1), np.diag(cov) def _grad_reconstruct_distribution_parameters( - self, variational_parameters: ArrayN - ) -> np.ndarray: + self, variational_parameters: ArrayNParams + ) -> Array1XNParams: """Gradient of the parameter reconstruction. Args: @@ -126,7 +139,7 @@ def _grad_reconstruct_distribution_parameters( grad_reconstruct_params = np.hstack((grad_mean, grad_std)) return grad_reconstruct_params - def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: + def draw(self, variational_parameters: ArrayNParams, n_draws: NSamples) -> ArrayNSamplesXNDims: """Draw *n_draw* samples from the variational distribution. Args: @@ -134,7 +147,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: n_draws: Number of samples to draw Returns: - Samples of shape (n_draws, n_dim) + Samples """ mean, cov = self.reconstruct_distribution_parameters(variational_parameters) samples = np.random.randn(n_draws, self.dimension) * np.sqrt(np.diag(cov)).reshape( @@ -142,7 +155,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: ) + mean.reshape(1, -1) return samples - def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def logpdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Log-PDF evaluated using the variational parameters at samples `x`. Args: @@ -163,7 +176,7 @@ def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: ) return logpdf.flatten() - def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def pdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """PDF of the variational distribution evaluated at samples *x*. First computes the log-PDF, which is numerically more stable for exponential distributions. @@ -178,7 +191,9 @@ def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def grad_params_logpdf( + self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims + ) -> ArrayNParamsXNSamples: """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. @@ -204,8 +219,10 @@ def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> n return score def total_grad_params_logpdf( - self, variational_parameters: ArrayN, standard_normal_sample_batch: np.ndarray - ) -> np.ndarray: + self, + variational_parameters: ArrayNParams, + standard_normal_sample_batch: ArrayNSamplesXNDims, + ) -> ArrayNSamplesXNParams: """Total log-PDF reparameterization gradient. Total log-PDF reparameterization gradient w.r.t. the variational parameters. @@ -222,8 +239,8 @@ def total_grad_params_logpdf( return total_grad def grad_sample_logpdf( - self, variational_parameters: ArrayN, sample_batch: np.ndarray - ) -> np.ndarray: + self, variational_parameters: ArrayNParams, sample_batch: ArrayNSamplesXNDims + ) -> ArrayNSamplesXNDims: """Computes the gradient of the log-PDF w.r.t. to the sample *x*. Args: @@ -242,20 +259,22 @@ def grad_sample_logpdf( ) return gradients_batch - def fisher_information_matrix(self, variational_parameters: ArrayN) -> np.ndarray: + def fisher_information_matrix( + self, variational_parameters: ArrayNParams + ) -> ArrayNParamsXNParams: r"""Compute the Fisher information matrix analytically. Args: variational_parameters: Variational parameters Returns: - Fisher information matrix of shape (n_parameters, n_parameters) + Fisher information matrix """ fisher_diag = np.exp(-2 * variational_parameters[self.dimension :]) fisher_diag = np.hstack((fisher_diag, 2 * np.ones(self.dimension))) return np.diag(fisher_diag) - def export_dict(self, variational_parameters: ArrayN) -> dict: + def export_dict(self, variational_parameters: ArrayNParams) -> dict: """Create a dict of the distribution based on the given parameters. Args: @@ -276,8 +295,8 @@ def export_dict(self, variational_parameters: ArrayN) -> dict: return export_dict def conduct_reparameterization( - self, variational_parameters: ArrayN, n_samples: int - ) -> tuple[np.ndarray, np.ndarray]: + self, variational_parameters: ArrayNParams, n_samples: NSamples + ) -> tuple[ArrayNSamplesXNDims, ArrayNSamplesXNDims]: """Conduct a reparameterization. Args: @@ -296,10 +315,10 @@ def conduct_reparameterization( def grad_params_reparameterization( self, - variational_parameters: ArrayN, - standard_normal_sample_batch: np.ndarray, - upstream_gradient: np.ndarray, - ) -> np.ndarray: + variational_parameters: ArrayNParams, + standard_normal_sample_batch: ArrayNSamplesXNDims, + upstream_gradient: ArrayNSamplesXNDims, + ) -> ArrayNSamplesXNParams: r"""Calculate the gradient of the reparameterization. Args: diff --git a/src/queens/variational_distributions/mixture_model.py b/src/queens/variational_distributions/mixture_model.py index de81b329f..ec52424e5 100644 --- a/src/queens/variational_distributions/mixture_model.py +++ b/src/queens/variational_distributions/mixture_model.py @@ -14,15 +14,30 @@ # """Mixture Model Variational Distribution.""" -from typing import Iterable +from typing import Generic, Iterable, TypeAlias, TypeVar import numpy as np -from queens.utils.type_hinting import Array1D, ArrayN, ArrayNxM -from queens.variational_distributions._variational_distribution import Variational - - -class MixtureModel(Variational): +from queens.variational_distributions._variational_distribution import ( + ArrayNParams, + ArrayNParamsComponent, + ArrayNParamsXNParams, + ArrayNParamsXNSamples, + ArrayNSamples, + ArrayNSamplesXNDims, + NDims, + NSamples, + Variational, +) + +NComponents = TypeVar("NComponents", bound=int) +V = TypeVar("V", bound=Variational) +ArrayNComponents: TypeAlias = np.ndarray[ # pylint: disable=invalid-name + tuple[NComponents], np.dtype[np.floating] +] + + +class MixtureModel(Variational, Generic[V]): r"""Mixture model variational distribution class. Every component is a member of the same distribution family. Uses the parameterization: @@ -40,7 +55,7 @@ class MixtureModel(Variational): n_parameters: Number of parameters used in the parameterization. """ - def __init__(self, base_distribution: Variational, dimension: int, n_components: int) -> None: + def __init__(self, base_distribution: V, dimension: NDims, n_components: NComponents) -> None: """Initialize mixture model. Args: @@ -52,7 +67,7 @@ def __init__(self, base_distribution: Variational, dimension: int, n_components: self.n_components = n_components self.base_distribution = base_distribution - def initialize_variational_parameters(self, random: bool = False) -> ArrayN: + def initialize_variational_parameters(self, random: bool = False) -> ArrayNParams: r"""Initialize variational parameters. Default weights initialization: @@ -68,7 +83,7 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: random: If True, a random initialization is used. Otherwise the default is selected Returns: - Variational parameters of shape (n_params,) + Variational parameters """ variational_parameters_components = ( self.base_distribution.initialize_variational_parameters(random) @@ -89,8 +104,8 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: return np.concatenate([variational_parameters_components, variational_parameters_weights]) def construct_variational_parameters( # pylint: disable=arguments-differ - self, parameters_per_component: list[Iterable[np.ndarray]], weights: Array1D - ) -> ArrayN: + self, parameters_per_component: list[Iterable[np.ndarray]], weights: ArrayNComponents + ) -> ArrayNParams: """Construct the variational parameters from the probabilities. Args: @@ -109,8 +124,8 @@ def construct_variational_parameters( # pylint: disable=arguments-differ return np.concatenate(variational_parameters) def _construct_component_variational_parameters( - self, variational_parameters: ArrayN - ) -> tuple[list, np.ndarray]: + self, variational_parameters: ArrayNParams + ) -> tuple[list[ArrayNParamsComponent], ArrayNComponents]: """Reconstruct the weights and parameters of the mixture components. Creates a list containing the variational parameters of the different components. @@ -137,8 +152,8 @@ def _construct_component_variational_parameters( return variational_parameters_list, weights def reconstruct_distribution_parameters( - self, variational_parameters: ArrayN - ) -> tuple[list, ArrayN]: + self, variational_parameters: ArrayNParams + ) -> tuple[list[tuple[list | np.ndarray]], ArrayNComponents]: """Reconstruct the weights and parameters of the mixture components. The list is nested, each entry correspond to the parameters of a component. @@ -165,7 +180,7 @@ def reconstruct_distribution_parameters( weights = weights / np.sum(weights) return distribution_parameters_list, weights - def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: + def draw(self, variational_parameters: ArrayNParams, n_draws: NSamples) -> ArrayNSamplesXNDims: """Draw *n_draw* samples from the variational distribution. Uses a two-step process: @@ -177,7 +192,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: n_draws: Number of samples to draw Returns: - Samples of shape (n_draws, n_dim) + Samples """ parameters, weights = self._construct_component_variational_parameters( variational_parameters @@ -192,7 +207,7 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: samples = np.concatenate(samples, axis=0) return samples - def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def logpdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Log-PDF evaluated using the variational parameters at samples *x*. Is a general implementation using the log-PDF function of the components. Uses the @@ -224,7 +239,7 @@ def logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: logpdf = np.log(logpdf) + max_logpdf return logpdf - def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def pdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Pdf evaluated using the variational parameters at given samples `x`. Args: @@ -237,7 +252,9 @@ def pdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: pdf = np.exp(self.logpdf(variational_parameters, x)) return pdf - def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> np.ndarray: + def grad_params_logpdf( + self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims + ) -> ArrayNParamsXNSamples: """Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. @@ -277,8 +294,8 @@ def grad_params_logpdf(self, variational_parameters: ArrayN, x: np.ndarray) -> n return score def fisher_information_matrix( - self, variational_parameters: ArrayN, n_samples: int = 10000 - ) -> np.ndarray: + self, variational_parameters: ArrayNParams, n_samples: int = 10000 + ) -> ArrayNParamsXNParams: """Approximate the Fisher information matrix using Monte Carlo. Args: @@ -286,17 +303,18 @@ def fisher_information_matrix( n_samples: Number of samples for a MC FIM estimation Returns: - Fisher information matrix (num parameters x num parameters) + Fisher information matrix """ samples = self.draw(variational_parameters, n_samples) scores = self.grad_params_logpdf(variational_parameters, samples) - fim = np.zeros((scores.shape[0], scores.shape[0])) + n_var_params = scores.shape[0] + fim = np.zeros((n_var_params, n_var_params)) for j in range(n_samples): fim = fim + np.outer(scores[:, j], scores[:, j]) fim = fim / n_samples return fim - def export_dict(self, variational_parameters: ArrayN) -> dict: + def export_dict(self, variational_parameters: ArrayNParams) -> dict: """Create a dict of the distribution based on the given parameters. Args: diff --git a/src/queens/variational_distributions/particle.py b/src/queens/variational_distributions/particle.py index 8cb22db34..6b39b3a8c 100644 --- a/src/queens/variational_distributions/particle.py +++ b/src/queens/variational_distributions/particle.py @@ -19,8 +19,15 @@ import numpy as np from queens.distributions.particle import Particle as ParticleDistribution -from queens.utils.type_hinting import ArrayN, ArrayNxM -from queens.variational_distributions._variational_distribution import Variational +from queens.variational_distributions._variational_distribution import ( + ArrayNParams, + ArrayNParamsXNParams, + ArrayNParamsXNSamples, + ArrayNSamples, + ArrayNSamplesXNDims, + NSamples, + Variational, +) class Particle(Variational): @@ -44,8 +51,8 @@ def __init__(self, sample_space: np.ndarray | Sequence[Sized]) -> None: super().__init__(self.particles_obj.dimension, n_parameters=len(sample_space)) def construct_variational_parameters( # pylint: disable=arguments-differ - self, probabilities: ArrayN, sample_space: np.ndarray - ) -> ArrayN: + self, probabilities: ArrayNParams, sample_space: np.ndarray | Sequence[Sized] + ) -> ArrayNParams: """Construct the variational parameters from the probabilities. Args: @@ -59,7 +66,7 @@ def construct_variational_parameters( # pylint: disable=arguments-differ variational_parameters = np.log(probabilities).flatten() return variational_parameters - def initialize_variational_parameters(self, random: bool = False) -> ArrayN: + def initialize_variational_parameters(self, random: bool = False) -> ArrayNParams: r"""Initialize variational parameters. Default initialization: @@ -73,7 +80,7 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: random: If True, a random initialization is used. Otherwise the default is selected Returns: - Variational parameters of shape (n_params,) + Variational parameters """ if random: variational_parameters = ( @@ -86,8 +93,8 @@ def initialize_variational_parameters(self, random: bool = False) -> ArrayN: return variational_parameters def reconstruct_distribution_parameters( - self, variational_parameters: ArrayN - ) -> tuple[ArrayN, np.ndarray]: + self, variational_parameters: ArrayNParams + ) -> tuple[ArrayNParams, np.ndarray]: """Reconstruct probabilities from the variational parameters. Args: @@ -102,7 +109,7 @@ def reconstruct_distribution_parameters( self.particles_obj = ParticleDistribution(probabilities, self.particles_obj.sample_space) return probabilities, self.particles_obj.sample_space - def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: + def draw(self, variational_parameters: ArrayNParams, n_draws: NSamples) -> ArrayNSamplesXNDims: """Draw *n_draws* samples from distribution. Args: @@ -110,17 +117,17 @@ def draw(self, variational_parameters: ArrayN, n_draws: int = 1) -> ArrayNxM: n_draws: Number of samples Returns: - Samples of shape (n_draws, n_dim) + Samples """ self.reconstruct_distribution_parameters(variational_parameters) return self.particles_obj.draw(n_draws) - def logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: + def logpdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Evaluate the natural logarithm of the PDF. Args: variational_parameters: Variational parameters of the distribution - x: Locations at which to evaluate the distribution of shape (n_samples, n_dim) + x: Locations at which to evaluate the distribution Returns: Log-PDF values at the locations x @@ -128,12 +135,12 @@ def logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: self.reconstruct_distribution_parameters(variational_parameters) return self.particles_obj.logpdf(x) - def pdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: + def pdf(self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims) -> ArrayNSamples: """Evaluate the probability density function (PDF). Args: variational_parameters: Variational parameters of the distribution - x: Locations at which to evaluate the distribution of shape (n_samples, n_dim) + x: Locations at which to evaluate the distribution Returns: Row vector of the PDF values @@ -141,7 +148,9 @@ def pdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: self.reconstruct_distribution_parameters(variational_parameters) return self.particles_obj.pdf(x) - def grad_params_logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np.ndarray: + def grad_params_logpdf( + self, variational_parameters: ArrayNParams, x: ArrayNSamplesXNDims + ) -> ArrayNParamsXNSamples: r"""Log-PDF gradient w.r.t. the variational parameters. Evaluated at samples *x*. Also known as the score function. @@ -151,7 +160,7 @@ def grad_params_logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np. Args: variational_parameters: Variational parameters of the distribution - x: Locations at which to evaluate the distribution of shape (n_samples, n_dim) + x: Locations at which to evaluate the distribution Returns: Score functions at the locations x @@ -172,7 +181,9 @@ def grad_params_logpdf(self, variational_parameters: ArrayN, x: ArrayNxM) -> np. # Get the samples return sample_scores[index].T - def fisher_information_matrix(self, variational_parameters: ArrayN) -> ArrayNxM: + def fisher_information_matrix( + self, variational_parameters: ArrayNParams + ) -> ArrayNParamsXNParams: r"""Compute the Fisher information matrix. For the given parameterization, the Fisher information yields: @@ -182,13 +193,13 @@ def fisher_information_matrix(self, variational_parameters: ArrayN) -> ArrayNxM: variational_parameters: Variational parameters of the distribution Returns: - Fisher information matrix of shape (n_parameters, n_parameters) + Fisher information matrix """ probabilities, _ = self.reconstruct_distribution_parameters(variational_parameters) fim = np.diag(probabilities) - np.outer(probabilities, probabilities) return fim - def export_dict(self, variational_parameters: ArrayN) -> dict: + def export_dict(self, variational_parameters: ArrayNParams) -> dict: """Create a dict of the distribution based on the given parameters. Args: