spharmgrid 0.1.0__py3-none-any.whl

spharmgrid/__init__.py

@@ -0,0 +1,63 @@
+"""xarray-first spherical-harmonic atmospheric operations using DUCC0.
+
+Importing :mod:`spharmgrid` registers the ``.sg`` accessors on xarray
+``DataArray`` and ``Dataset`` objects.  DUCC0 supplies all numerical
+spherical-harmonic transforms; spharmgrid supplies the GL/CC, xarray, and CF
+operations layer around them.
+"""
+
+# Import solely for xarray accessor registration after direct functions exist.
+from . import accessors as _accessors  # noqa: F401
+from .grids import Grid, clenshaw_curtis_grid, detect_grid, gaussian_grid
+from .kinematics import (
+    divergence,
+    divergent_wind,
+    helmholtz,
+    inverse_vector_laplacian,
+    kinematics,
+    potentials,
+    rotational_wind,
+    streamfunction,
+    vector_laplacian,
+    velocity_potential,
+    vorticity,
+    wind,
+)
+from .operators import (
+    EARTH_RADIUS_M,
+    gradient,
+    inverse_gradient,
+    inverse_laplacian,
+    laplacian,
+)
+from .regrid import regrid, regrid_vector
+from .spectral import SpectralRange, filter, parse_spectral
+
+__all__ = [
+    "EARTH_RADIUS_M",
+    "Grid",
+    "SpectralRange",
+    "clenshaw_curtis_grid",
+    "detect_grid",
+    "divergent_wind",
+    "divergence",
+    "filter",
+    "gaussian_grid",
+    "gradient",
+    "helmholtz",
+    "inverse_gradient",
+    "inverse_laplacian",
+    "inverse_vector_laplacian",
+    "kinematics",
+    "laplacian",
+    "parse_spectral",
+    "potentials",
+    "regrid",
+    "regrid_vector",
+    "rotational_wind",
+    "streamfunction",
+    "velocity_potential",
+    "vector_laplacian",
+    "vorticity",
+    "wind",
+]

spharmgrid/_ducc.py

