Skip to content

Documentation for levy-stable-jax

levy_stable_jax

Implementation of Levy Stable distributions in JAX.

Params

The different parametrizations that are allowed for the stable distribution:

  • N0: The main parametrization. It is the most favorable from a numerical perspective (the pdf is continuous in the parameters).

  • N1: The parametrization used by scipy by default. The most intuitive for users, since the loc parameter is equal to the mean of the distribution in this parametrization.

Source code in src/levy_stable_jax/_typing.py 
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
class Params:
    """
    The different parametrizations that are allowed for the stable
    distribution:

    - N0: The main parametrization. It is the most favorable from a numerical
    perspective (the pdf is continuous in the parameters).

    - N1: The parametrization used by scipy by default. The most intuitive
    for users, since the `loc` parameter is equal to the mean of the distribution
    in this parametrization.

    """

    N0: Literal["N0"] = "N0"
    N1: Literal["N1"] = "N1"

logpdf(x, alpha, beta, loc=0.0, scale=1.0, param=Params.N0)

The probability density function for the Levy-Stable distribution.

Parameters:

Name Type Description Default
x INPUT_TYPE

the values at which to evaluate the PDF

 required
alpha INPUT_TYPE

the alpha parameter of the distribution

 required
beta INPUT_TYPE

the beta parameter of the distribution

 required
loc INPUT_TYPE

the location parameter of the distribution (default 0.0, meaning depends on the parametrization)

 0.0
scale INPUT_TYPE

the scale parameter of the distribution (default 1.0, meaning depends on the parametrization)

 1.0
param Param

the parametrization of the distribution (see Params)

 N0

Examples:

>>> logpdf(0.0, 1.5, 0.0, 1.0) # doctest: +ELLIPSIS
-1.603550...
Source code in src/levy_stable_jax/distribution.py 
124
125
126
127
128
129
130
131
132
133
134
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
def logpdf(
    x: INPUT_TYPE,
    alpha: INPUT_TYPE,
    beta: INPUT_TYPE,
    loc: INPUT_TYPE = 0.0,
    scale: INPUT_TYPE = 1.0,
    param: Param = Params.N0,
) -> JArray | float:
    """
    The probability density function for the Levy-Stable distribution.

    Args:
        x: the values at which to evaluate the PDF
        alpha: the alpha parameter of the distribution
        beta: the beta parameter of the distribution
        loc: the location parameter of the distribution
            (default 0.0, meaning depends on the parametrization)
        scale: the scale parameter of the distribution
            (default 1.0, meaning depends on the parametrization)
        param: the parametrization of the distribution (see Params)

    Examples:

    ```py
    >>> logpdf(0.0, 1.5, 0.0, 1.0) # doctest: +ELLIPSIS
    -1.603550...

    ```

    """
    (alpha_, beta_, x_, loc_, scale_), is_scalar = _canonicalize_input(
        [alpha, beta, x, loc, scale]
    )

    # TODO: implement the other parametrizations
    assert param == Params.N0, param
    res = _logpdf_n0(x_, alpha_, beta_, loc_, scale_)
    if is_scalar:
        res = res.item()
    return res

param_convert(alpha, beta, loc, scale, param_from, param_to)

Shifts the loc and scale of the alpha-stable distribution from one parametrization to another.

Parameters:

Name Type Description Default
alpha INPUT

The stability parameter of the stable distribution (0-2.0].

 required
beta INPUT

The skewness parameter of the stable distribution.

 required
loc INPUT

The location parameter of the stable distribution.

 required
scale INPUT

The scale parameter of the stable distribution.

 required
param_from Param

The initial parametrization of the stable distribution.

 required
param_to Param

The requested parametrization

 required

Returns:

Type Description
Tuple[Array | float, Array | float]

