Skip to content

Bayes¤

The bayesian module is designed to ease the calculation of things like convariance and fisher matrices in differentiable ways. It implements two likelihood functions poiss_loglike and chi2_loglike. They both take in a pytree and data and return a scalar log likelihood. The poiss_loglike function assumes the data is poisson distributed and the chi2_loglike function assumes the data is normally distributed. To use these functions the input pytree must have a .model() function.

There are also four functions used to calcualte fisher and covariances matrices: fisher_matrix, covariance_matrix, self_fisher_matrix, self_covariance_matrix. The fisher_matrix and covariance_matrix functions take in a pytree, parameters, a log likelihood function and data. They return the fisher and covariance matrices respectively. The self_fisher_matrix and self_covariance_matrix functions take in a pytree, parameters and a log likelihood function. They return the fisher and covariance matrices respectively, but the data is generated from the model itself.

Full API

__all__ = ['poiss_loglike', 'chi2_loglike', 'covaraince_entropy', 'fisher_matrix', 'covariance_matrix', 'self_fisher_matrix', 'self_covariance_matrix'] module-attribute ¤

Assumes all pytrees have a .model() function TODO: Allow args and kwargs to be passed to pytree.model() -> Actually kind of hard, since we already have * args and * kwargs for the loglike function

chi2_loglike(pytree, data, noise=1) ¤

Chi2 log likelihood of the pytree given the data. Assumes the pytree has a .model() function.

Parameters:

Name Type Description Default
pytree Base

Pytree with a .model() function.

 required
data Array

Data to compare the model to.

 required
noise float

The 'scale' parameter representing the noise in the data.

 1

Returns:

Name Type Description
log_likelihood Array

Log likelihood of the pytree given the data.
Source code in zodiax/bayes.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def chi2_loglike(pytree : Base(), data : Array, noise : float = 1) -> float:
    """
    Chi2 log likelihood of the pytree given the data. Assumes the pytree has a
    .model() function.

    Parameters
    ----------
    pytree : Base
        Pytree with a .model() function.
    data : Array
        Data to compare the model to.
    noise : float = 1
        The 'scale' parameter representing the noise in the data.

    Returns
    -------
    log_likelihood : Array
        Log likelihood of the pytree given the data.
    """
    return jsp.stats.chi2.logpdf(pytree.model(), data, scale=noise).sum()

covaraince_entropy(covariance_matrix) ¤

Calculates the entropy of a covariance matrix.

Parameters:

Name Type Description Default
covariance_matrix Array

The covariance matrix to calculate the entropy of.

 required

Returns:

Name Type Description
entropy Array

The entropy of the covariance matrix.
Source code in zodiax/bayes.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def covaraince_entropy(covariance_matrix : Array) -> Array:
    """
    Calculates the entropy of a covariance matrix.

    Parameters
    ----------
    covariance_matrix : Array
        The covariance matrix to calculate the entropy of.

    Returns
    -------
    entropy : Array
        The entropy of the covariance matrix.
    """
    sign, logdet = np.linalg.slogdet(covariance_matrix)
    return 0.5 * (np.log(2 * np.pi * np.e) + (sign * logdet))

covariance_matrix(pytree, parameters, loglike_fn, *loglike_args, shape_dict={}, **loglike_kwargs) ¤

Calculates the covariance matrix of the pytree parameters. The shaped_dict parameter is used to specify the shape of the differentiated vector for specific parameters. For example, if the parameter param is a 1D array of shape (5,) and we wanted to calculate the covariance matrix of the mean, we can pass in shape_dict={'param': (1,)}. This will differentiate the log likelihood with respect to the mean of the parameters.

Parameters:

Name Type Description Default
pytree Base

Pytree with a .model() function.

 required
parameters Union[str, list]

A path or list of paths or list of nested paths.

 required
loglike_fn callable

The log likelihood function to differentiate.

 required
*loglike_args Any

The args to pass to the log likelihood function.

 ()
shape_dict dict

A dictionary specifying the shape of the differentiated vector for specific parameters.

 {}
**loglike_kwargs Any

The kwargs to pass to the log likelihood function.

 {}

Returns:

Name Type Description
covariance_matrix Array

The covariance matrix of the pytree parameters.
Source code in zodiax/bayes.py 
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
def covariance_matrix(
        pytree           : Base(), 
        parameters       : Params,
        loglike_fn       : callable,
        *loglike_args    : Any,
        shape_dict       : dict = {},
        **loglike_kwargs : Any
        ) -> Array:
    """
    Calculates the covariance matrix of the pytree parameters. The
    `shaped_dict` parameter is used to specify the shape of the differentiated
    vector for specific parameters. For example, if the parameter `param` is
    a 1D array of shape (5,) and we wanted to calculate the covariance matrix
    of the mean, we can pass in `shape_dict={'param': (1,)}`. This will
    differentiate the log likelihood with respect to the mean of the parameters.

    Parameters
    ----------
    pytree : Base
        Pytree with a .model() function.
    parameters : Union[str, list]
        A path or list of paths or list of nested paths.
    loglike_fn : callable
        The log likelihood function to differentiate.
    *loglike_args : Any
        The args to pass to the log likelihood function.
    shape_dict : dict = {}
        A dictionary specifying the shape of the differentiated vector for
        specific parameters.
    **loglike_kwargs : Any
        The kwargs to pass to the log likelihood function.

    Returns
    -------
    covariance_matrix : Array
        The covariance matrix of the pytree parameters.
    """
    return -np.linalg.inv(fisher_matrix(pytree, parameters, loglike_fn, 
        *loglike_args, shape_dict=shape_dict, **loglike_kwargs))

