Dev healpix

.github/workflows/test.yml

@@ -18,21 +18,10 @@ jobs:
       uses: actions/setup-python@v3
       with:
         python-version: '3.10'
-    - name: Add conda to system path
-      run: |
-        # $CONDA is an environment variable pointing to the root of the miniconda directory
-        echo $CONDA/bin >> $GITHUB_PATH
-    - name: Install dependencies
-      run: |
-        conda env update --file environment.yml --name base
-    - name: Lint with flake8
-      run: |
-        conda install flake8
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+    - uses: mamba-org/setup-micromamba@v2
+      with:
+        environment-file: environment.yml
+        generate-run-shell: true   
     - name: Test with pytest
-      run: |
-        conda install pytest
-        pytest
+      run: pytest
+      shell: micromamba-shell {0}

.gitignore

@@ -17,6 +17,7 @@ Untitled*
 
 # python package
 *.egg-info
+*build
 
 # Jupyterhub checkpoints
 *.ipynb_checkpoints*

environment.yml

@@ -2,20 +2,20 @@ name: pywk99
 channels:
   - conda-forge
 dependencies:
-  - python>=3.9
-  - pytest
-  - numpy
-  - pandas
-  - xarray
-  - netcdf4
-  - shapely
-  - matplotlib
   - cartopy
   - cmocean
-  - ipython
+  - dask
+  - healpy
   - ipykernel
+  - ipython
+  - matplotlib
+  - netcdf4
+  - numpy
+  - pandas
   - pip
-  - dask
+  - pytest
+  - python>=3.9
+  - shapely
+  - xarray
   - pip:
-    - -e .
-
+      - -e .

pyproject.toml

@@ -8,12 +8,13 @@ authors = [
     {name = "MPI-M"},
 ]
 dependencies = [
-    "numpy",
-    "xarray",
-    "netcdf4",
+    "cartopy",
     "matplotlib",
+    "netcdf4",
+    "numpy",
+    "pandas",
     "shapely",
-    "cartopy",
+    "xarray",
 ]
 description = "Wheeler and Kiladis Wavenumber-Frequency Analysis in Python"
 requires-python = ">=3.10"
