kpm.py 21.1 KB
Newer Older
Pablo Piskunow's avatar
Pablo Piskunow committed
1
2
3
4
5
6
7
8
9
# -*- coding: utf-8 -*-
# Copyright 2011-2016 Kwant authors.
#
# This file is part of Kwant.  It is subject to the license terms in the file
# LICENSE.rst found in the top-level directory of this distribution and at
# http://kwant-project.org/license.  A list of Kwant authors can be found in
# the file AUTHORS.rst at the top-level directory of this distribution and at
# http://kwant-project.org/authors.

10
import math
Pablo Piskunow's avatar
Pablo Piskunow committed
11
12
13
14
15
16
import numpy as np
import scipy
import scipy.sparse.linalg as sla
import scipy.fftpack as fft

from . import system
Joseph Weston's avatar
Joseph Weston committed
17
from ._common import ensure_rng
18
from .operator import _LocalOperator
Pablo Piskunow's avatar
Pablo Piskunow committed
19

20
21
22
__all__ = ['SpectralDensity']


Pablo Piskunow's avatar
Pablo Piskunow committed
23
24
25
class SpectralDensity:
    """Calculate the spectral density of an operator.

Pablo Piskunow's avatar
Pablo Piskunow committed
26
    This class makes use of the kernel polynomial
Pablo Piskunow's avatar
Pablo Piskunow committed
27
28
29
30
31
32
33
34
    method (KPM), presented in [1]_, to obtain the spectral density
    :math:`ρ_A(e)`, as a function of the energy :math:`e`, of some
    operator :math:`A` that acts on a kwant system or a Hamiltonian.
    In general

    .. math::
       ρ_A(e) = ρ(e) A(e),

35
36
37
    where :math:`ρ(e) = \sum_{k=0}^{D-1} δ(E-E_k)` is the density of
    states, and :math:`A(e)` is the expectation value of :math:`A` for
    all the eigenstates with energy :math:`e`.
Pablo Piskunow's avatar
Pablo Piskunow committed
38
39
40

    Parameters
    ----------
41
    hamiltonian : `~kwant.system.FiniteSystem` or matrix Hamiltonian
Pablo Piskunow's avatar
Pablo Piskunow committed
42
        If a system is passed, it should contain no leads.
43
44
    params : dict, optional
        Additional parameters to pass to the Hamiltonian and operator.
Pablo Piskunow's avatar
Pablo Piskunow committed
45
46
47
48
    operator : operator, dense matrix, or sparse matrix, optional
        Operator for which the spectral density will be evaluated. If
        it is callable, the ``densities`` at each energy will have the
        dimension of the result of `operator(bra, ket)`. If it has a
Pablo Piskunow's avatar
Pablo Piskunow committed
49
50
        ``dot`` method, such as ``numpy.ndarray`` and
        ``scipy.sparse.matrices``, the densities will be scalars.
51
52
        If no operator is provided, the density of states is calculated
        with a faster algorithm.
53
    num_vectors : positive int, default: 10
Pablo Piskunow's avatar
Pablo Piskunow committed
54
        Number of random vectors for the KPM method.
55
56
57
58
59
60
    num_moments : positive int, default: 100
        Number of moments, order of the KPM expansion. Mutually exclusive
        with 'energy_resolution'.
    energy_resolution : positive float, optional
        The resolution in energy of the KPM approximation to the spectral
        density. Mutually exclusive with 'num_moments'.
Pablo Piskunow's avatar
Pablo Piskunow committed
61
62
63
64
65
66
    vector_factory : function, optional
        The user defined function ``f(n)`` generates random vectors of
        length ``n`` that will be used in the algorithm.
        If not provided, random phase vectors are used.
        The default random vectors are optimal for most cases, see the
        discussions in [1]_ and [2]_.
67
68
69
    bounds : pair of floats, optional
        Lower and upper bounds for the eigenvalue spectrum of the system.
        If not provided, they are computed.
70
    eps : positive float, default: 0.05
71
        Parameter to ensure that the rescaled spectrum lies in the
Pablo Piskunow's avatar
Pablo Piskunow committed
72
        interval ``(-1, 1)``; required for stability.
Pablo Piskunow's avatar
Pablo Piskunow committed
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
    rng : seed, or random number generator, optional
        Random number generator used by ``vector_factory``.
        If not provided, numpy's rng will be used; if it is an Integer,
        it will be used to seed numpy's rng, and if it is a random
        number generator, this is the one used.

    Notes
    -----
    When passing an ``operator`` defined in `~kwant.operator`, the
    result of ``operator(bra, ket)`` depends on the attribute ``sum``
    of such operator. If ``sum=True``, densities will be scalars, that
    is, total densities of the system. If ``sum=False`` the densities
    will be arrays of the length of the system, that is, local
    densities.

    .. [1] `Rev. Mod. Phys., Vol. 78, No. 1 (2006)
       <https://arxiv.org/abs/cond-mat/0504627>`_.
    .. [2] `Phys. Rev. E 69, 057701 (2004)
       <https://arxiv.org/abs/cond-mat/0401202>`_

    Examples
    --------
    In the following example, we will obtain the density of states of a
    graphene sheet, defined as a honeycomb lattice with first nearest
    neighbors coupling.

    We start by importing kwant and defining a
    `~kwant.system.FiniteSystem`,

    >>> import kwant
    ...
    >>> def circle(pos):
    ...     x, y = pos
    ...     return x**2 + y**2 < 100
    ...
    >>> lat = kwant.lattice.honeycomb()
    >>> syst = kwant.Builder()
    >>> syst[lat.shape(circle, (0, 0))] = 0
    >>> syst[lat.neighbors()] = 1

    and after finalizing the system, create an instance of
Viacheslav Ostroukh's avatar
Viacheslav Ostroukh committed
114
    `~kwant.kpm.SpectralDensity`
Pablo Piskunow's avatar
Pablo Piskunow committed
115
116
117
118

    >>> fsyst = syst.finalized()
    >>> rho = kwant.kpm.SpectralDensity(fsyst)

Pablo Piskunow's avatar
Pablo Piskunow committed
119
    The ``energies`` and ``densities`` can be accessed with
Pablo Piskunow's avatar
Pablo Piskunow committed
120
121
122
123
124
125
126
127
128
129

    >>> energies, densities = rho()

    or

    >>> energies, densities = rho.energies, rho.densities

    Attributes
    ----------
    energies : array of floats
130
        Array of sampling points with length ``2 * num_moments`` in
Pablo Piskunow's avatar
Pablo Piskunow committed
131
132
133
134
135
        the range of the spectrum.
    densities : array of floats
        Spectral density of the ``operator`` evaluated at the energies.
    """