fisher_matrix(pytree, parameters, loglike_fn, *loglike_args, shape_dict={}, **loglike_kwargs) ¤

Calculates the Fisher information matrix of the pytree parameters. The shaped_dict parameter is used to specify the shape of the differentiated vector for specific parameters. For example, if the parameter param is a 1D array of shape (5,) and we wanted to calculate the Fisher information of the mean, we can pass in shape_dict={'param': (1,)}. This will differentiate the log likelihood with respect to the mean of the parameters.

Parameters:

Name Type Description Default
pytree Base

Pytree with a .model() function.

 required
parameters Union[str, list]

A path or list of paths or list of nested paths.

 required
loglike_fn callable

The log likelihood function to differentiate.

 required
*loglike_args Any

The args to pass to the log likelihood function.

 ()
shape_dict dict

A dictionary specifying the shape of the differentiated vector for specific parameters.

 {}
**loglike_kwargs Any

The kwargs to pass to the log likelihood function.

 {}

Returns:

Name Type Description
fisher_matrix Array

The Fisher information matrix of the pytree parameters.
Source code in zodiax/bayes.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
def fisher_matrix(
        pytree           : Base(), 
        parameters       : Params,
        loglike_fn       : callable,
        *loglike_args    : Any,
        shape_dict       : dict = {},
        **loglike_kwargs : Any,
        ) -> Array:
    """
    Calculates the Fisher information matrix of the pytree parameters. The
    `shaped_dict` parameter is used to specify the shape of the differentiated
    vector for specific parameters. For example, if the parameter `param` is
    a 1D array of shape (5,) and we wanted to calculate the Fisher information
    of the mean, we can pass in `shape_dict={'param': (1,)}`. This will
    differentiate the log likelihood with respect to the mean of the parameters.

    Parameters
    ----------
    pytree : Base
        Pytree with a .model() function.
    parameters : Union[str, list]
        A path or list of paths or list of nested paths.
    loglike_fn : callable
        The log likelihood function to differentiate.
    *loglike_args : Any
        The args to pass to the log likelihood function.
    shape_dict : dict = {}
        A dictionary specifying the shape of the differentiated vector for
        specific parameters.
    **loglike_kwargs : Any
        The kwargs to pass to the log likelihood function.

    Returns
    -------
    fisher_matrix : Array
        The Fisher information matrix of the pytree parameters.
    """
    # Build X vec
    pytree = zodiax.tree.set_array(pytree, parameters)
    shapes, lengths = _shapes_and_lengths(pytree, parameters, shape_dict)
    X = np.zeros(_lengths_to_N(lengths))

    # Build function to calculate FIM and calculate
    @hessian
    def calc_fim(X):
        parametric_pytree = _perturb(X, pytree, parameters, shapes, lengths)
        return loglike_fn(parametric_pytree, *loglike_args, **loglike_kwargs)
    return calc_fim(X)

poiss_loglike(pytree, data) ¤

Poissonian log likelihood of the pytree given the data. Assumes the pytree has a .model() function.

Parameters:

Name Type Description Default
pytree Base

Pytree with a .model() function.

 required
data Array

Data to compare the model to.

 required

Returns:

Name Type Description
log_likelihood Array

Log likelihood of the pytree given the data.
Source code in zodiax/bayes.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def poiss_loglike(pytree : Base(), data : Array) -> float:
    """
    Poissonian log likelihood of the pytree given the data. Assumes the pytree
    has a .model() function.

    Parameters
    ----------
    pytree : Base
        Pytree with a .model() function.
    data : Array
        Data to compare the model to.

    Returns
    -------
    log_likelihood : Array
        Log likelihood of the pytree given the data.
    """
    return jsp.stats.poisson.logpmf(pytree.model(), data).sum()

self_covariance_matrix(pytree, parameters, loglike_fn, *loglike_args, shape_dict={}, **loglike_kwargs) ¤

Calculates the covariance matrix of the pytree parameters with respect to itself. This can be used to optimise the fisher information of the pytree model. Simply calcluates the data from the .model() method and passes it to the covariance_matrix function. The shaped_dict parameter is used to specify the shape of the differentiated vector for specific parameters. For example, if the parameter param is a 1D array of shape (5,) and we wanted to calculate the covariance matrix of the mean, we can pass in shape_dict={'param': (1,)}. This will differentiate the log likelihood with respect to the mean of the parameters.

Parameters:

Name Type Description Default
pytree Base

Pytree with a .model() function.

 required
parameters Union[str, list]

A path or list of paths or list of nested paths.

 required