@@ -0,0 +1,202 @@
+"""Small internal adapters around DUCC0's spherical-harmonic functions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import cache
+from typing import Literal, cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+from .grids import Grid
+
+Geometry = Literal["CC", "GL"]
+
+DUCC_THREADS = 4
+
+
+@dataclass(frozen=True, slots=True)
+class TransformSpec:
+    """A scalar/vector transform bandwidth supported by its source geometry."""
+
+    lmax: int
+    mmax: int
+
+    def __post_init__(self) -> None:
+        if self.lmax < 0 or self.mmax < 0 or self.mmax > self.lmax:
+            raise ValueError("transform limits must satisfy 0 <= mmax <= lmax")
+
+
+def geometry_for(grid: Grid) -> Geometry:
+    """Map the public lowercase grid label to DUCC's geometry label."""
+    return "GL" if grid.kind == "gl" else "CC"
+
+
+@cache
+def alm_degrees(lmax: int, mmax: int) -> NDArray[np.int64]:
+    """Return total degree for each contiguous healpy-ordered coefficient."""
+    degrees = np.concatenate(
+        [np.arange(m, lmax + 1, dtype=np.int64) for m in range(mmax + 1)]
+    )
+    degrees.setflags(write=False)
+    return degrees
+
+
+def scalar_analysis(
+    frame: NDArray[np.generic],
+    *,
+    spec: TransformSpec,
+    geometry: Geometry,
+    phi0: float,
+) -> NDArray[np.complex128]:
+    """Analyze one north-to-south, cyclic-eastward scalar map."""
+    _require_frame(frame)
+    import ducc0
+
+    values = np.ascontiguousarray(frame, dtype=np.float64)
+    result = ducc0.sht.analysis_2d(
+        map=values[np.newaxis, :, :],
+        spin=0,
+        lmax=spec.lmax,
+        mmax=spec.mmax,
+        geometry=geometry,
+        phi0=phi0,
+        nthreads=DUCC_THREADS,
+    )
+    return cast(NDArray[np.complex128], result)
+
+
+def scalar_synthesis(
+    alm: NDArray[np.complexfloating],
+    *,
+    spec: TransformSpec,
+    geometry: Geometry,
+    ntheta: int,
+    nphi: int,
+    phi0: float,
+) -> NDArray[np.float64]:
+    """Synthesize one scalar map on a north-to-south DUCC geometry."""
+    import ducc0
+
+    values = np.ascontiguousarray(alm)
+    result = ducc0.sht.synthesis_2d(
+        alm=values,
+        spin=0,
+        lmax=spec.lmax,
+        mmax=spec.mmax,
+        geometry=geometry,
+        ntheta=ntheta,
+        nphi=nphi,
+        phi0=phi0,
+        nthreads=DUCC_THREADS,
+    )
+    return cast(NDArray[np.float64], result[0])
+
+
+def scalar_derivative_synthesis(
+    alm: NDArray[np.complexfloating],
+    *,
+    spec: TransformSpec,
+    geometry: Geometry,
+    ntheta: int,
+    nphi: int,
+    phi0: float,
+) -> NDArray[np.float64]:
+    """Synthesize theta and eastward angular derivatives of a scalar field.
+
+    DUCC's ``DERIV1`` mode returns ``(d/dtheta, (1/sin(theta)) d/dphi)``.
+    The caller maps the first component from southward theta to northward
+    latitude and applies the physical-radius factor.
+    """
+    import ducc0
+
+    result = ducc0.sht.synthesis_2d(
+        alm=np.ascontiguousarray(alm),
+        spin=1,
+        lmax=spec.lmax,
+        mmax=spec.mmax,
+        geometry=geometry,
+        ntheta=ntheta,
+        nphi=nphi,
+        phi0=phi0,
+        nthreads=DUCC_THREADS,
+        mode="DERIV1",
+    )
+    return cast(NDArray[np.float64], result)
+
+
+def vector_analysis(
+    u: NDArray[np.generic],
+    v: NDArray[np.generic],
+    *,
+    spec: TransformSpec,
+    geometry: Geometry,
+    phi0: float,
+) -> NDArray[np.complex128]:
+    """Analyze geographic eastward/northward wind into DUCC E/B coefficients.
+
+    DUCC's spin-1 map components are ``(v_theta, v_phi)``.  Geographic wind
+    uses northward ``v``, while increasing theta points south, so spharmgrid
+    maps the components as ``(-v, u)``.
+    """
+    _require_frame(u)
+    _require_frame(v)
+    if u.shape != v.shape:
+        raise ValueError(
+            f"wind frames have incompatible shapes: {u.shape} and {v.shape}"
+        )
+    import ducc0
+
+    eastward = np.asarray(u, dtype=np.float64)
+    northward = np.asarray(v, dtype=np.float64)
+    vector_map = np.stack((-northward, eastward), axis=0)
+    result = ducc0.sht.analysis_2d(
+        map=np.ascontiguousarray(vector_map),
+        spin=1,
+        lmax=spec.lmax,
+        mmax=spec.mmax,
+        geometry=geometry,
+        phi0=phi0,
+        nthreads=DUCC_THREADS,
+    )
+    return cast(NDArray[np.complex128], result)
+
+
+def vector_synthesis(
+    alm_e: NDArray[np.complexfloating],
+    alm_b: NDArray[np.complexfloating],
+    *,
+    spec: TransformSpec,
+    geometry: Geometry,
+    ntheta: int,
+    nphi: int,
+    phi0: float,
+) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+    """Synthesize DUCC E/B coefficients into geographic eastward/northward wind."""
+    import ducc0
+
+    coefficients = np.stack((alm_e, alm_b), axis=0)
+    result = ducc0.sht.synthesis_2d(
+        alm=np.ascontiguousarray(coefficients),
+        spin=1,
+        lmax=spec.lmax,
+        mmax=spec.mmax,
+        geometry=geometry,
+        ntheta=ntheta,
+        nphi=nphi,
+        phi0=phi0,
+        nthreads=DUCC_THREADS,
+    )
+    # This is the inverse of ``(-v, u)`` in ``vector_analysis``.
+    return (
+        cast(NDArray[np.float64], result[1]),
+        cast(NDArray[np.float64], -result[0]),
+    )
+
+
+def _require_frame(frame: NDArray[np.generic]) -> None:
+    if frame.ndim != 2:
+        raise ValueError(
+            f"DUCC kernel requires a two-dimensional frame, got {frame.shape}"
+        )