136
    def __init__(self, hamiltonian, params=None, operator=None,
137
                 num_vectors=10, num_moments=None, energy_resolution=None,
138
                 vector_factory=None, bounds=None, eps=0.05, rng=None):
139
140
141
142
143

        if num_moments and energy_resolution:
            raise TypeError("either 'num_moments' or 'energy_resolution' "
                            "must be provided.")

Pablo Piskunow's avatar
Pablo Piskunow committed
144
        rng = ensure_rng(rng)
145
        # self.eps ensures that the rescaled Hamiltonian has a
Pablo Piskunow's avatar
Pablo Piskunow committed
146
        # spectrum strictly in the interval (-1,1).
147
        self.eps = eps
148
149

        # Normalize the format of 'ham'
150
151
152
        if isinstance(hamiltonian, system.System):
            hamiltonian = hamiltonian.hamiltonian_submatrix(params=params,
                                                            sparse=True)
Pablo Piskunow's avatar
Pablo Piskunow committed
153
        try:
154
            hamiltonian = scipy.sparse.csr_matrix(hamiltonian)
155
        except Exception:
156
157
            raise ValueError("'hamiltonian' is neither a matrix "
                             "nor a Kwant system.")
158
159

        # Normalize 'operator' to a common format.
Pablo Piskunow's avatar
Pablo Piskunow committed
160
161
        if operator is None:
            self.operator = None
162
163
164
165
166
167
168
        elif isinstance(operator, _LocalOperator):
            self.operator = operator.bind(params=params)
        elif callable(operator):
            self.operator = operator
        elif hasattr(operator, 'dot'):
            operator = scipy.sparse.csr_matrix(operator)
            self.operator = lambda bra, ket: np.vdot(bra, operator.dot(ket))
Pablo Piskunow's avatar
Pablo Piskunow committed
169
        else:
170
171
172
            raise ValueError('Parameter `operator` has no `.dot` '
                             'attribute and is not callable.')

Pablo Piskunow's avatar
Pablo Piskunow committed
173
174
175
        self._vector_factory = vector_factory or \
            (lambda n: np.exp(2j * np.pi * rng.random_sample(n)))
        # store this vector for reproducibility