(loc, scale) in the requested parametrization.
Source code in src/levy_stable_jax/_utils.py 
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
def param_convert(
    alpha: INPUT,
    beta: INPUT,
    loc: INPUT,
    scale: INPUT,
    param_from: Param,
    param_to: Param,
) -> Tuple[JArray | float, JArray | float]:
    """
    Shifts the loc and scale of the alpha-stable distribution from
    one parametrization to another.

    Args:
        alpha: The stability parameter of the stable distribution (0-2.0].
        beta: The skewness parameter of the stable distribution.
        loc: The location parameter of the stable distribution.
        scale: The scale parameter of the stable distribution.
        param_from: The initial parametrization of the stable distribution.
        param_to: The requested parametrization

    Returns:
        (loc, scale) in the requested parametrization.
    """

    # For now, we only need to shift the location. The conversion of the scale
    # will only be needed for N2/N3 parametrization.
    def _phi() -> JArray:
        return jnp.tan(jnp.pi * alpha / 2)

    if param_from == param_to:
        return (loc, scale)
    # All N0 -> Nx
    if param_from == Params.N0:
        if param_to == Params.N1:
            # N0 -> N1
            loc_to = jnp.where(
                alpha == 1,
                loc - beta * 2 / jnp.pi * _phi(),
                loc - beta * scale * _phi(),
            )
        return (loc_to, scale)
    # All Nx -> N0
    elif param_from == Params.N1 and param_to == Params.N0:
        # N1 -> N0
        loc_to = jnp.where(
            alpha == 1, loc + beta * 2 / jnp.pi * _phi(), loc + beta * scale * _phi()
        )
        return (loc_to, scale)
    else:
        # For other combinations (which do not exist right now),
        # convert first to N0 and then from N0 to the requested parametrization.
        loc_n0, scale_n0 = param_convert(alpha, beta, loc, scale, param_from, Params.N0)
        return param_convert(alpha, beta, loc_n0, scale_n0, Params.N0, param_to)

pdf(x, alpha, beta, loc=0.0, scale=1.0, param=Params.N0)

The probability density function for the Levy-Stable distribution.

Parameters:

Name Type Description Default
x INPUT_TYPE

the values at which to evaluate the PDF

 required
alpha INPUT_TYPE

the alpha parameter of the distribution

 required
beta INPUT_TYPE

the beta parameter of the distribution

 required
loc INPUT_TYPE

the location parameter of the distribution (default 0.0, meaning depends on the parametrization)

 0.0
scale INPUT_TYPE

the scale parameter of the distribution (default 1.0, meaning depends on the parametrization)

 1.0
param Param

the parametrization of the distribution

 N0

Returns:

Type Description
Array | float

the value of the PDF at the given point(s)

Examples:

Evaluation of the unit Levy-stable distribution at 0.0, with alpha=1.5, beta=0.0. 

>>> pdf(0.0, 1.5, 0.0) # doctest: +ELLIPSIS
0.28568...
>>> pdf(0.0, 1.5, 0.0,param=Params.N0) # doctest: +ELLIPSIS
0.28568...

Source code in src/levy_stable_jax/distribution.py 
 77
 78
 79
 80
 81
 82
 83
 84
 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
def pdf(
    x: INPUT_TYPE,
    alpha: INPUT_TYPE,
    beta: INPUT_TYPE,
    loc: INPUT_TYPE = 0.0,
    scale: INPUT_TYPE = 1.0,
    param: Param = Params.N0,
) -> JArray | float:
    """
    The probability density function for the Levy-Stable distribution.

    Args:
        x: the values at which to evaluate the PDF
        alpha: the alpha parameter of the distribution
        beta: the beta parameter of the distribution
        loc: the location parameter of the distribution
            (default 0.0, meaning depends on the parametrization)
        scale: the scale parameter of the distribution
            (default 1.0, meaning depends on the parametrization)
        param: the parametrization of the distribution

    Returns:
        the value of the PDF at the given point(s)

    Examples:

    Evaluation of the unit Levy-stable distribution at 0.0, with alpha=1.5, beta=0.0.
    ```py
    >>> pdf(0.0, 1.5, 0.0) # doctest: +ELLIPSIS
    0.28568...
    >>> pdf(0.0, 1.5, 0.0,param=Params.N0) # doctest: +ELLIPSIS
    0.28568...

    ```
    """
    # TODO: implement the other parametrizations
    (alpha_, beta_, x_, loc_, scale_), is_scalar = _canonicalize_input(
        [alpha, beta, x, loc, scale]
    )

    assert param == Params.N0, param
    res = _pdf_n0(x_, alpha_, beta_, loc_, scale_)
    if is_scalar:
        res = res.item()
    return res