loglike_fn callable

The log likelihood function to differentiate.

 required
*loglike_args Any

The args to pass to the log likelihood function.

 ()
shape_dict dict

A dictionary specifying the shape of the differentiated vector for specific parameters.

 {}
**loglike_kwargs Any

The kwargs to pass to the log likelihood function.

 {}

Returns:

Name Type Description
covariance_matrix Array

The covariance matrix of the pytree parameters.
Source code in zodiax/bayes.py 
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
def self_covariance_matrix(
        pytree           : Base(), 
        parameters       : Params, 
        loglike_fn       : callable,
        *loglike_args    : Any,
        shape_dict       : dict = {},
        **loglike_kwargs : Any
        ) -> Array:
    """
    Calculates the covariance matrix of the pytree parameters with respect to
    itself. This can be used to optimise the fisher information of the pytree
    model. Simply calcluates the data from the `.model()` method and passes it
    to the `covariance_matrix` function. The `shaped_dict` parameter is used to
    specify the shape of the differentiated vector for specific parameters. For
    example, if the parameter `param` is a 1D array of shape (5,) and we wanted
    to calculate the covariance matrix of the mean, we can pass in
    `shape_dict={'param': (1,)}`. This will differentiate the log likelihood
    with respect to the mean of the parameters.

    Parameters
    ----------
    pytree : Base
        Pytree with a .model() function.
    parameters : Union[str, list]
        A path or list of paths or list of nested paths.
    loglike_fn : callable
        The log likelihood function to differentiate.
    *loglike_args : Any
        The args to pass to the log likelihood function.
    shape_dict : dict = {}
        A dictionary specifying the shape of the differentiated vector for
        specific parameters.
    **loglike_kwargs : Any
        The kwargs to pass to the log likelihood function.

    Returns
    -------
    covariance_matrix : Array
        The covariance matrix of the pytree parameters.
    """
    data = pytree.model()
    loglike_args = [pytree.model()] + list(loglike_args)
    return covariance_matrix(pytree, parameters, loglike_fn, *loglike_args, 
        shape_dict=shape_dict, **loglike_kwargs)

self_fisher_matrix(pytree, parameters, loglike_fn, *loglike_args, shape_dict={}, **loglike_kwargs) ¤

Calculates the Fisher information matrix of the pytree parameters with respect to itself. This can be used to optimise the fisher information of the pytree model. Simply calcluates the data from the .model() method and passes it to the fisher_matrix function. The shaped_dict parameter is used to specify the shape of the differentiated vector for specific parameters. For example, if the parameter param is a 1D array of shape (5,) and we wanted to calculate the covariance matrix of the mean, we can pass in shape_dict={'param': (1,)}. This will differentiate the log likelihood with respect to the mean of the parameters.

Parameters:

Name Type Description Default
pytree Base

Pytree with a .model() function.

 required
parameters Union[str, list]

A path or list of paths or list of nested paths.

 required
loglike_fn callable

The log likelihood function to differentiate.

 required
*loglike_args Any

The args to pass to the log likelihood function.

 ()
shape_dict dict

A dictionary specifying the shape of the differentiated vector for specific parameters.

 {}
**loglike_kwargs Any

The kwargs to pass to the log likelihood function.

 {}

Returns:

Name Type Description
fisher_matrix Array

The Fisher information matrix of the pytree parameters.
Source code in zodiax/bayes.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
def self_fisher_matrix(
        pytree           : Base(), 
        parameters       : Params, 
        loglike_fn       : callable,
        *loglike_args    : Any,
        shape_dict       : dict = {},
        **loglike_kwargs : Any
        ) -> Array:
    """
    Calculates the Fisher information matrix of the pytree parameters with
    respect to itself. This can be used to optimise the fisher information of
    the pytree model. Simply calcluates the data from the `.model()` method and
    passes it to the `fisher_matrix` function. The `shaped_dict` parameter is
    used to specify the shape of the differentiated vector for specific
    parameters. For example, if the parameter `param` is a 1D array of shape
    (5,) and we wanted to calculate the covariance matrix of the mean, we can
    pass in `shape_dict={'param': (1,)}`. This will differentiate the log
    likelihood with respect to the mean of the parameters.

    Parameters
    ----------
    pytree : Base
        Pytree with a .model() function.
    parameters : Union[str, list]
        A path or list of paths or list of nested paths.
    loglike_fn : callable
        The log likelihood function to differentiate.
    *loglike_args : Any
        The args to pass to the log likelihood function.
    shape_dict : dict = {}
        A dictionary specifying the shape of the differentiated vector for
        specific parameters.
    **loglike_kwargs : Any
        The kwargs to pass to the log likelihood function.

    Returns
    -------
    fisher_matrix : Array
        The Fisher information matrix of the pytree parameters.
    """
    # Build X vec and get data
    data = pytree.model()
    loglike_args = [pytree.model()] + list(loglike_args)
    return fisher_matrix(pytree, parameters, loglike_fn, *loglike_args, 
        shape_dict=shape_dict, **loglike_kwargs)

¤

¤

¤

¤

¤

¤

¤

¤

¤