176
        self._v0 = np.exp(2j * np.pi * rng.random_sample(hamiltonian.shape[0]))
Pablo Piskunow's avatar
Pablo Piskunow committed
177
        self._rand_vect_list = []
178
        # Hamiltonian rescaled as in Eq. (24)
179
180
181
182
        self.hamiltonian, (self._a, self._b) = _rescale(hamiltonian,
                                                        eps=self.eps,
                                                        v0=self._v0,
                                                        bounds=bounds)
183
        self.bounds = (self._b - self._a, self._b + self._a)
184

185
186
187
188
189
190
191
192
193
194
195
196
197
        if energy_resolution:
            num_moments = math.ceil((1.6 * self._a) / energy_resolution)
        elif num_moments is None:
            num_moments = 100

        must_be_positive_int = ['num_vectors', 'num_moments']
        for var in must_be_positive_int:
            val = locals()[var]
            if val <= 0 or val != int(val):
                raise ValueError('{} must be a positive integer'.format(var))
        if eps <= 0:
            raise ValueError('eps must be positive')

198
199
200
201
202
        for r in range(num_vectors):
            self._rand_vect_list.append(
                self._vector_factory(self.hamiltonian.shape[0]))
        self._last_two_alphas = [0.] * num_vectors
        self._moments_list = [0.] * num_vectors
203

204
205
        self.num_moments = num_moments
        self.num_vectors = 0  # new random vectors will be used
206
207
        self._update_moments_list(self.num_moments, num_vectors)
        self.num_vectors = num_vectors
208

209
        moments = self._moments()
Pablo Piskunow's avatar
Pablo Piskunow committed
210
        xk_rescaled, rho, self._gammas = _calc_fft_moments(
211
212
            moments, 2 * self.num_moments)
        self.energies = xk_rescaled * self._a + self._b
Pablo Piskunow's avatar
Pablo Piskunow committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
        self.densities = rho

    def __call__(self, energy=None):
        """Return the spectral density evaluated at ``energy``.

        Parameters
        ----------
        energy : float or sequence of float, optional

        Returns
        -------
        float, if ``energy`` is float, array of float if ``energy``
        is a sequence, a tuple (energies, densities) if
        ``energy`` was not provided.

        Notes
        -----
230
231
        If ``energy`` is not provided, then the densities are obtained
        by Fast Fourier Transform of the Chebyshev moments.
Pablo Piskunow's avatar
Pablo Piskunow committed
232
233
234
235
236
237
        """
        if energy is None:
            return self.energies, self.densities
        else:
            energy = np.asarray(energy, dtype=complex)
            rescaled_energy = (energy - self._b) / self._a
238
239
            g_e = (np.pi * np.sqrt(1 - rescaled_energy)
                   * np.sqrt(1 + rescaled_energy))
Pablo Piskunow's avatar
Pablo Piskunow committed
240

241
            moments = self._moments()
Pablo Piskunow's avatar
Pablo Piskunow committed
242
243
244
245
246
247
248
249
250

            m = np.arange(self.num_moments)

            kernel = ((self.num_moments - m + 1) *
                      np.cos(np.pi * m/(self.num_moments + 1)) +
                      np.sin(np.pi * m/(self.num_moments + 1)) /
                      np.tan(np.pi/(self.num_moments + 1)))
            kernel = kernel / (self.num_moments + 1)

251
252
253
            # transposes handle the case where operators have vector outputs
            coef_cheb = np.transpose(moments.transpose() * kernel)
            coef_cheb[1:] = 2 * coef_cheb[1:]
Pablo Piskunow's avatar
Pablo Piskunow committed
254

255
            return np.transpose(np.polynomial.chebyshev.chebval(
Pablo Piskunow's avatar
Pablo Piskunow committed
256
257
                    rescaled_energy, coef_cheb) / g_e).real

258
    def integrate(self, distribution_function=None):
Pablo Piskunow's avatar
Pablo Piskunow committed
259
260
261
        """Returns the total spectral density.

        Returns the integral over the whole spectrum with an optional
262
263
        distribution function. ``distribution_function`` should be able
        to take arrays as input. Defined using Gauss-Chebyshev
Pablo Piskunow's avatar
Pablo Piskunow committed
264
265
266
        integration.
        """
        # This factor divides the sum to normalize the Gauss integral
267
268
        # and rescales the integral back with ``self._a`` to normal
        # scale.
269
        factor = self._a / (2 * self.num_moments)