spharmgrid/_vector.py

@@ -0,0 +1,248 @@
+"""Private xarray adapters for paired tangent-vector SHT operations."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import cast
+
+import numpy as np
+import xarray as xr
+from numpy.typing import NDArray
+
+from ._xarray import (
+    FieldLayout,
+    apply_ufunc_options,
+    exact_align,
+    field_layout,
+    require_dataarray,
+    restore_output,
+)
+from .grids import grids_equivalent
+
+VectorScalarKernel = Callable[
+    [NDArray[np.generic], NDArray[np.generic]], NDArray[np.float64]
+]
+VectorPairKernel = Callable[
+    [NDArray[np.generic], NDArray[np.generic]],
+    tuple[NDArray[np.float64], NDArray[np.float64]],
+]
+VectorQuadKernel = Callable[
+    [NDArray[np.generic], NDArray[np.generic]],
+    tuple[
+        NDArray[np.float64],
+        NDArray[np.float64],
+        NDArray[np.float64],
+        NDArray[np.float64],
+    ],
+]
+
+
+def vector_inputs(
+    u: xr.DataArray, v: xr.DataArray
+) -> tuple[FieldLayout, xr.DataArray, xr.DataArray]:
+    """Validate, canonicalize, and exactly align a geographic vector pair."""
+    u = require_dataarray(u, "u")
+    v = require_dataarray(v, "v")
+    layout = field_layout(u)
+    v_layout = field_layout(v)
+    _require_matching_layouts(layout, v_layout, "u", "v")
+    canonical_u = layout.canonicalize(u)
+    canonical_v = v_layout.canonicalize(v)
+    return layout, *exact_align(canonical_u, canonical_v, names=("u", "v"))
+
+
+def vector_scalar_transform(
+    original: xr.DataArray,
+    source: FieldLayout,
+    canonical_u: xr.DataArray,
+    canonical_v: xr.DataArray,
+    transform: VectorScalarKernel,
+) -> xr.DataArray:
+    """Apply one paired-vector kernel that returns a scalar map."""
+    dask_field = _dask_input(canonical_u, canonical_v)
+
+    def kernel(
+        frame_u: NDArray[np.generic], frame_v: NDArray[np.generic]
+    ) -> NDArray[np.float64]:
+        return transform(frame_u, frame_v)
+
+    output = xr.apply_ufunc(
+        kernel,
+        canonical_u,
+        canonical_v,
+        input_core_dims=[
+            [source.latitude_dim, source.longitude_dim],
+            [source.latitude_dim, source.longitude_dim],
+        ],
+        output_core_dims=[[source.latitude_dim, source.longitude_dim]],
+        vectorize=True,
+        output_dtypes=[np.float64],
+        **apply_ufunc_options(dask_field),
+    )
+    return restore_output(
+        cast(xr.DataArray, output),
+        source=source,
+        target=source,
+        original_dims=original.dims,
+    )
+
+
+def vector_pair_transform(
+    original: xr.DataArray,
+    source: FieldLayout,
+    target: FieldLayout,
+    canonical_u: xr.DataArray,
+    canonical_v: xr.DataArray,
+    transform: VectorPairKernel,
+) -> tuple[xr.DataArray, xr.DataArray]:
+    """Apply one paired-vector kernel that returns two target-grid maps."""
+    dask_field = _dask_input(canonical_u, canonical_v)
+    output_changes_shape = _output_changes_shape(source, target)
+    output_sizes = (
+        {target.latitude_dim: target.grid.nlat, target.longitude_dim: target.grid.nlon}
+        if output_changes_shape
+        else None
+    )
+
+    def kernel(
+        frame_u: NDArray[np.generic], frame_v: NDArray[np.generic]
+    ) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
+        return transform(frame_u, frame_v)
+
+    outputs = xr.apply_ufunc(
+        kernel,
+        canonical_u,
+        canonical_v,
+        input_core_dims=[
+            [source.latitude_dim, source.longitude_dim],
+            [source.latitude_dim, source.longitude_dim],
+        ],
+        output_core_dims=[
+            [target.latitude_dim, target.longitude_dim],
+            [target.latitude_dim, target.longitude_dim],
+        ],
+        vectorize=True,
+        output_dtypes=[np.float64, np.float64],
+        exclude_dims=(
+            {source.latitude_dim, source.longitude_dim}
+            if output_changes_shape
+            else set()
+        ),
+        **apply_ufunc_options(dask_field, output_sizes=output_sizes),
+    )
+    first, second = cast(tuple[xr.DataArray, xr.DataArray], outputs)
+    return (
+        restore_output(
+            first,
+            source=source,
+            target=target,
+            original_dims=original.dims,
+        ),
+        restore_output(
+            second,
+            source=source,
+            target=target,
+            original_dims=original.dims,
+        ),
+    )
+
+
+def vector_quad_transform(
+    original: xr.DataArray,
+    source: FieldLayout,
+    canonical_u: xr.DataArray,
+    canonical_v: xr.DataArray,
+    transform: VectorQuadKernel,
+) -> tuple[xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray]:
+    """Apply one paired-vector kernel that returns four source-grid maps."""
+    dask_field = _dask_input(canonical_u, canonical_v)
+
+    def kernel(
+        frame_u: NDArray[np.generic], frame_v: NDArray[np.generic]
+    ) -> tuple[
+        NDArray[np.float64],
+        NDArray[np.float64],
+        NDArray[np.float64],
+        NDArray[np.float64],
+    ]:
+        return transform(frame_u, frame_v)
+
+    outputs = xr.apply_ufunc(
+        kernel,
+        canonical_u,
+        canonical_v,
+        input_core_dims=[
+            [source.latitude_dim, source.longitude_dim],
+            [source.latitude_dim, source.longitude_dim],
+        ],
+        output_core_dims=[
+            [source.latitude_dim, source.longitude_dim],
+            [source.latitude_dim, source.longitude_dim],
+            [source.latitude_dim, source.longitude_dim],
+            [source.latitude_dim, source.longitude_dim],
+        ],
+        vectorize=True,
+        output_dtypes=[np.float64, np.float64, np.float64, np.float64],
+        **apply_ufunc_options(dask_field),
+    )
+    first, second, third, fourth = cast(
+        tuple[xr.DataArray, xr.DataArray, xr.DataArray, xr.DataArray], outputs
+    )
+    return (
+        restore_output(
+            first,
+            source=source,
+            target=source,
+            original_dims=original.dims,
+        ),
+        restore_output(
+            second,
+            source=source,
+            target=source,
+            original_dims=original.dims,
+        ),
+        restore_output(
+            third,
+            source=source,
+            target=source,
+            original_dims=original.dims,
+        ),
+        restore_output(
+            fourth,
+            source=source,
+            target=source,
+            original_dims=original.dims,
+        ),
+    )
+
+
+def _dask_input(first: xr.DataArray, second: xr.DataArray) -> xr.DataArray:
+    return first if first.chunks is not None else second
+
+
+def _output_changes_shape(source: FieldLayout, target: FieldLayout) -> bool:
+    return (
+        source.latitude_dim != target.latitude_dim
+        or source.longitude_dim != target.longitude_dim
+        or source.grid.nlat != target.grid.nlat
+        or source.grid.nlon != target.grid.nlon
+    )
+
+
+def _require_matching_layouts(
+    first: FieldLayout, second: FieldLayout, first_name: str, second_name: str
+) -> None:
+    if (
+        first.latitude_dim != second.latitude_dim
+        or first.longitude_dim != second.longitude_dim
+        or first.coordinates.latitude_name != second.coordinates.latitude_name
+        or first.coordinates.longitude_name != second.coordinates.longitude_name
+    ):
+        raise ValueError(
+            f"{first_name} and {second_name} must use the same horizontal "
+            "coordinate names and dimensions"
+        )
+    if not grids_equivalent(first.grid, second.grid):
+        raise ValueError(
+            f"{first_name} and {second_name} must describe equivalent supported grids"
+        )