rvs(alpha, beta, prng, loc=0.0, scale=1.0, shape=(), param=Params.N0)

Generate random samples from the Levy-Stable distribution.

Parameters:

Name Type Description Default
alpha INPUT_TYPE

the alpha parameter of the distribution

 required
beta INPUT_TYPE

the beta parameter of the distribution

 required
prng ArrayLike

the pseudo-random number generator key

 required
loc INPUT_TYPE

the location parameter of the distribution

 0.0
scale INPUT_TYPE

the scale parameter of the distribution

 1.0
shape Shape

the shape of the output array

 ()
param Param

the parametrization of the distribution

 N0

Returns: - the generated samples

Examples:

>>> import jax
>>> prng = jax.random.PRNGKey(1)
>>> rvs(alpha=1.5, beta=0.0, loc=0.0, scale=1.0, param=Params.N1,
...  shape=(10,), prng=prng) # doctest: +ELLIPSIS
Array([-0.750..., -0.495..., ...], ...)
Source code in src/levy_stable_jax/distribution.py 
166
167
168
169
170
171
172
173
174
175
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
def rvs(
    alpha: INPUT_TYPE,
    beta: INPUT_TYPE,
    prng: jax.typing.ArrayLike,
    loc: INPUT_TYPE = 0.0,
    scale: INPUT_TYPE = 1.0,
    shape: Shape = (),
    param: Param = Params.N0,
) -> JArray:
    """
    Generate random samples from the Levy-Stable distribution.

    Args:
        alpha: the alpha parameter of the distribution
        beta: the beta parameter of the distribution
        prng: the pseudo-random number generator key
        loc: the location parameter of the distribution
        scale: the scale parameter of the distribution
        shape: the shape of the output array
        param: the parametrization of the distribution

    Returns:
    - the generated samples

    Examples:

    ```py
    >>> import jax
    >>> prng = jax.random.PRNGKey(1)
    >>> rvs(alpha=1.5, beta=0.0, loc=0.0, scale=1.0, param=Params.N1,
    ...  shape=(10,), prng=prng) # doctest: +ELLIPSIS
    Array([-0.750..., -0.495..., ...], ...)

    ```
    """
    # The sampling algorithm is for N1 parametrization.
    unit_n1 = _sample_unit_n1(jnp.asarray(alpha), jnp.asarray(beta), prng, shape)
    return _values_n1_to_param(alpha, beta, unit_n1, loc, scale, param)  # type:ignore

set_stable(p)

Manages the parametrization of scipy's levy_stable distribution.

Since the parametrization of scipy is at the module level, this function provides a way to temporarily change the parametrization of the distribution.

Example:

from scipy.stats import levy_stable as sp_levy_stable
with set_stable(Params.N0):
    # Will be parametrized as N0 / scipy's S0 instead of default S1.
    _ = sp_levy_stable.rvs(alpha=1.5, beta=0.0, size=10)

TODO: turn into doctest, there is an issue with indentation

