Merge branch 'devel' into devel

Author: falconstryker

.github/workflows/marsfiles_test.yml

@@ -59,9 +59,42 @@ jobs:
           pip install numpy netCDF4 xarray scipy matplotlib
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
 
-      # Install the package in editable mode
-      - name: Install package
-        run: pip install -e .
+      # Install pyshtools for spatial analysis capabilities
+      - name: Install pyshtools and spectral dependencies (Ubuntu)
+        if: runner.os == 'Linux'
+        shell: bash
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y libfftw3-dev
+          pip install pyshtools
+
+      - name: Install pyshtools and spectral dependencies (macOS)
+        if: runner.os == 'macOS'
+        shell: bash
+        run: |
+          pip install pyshtools
+
+      - name: Install pyshtools and spectral dependencies (Windows)
+        if: runner.os == 'Windows'
+        shell: pwsh
+        run: |
+          pip install pyshtools
+
+      # Install the package with spectral extension
+      - name: Install package with spectral capabilities
+        run: |
+          pip install -e .
+      
+      # Create a test profile if needed
+      - name: Create amescap profile
+        shell: bash
+        run: |
+          mkdir -p $HOME
+          mkdir -p mars_templates
+          echo "# AmesCAP profile" > mars_templates/amescap_profile
+          echo "export PYTHONPATH=$PYTHONPATH:$(pwd)" >> mars_templates/amescap_profile
+          cp mars_templates/amescap_profile $HOME/.amescap_profile
+          echo "Created profile at $HOME/.amescap_profile"
 
       # Create a patch for the test file to fix Windows path issues
       - name: Create test_marsfiles.py path fix for Windows

README.rst

@@ -39,6 +39,10 @@ For reproducible analysis, we recommend installing CAP in a dedicated virtual en
 
     # Install CAP and its dependencies
     pip install git+https://github.com/NASA-Planetary-Science/AmesCAP.git
+    # OR install a specific branch with:
+    pip install git+https://github.com/NASA-Planetary-Science/AmesCAP.git@devel
+
+    # For spectral analysis capabilities, please follow the installation instructions.
 
     # Copy amescap_profile to your home directory, which varies by OS, shell, and package manager:
     # pip + Unix/MacOS (bash, csh, tcsh, zsh) OR Windows Cygwin:
@@ -54,6 +58,8 @@ For reproducible analysis, we recommend installing CAP in a dedicated virtual en
 
 This ensures consistent package versions across different systems.
 
+For spectral analysis capabilities, please follow the `installation instructions <https://amescap.readthedocs.io/en/latest/installation.html>`_.
+
 Available Commands
 ^^^^^^^^^^^^^^^
 After installation, the following commands will be available:

amescap/Spectral_utils.py

@@ -31,6 +31,28 @@
 # ======================================================================
 #                           DEFINITIONS
 # ======================================================================