@@ -26,4 +27,4 @@ classifiers=[
 dynamic = ["version"]
 
 [tool.setuptools.dynamic]
-version = {attr = "pywk99.__version__"}
+version = {attr = "pywk99.__version__"}

pywk99/__init__.py

@@ -3,6 +3,9 @@
 
 Changelog
 ---------
+Apr 8, 2025: Version 0.4.3
+    - Adding support for healpix
+
 Mar 7, 2025: Version 0.4.2
     - Organizing first github public version.
 
@@ -67,4 +70,4 @@
     - Several bugs where detected and corrected with the tests.
 """
 
-__version__ = "0.4.2"
+__version__ = "0.4.3"

pywk99/filter/filter.py

@@ -1,6 +1,6 @@
 """Filter for equatorial wave bands following Wheeler and Kiladis 1999."""
 
-from typing import Union
+from typing import Optional, Union, Tuple
 import xarray as xr
 import numpy as np
 
@@ -11,63 +11,74 @@
 from pywk99.timeseries.timeseries import remove_seasonal_cycle
 from pywk99.timeseries.timeseries import check_variable_coordinates_are_sorted
 from pywk99.filter.window import FilterPoint, FilterWindow
-
-
-def filter_variable(variable: xr.DataArray,
-                    filter_windows: Union[FilterWindow, list[FilterWindow]],
-                    taper: bool = True,
-                    taper_alpha: float = 0.5,
-                    rm_seasonal_cycle: bool = True,
-                    rm_linear_trend: bool = True) -> xr.Dataset:
-    modified_variable = _preprocess(variable,
-                                    taper, taper_alpha,
-                                    rm_seasonal_cycle,
-                                    rm_linear_trend)
-    spectrum = fourier_transform(modified_variable)
+from pywk99.grid.grid import dataarray_to_equatorial_latlon_grid
+
+
+def filter_variable(
+    variable: xr.DataArray,
+    filter_windows: Union[FilterWindow, list[FilterWindow]],
+    taper: bool = True,
+    taper_alpha: float = 0.5,
+    rm_seasonal_cycle: bool = True,
+    rm_linear_trend: bool = True,
+    grid_type: str = "latlon",
+    grid_dict: Optional[dict] = None,
+) -> xr.Dataset:
+    new_variable = dataarray_to_equatorial_latlon_grid(
+        variable, grid_type, grid_dict
+    )
+    new_variable = _preprocess(
+        new_variable, taper, taper_alpha, rm_seasonal_cycle, rm_linear_trend
+    )
+    spectrum = fourier_transform(new_variable)
     if isinstance(filter_windows, FilterWindow):
         filter_windows = [filter_windows]
     data_vars = dict()
     seen_window_names = []
     for window in filter_windows:
         field_name = _set_field_name(window.name, seen_window_names)
         seen_window_names.append(window.name)
-        data_vars[field_name] = _filter(spectrum, window, variable.coords)
-    filtered_variables =  xr.Dataset(data_vars=data_vars)
+        data_vars[field_name] = _filter(spectrum, window, new_variable.coords)
+    filtered_variables = xr.Dataset(data_vars=data_vars)
     return filtered_variables
 
 
-def modify_spectrum(spectrum: xr.DataArray,
-                    filter_window: FilterWindow,
-                    action: str = "filter") -> xr.DataArray:
+def modify_spectrum(
+    spectrum: xr.DataArray, filter_window: FilterWindow, action: str = "filter"
+) -> xr.DataArray:
     """Filter the spectrum with the wave filter window."""
     mask = _get_window_mask(spectrum, filter_window, action)
     modified_spectrum = mask * spectrum
     return modified_spectrum
 
 
-def _filter(spectrum: xr.DataArray,
-            filter_window: FilterWindow,
-            xarray_coords) -> xr.DataArray:
+def _filter(
+    spectrum: xr.DataArray, filter_window: FilterWindow, xarray_coords
+) -> xr.DataArray:
     masked_spectrum = modify_spectrum(spectrum, filter_window, "filter")
-    filtered_variable = inverse_fourier_transform(masked_spectrum, xarray_coords)
+    filtered_variable = inverse_fourier_transform(
+        masked_spectrum, xarray_coords
+    )
     return filtered_variable
 
 
-def _preprocess(variable: Union[xr.DataArray, xr.Dataset],
-                taper: bool,
-                taper_alpha: float,
-                rm_seasonal_cycle: bool,
-                rm_linear_trend: bool
-                ) -> Union[xr.DataArray, xr.Dataset]:
+def _preprocess(
+    variable: Union[xr.DataArray, xr.Dataset],
+    taper: bool,
+    taper_alpha: float,
+    rm_seasonal_cycle: bool,
+    rm_linear_trend: bool,
+) -> Union[xr.DataArray, xr.Dataset]:
     check_variable_coordinates_are_sorted(variable)
     modified_variable = variable.transpose("time", "lon", ...)
     if rm_linear_trend:
         modified_variable = remove_linear_trend(modified_variable)
     if rm_seasonal_cycle:
         modified_variable = remove_seasonal_cycle(modified_variable)
     if taper:
-        modified_variable = taper_variable_time_ends(modified_variable,
-                                                     taper_alpha)
+        modified_variable = taper_variable_time_ends(
+            modified_variable, taper_alpha
+        )
     return modified_variable
 
 
@@ -79,27 +90,31 @@ def _set_field_name(window_name: str, seen_window_names: list[str]) -> str:
         return f"{window_name}{name_count + 1}"
 
 
-def _get_window_mask(spectrum: xr.DataArray,
-                     filter_window: FilterWindow,
-                     action: str = "filter") -> xr.DataArray:
+def _get_window_mask(
+    spectrum: xr.DataArray, filter_window: FilterWindow, action: str = "filter"
+) -> xr.DataArray:
     """Get a wavenumber-frequency mask corresponding to the filter window."""
-    bbox_wavenumbers, bbox_frequencies = \
-        _window_bbox_wavenumber_and_frequencies(
-            spectrum, filter_window
-        )
+    bbox_wavenumbers, bbox_frequencies = (
+        _window_bbox_wavenumber_and_frequencies(spectrum, filter_window)
+    )
     mask = _get_mask_base(spectrum, action)
     include_fft_reflection = bool(np.any(spectrum.frequency.values < 0))
     for wavenumber in bbox_wavenumbers.values:
         for frequency in bbox_frequencies.values:
-            mask = _modify_mask_value_at_point(mask, filter_window, action,
-                                               wavenumber, frequency,
-                                               include_fft_reflection)
+            mask = _modify_mask_value_at_point(
+                mask,
+                filter_window,
+                action,
+                wavenumber,
+                frequency,
+                include_fft_reflection,
+            )
     return mask
 
 
 def _window_bbox_wavenumber_and_frequencies(
-        spectrum: xr.DataArray,
-        filter_window: FilterWindow) -> tuple[np.ndarray, np.ndarray]:
+    spectrum: xr.DataArray, filter_window: FilterWindow
+) -> Tuple[np.ndarray, np.ndarray]:
     """Get the wavenumbers and frequencies of the window bounding box."""
     k_wmin, w_wmin, k_wmax, w_wmax = filter_window.bounds
     if not np.all(np.diff(spectrum.wavenumber.values) > 0):
@@ -121,12 +136,13 @@ def _get_mask_base(spectrum: xr.DataArray, action: str) -> xr.DataArray:
 
 
 def _modify_mask_value_at_point(
-        mask: xr.DataArray,
-        filter_window: FilterWindow,
-        action: str,
-        wavenumber: float,
-        frequency: float,
-        include_fft_reflection: bool = True) -> xr.DataArray:
+    mask: xr.DataArray,
+    filter_window: FilterWindow,
+    action: str,
+    wavenumber: float,
+    frequency: float,
+    include_fft_reflection: bool = True,
+) -> xr.DataArray:
     ACTION_VALUE = {"filter": True, "substract": False}
     point = FilterPoint(wavenumber, frequency)
     point_is_contained = filter_window.covers(point)
@@ -135,12 +151,15 @@ def _modify_mask_value_at_point(
         mask.loc[point_loc_dict1] = ACTION_VALUE[action]
         if include_fft_reflection:
             # rounding errors in the index make the following necessary
-            approx_point_loc_dict2 = dict(wavenumber=-wavenumber,
-                                          frequency=-frequency)
-            reflection_point = mask.sel(approx_point_loc_dict2,
-                                        method='nearest')
+            approx_point_loc_dict2 = dict(
+                wavenumber=-wavenumber, frequency=-frequency
+            )
+            reflection_point = mask.sel(
+                approx_point_loc_dict2, method="nearest"
+            )
             point_loc_dict2 = dict(
                 wavenumber=reflection_point.wavenumber.values,
-                frequency=reflection_point.frequency.values)
+                frequency=reflection_point.frequency.values,
+            )
             mask.loc[point_loc_dict2] = ACTION_VALUE[action]
     return mask

pywk99/filter/window.py

@@ -1,7 +1,7 @@
 """Define filtering windows for various waves following Wheeler and Kiladis."""
 
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Optional, Tuple, Union
 import numpy as np
 from shapely.geometry import Point, Polygon, MultiPolygon
 
@@ -18,7 +18,7 @@ class FilterWindow:
     polygon: Union[Polygon, MultiPolygon]
 
     @property
-    def bounds(self) -> tuple[float, float, float, float]:
+    def bounds(self) -> Tuple[float, float, float, float]:
         return self.polygon.bounds
 
     def union(self, other) -> "FilterWindow":

pywk99/grid/__init__.py

@@ -0,0 +1,5 @@
+"""
+Transform grids to latlon
+"""
+
+from pywk99.grid.grid import dataset_to_equatorial_latlon_grid

pywk99/grid/grid.py

@@ -0,0 +1,34 @@
+"""Convert datasets to latlon grid"""
+
+from typing import Optional
+import xarray as xr
+
+
+from pywk99.grid.healpix import dataarray_healpix_to_equatorial_latlon
+from pywk99.grid.healpix import dataset_healpix_to_equatorial_latlon
+
+
+def dataset_to_equatorial_latlon_grid(
+    dataset: xr.Dataset, grid_type: str, grid_dict: Optional[dict]
+) -> xr.Dataset:
+    if grid_type == "latlon":
+        return dataset
+    elif grid_type == "healpix":
+        if grid_dict is None:
+            raise ValueError("No grid_dict provided for healpix conversion.")
+        return dataset_healpix_to_equatorial_latlon(dataset, **grid_dict)
+    else:
+        raise ValueError("Grid type not found.")
+
+
+def dataarray_to_equatorial_latlon_grid(
+    dataarray: xr.DataArray, grid_type: str, grid_dict: Optional[dict]
+) -> xr.DataArray:
+    if grid_type == "latlon":
+        return dataarray
+    elif grid_type == "healpix":
+        if grid_dict is None:
+            raise ValueError("No grid_dict provided for healpix conversion.")
+        return dataarray_healpix_to_equatorial_latlon(dataarray, **grid_dict)
+    else:
+        raise ValueError("Grid type not found.")