Source code in src/levy_stable_jax/_utils.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
@contextmanager
def set_stable(p: Param) -> Generator[Any, Any, Any]:
    """
    Manages the parametrization of scipy's levy_stable distribution.

    Since the parametrization of scipy is at the module level, this function
    provides a way to temporarily change the parametrization of the distribution.

    Example:

    ```python
    from scipy.stats import levy_stable as sp_levy_stable
    with set_stable(Params.N0):
        # Will be parametrized as N0 / scipy's S0 instead of default S1.
        _ = sp_levy_stable.rvs(alpha=1.5, beta=0.0, size=10)
    ```

    TODO: turn into doctest, there is an issue with indentation
    """
    curr = sp_levy_stable.parameterization
    if p == Params.N0:
        sp_levy_stable.parameterization = "S0"
    elif p == Params.N1:
        sp_levy_stable.parameterization = "S1"
    else:
        # TODO: proper error
        raise ValueError(f"Invalid parametrization {p}")
    try:
        yield
    finally:
        sp_levy_stable.parameterization = curr

shift_scale(alpha, beta, loc, scale, a, b, param)

Given an alpha-stable distribution X with parameters alpha, beta, loc, scale, returns the distribution Y = a * X + b, in the same parametrization. Since the factor alpha is the same, it returns a tuple of (beta, loc, scale)

Parameters:

Name Type Description Default
alpha INPUT

The stability parameter of the stable distribution (0-2.0].

 required
beta INPUT

The skewness parameter of the stable distribution.

 required
loc INPUT

The location parameter of the stable distribution.

 required
scale INPUT

The scale parameter of the stable distribution.

 required
a INPUT

The scaling factor.

 required
b INPUT

The shift factor.

 required
param Param

The parametrization of the stable distribution.

 required

Example: 

import jax.numpy as jnp
>>> shift_scale(2.0, 0.0, 0.0, 1.0, a=2.0, b=1.0, param=Params.N1) # doctest: +ELLIPSIS
(Array(0., dtype=...), Array(1., dtype=...), Array(2., dtype=...))

Source code in src/levy_stable_jax/_utils.py 
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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
def shift_scale(
    alpha: INPUT,
    beta: INPUT,
    loc: INPUT,
    scale: INPUT,
    a: INPUT,
    b: INPUT,
    param: Param,
) -> Tuple[JArray | float, JArray | float, JArray]:
    """
    Given an alpha-stable distribution X with parameters alpha, beta, loc, scale,
    returns the distribution Y = a * X + b, in the same parametrization.
    Since the factor alpha is the same, it returns a tuple of (beta, loc, scale)

    Args:
        alpha: The stability parameter of the stable distribution (0-2.0].
        beta: The skewness parameter of the stable distribution.
        loc: The location parameter of the stable distribution.
        scale: The scale parameter of the stable distribution.
        a: The scaling factor.
        b: The shift factor.
        param: The parametrization of the stable distribution.

    Example:
    ```python
    import jax.numpy as jnp
    >>> shift_scale(2.0, 0.0, 0.0, 1.0, a=2.0, b=1.0, param=Params.N1) # doctest: +ELLIPSIS
    (Array(0., dtype=...), Array(1., dtype=...), Array(2., dtype=...))

    ```

    """

    beta2 = jnp.sign(a) * beta
    scale2 = jnp.abs(a) * scale
    if param == Params.N0:
        return (beta2, loc * a + b, scale2)
    elif param == Params.N1:
        # This code will not support a == 0, but then we have other degenerate distribution problems
        # anyway...
        loc2 = jnp.where(
            alpha == 1,
            a * loc + b - (2 / jnp.pi) * (a * jnp.log(jnp.abs(a))) * (beta * scale),
            a * loc + b,
        )
        return (beta2, loc2, scale2)

sum(alpha, beta, loc=0.0, scale=1.0, param=Params.N0, axis=None)

Computes the sum of independent stable random distributions.

Parameters:

Name Type Description Default
alpha INPUT

The stability parameter of the stable distribution (0-2.0]. It is fixed for all distributions.

 required
beta INPUT

The skewness parameter of the stable distribution.

 required
loc INPUT

The location parameter of the stable distribution (exact definition depending on the choice of pametrization)

 0.0
scale INPUT

The scale parameter of the stable distribution.

 1.0
param Param

The parametrization of the stable distribution.

 N0

Returns: A tuple of (beta_sum, loc_sum, scale_sum)

