grib2io.utils.gauss_grid

 1"""
 2Tools for working with Gaussian grids.
 3
 4Adopted from: https://gist.github.com/ajdawson/b64d24dfac618b91974f
 5"""
 6
 7from __future__ import absolute_import, division, print_function
 8
 9import functools
10
11import numpy as np
12import numpy.linalg as la
13from numpy.polynomial.legendre import legcompanion, legder, legval
14
15
16def __single_arg_fast_cache(func):
17    """Caching decorator for functions of one argument."""
18
19    class CachingDict(dict):
20        def __missing__(self, key):
21            result = self[key] = func(key)
22            return result
23
24        @functools.wraps(func)
25        def __getitem__(self, *args, **kwargs):
26            return super(CachingDict, self).__getitem__(*args, **kwargs)
27
28    return CachingDict().__getitem__
29
30
31@__single_arg_fast_cache
32def gaussian_latitudes(nlat: int):
33    """
34    Construct latitudes for a Gaussian grid.
35
36    Parameters
37    ----------
38    nlat
39        The number of latitudes in the Gaussian grid.
40
41    Returns
42    -------
43    latitudes
44        `numpy.ndarray` of latitudes (in degrees) with a length of `nlat`.
45    """
46    if abs(int(nlat)) != nlat:
47        raise ValueError("nlat must be a non-negative integer")
48    # Create the coefficients of the Legendre polynomial and construct the
49    # companion matrix:
50    cs = np.array([0] * nlat + [1], dtype=int)
51    cm = legcompanion(cs)
52    # Compute the eigenvalues of the companion matrix (the roots of the
53    # Legendre polynomial) taking advantage of the fact that the matrix is
54    # symmetric:
55    roots = la.eigvalsh(cm)
56    roots.sort()
57    # Improve the roots by one application of Newton's method, using the
58    # solved root as the initial guess:
59    fx = legval(roots, cs)
60    fpx = legval(roots, legder(cs))
61    roots -= fx / fpx
62    # The roots should exhibit symmetry, but with a sign change, so make sure
63    # this is the case:
64    roots = (roots - roots[::-1]) / 2.0
65    # Convert the roots from the interval [-1, 1] to latitude values on the
66    # interval [-90, 90] degrees:
67    latitudes = np.rad2deg(np.arcsin(roots))
68    # Flip latitudes such that it is oriented from North to South [90, -90]
69    latitudes = np.flip(latitudes)
70    return latitudes