Pablo Piskunow's avatar
Pablo Piskunow committed
270
271
272
273
274
275
276
277
278
        if distribution_function is None:
            rho = self._gammas
        else:
            # The evaluation of the distribution function should be at
            # the energies without rescaling.
            distribution_array = distribution_function(self.energies)
            rho = np.transpose(self._gammas.transpose() * distribution_array)
        return factor * np.sum(rho, axis=0)

279
280
    def add_moments(self, num_moments=None, *, energy_resolution=None):
        """Increase the number of Chebyshev moments.
Pablo Piskunow's avatar
Pablo Piskunow committed
281

282
283
284
285
286
287
288
289
290
        Parameters
        ----------
        num_moments: positive int
            The number of Chebyshev moments to add. Mutually
            exclusive with 'energy_resolution'.
        energy_resolution: positive float, optional
            Features wider than this resolution are visible
            in the spectral density. Mutually exclusive with
            'num_moments'.
Pablo Piskunow's avatar
Pablo Piskunow committed
291
        """
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
        if not ((num_moments is None) ^ (energy_resolution is None)):
            raise TypeError("either 'num_moments' or 'energy_resolution' "
                            "must be provided.")

        if energy_resolution:
            if energy_resolution <= 0:
                raise ValueError("'energy_resolution' must be positive"
                                 .format(energy_resolution))
            # factor of 1.6 comes from the fact that we use the
            # Jackson kernel when calculating the FFT, which has
            # maximal slope π/2. Rounding to 1.6 ensures that the
            # energy resolution is sufficient.
            present_resolution = self._a * 1.6 / self.num_moments
            if present_resolution < energy_resolution:
                raise ValueError('Energy resolution is already smaller '
                                 'than the requested resolution')
            num_moments = math.ceil((1.6 * self._a) / energy_resolution)

        if (num_moments is None or num_moments <= 0
            or num_moments != int(num_moments)):
            raise ValueError("'num_moments' must be a positive integer")

        self._update_moments_list(self.num_moments + num_moments,
                                  self.num_vectors)
        self.num_moments += num_moments

        # recalculate quantities derived from the moments
        moments = self._moments()
        xk_rescaled, rho, self._gammas = _calc_fft_moments(
            moments, 2 * self.num_moments)
        self.energies = xk_rescaled * self._a + self._b
        self.densities = rho

    def add_vectors(self, num_vectors):
        """Increase the number of random vectors.
Pablo Piskunow's avatar
Pablo Piskunow committed
327
328
329

        Parameters
        ----------
330
331
        num_vectors: positive int
            The number of random vectors to add.
Pablo Piskunow's avatar
Pablo Piskunow committed
332
        """
333
334
335
        if num_vectors <= 0 or num_vectors != int(num_vectors):
            raise ValueError("'num_vectors' must be a positive integer")
        for r in range(num_vectors):
336
337
            self._rand_vect_list.append(
                self._vector_factory(self.hamiltonian.shape[0]))
338
339
340
341
342
343
344
345
346
347
348
349
        self._moments_list.extend([0.] * num_vectors)
        self._last_two_alphas.extend([0.] * num_vectors)
        self._update_moments_list(self.num_moments,
                                  self.num_vectors + num_vectors)
        self.num_vectors += num_vectors

        # recalculate quantities derived from the moments
        moments = self._moments()
        xk_rescaled, rho, self._gammas = _calc_fft_moments(
            moments, 2 * self.num_moments)
        self.energies = xk_rescaled * self._a + self._b
        self.densities = rho
Pablo Piskunow's avatar
Pablo Piskunow committed
350

351
    def _moments(self):
Pablo Piskunow's avatar
Pablo Piskunow committed
352
        # sum moments of all random vectors
353
        moments = np.sum(np.asarray(self._moments_list).real, axis=0)
Pablo Piskunow's avatar
Pablo Piskunow committed
354
        # divide by the number of random vectors
355
        moments /= self.num_vectors
Pablo Piskunow's avatar
Pablo Piskunow committed
356
        # divide by scale factor to reflect the integral rescaling
357
        return moments / self._a
Pablo Piskunow's avatar
Pablo Piskunow committed
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374

    def _update_moments_list(self, n_moments, n_rand):
        """Calculate the Chebyshev moments of an operator's spectral
        density.

        The algorithm is based on the KPM method as depicted in `Rev.
        Mod. Phys., Vol. 78, No. 1 (2006)
        <https://arxiv.org/abs/cond-mat/0504627>`_.

        Parameters
        ----------
        n_moments : integer
            Number of Chebyshev moments.
        n_rand : integer
            Number of random vectors used for sampling.
        """