Example: the sum of two centered Gaussian distributions is a centered Gaussian distribution with a standard deviation of sqrt(2). 

import jax.numpy as jnp
>>> sum(2.0, 0.0, 0.0, jnp.asarray([1.0, 1.0]), Params.N1) # doctest: +ELLIPSIS
(Array(0., dtype=float64), Array(0., dtype=float64), Array(1.4142135...,...))

Example: dot product. Say that x1 ~ S(2, 0, 1, 1) and x2 ~ S(2, 0, 1.1, 1) are independent. We want to compute the distribution z = A . x + b where x = [x1, x2], A is a 3 x 2 matrix, b is a 3-vector, and . is the matrix-vector product.

This can be written by a summing a scale operation:

>>> locs = jnp.asarray([1.0, 1.1])
>>> scales = 2.0
>>> A = jnp.asarray([[1,2],[3,4],[5,6]])
>>> b = jnp.asarray([[2,2,2]]).T
>>> (beta1, loc1, scale1) = shift_scale(2.0, 0.0,locs, scales, A,b,"N1")
>>> print(beta1)
[[0. 0.]
 [0. 0.]
 [0. 0.]]
>>> print(loc1)
[[3.  4.2]
 [5.  6.4]
 [7.  8.6]]
>>> print(scale1)
[[ 2.  4.]
 [ 6.  8.]
 [10. 12.]]
>>> sum(2.0, beta1, loc1, scale1, param="N1", axis=0) # doctest: +ELLIPSIS
(Array([0., 0.], ...), Array([15. , 19.2], ...), Array([11.83215957, 14.96662955], ...))
Source code in src/levy_stable_jax/_utils.py 
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
174
175
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
221
222
223
224
225
226
227
228
229
230
231
232
def sum(
    alpha: INPUT,
    beta: INPUT,
    loc: INPUT = 0.0,
    scale: INPUT = 1.0,
    param: Param = Params.N0,
    axis: Union[int, Sequence[int], None] = None,
) -> Tuple[JArray, JArray, JArray]:
    """
    Computes the sum of independent stable random distributions.

    Args:
        alpha: The stability parameter of the stable distribution (0-2.0]. It is fixed
            for all distributions.
        beta: The skewness parameter of the stable distribution.
        loc: The location parameter of the stable distribution (exact definition
            depending on the choice of pametrization)
        scale: The scale parameter of the stable distribution.
        param: The parametrization of the stable distribution.

    Returns: A tuple of (beta_sum, loc_sum, scale_sum)

    Example: the sum of two centered Gaussian distributions is a centered Gaussian distribution with
    a standard deviation of sqrt(2).
    ```python
    import jax.numpy as jnp
    >>> sum(2.0, 0.0, 0.0, jnp.asarray([1.0, 1.0]), Params.N1) # doctest: +ELLIPSIS
    (Array(0., dtype=float64), Array(0., dtype=float64), Array(1.4142135...,...))

    ```

    Example: dot product.
    Say that x1 ~ S(2, 0, 1, 1) and x2 ~ S(2, 0, 1.1, 1) are independent.
    We want to compute the distribution z = A . x + b where
    x = [x1, x2], A is a 3 x 2 matrix, b is a 3-vector, and . is the matrix-vector product.

    This can be written by a summing a scale operation:

    ```python
    >>> locs = jnp.asarray([1.0, 1.1])
    >>> scales = 2.0
    >>> A = jnp.asarray([[1,2],[3,4],[5,6]])
    >>> b = jnp.asarray([[2,2,2]]).T
    >>> (beta1, loc1, scale1) = shift_scale(2.0, 0.0,locs, scales, A,b,"N1")
    >>> print(beta1)
    [[0. 0.]
     [0. 0.]
     [0. 0.]]
    >>> print(loc1)
    [[3.  4.2]
     [5.  6.4]
     [7.  8.6]]
    >>> print(scale1)
    [[ 2.  4.]
     [ 6.  8.]
     [10. 12.]]
    >>> sum(2.0, beta1, loc1, scale1, param="N1", axis=0) # doctest: +ELLIPSIS
    (Array([0., 0.], ...), Array([15. , 19.2], ...), Array([11.83215957, 14.96662955], ...))

    ```

    """
    # Adopting the notation of Nolan 2020.
    gamma = scale
    delta = loc
    phi = jnp.tan(jnp.pi * alpha / 2)
    gamma_alpha = jnp.sum(gamma**alpha, axis=axis)
    gamma_sum = gamma_alpha ** (1.0 / alpha)
    # TODO: this is not robust to gamma == 0
    glg = gamma * jnp.log(gamma)
    beta_sum = jnp.sum(beta * gamma**alpha, axis=axis) / gamma_alpha
    # Proposition 1.3 in Nolan 2020.
    if param == Params.N0:
        s = jnp.sum(delta, axis=axis)
        s1 = jnp.sum(beta * gamma, axis=axis)
        s1_a1 = jnp.sum(beta * glg, axis=axis)
        delta_sum = s + phi * (beta_sum * gamma - s1)
        delta_sum_a1 = s + 2.0 / jnp.pi * (
            beta_sum * gamma_sum * jnp.log(gamma_sum) - s1_a1
        )
        delta_sum = jnp.where(alpha == 1, delta_sum_a1, delta_sum)
    elif param == Params.N1:
        delta_sum = jnp.sum(delta, axis=axis)
    return (beta_sum, delta_sum, gamma_sum)