+# Try to import pyshtools with proper error handling
+try:
+    import pyshtools
+    PYSHTOOLS_AVAILABLE = True
+except ImportError:
+    PYSHTOOLS_AVAILABLE = False
+    print(
+        f"{Yellow}__________________\n"
+        f"Zonal decomposition relies on the pyshtools library, "
+        f"referenced at:\n\n"
+        f"Mark A. Wieczorek and Matthias Meschede (2018). "
+        f"SHTools - Tools for working with spherical harmonics,"
+        f"Geochemistry, Geophysics, Geosystems, 2574-2592, "
+        f"doi:10.1029/2018GC007529\n\nPlease consult pyshtools  "
+        f"documentation at:\n"
+        f" {Cyan}https://pypi.org/project/pyshtools\n"
+        f"{Yellow}And installation instructions for CAP with pyshtools:\n"
+        f" {Cyan}https://amescap.readthedocs.io/en/latest/installation."
+        f"html#_spectral_analysis{Yellow}\n"
+        f"__________________{Nclr}\n\n"
+    )
+
 def diurn_extract(VAR, N, tod, lon):
     """
     Extract the diurnal component of a field. Original code by John
@@ -414,3 +436,122 @@ def zeroPhi_filter(VAR, btype, low_highcut, fs, axis=0, order=4,
         return VAR_trend + VAR_f
     else:
         return VAR_f
+
+def zonal_decomposition(VAR):
+    """
+    Decomposition into spherical harmonics. [A. Kling, 2020]
+
+    :param VAR: Detrend variable for decomposition. Lat is SECOND to
+        LAST dimension and lon is LAST (e.g., ``[time,lat,lon]`` or
+        ``[time,lev,lat,lon]``)
+
+    :return: (COEFFS) coefficient for harmonic decomposion, shape is
+        flattened (e.g., ``[time, 2, lat/2, lat/2]``
+        ``[time x lev, 2, lat/2, lat/2]``);
+        (power_per_l) power spectral density, shape is re-organized
+        (e.g., [time, lat/2] or [time, lev, lat/2])
+
+    .. NOTE:: Output size is (``[...,lat/2, lat/2]``) as lat is the
+        smallest dimension. This matches the Nyquist frequency.
+        
+    """
+    if not PYSHTOOLS_AVAILABLE:
+        raise ImportError(
+            "This function requires pyshtools. Install with:\n"
+            "conda install -c conda-forge pyshtools\n"
+            "or\n"
+            "pip install amescap[spectral]"
+        )
+
+    var_shape = np.array(VAR.shape)
+
+    # Flatten array (e.g., [10, 36, lat, lon] -> [360, lat, lon])
+    nflatten = int(np.prod(var_shape[:-2]))
+    reshape_flat = np.append(nflatten, var_shape[-2:])
+    VAR = VAR.reshape(reshape_flat)
+
+    coeff_out_shape = np.append(var_shape[0:-2],
+                                [2, var_shape[-2]//2, var_shape[-2]//2])
+    psd_out_shape = np.append(var_shape[0:-2], var_shape[-2]//2)
+    coeff_flat_shape = np.append(nflatten, coeff_out_shape[-3:])
+    COEFFS = np.zeros(coeff_flat_shape)
+
+    psd = np.zeros((nflatten,var_shape[-2]//2))
+
+    for ii in range(0,nflatten):
+        COEFFS[ii,...] = pyshtools.expand.SHExpandDH(VAR[ii,...], sampling = 2)
+        psd[ii,:] = pyshtools.spectralanalysis.spectrum(COEFFS[ii,...])
+
+    return  COEFFS, psd.reshape(psd_out_shape)
+
+def zonal_construct(COEFFS_flat, VAR_shape, btype=None, low_highcut=None):
+    """
+    Recomposition into spherical harmonics
+
+    :param COEFFS_flat: Spherical harmonic coefficients as a flattened
+        array, (e.g., ``[time, lat, lon]`` or
+        ``[time x lev, 2, lat, lon]``)
+    :type COEFFS_flat: array
+    
+    :param VAR_shape: shape of the original variable
+    :type VAR_shape: tuple
+    
+    :param btype: filter type: "low", "high", or "band". If None,
+        returns reconstructed array using all zonal wavenumbers
+    :type btype: str or None
+    
+    :param low_high_cut: low, high or [low, high] zonal wavenumber
+    :type low_high_cut: int
+
+    :return: detrended variable reconstructed to original size
+        (e.g., [time, lev, lat, lon])
+
+    .. NOTE:: The minimum and maximum wavelenghts in [km] are computed::
+        dx = 2*np.pi * 3400
+        L_min = (1./kmax) * dx
+        L_max = 1./max(kmin, 1.e-20) * dx
+        if L_max > 1.e20:
+            L_max = np.inf
+        print("(kmin,kmax) = ({kmin}, {kmax})
+              -> dx min = {L_min} km, dx max = {L_max} km")
+              
+    """
+    if not PYSHTOOLS_AVAILABLE:
+        raise ImportError(
+            "This function requires pyshtools. Install with:\n"
+            "conda install -c conda-forge pyshtools\n"
+            "or\n"
+            "pip install amescap[spectral]"
+        )
+
+    # Initialization
+    nflatten = COEFFS_flat.shape[0]
+    kmin = 0
+    kmax = COEFFS_flat.shape[-1]
+
+    VAR = np.zeros((nflatten, VAR_shape[-2], VAR_shape[-1]))
+
+    if btype == "low":
+        kmax= int(low_highcut)
+    if btype == "high":
+        kmin= int(low_highcut)
+    if btype == "band":
+        kmin, kmax= int(low_highcut[0]), int(low_highcut[1])
+
+    for ii in range(0, nflatten):
+        # Filtering
+        COEFFS_flat[ii, :, :kmin, :] = 0.
+        COEFFS_flat[ii, :, kmax:, :] = 0.
+        VAR[ii, :, :] = pyshtools.expand.MakeGridDH(COEFFS_flat[ii, :, :],
+                                                  sampling = 2)
+    return  VAR.reshape(VAR_shape)
+
+
+def init_shtools():
+    """
+    Check if pyshtools is available and return its availability status
+    
+    :return: True if pyshtools is available, False otherwise
+    :rtype: bool
+    """
+    return PYSHTOOLS_AVAILABLE

bin/MarsCalendar.py

@@ -29,13 +29,12 @@
 import sys          # System commands
 import argparse     # Parse arguments
 import numpy as np
+import functools    # For function decorators
+import traceback    # For printing stack traces
 
 # Load amesCAP modules
 from amescap.FV3_utils import (sol2ls, ls2sol)
 
-import functools
-import traceback
-
 def debug_wrapper(func):
     """
     A decorator that wraps a function with error handling based on the