375
        if self.num_vectors == n_rand:
Pablo Piskunow's avatar
Pablo Piskunow committed
376
            r_start = 0
377
            new_rand_vect = 0
378
379
380
        elif self.num_vectors < n_rand:
            r_start = self.num_vectors
            new_rand_vect = n_rand - self.num_vectors
Pablo Piskunow's avatar
Pablo Piskunow committed
381
        else:
382
            raise ValueError('Cannot decrease number of random vectors')
Pablo Piskunow's avatar
Pablo Piskunow committed
383
384

        if n_moments == self.num_moments:
385
            m_start = 2
Pablo Piskunow's avatar
Pablo Piskunow committed
386
            new_moments = 0
387
388
389
            if new_rand_vect == 0:
                # nothing new to calculate
                return
Pablo Piskunow's avatar
Pablo Piskunow committed
390
391
        else:
            new_moments = n_moments - self.num_moments
392
393
            m_start = self.num_moments
            if new_moments < 0:
394
                raise ValueError('Cannot decrease number of moments')
Pablo Piskunow's avatar
Pablo Piskunow committed
395

396
397
398
            if new_rand_vect != 0:
                raise ValueError("Only 'num_moments' *or* 'num_vectors' "
                                 "may be updated at a time.")
399
400

        for r in range(r_start, n_rand):
Pablo Piskunow's avatar
Pablo Piskunow committed
401
402
403
            alpha_zero = self._rand_vect_list[r]

            one_moment = [0.] * n_moments
404
            if new_rand_vect > 0:
Pablo Piskunow's avatar
Pablo Piskunow committed
405
                alpha = alpha_zero
406
                alpha_next = self.hamiltonian.matvec(alpha)
Pablo Piskunow's avatar
Pablo Piskunow committed
407
408
409
410
411
412
413
                if self.operator is None:
                    one_moment[0] = np.vdot(alpha_zero, alpha_zero)
                    one_moment[1] = np.vdot(alpha_zero, alpha_next)
                else:
                    one_moment[0] = self.operator(alpha_zero, alpha_zero)
                    one_moment[1] = self.operator(alpha_zero, alpha_next)

414
            if new_moments > 0:
Pablo Piskunow's avatar
Pablo Piskunow committed
415
416
417
418
419
420
421
422
423
424
425
426
427
                (alpha, alpha_next) = self._last_two_alphas[r]
                one_moment[0:self.num_moments] = self._moments_list[r]
            # Iteration over the moments
            # Two cases can occur, depicted in Eq. (28) and in Eq. (29),
            # respectively.
            # ----
            # In the first case, self.operator is None and we can use
            # Eqs. (34) and (35) to obtain the density of states, with
            # two moments ``one_moment`` for every new alpha.
            # ----
            # In the second case, the operator is not None and a matrix
            # multiplication should be used.
            if self.operator is None:
428
                for n in range(m_start//2, n_moments//2):
Pablo Piskunow's avatar
Pablo Piskunow committed
429
                    alpha_save = alpha_next
430
431
                    alpha_next = (2 * self.hamiltonian.matvec(alpha_next)
                                  - alpha)
Pablo Piskunow's avatar
Pablo Piskunow committed
432
433
                    alpha = alpha_save
                    # Following Eqs. (34) and (35)
434
435
436
437
438
439
                    one_moment[2*n] = (2 * np.vdot(alpha, alpha)
                                       - one_moment[0])
                    one_moment[2*n+1] = (2 * np.vdot(alpha_next, alpha)
                                         - one_moment[1])
                if n_moments % 2:
                    # odd moment
440
441
                    one_moment[n_moments - 1] = (
                        2 * np.vdot(alpha_next, alpha_next) - one_moment[0])
Pablo Piskunow's avatar
Pablo Piskunow committed
442
            else:
443
                for n in range(m_start, n_moments):
Pablo Piskunow's avatar
Pablo Piskunow committed
444
                    alpha_save = alpha_next
445
446
                    alpha_next = (2 * self.hamiltonian.matvec(alpha_next)
                                  - alpha)
Pablo Piskunow's avatar
Pablo Piskunow committed
447
448
449
                    alpha = alpha_save
                    one_moment[n] = self.operator(alpha_zero, alpha_next)

450
451
            self._last_two_alphas[r] = (alpha, alpha_next)
            self._moments_list[r] = one_moment[:]
Pablo Piskunow's avatar
Pablo Piskunow committed
452
453
454
455
456


# ### Auxiliary functions


457
def _rescale(hamiltonian, eps, v0, bounds):
Pablo Piskunow's avatar
Pablo Piskunow committed
458
459
460
461
    """Rescale a Hamiltonian and return a LinearOperator

    Parameters
    ----------
462
    hamiltonian : 2D array
Pablo Piskunow's avatar
Pablo Piskunow committed
463
        Hamiltonian of the system.
464
    eps : scalar
Pablo Piskunow's avatar
Pablo Piskunow committed
465
        Ensures that the bounds 'a' and 'b' are strict.
466
    v0 : random vector, or None
Pablo Piskunow's avatar
Pablo Piskunow committed
467
468
        Used as the initial residual vector for the algorithm that
        finds the lowest and highest eigenvalues.
469
    bounds : tuple, or None
Pablo Piskunow's avatar
Pablo Piskunow committed
470
471
472
        Boundaries of the spectrum. If not provided the maximum and
        minimum eigenvalues are calculated.
    """
473
    # Relative tolerance to which to calculate eigenvalues.  Because after
474
475
476
    # rescaling we will add eps / 2 to the spectral bounds, we don't need
    # to know the bounds more accurately than eps / 2.
    tol = eps / 2
Pablo Piskunow's avatar
Pablo Piskunow committed
477
478

    if bounds:
479
        lmin, lmax = bounds
Pablo Piskunow's avatar
Pablo Piskunow committed
480
    else:
481
482
483
484
        lmax = float(sla.eigsh(hamiltonian, k=1, which='LA',
                               return_eigenvectors=False, tol=tol, v0=v0))
        lmin = float(sla.eigsh(hamiltonian, k=1, which='SA',
                               return_eigenvectors=False, tol=tol, v0=v0))
485

486
    a = np.abs(lmax-lmin) / (2. - eps)
Pablo Piskunow's avatar
Pablo Piskunow committed
487
488
    b = (lmax+lmin) / 2.

489
    if lmax - lmin <= abs(lmax + lmin) * tol / 2:
Pablo Piskunow's avatar
Pablo Piskunow committed
490
491
492
493
494
        raise ValueError(
            'The Hamiltonian has a single eigenvalue, it is not possible to '
            'obtain a spectral density.')

    def rescaled(v):
495
        return (hamiltonian.dot(v) - b * v) / a
Pablo Piskunow's avatar
Pablo Piskunow committed
496

497
    rescaled_ham = sla.LinearOperator(shape=hamiltonian.shape, matvec=rescaled)
498
499

    return rescaled_ham, (a, b)
Pablo Piskunow's avatar
Pablo Piskunow committed
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534


def _calc_fft_moments(moments, n_sampling):
    """This function takes the normalised moments and returns an array
    of points and an array of the evaluated function at those points.
    """
    moments = np.asarray(moments)
    # extra_shape handles the case where operators have vector outputs
    n_moments, *extra_shape = moments.shape
    moments_ext = np.zeros([n_sampling] + extra_shape)

    # Jackson kernel, as in Eq. (71), and kernel improved moments,
    # as in Eq. (81).
    m = np.arange(n_moments)
    kernel = ((n_moments - m + 1) * np.cos(np.pi * m / (n_moments + 1)) +
              np.sin(np.pi * m / (n_moments + 1)) /
              np.tan(np.pi / (n_moments + 1))) / (n_moments + 1)

    # special points at the abscissas of Chebyshev integration
    k = np.arange(0, n_sampling)
    xk_rescaled = np.cos(np.pi * (k + 0.5) / n_sampling)
    # prefactor in Eq. (80)
    gk = np.pi * np.sqrt(1 - xk_rescaled ** 2)

    # transposes handle the case where operators have vector outputs
    moments_ext[:n_moments] = np.transpose(moments.transpose() * kernel)
    # The function evaluated in these special data points is the FFT of
    # the moments as in Eq. (83).
    gammas = np.transpose(fft.dct(moments_ext.transpose(), type=3))

    # Element-wise division of moments_FFT over gk, as in Eq. (83).
    rho = np.transpose(np.divide(gammas.transpose(), gk))

    # Reverse energies and densities to set ascending order.
    return xk_rescaled[::-1], rho[::-1], gammas[::-1]