levy_stable_jax.estimation

Different algorithms to estimate the parameters of a stable distribution.

These algorithms should be all considered as experimental and not part of the standard API: their interface is not likely to change but they may not be as numerically stable.

fit_ll(samples, param, weights=None, alpha=None, beta=None)

Maximum likelihood evaluation for univariate Lévy-Stable distributions.

Parameters:

Name Type Description Default
param Param

the parametrization of the returned distribution

 required
samples Array

a 1-d array with the observed samples.

 required
weights Optional[Array]

optional, the weights on each of the samples.

 None
alpha Optional[float]

optional, the value of alpha to be used in the optimization.

 None
beta Optional[float]

optional, the value of beta to be used in the optimization.

 None

Examples:

>>> samples = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
>>> fit_quantiles(samples, Params.N0) # doctest: +ELLIPSIS
Array([2.        , 0.        , 3.        , 1.048...], dtype=float64)

The first value is alpha, then beta, then the location parameter and then the scale parameter.

Source code in src/levy_stable_jax/estimation.py 
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 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
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def fit_ll(
    samples: JArray,
    param: Param,
    weights: Optional[JArray] = None,
    alpha: Optional[float] = None,
    beta: Optional[float] = None,
) -> JArray:
    """
    Maximum likelihood evaluation for univariate Lévy-Stable distributions.

    Args:
        param: the parametrization of the returned distribution
        samples: a 1-d array with the observed samples.
        weights: optional, the weights on each of the samples.
        alpha: optional, the value of alpha to be used in the optimization.
        beta: optional, the value of beta to be used in the optimization.

    Examples:

    ```py
    >>> samples = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
    >>> fit_quantiles(samples, Params.N0) # doctest: +ELLIPSIS
    Array([2.        , 0.        , 3.        , 1.048...], dtype=float64)

    ```

    The first value is alpha, then beta, then the location parameter and then
     the scale parameter.
    """
    assert samples.ndim == 1, samples
    assert samples.size > 4, samples.size

    def _unpack(theta: JArray) -> Tuple[JArray, JArray, JArray, JArray]:
        idx = 0
        if alpha is None:
            alpha_ = theta[idx]
            idx += 1
        else:
            alpha_ = jnp.asarray(alpha)
        if beta is None:
            beta_ = theta[idx]
            idx += 1
        else:
            beta_ = jnp.asarray(beta)
        loc_ = theta[idx]
        scale_ = theta[idx + 1]
        return (alpha_, beta_, loc_, scale_)

    def _pack(alpha_: JArray, beta_: JArray, loc_: JArray, scale_: JArray) -> JArray:
        vals = [
            [alpha_] if alpha is None else [],
            [beta_] if beta is None else [],
            [loc_],
            [scale_],
        ]
        return jnp.array([x for val in vals for x in val])

    def _log_likelikhood_n0(theta: JArray) -> JArray:
        (alpha_, beta_, loc_, scale_) = _unpack(theta)
        logs = logpdf(
            samples, alpha=alpha_, beta=beta_, loc=loc_, scale=scale_, param=Params.N0
        )
        if weights is None:
            res = -jnp.mean(logs)
        else:
            res = -jnp.mean(logs * weights)
        return res

    solver = jaxopt.ScipyBoundedMinimize(
        fun=_log_likelikhood_n0, method="l-bfgs-b", tol=1e-6
    )
    # Get a crude start
    # TODO: this is not using the weights
    # All the work is done in the N0 parametrization.
    start = fit_quantiles(samples, Params.N0)
    theta_start = _pack(*start)

    # Build the bounds for the solver:
    def _bounds() -> Tuple[JArray, JArray]:
        lower = []
        upper = []
        idx = 0
        if alpha is None:
            lower.append(1.1)
            upper.append(2.0)
            idx += 1
        if beta is None:
            lower.append(-1.0)
            upper.append(1.0)
            idx += 1
        # Loc
        lower.append(-1e4)
        upper.append(1e4)
        # Scale
        lower.append(1e-3)
        upper.append(1e2)
        return (jnp.array(lower), jnp.array(upper))

    theta_bounds = _bounds()
    sres = solver.run(theta_start, bounds=theta_bounds)
    (alpha0, beta0, loc0, scale0) = _unpack(sres.params)
    # Convert to the requested parametrization
    (loc_x, scale_x) = param_convert(alpha0, beta0, loc0, scale0, Params.N0, param)
    return jnp.array([alpha0, beta0, loc_x, scale_x])

