1"""
2Generate Gauss-Kronrod quadrature abscissas and weights.
3
4Example usage
5-------------
6
7Generate Gauss-Kronrod quadrature rules::
8
9    >>> distribution = chaospy.Beta(2, 2, lower=-1, upper=1)
10    >>> for order in range(5):  # doctest: +NORMALIZE_WHITESPACE
11    ...     abscissas, weights = chaospy.generate_quadrature(
12    ...         order, distribution, rule="kronrod")
13    ...     print(abscissas.round(2), weights.round(2))
14    [[0.]] [1.]
15    [[-0.65 -0.    0.65]] [0.23 0.53 0.23]
16    [[-0.82 -0.45  0.    0.45  0.82]] [0.07 0.26 0.34 0.26 0.07]
17    [[-0.89 -0.65 -0.34 -0.    0.34  0.65  0.89]]
18     [0.03 0.12 0.22 0.26 0.22 0.12 0.03]
19    [[-0.93 -0.77 -0.54 -0.29 -0.    0.29  0.54  0.77  0.93]]
20     [0.01 0.06 0.13 0.19 0.22 0.19 0.13 0.06 0.01]
21
22Compare Gauss-Kronrod builds on top of Gauss-Legendre quadrature to pure
23Gauss-Legendre::
24
25    >>> distribution = chaospy.Uniform(-1, 1)
26    >>> for order in range(5):
27    ...     abscissas1, weights1 = chaospy.generate_quadrature(
28    ...         order, distribution, rule="gaussian")
29    ...     abscissas2, weights2 = chaospy.generate_quadrature(
30    ...         order+1, distribution, rule="kronrod")
31    ...     print(abscissas1.round(2), abscissas2[:, 1::2].round(2))
32    [[0.]] [[-0.]]
33    [[-0.58  0.58]] [[-0.58  0.58]]
34    [[-0.77 -0.    0.77]] [[-0.77 -0.    0.77]]
35    [[-0.86 -0.34  0.34  0.86]] [[-0.86 -0.34  0.34  0.86]]
36    [[-0.91 -0.54  0.    0.54  0.91]] [[-0.91 -0.54 -0.    0.54  0.91]]
37
38Gauss-Kronrod build on top of Gauss-Hermite quadrature::
39
40    >>> distribution = chaospy.Normal(0, 1)
41    >>> for order in range(2):
42    ...     abscissas1, weights1 = chaospy.generate_quadrature(
43    ...         order, distribution, rule="gaussian")
44    ...     abscissas2, weights2 = chaospy.generate_quadrature(
45    ...         order+1, distribution, rule="kronrod")
46    ...     print(abscissas1.round(2), abscissas2.round(2))
47    [[0.]] [[-1.73  0.    1.73]]
48    [[-1.  1.]] [[-2.45 -1.    0.    1.    2.45]]
49
50Applying Gauss-Kronrod to Gauss-Hermite quadrature, for an order known to not
51exist::
52
53    >>> chaospy.generate_quadrature(  # doctest: +IGNORE_EXCEPTION_DETAIL
54    ...     5, distribution, rule="kronrod")
55    Traceback (most recent call last):
56        ...
57    numpy.linalg.LinAlgError: \
58Invalid recurrence coefficients can not be used for constructing Gaussian quadrature rule
59
60Multivariate support::
61
62    >>> distribution = chaospy.J(
63    ...     chaospy.Uniform(0, 1), chaospy.Beta(4, 5))
64    >>> abscissas, weights = chaospy.generate_quadrature(
65    ...     2, distribution, rule="kronrod")
66    >>> abscissas.round(3)
67    array([[0.037, 0.037, 0.037, 0.037, 0.037, 0.211, 0.211, 0.211, 0.211,
68            0.211, 0.5  , 0.5  , 0.5  , 0.5  , 0.5  , 0.789, 0.789, 0.789,
69            0.789, 0.789, 0.963, 0.963, 0.963, 0.963, 0.963],
70           [0.144, 0.297, 0.444, 0.612, 0.796, 0.144, 0.297, 0.444, 0.612,
71            0.796, 0.144, 0.297, 0.444, 0.612, 0.796, 0.144, 0.297, 0.444,
72            0.612, 0.796, 0.144, 0.297, 0.444, 0.612, 0.796]])
73    >>> weights.round(3)
74    array([0.006, 0.027, 0.035, 0.026, 0.004, 0.016, 0.067, 0.086, 0.065,
75           0.011, 0.02 , 0.085, 0.11 , 0.083, 0.014, 0.016, 0.067, 0.086,
76           0.065, 0.011, 0.006, 0.027, 0.035, 0.026, 0.004])
77"""
78from __future__ import division
79import math
80
81import numpy
82import chaospy
83
84from .utils import combine_quadrature
85
86
87def kronrod(
88        order,
89        dist,
90        recurrence_algorithm="stieltjes",
91        rule="clenshaw_curtis",
92        tolerance=1e-10,
93        scaling=3,
94        n_max=5000,
95):
96    """
97    Generate Gauss-Kronrod quadrature abscissas and weights.
98
99    Gauss-Kronrod quadrature is an adaptive method for Gaussian quadrature
100    rule. It builds on top of other quadrature rules by extending "pure"
101    Gaussian quadrature rules with extra abscissas and new weights such that
102    already used abscissas can be reused. For more details, see `Wikipedia
103    article
104    <https://en.wikipedia.org/wiki/Gauss%E2%80%93Kronrod_quadrature_formula>`_.
105
106    For each order ``N`` taken with ordinary Gaussian quadrature, Gauss-Kronrod
107    will create ``2N+1`` abscissas where all of the ``N`` "old" abscissas are
108    all interlaced between the "new" ones.
109
110    The algorithm is well suited for any Jacobi scheme, i.e. quadrature
111    involving Uniform or Beta distribution, and might work on others as well.
112    However, it will not work everywhere. For example `Kahaner and Monegato
113    <https://link.springer.com/article/10.1007/BF01590820>`_ showed that higher
114    order Gauss-Kronrod quadrature for Gauss-Hermite and Gauss-Laguerre does
115    not exist.
116
117    Args:
118        order (int):
119            The order of the quadrature.
120        dist (chaospy.distributions.baseclass.Distribution):
121            The distribution which density will be used as weight function.
122        recurrence_algorithm (str):
123            Name of the algorithm used to generate abscissas and weights. If
124            omitted, ``analytical`` will be tried first, and ``stieltjes`` used
125            if that fails.
126        rule (str):
127            In the case of ``lanczos`` or ``stieltjes``, defines the
128            proxy-integration scheme.
129        tolerance (float):
130            The allowed relative error in norm between two quadrature orders
131            before method assumes convergence.
132        scaling (float):
133            A multiplier the adaptive order increases with for each step
134            quadrature order is not converged. Use 0 to indicate unit
135            increments.
136        n_max (int):
137            The allowed number of quadrature points to use in approximation.
138
139    Returns:
140        (numpy.ndarray, numpy.ndarray):
141            abscissas:
142                The quadrature points for where to evaluate the model function
143                with ``abscissas.shape == (len(dist), N)`` where ``N`` is the
144                number of samples.
145            weights:
146                The quadrature weights with ``weights.shape == (N,)``.
147
148    Raises:
149        ValueError:
150            Error raised if Kronrod algorithm results in negative recurrence
151            coefficients.
152
153    Notes:
154        Code is adapted from `quadpy <https://github.com/nschloe/quadpy>`_,
155        which adapted his code from `W. Gautschi
156        <https://www.cs.purdue.edu/archives/2002/wxg/codes/OPQ.html>`_.
157        Algorithm for calculating Kronrod-Jacobi matrices was first published
158        Algorithm as proposed by Laurie :cite:`laurie_calculation_1997`.
159
160    Example:
161        >>> distribution = chaospy.Uniform(-1, 1)
162        >>> abscissas, weights = chaospy.quadrature.kronrod(4, distribution)
163        >>> abscissas.round(2)
164        array([[-0.98, -0.86, -0.64, -0.34,  0.  ,  0.34,  0.64,  0.86,  0.98]])
165        >>> weights.round(3)
166        array([0.031, 0.085, 0.133, 0.163, 0.173, 0.163, 0.133, 0.085, 0.031])
167
168    """
169    length = int(numpy.ceil(3*(order+1) / 2.0))
170    coefficients = chaospy.construct_recurrence_coefficients(
171        order=length,
172        dist=dist,
173        recurrence_algorithm=recurrence_algorithm,
174        rule=rule,
175        tolerance=tolerance,
176        scaling=scaling,
177        n_max=n_max,
178    )
179
180    coefficients = [kronrod_jacobi(order, coeffs) for coeffs in coefficients]
181
182    abscissas, weights = chaospy.coefficients_to_quadrature(coefficients)
183    return combine_quadrature(abscissas, weights)
184
185
186def kronrod_jacobi(order, coeffs):
187    """
188    Create the three-terms-recursion coefficients resulting from the
189    Kronrod-Jacobi matrix.
190
191    Augment three terms recurrence coefficients to add extra Gauss-Kronrod
192    terms.
193
194    Args:
195        order (int):
196            Order of the Gaussian quadrature rule.
197        coeffs (numpy.ndarray):
198            Three terms recurrence coefficients of the Gaussian quadrature
199            rule.
200
201    Returns:
202        Three terms recurrence coefficients of the Gauss-Kronrod quadrature
203        rule.
204    """
205    if not order:
206        return kronrod_jacobi(1, coeffs)[:, :1]
207
208    bound = int(math.floor(3*order/2.0))+1
209    coeffs_a = numpy.zeros(2*order+1)
210    coeffs_a[:bound] = coeffs[0][:bound]
211
212    bound = int(math.ceil(3*order/2.0))+1
213    coeffs_b = numpy.zeros(2*order+1)
214    coeffs_b[:bound] = coeffs[1][:bound]
215
216    sigma = numpy.zeros((2, order//2+2))
217    sigma[1, 1] = coeffs_b[order+1]
218
219    for idx in range(order-1):
220        idy = numpy.arange((idx+1)//2, -1, -1)
221        sigma[0, idy+1] = numpy.cumsum(
222            (coeffs_a[idy+order+1]-coeffs_a[idx-idy])*sigma[1, idy+1]+
223            coeffs_b[idy+order+1]*sigma[0, idy]-
224            coeffs_b[idx-idy]*sigma[0, idy+1]
225        )
226        sigma = numpy.roll(sigma, 1, axis=0)
227
228    sigma[0, 1:order//2+2] = sigma[0, :order//2+1]
229    for idx in range(order-1, 2*order-2):
230        idy = numpy.arange(idx-order+1, (idx-1)//2+1)
231        j = order-1-idx+idy
232        sigma[0, j+1] = numpy.cumsum(
233            -(coeffs_a[idy+order+1]-coeffs_a[idx-idy])*sigma[1, j+1]-
234            coeffs_b[idy+order+1]*sigma[0, j+1]+
235            coeffs_b[idx-idy]*sigma[0, j+2]
236        )
237        j = j[-1]
238        idy = (idx+1)//2
239        if idx%2 == 0:
240            coeffs_a[idy+order+1] = (
241                coeffs_a[idy]+
242                (sigma[0, j+1]-coeffs_b[idy+order+1]*sigma[0, j+2])/
243                sigma[1, j+2]
244            )
245        else:
246            coeffs_b[idy+order+1] = sigma[0, j+1]/sigma[0, j+2]
247        sigma = numpy.roll(sigma, 1, axis=0)
248
249    coeffs_a[2*order] = (coeffs_a[order-1]-
250                         coeffs_b[2*order]*sigma[0, 1]/sigma[1, 1])
251    return numpy.asfarray([coeffs_a, coeffs_b])
252