fit_quantiles(samples, param)

A rough approximation of the distribution parameters based on quantiles.

For now, it is just wrapping the scipy code.

Examples:

>>> samples = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0, 30.0])
>>> fit_quantiles(samples, Params.N0)  # doctest: +ELLIPSIS
Array([0.76..., 0.89..., 3.07..., 0.64...], dtype=float64)

The first value is alpha, then beta, then the location parameter and then the scale parameter.

Source code in src/levy_stable_jax/estimation.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
def fit_quantiles(samples: JArray, param: Param) -> JArray:
    """
    A rough approximation of the distribution parameters based on quantiles.

    For now, it is just wrapping the scipy code.

    Examples:

    ```py
    >>> samples = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0, 30.0])
    >>> fit_quantiles(samples, Params.N0)  # doctest: +ELLIPSIS
    Array([0.76..., 0.89..., 3.07..., 0.64...], dtype=float64)

    ```

    The first value is alpha, then beta, then the location parameter and then
     the scale parameter.

    """
    # TODO: it would be nice to have pure jax code, but the scipy code
    # depends on the scipy.interpolate module, which is not going to be
    # implemented in jax.
    assert samples.ndim == 1, samples
    assert samples.size > 4, samples.size
    if param == Params.N1:
        res = scipy.stats._levy_stable._fitstart_S1(samples)
    else:
        res = scipy.stats._levy_stable._fitstart_S0(samples)
    return jnp.array(res)

levy_stable_jax.pymc

Monte Carlo sampling of Levy-stable distributions using PyMC.

As far as users are concerned, the only relevant function should be LevyStableN0.

A full example notebook is available here: https://github.com/tjhunter/levy-stable-jax/blob/main/notebooks/pymc_levy.ipynb

LevyStableN0(name, alpha, beta, loc, scale, observed=None)

A Lévy-stable distribution. The parametrization follows the "0" notation from Nolan (2022). It is also known as the "S0" parametrization in scipy.

Parameters:

Name Type Description Default
alpha(tensor)

the stability parameter. Must be in (1.1, 2].

 required
beta(tensor)

the skewness parameter. Must be in [-1, 1].

 required
loc(tensor)

the location parameter (delta in Nolan's notation).

 required
scale(tensor)

the scale parameter (gamma in Nolan's notation).

 required

This distribution is explicitly implemented and parametrized for the N0 parametrization, because the N1 parametrization is not continuous for alpha close to 1, and is not recommended in general for numerical work. Use the shift_scale function if you wish to convert to other parametrizations.

Practical tips

It is currently only implemented for alpha in (1.1 - 2.0]. Smaller values will get trimmed.

In the context of MCMC, it is highly recommended to keep beta in (-0.8, 0.8). The Lévy-stable distribution changes abruptly around values of beta close to 1 or -1. As a result, the sampler will reject many values for small changes around these values. The symptoms will be a high number of divergences and slow progress. It is better to trim the range of accepted values of beta and inspect if the posterior distribution is highly skewed towards the extremes.

You cannot really infer a good value for |beta| > 0.8, unless you have an extremely large amount of points. In addition, all other parameters being equal, the log-likelihood as a function of beta is very flat around |beta| ~ (0.8-0.99) and alpha > 1.5: a very large number of observations would be necessary to observed a rare event in the non-heavy tail that would allow to distinguish between these values.

Because of the rather complex and non-convex shape of the log-likelihood, the sampler may struggle to initially reach a region of high probability. I have found that tune=6000 really helps further convergence.

Since all the code is implemented in JAX, a JAX-based sampler will be much faster. It is recommended to use pymc.sampling_jax.sample_numpyro_nuts

TODO: turn into a proper distribution.

Source code in src/levy_stable_jax/pymc.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def LevyStableN0(name, alpha, beta, loc, scale, observed=None):
    """
    A Lévy-stable distribution. The parametrization follows the "0" notation
    from Nolan (2022). It is also known as the "S0" parametrization in scipy.

    Args:
        alpha(tensor): the stability parameter. Must be in (1.1, 2].
        beta(tensor): the skewness parameter. Must be in [-1, 1].
        loc(tensor): the location parameter (delta in Nolan's notation).
        scale(tensor): the scale parameter (gamma in Nolan's notation).

    This distribution is explicitly implemented and parametrized for the N0
    parametrization, because the N1 parametrization is not continuous
    for alpha close to 1, and is not recommended in general for numerical work.
    Use the `shift_scale` function if you wish to convert to other parametrizations.

    **Practical tips**

    It is currently only implemented for alpha in (1.1 - 2.0]. Smaller values
    will get trimmed.

    In the context of MCMC, it is highly recommended to keep beta in (-0.8, 0.8).
    The Lévy-stable distribution changes abruptly around values of beta close to
    1 or -1. As a result, the sampler will reject many values for small changes
    around these values. The symptoms will be a high number of divergences and slow
    progress. It is better to trim the range of accepted values of beta and inspect
    if the posterior distribution is highly skewed towards the extremes.

    You cannot really infer a good value for |beta| > 0.8, unless you have an
    extremely large amount of points.
    In addition, all other parameters being equal, the log-likelihood as a function of beta
    is very flat around |beta| ~ (0.8-0.99) and alpha > 1.5: a very large number
    of observations would be necessary to observed a rare event in the non-heavy tail that
    would allow to distinguish between these values.

    Because of the rather complex and non-convex shape of the log-likelihood,
    the sampler may struggle to initially reach a region of high probability.
    I have found that tune=6000 really helps further convergence.

    Since all the code is implemented in JAX, a JAX-based sampler will be much faster.
    It is recommended to use `pymc.sampling_jax.sample_numpyro_nuts`

    TODO: turn into a proper distribution.

    """
    return pm.CustomDist(
        name,
        alpha,
        beta,
        loc,
        scale,
        logp=_levy_stable_logp_op,
        random=_sample,
        observed=observed,
    )