diff --git a/.github/workflows/marscalendar_test.yml b/.github/workflows/marscalendar_test.yml index 11d4376..934e12a 100644 --- a/.github/workflows/marscalendar_test.yml +++ b/.github/workflows/marscalendar_test.yml @@ -20,7 +20,6 @@ on: - 'bin/MarsCalendar.py' - 'tests/test_marscalendar.py' - '.github/workflows/marscalendar_test.yml' - jobs: test: # Run on multiple OS and Python versions for comprehensive testing @@ -32,11 +31,13 @@ jobs: steps: # Checkout the repository - uses: actions/checkout@v3 + # Set up the specified Python version - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} + # Install dependencies - name: Install dependencies shell: bash @@ -44,10 +45,12 @@ jobs: python -m pip install --upgrade pip pip install numpy if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + # Install the package in editable mode - name: Install package run: pip install -e . - # Run the tests + + # Run the tests - name: Run MarsCalendar tests run: | cd tests diff --git a/.github/workflows/marsfiles_test.yml b/.github/workflows/marsfiles_test.yml index fc32478..ec5c641 100644 --- a/.github/workflows/marsfiles_test.yml +++ b/.github/workflows/marsfiles_test.yml @@ -31,13 +31,16 @@ jobs: fail-fast: false runs-on: ${{ matrix.os }} steps: + # Checkout the repository - uses: actions/checkout@v3 + # Set up the specified Python version - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - + + # Cache pip dependencies - name: Cache pip dependencies uses: actions/cache@v3 with: @@ -62,13 +65,15 @@ jobs: sudo apt-get update sudo apt-get install -y libfftw3-dev pip install pyshtools - + + # Install pyshtools for spatial analysis capabilities (macos) - name: Install pyshtools and spectral dependencies (macOS) if: runner.os == 'macOS' shell: bash run: | pip install pyshtools + # Install pyshtools for spatial analysis capabilities (Windows) - name: Install pyshtools and spectral dependencies (Windows) if: runner.os == 'Windows' shell: pwsh diff --git a/.github/workflows/marsformat_test.yml b/.github/workflows/marsformat_test.yml index b8cc7ac..b9503a5 100644 --- a/.github/workflows/marsformat_test.yml +++ b/.github/workflows/marsformat_test.yml @@ -22,7 +22,6 @@ on: - '.github/workflows/marsformat_test.yml' jobs: test: - # Run on multiple OS and Python versions for comprehensive testing strategy: matrix: os: [ubuntu-latest, macos-latest, windows-latest] @@ -31,33 +30,39 @@ jobs: steps: # Checkout the repository - uses: actions/checkout@v3 + # Set up the specified Python version - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - # Install dependencies + + # Install dependencies - name: Install dependencies shell: bash run: | python -m pip install --upgrade pip pip install numpy netCDF4 xarray if [ -f requirements.txt ]; then pip install -r requirements.txt; fi - # Install the package in editable mode + + # Install the package in editable mode - name: Install package run: pip install -e . - # Set HOME for Windows since it might be used by the script + + # Set HOME for Windows since it might be used by the script - name: Set HOME environment variable for Windows if: runner.os == 'Windows' shell: pwsh run: | echo "HOME=$env:USERPROFILE" >> $env:GITHUB_ENV + # Set up AmesCAP configuration - handle platform differences - name: Set up AmesCAP configuration shell: bash run: | mkdir -p $HOME/.amescap cp mars_templates/amescap_profile $HOME/.amescap_profile + # Print out environment info - name: Show environment info run: | @@ -65,6 +70,7 @@ jobs: echo "Working directory: $(pwd)" echo "Home directory: $HOME" echo "Environment variables: $(env)" + # Free up disk space - name: Free up disk space if: runner.os == 'Linux' @@ -73,23 +79,28 @@ jobs: sudo rm -rf /usr/share/dotnet sudo rm -rf /opt/ghc sudo rm -rf /usr/local/share/boost + # Create temporary directory for tests - name: Create temporary test directory shell: bash run: | mkdir -p $HOME/marsformat_tests echo "TMPDIR=$HOME/marsformat_tests" >> $GITHUB_ENV + # Run the integration tests with cleanup between tests - name: Run MarsFormat tests run: | cd tests python -m unittest -v test_marsformat.py + # Clean up temporary files to avoid disk space issues - OS specific - name: Clean up temp files (Unix) if: runner.os != 'Windows' && always() shell: bash run: | rm -rf $HOME/marsformat_tests || true + + # Clean up temporary files (Windows) - name: Clean up temp files (Windows) if: runner.os == 'Windows' && always() shell: pwsh diff --git a/.github/workflows/marsinterp_test.yml b/.github/workflows/marsinterp_test.yml index 0c7358a..cf91570 100644 --- a/.github/workflows/marsinterp_test.yml +++ b/.github/workflows/marsinterp_test.yml @@ -31,13 +31,16 @@ jobs: fail-fast: true runs-on: ${{ matrix.os }} steps: + # Checkout the repository - uses: actions/checkout@v3 + # Set up the specified Python version - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} + # Cache pip dependencies - name: Cache pip dependencies uses: actions/cache@v3 with: diff --git a/.github/workflows/marsplot_test.yml b/.github/workflows/marsplot_test.yml index 6b00aa4..86401ca 100644 --- a/.github/workflows/marsplot_test.yml +++ b/.github/workflows/marsplot_test.yml @@ -22,7 +22,6 @@ on: - '.github/workflows/marsplot_test.yml' jobs: test: - # Run on multiple OS and Python versions for comprehensive testing strategy: matrix: os: [ubuntu-latest, macos-latest, windows-latest] @@ -31,33 +30,39 @@ jobs: steps: # Checkout the repository - uses: actions/checkout@v3 + # Set up the specified Python version - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - # Install dependencies + + # Install dependencies - name: Install dependencies shell: bash run: | python -m pip install --upgrade pip pip install numpy netCDF4 xarray if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + # Install the package in editable mode - name: Install package run: pip install -e . + # Set HOME for Windows since it might be used by the script - name: Set HOME environment variable for Windows if: runner.os == 'Windows' shell: pwsh run: | echo "HOME=$env:USERPROFILE" >> $env:GITHUB_ENV + # Set up AmesCAP configuration - handle platform differences - name: Set up AmesCAP configuration shell: bash run: | mkdir -p $HOME/.amescap cp mars_templates/amescap_profile $HOME/.amescap_profile + # Print out environment info - name: Show environment info run: | @@ -65,6 +70,7 @@ jobs: echo "Working directory: $(pwd)" echo "Home directory: $HOME" echo "Environment variables: $(env)" + # Free up disk space - name: Free up disk space if: runner.os == 'Linux' @@ -73,23 +79,28 @@ jobs: sudo rm -rf /usr/share/dotnet sudo rm -rf /opt/ghc sudo rm -rf /usr/local/share/boost + # Create temporary directory for tests - name: Create temporary test directory shell: bash run: | mkdir -p $HOME/marsplot_tests echo "TMPDIR=$HOME/marsplot_tests" >> $GITHUB_ENV + # Run the integration tests with cleanup between tests - name: Run MarsPlot tests run: | cd tests python -m unittest -v test_marsplot.py - # Clean up temporary files to avoid disk space issues - OS specific + + # Clean up temporary files to avoid disk space issues - unix - name: Clean up temp files (Unix) if: runner.os != 'Windows' && always() shell: bash run: | rm -rf $HOME/marsplot_tests || true + + # Clean up temporary files to avoid disk space issues - windows - name: Clean up temp files (Windows) if: runner.os == 'Windows' && always() shell: pwsh diff --git a/.github/workflows/marspull_test.yml b/.github/workflows/marspull_test.yml index 3ddd52a..8c1a0ec 100644 --- a/.github/workflows/marspull_test.yml +++ b/.github/workflows/marspull_test.yml @@ -23,7 +23,6 @@ on: jobs: test: - # Run on multiple OS and Python versions for comprehensive testing strategy: matrix: os: [ubuntu-latest, macos-latest, windows-latest] diff --git a/.github/workflows/marsvars_test.yml b/.github/workflows/marsvars_test.yml index d16d71d..237c644 100644 --- a/.github/workflows/marsvars_test.yml +++ b/.github/workflows/marsvars_test.yml @@ -35,13 +35,16 @@ jobs: fail-fast: true runs-on: ${{ matrix.os }} steps: + # Checkout the repository - uses: actions/checkout@v3 + # Set up the specified Python version - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v3 with: python-version: ${{ matrix.python-version }} - + + # Cache pip dependencies - name: Cache pip dependencies uses: actions/cache@v3 with: diff --git a/amescap/FV3_utils.py b/amescap/FV3_utils.py index f75777c..b53396a 100644 --- a/amescap/FV3_utils.py +++ b/amescap/FV3_utils.py @@ -1,9 +1,9 @@ # !/usr/bin/env python3 """ -FV3_utils contains internal Functions for processing data in MGCM +FV3_utils contains internal Functions for processing data in MGCM output files such as vertical interpolation. -These functions can be used on their own outside of CAP if they are +These functions can be used on their own outside of CAP if they are imported as a module:: from /u/path/FV3_utils import fms_press_calc @@ -13,37 +13,34 @@ * ``numpy`` * ``warnings`` * ``scipy`` - """ # Load generic Python modules import warnings # suppress errors triggered by NaNs import numpy as np from scipy import optimize +from scipy.spatial import cKDTree # p_half = half-level = layer interfaces # p_full = full-level = layer midpoints + def fms_press_calc(psfc, ak, bk, lev_type='full'): """ Returns the 3D pressure field from the surface pressure and the ak/bk coefficients. - :param psfc: the surface pressure [Pa] or an array of surface + :param psfc: the surface pressure [Pa] or an array of surface pressures (1D, 2D, or 3D if time dimension) - :type psfc: array - + :type psfc: array :param ak: 1st vertical coordinate parameter - :type ak: array - + :type ak: array :param bk: 2nd vertical coordinate parameter - :type bk: array: - - :param lev_type: "full" (layer midpoints) or "half" + :type bk: array: + :param lev_type: "full" (layer midpoints) or "half" (layer interfaces). Defaults to "full." - :type lev_type: str - - :return: the 3D pressure field at the full levels + :type lev_type: str + :return: the 3D pressure field at the full levels ``PRESS_f(Nk-1:,:,:)`` or half-levels ``PRESS_h(Nk,:,:,)`` [Pa] Calculation:: @@ -56,12 +53,12 @@ def fms_press_calc(psfc, ak, bk, lev_type='full'): ---Nk-1--- -------- p_full --- Nk --- SFC ======== p_half / / / / / - + .. note:: Some literature uses pk (pressure) instead of ak with ``p3d = ps * bk + P_ref * ak`` instead of ``p3d = ps * bk + ak`` - """ + Nk = len(ak) # If ``psfc`` is a float (e.g., ``psfc = 700.``), make it a # 1-element array (e.g., ``psfc = [700]``) @@ -86,7 +83,8 @@ def fms_press_calc(psfc, ak, bk, lev_type='full'): # Pressure at layer midpoints. The size of Z axis is ``Nk-1`` PRESS_f = np.zeros((len(psfc_flat), Nk-1)) - if ak[0] == 0 and bk[0] == 0: # Top layer (1st element is ``i = 0`` in Python) + if ak[0] == 0 and bk[0] == 0: + # Top layer (1st element is ``i = 0`` in Python) PRESS_f[:, 0] = 0.5 * (PRESS_h[:, 0]+PRESS_h[:, 1]) else: PRESS_f[:, 0] = ( @@ -105,43 +103,38 @@ def fms_press_calc(psfc, ak, bk, lev_type='full'): #TODO if lev_type == "full": new_dim_f = np.append(Nk-1, psfc.shape) - return PRESS_f.T.reshape(new_dim_f) # (dz = dp/(-rho g)) and # (rho = p/(r T)) => (dz = rT/g * (-dp/p)) @@ -178,12 +171,12 @@ def fms_Z_calc(psfc, ak, bk, T, topo=0., lev_type="full"): Z_half calculation:: - # Hydrostatic relation within the layer > (P(k+1)/P(k) = + # Hydrostatic relation within the layer > (P(k+1)/P(k) = # exp(-DZ(k)/H)) - + # layer thickness: DZ(k) = rT/g * -(du) - + # previous layer altitude + thickness of layer: Z_h k) = Z_h(k+1) +DZ_h(h) @@ -192,17 +185,17 @@ def fms_Z_calc(psfc, ak, bk, T, topo=0., lev_type="full"): # previous altitude + half the thickness of previous layer and # half of current layer Z_f(k) = Z_f(k+1) + (0.5 DZ(k) + 0.5 DZ(k+1)) - + # Add ``+0.5 DZ(k)-0.5 DZ(k)=0`` and re-organiz the equation Z_f(k) = Z_f(k+1) + DZ(k) + 0.5 (DZ(k+1) - DZ(k)) Z_f(k) = Z_h(k+1) + 0.5 (DZ(k+1) - DZ(k)) The specific heat ratio: ``γ = cp/cv (cv = cp-R)`` => ``γ = cp/(cp-R)`` Also ``(γ-1)/γ=R/cp`` - + The dry adiabatic lapse rate: ``Γ = g/cp`` => ``Γ = (gγ)/R`` - + The isentropic relation: ``T2 = T1(p2/p1)**(R/cp)`` @@ -224,16 +217,16 @@ def fms_Z_calc(psfc, ak, bk, T, topo=0., lev_type="full"): Line 7: ``Z_full[k] = Z_half[k] + (T_half[k+1]-T_full[k])/Γ`` Pulling out ``Tfull`` from above equation and using ``Γ = (gγ)/R``:: - + Z_full[k] = (Z_half[k+1] + (R Tfull[k]) / (gγ)(T_half[k+1] / T_full[k] - 1)) Using the isentropic relation above:: - + Z_full = (Z_half[k+1] + (R Tfull[k]) / (gγ)(p_half[k+1] / p_full[k])**(R/Cp)-1)) - """ + g = 3.72 # acc. m/s2 r_co2 = 191.00 # kg/mol Nk = len(ak) @@ -287,6 +280,7 @@ def fms_Z_calc(psfc, ak, bk, T, topo=0., lev_type="full"): raise Exception("Altitude level type not recognized: use 'full' " "or 'half'") + # TODO: delete: Former version of find_n(): only provides 1D > 1D and # ND > 1D mapping @@ -298,16 +292,13 @@ def find_n0(Lfull_IN, Llev_OUT, reverse_input=False): :param Lfull_IN: Input pressure [Pa] or altitude [m] at layer midpoints. ``Level`` dimension is FIRST - :type Lfull_IN: array - + :type Lfull_IN: array :param Llev_OUT: Desired level type for interpolation [Pa] or [m] - :type Llev_OUT: float or 1D array - + :type Llev_OUT: float or 1D array :param reverse_input: Reverse array (e.g., if ``z(0) = 120 km``, ``z(N) = 0km`` -- which is typical -- or if input data is ``p(0) = 1000Pa``, ``p(N) = 0Pa``) - :type reverse_input: bool - + :type reverse_input: bool :return: ``n`` index for the level(s) where the pressure is just below ``plev`` @@ -319,8 +310,8 @@ def find_n0(Lfull_IN, Llev_OUT, reverse_input=False): If ``Lfull_IN`` is ND ``[lev, time, lat, lon]`` and ``Llev_OUT`` is a 1D array of size ``klev`` then ``n`` is an array of size ``[klev, Ndim]`` with ``Ndim = [time, lat, lon]`` - """ + # Number of original layers Lfull_IN = np.array(Lfull_IN) Nlev = len(np.atleast_1d(Llev_OUT)) @@ -350,32 +341,33 @@ def find_n0(Lfull_IN, Llev_OUT, reverse_input=False): n[i, j] = n[i, j] - 1 return n + def find_n(X_IN, X_OUT, reverse_input=False, modulo=None): """ Maps the closest index from a 1D input array to a ND output array just below the input values. :param X_IN: Source level [Pa] or [m] - :type X_IN: float or 1D array - + :type X_IN: float or 1D array :param X_OUT: Desired pressure [Pa] or altitude [m] at layer midpoints. Level dimension is FIRST - :type X_OUT: array - + :type X_OUT: array :param reverse_input: If input array is decreasing (e.g., if z(0) = 120 km, z(N) = 0 km, which is typical, or if data is p(0) = 1000 Pa, p(N) = 0 Pa, which is uncommon) - :type reverse_input: bool - + :type reverse_input: bool :return: The index for the level(s) where the pressure < ``plev`` - """ - if type(X_IN) != np.ndarray: # Number of original layers + + if type(X_IN) != np.ndarray: + # Number of original layers X_IN = np.array(X_IN) - if len(np.atleast_1d(X_OUT)) == 1: # If one value is requested, convert float to array + if len(np.atleast_1d(X_OUT)) == 1: + # If one value is requested, convert float to array X_OUT = np.array([X_OUT]) - elif type(X_OUT) != np.ndarray: # Convert list to numpy array as needed + elif type(X_OUT) != np.ndarray: + # Convert list to numpy array as needed X_OUT = np.array(X_OUT) # Get input variable dimensions @@ -410,7 +402,8 @@ def find_n(X_IN, X_OUT, reverse_input=False, modulo=None): # out of the larger loop over all of the array elements if len(dimsIN) == 1: for i in range(N_OUT): - for j in range(Ndim): # Handle the case where j might be out of bounds + for j in range(Ndim): + # Handle the case where j might be out of bounds if j < NdimsOUT: n[i, j] = np.argmin(np.abs(X_OUT[i, j] - X_IN[:])) if X_IN[n[i, j]] > X_OUT[i, j]: @@ -420,7 +413,8 @@ def find_n(X_IN, X_OUT, reverse_input=False, modulo=None): n[i, j] = 0 elif len(dimsOUT) == 1: for i in range(N_OUT): - for j in range(Ndim): # Handle the case where j might be out of bounds for X_IN + for j in range(Ndim): + # Handle the case where j might be out of bounds for X_IN if j < NdimsIN: n[i, j] = np.argmin(np.abs(X_OUT[i] - X_IN[:, j])) if X_IN[n[i, j], j] > X_OUT[i]: @@ -430,7 +424,8 @@ def find_n(X_IN, X_OUT, reverse_input=False, modulo=None): n[i, j] = 0 else: for i in range(N_OUT): - for j in range(Ndim): # Handle the case where j might be out of bounds for either array + for j in range(Ndim): + # Handle the case where j might be out of bounds for either array if j < NdimsIN and j < NdimsOUT: n[i, j] = np.argmin(np.abs(X_OUT[i, j] - X_IN[:, j])) if X_IN[n[i, j], j] > X_OUT[i, j]: @@ -443,24 +438,22 @@ def find_n(X_IN, X_OUT, reverse_input=False, modulo=None): n = np.squeeze(n) return n + def expand_index(Nindex, VAR_shape_axis_FIRST, axis_list): """ Repeat interpolation indices along an axis. :param Nindex: Interpolation indices, size is (``n_axis``, ``Nfull = [time, lat, lon]``) - :type Nindex: idx - + :type Nindex: idx :param VAR_shape_axis_FIRST: Shape for the variable to interpolate with interpolation axis first (e.g., ``[tod, time, lev, lat, lon]``) - :type VAR_shape_axis_FIRST: tuple - + :type VAR_shape_axis_FIRST: tuple :param axis_list: Position or list of positions for axis to insert (e.g., ``2`` for ``lev`` in ``[tod, time, lev, lat, lon]``, ``[2, 4]`` for ``lev`` and ``lon``). The axis positions are those for the final shape (``VAR_shape_axis_FIRST``) and must be INCREASING - :type axis_list: int or list - + :type axis_list: int or list :return: ``LFULL`` a 2D array (size ``n_axis``, ``NfFULL = [time, lev, lat, lon]``) with the indices expanded along the ``lev`` dimension and flattened @@ -473,8 +466,8 @@ def expand_index(Nindex, VAR_shape_axis_FIRST, axis_list): This routine expands the indices from ``[tod, time, lat, lon]`` to ``[tod, time, lev, lat, lon]`` with ``Nfull = [time, lev, lat, lon]`` for use in interpolation. - """ + # If one element, turn axis to list if len(np.atleast_1d(axis_list)) == 1: axis_list = [axis_list] @@ -508,43 +501,37 @@ def expand_index(Nindex, VAR_shape_axis_FIRST, axis_list): # Return the new, flattened version of ``Nindex`` return Nindex.reshape(dimsOUT_flat) + def vinterp(varIN, Lfull, Llev, type_int="log", reverse_input=False, masktop=True, index=None): """ Vertical linear or logarithmic interpolation for pressure or altitude. - + :param varIN: Variable to interpolate (VERTICAL AXIS FIRST) - :type varIN: ND array - + :type varIN: ND array :param Lfull: Pressure [Pa] or altitude [m] at full layers, same dimensions as ``varIN`` - :type Lfull: array - + :type Lfull: array :param Llev: Desired level for interpolation [Pa] or [m]. May be increasing or decreasing as the output levels are processed one at a time - :type Llev: 1D array - + :type Llev: 1D array :param type_int: "log" for logarithmic (typically pressure), "lin" for linear (typically altitude) - :type type_int: str - + :type type_int: str :param reverse_input: Reverse input arrays. e.g., if ``zfull[0]`` = 120 km then ``zfull[N]`` = 0km (typical) or if input data is ``pfull[0]``=1000 Pa, ``pfull[N]``=0 Pa - :type reverse_input: bool - + :type reverse_input: bool :param masktop: Set to NaN values if above the model top - :type masktop: bool - + :type masktop: bool :param index: Indices for the interpolation, already processed as ``[klev, Ndim]``. Indices calculated if not provided - :type index: None or array - + :type index: None or array :return: ``varOUT`` variable interpolated on the ``Llev`` pressure or altitude levels - .. note:: + .. note:: This interpolation assumes pressure decreases with height:: -- 0 -- TOP [0 Pa] : [120 km]| X_OUT = Xn*A + (1-A)*Xn + 1 @@ -560,8 +547,8 @@ def vinterp(varIN, Lfull, Llev, type_int="log", reverse_input=False, with ``A = log(Llev/pn + 1) / log(pn/pn + 1)`` in "log" mode or ``A = (zlev-zn + 1) / (zn-zn + 1)`` in "lin" mode - """ + Nlev = len(np.atleast_1d(Llev)) if Nlev == 1: # Special case where only 1 layer is requested @@ -639,6 +626,7 @@ def vinterp(varIN, Lfull, Llev, type_int="log", reverse_input=False, + (1-alpha) * varIN.flatten()[nindexp1]) return np.reshape(varOUT, dimsOUT) + def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", modulo=None): """ @@ -646,40 +634,32 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", :param var_IN: Variable on a REGULAR grid (e.g., ``[lev, lat, lon]`` or ``[time, lev, lat, lon]``) - :type var_IN: ND array - + :type var_IN: ND array :param x: Original position array (e.g., ``time``) - :type x: 1D array - + :type x: 1D array :param xi: Target array to interpolate the array on - :type xi: 1D array - + :type xi: 1D array :param axis: Position of the interpolation axis (e.g., ``0`` for a temporal interpolation on ``[time, lev, lat, lon]``) - :type axis: int - + :type axis: int :param reverse_input: Reverse input arrays (e.g., if ``zfull(0)``= 120 km, ``zfull(N)``= 0 km, which is typical) - :type reverse_input: bool - + :type reverse_input: bool :param type_int: "log" for logarithmic (typically pressure), "lin" for linear - :type type_int: str - + :type type_int: str :param modulo: For "lin" interpolation only, use cyclic input (e.g., when using ``modulo = 24`` for time of day, 23.5 and 00 am are considered 30 min apart, not 23.5 hr apart) - :type modulo: float - + :type modulo: float :return: ``VAR_OUT`` interpolated data on the requested axis - .. note:: + .. note:: This routine is similar but simpler than the vertical interpolation ``vinterp()`` as the interpolation axis is assumed to be fully defined by a 1D array such as ``time``, ``pstd`` or ``zstd`` rather than 3D arrays like ``pfull`` or ``zfull``. - For lon/lat interpolation, consider using ``interp_KDTree()``. Calculation:: @@ -687,8 +667,8 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", X_OUT = Xn*A + (1-A)*Xn + 1 with ``A = log(xi/xn + 1) / log(xn/xn + 1)`` in "log" mode or ``A = (xi-xn + 1)/(xn-xn + 1)`` in "lin" mode - """ + # Convert list to numpy array as needed if type(var_IN) != np.ndarray: var_IN = np.array(var_IN) @@ -704,18 +684,21 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", dimsIN = var_IN.shape dimsOUT = tuple(np.append(len(xi), dimsIN[1:])) var_OUT = np.zeros(dimsOUT) - + for k in range(0, len(index)): n = index[k] np1 = n+1 - # Treatment of edge cases where the interpolated value is outside the domain - # (i.e. ``n`` is the last element and ``n+1`` does not exist) + # Treatment of edge cases where the interpolated value is + # outside the domain (i.e. ``n`` is the last element and ``n+1`` + # does not exist) if np1 >= len(x): if modulo is not None: - # If looping around (e.g., longitude, time of day...) replace ``n+1`` by the first element + # If looping around (e.g., longitude, time of day...) + # replace ``n+1`` by the first element np1 = 0 else: - # Sets the interpolated value to NaN in ``xi`` as last value as ``x[n] - x[np1] = 0`` + # Sets the interpolated value to NaN in ``xi`` as last + # value as ``x[n] - x[np1] = 0`` np1 -= 1 if n == -1 and modulo is None: @@ -723,8 +706,9 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", # Also set ``n = n+1`` (which results in NaN) if ``n = -1`` # (i.e., if the requested value is samller than first # element array) and the values are NOT cyclic. - - if type_int == "log": # Add error handling to avoid logarithm and division issues + + if type_int == "log": + # Add error handling to avoid logarithm and division issues if x[np1] <= 0 or xi[k] <= 0 or x[n] <= 0 or x[n] == x[np1]: alpha = 0 # Default to 0 if we can't compute logarithm var_OUT[k, :] = var_IN[np1, ...] # Use nearest value @@ -738,7 +722,8 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", var_OUT[k, :] = var_IN[np1, ...] elif type_int == "lin": if modulo is None: - if x[n] == x[np1]: # Avoid division by zero + if x[n] == x[np1]: + # Avoid division by zero alpha = 0 var_OUT[k, :] = var_IN[np1, ...] else: @@ -747,7 +732,8 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", else: # Handle modulo case with similar error checking denom = np.mod(x[n]-x[np1] + modulo, modulo) - if denom == 0: # Avoid division by zero + if denom == 0: + # Avoid division by zero alpha = 0 var_OUT[k, :] = var_IN[np1, ...] else: @@ -756,25 +742,24 @@ def axis_interp(var_IN, x, xi, axis, reverse_input=False, type_int="lin", return np.moveaxis(var_OUT, 0, axis) + def layers_mid_point_to_boundary(pfull, sfc_val): """ A general description for the layer boundaries is:: p_half = ps*bk + pk - This routine converts the coordinate of the layer MIDPOINTS, - ``p_full`` or ``bk``, into the coordinate of the layer BOUNDARIES + This routine converts the coordinate of the layer MIDPOINTS, + ``p_full`` or ``bk``, into the coordinate of the layer BOUNDARIES ``p_half``. The surface value must be provided. :param p_full: Pressure/sigma values for the layer MIDPOINTS, INCREASING with ``N`` (e.g., [0.01 -> 720] or [0.001 -> 1]) - :type p_full: 1D array - + :type p_full: 1D array :param sfc_val: The surface value for the lowest layer's boundary ``p_half[N]`` (e.g., ``sfc_val`` = 720 Pa or ``sfc_val`` = 1 in sigma coordinates) - :type sfc_val: float - + :type sfc_val: float :return: ``p_half`` the pressure at the layer boundaries (size = ``N+1``) @@ -788,7 +773,7 @@ def layers_mid_point_to_boundary(pfull, sfc_val): ---Nk-1--- -------- p_full --- Nk --- SFC ======== p_half / / / / / - + We have:: pfull[N] = ((phalf[N]-phalf[N-1]) / np.log(phalf[N]/phalf[N-1])) @@ -796,21 +781,20 @@ def layers_mid_point_to_boundary(pfull, sfc_val): = phalf[N] - pfull[N] log(phalf[N]) We want to solve for ``phalf[N-1] = X``:: - + v v v X - pfull[N] log(X) = B ``=> X= -pfull[N] W{-exp(-B/pfull[N])/pfull[N]}`` - + with ``B = phalf[N] - pfull[N] log(phalf[N])`` (known at N) and - + ``W`` is the product-log (Lambert) function. This was tested on an L30 simulation: The values of ``phalf`` are reconstructed from ``pfull`` with a max error of: - + ``100*(phalf - phalf_reconstruct)/phalf < 0.4%`` at the top. - """ def lambertW_approx(x): @@ -835,28 +819,24 @@ def lambertW_approx(x): / pfull[i-1])) return phalf + def polar2XYZ(lon, lat, alt, Re=3400*10**3): """ Spherical to cartesian coordinate transformation. :param lon: Longitude in radians - :type lon: ND array - + :type lon: ND array :param lat: Latitude in radians - :type lat: ND array - + :type lat: ND array :param alt: Altitude [m] - :type alt: ND array - + :type alt: ND array :param Re: Planetary radius [m], defaults to 3400*10^3 - :type Re: float - + :type Re: float :return: ``X``, ``Y``, ``Z`` in cartesian coordinates [m] .. note:: This is a classic polar coordinate system with ``colatitude = pi/2 - lat`` where ``cos(colat) = sin(lat)`` - """ lon = np.array(lon) lat = np.array(lat) @@ -869,6 +849,7 @@ def polar2XYZ(lon, lat, alt, Re=3400*10**3): Z = R * np.sin(lat)*np.ones_like(lon) return X, Y, Z + def interp_KDTree(var_IN, lat_IN, lon_IN, lat_OUT, lon_OUT, N_nearest=10): """ Inverse distance-weighted interpolation using nearest neighboor for @@ -877,33 +858,26 @@ def interp_KDTree(var_IN, lat_IN, lon_IN, lat_OUT, lon_OUT, N_nearest=10): :param var_IN: ND variable to regrid (e.g., ``[lev, lat, lon]``, ``[time, lev, lat, lon]`` with ``[lat, lon]`` dimensions LAST [°]) - :type var_IN: ND array - + :type var_IN: ND array :param lat_IN: latitude [°] (``LAT[y, x]`` array for irregular grids) - :type lat_IN: 1D or 2D array - + :type lat_IN: 1D or 2D array :param lon_IN: latitude [°] (``LAT[y, x]`` array for irregular grids) - :type lon_IN: 1D or 2D array - + :type lon_IN: 1D or 2D array :param lat_OUT: latitude [°] for the TARGET grid structure (or ``LAT1[y,x]`` for irregular grids) - :type lat_OUT: 1D or 2D array - + :type lat_OUT: 1D or 2D array :param lon_OUT: longitude [°] for the TARGET grid structure (or ``LON1[y,x]`` for irregular grids) - :type lon_OUT: 1D or 2D array - + :type lon_OUT: 1D or 2D array :param N_nearest: number of nearest neighbours for the search - :type N_nearest: int - + :type N_nearest: int :return: ``VAR_OUT`` interpolated data on the target grid .. note:: This implementation is much FASTER than ``griddata`` and it supports unstructured grids like an MGCM tile. - The nearest neighbour interpolation is only done on the lon/lat axis (not level). Although this interpolation works well on the 3D field [x, y, z], this is typically not what is expected. In @@ -911,12 +885,7 @@ def interp_KDTree(var_IN, lat_IN, lon_IN, lat_OUT, lon_OUT, N_nearest=10): on the target grid are 100's of km away while the closest points in the vertical are a few 10's -100's meter in the PBL. This would result in excessive weighting in the vertical. - """ - - # TODO Import called each time. Should be moved out of the routine - # if ``scipy`` is a requirement for the pipeline - from scipy.spatial import cKDTree dimsIN = var_IN.shape nlon_IN = dimsIN[-1] @@ -964,23 +933,21 @@ def interp_KDTree(var_IN, lat_IN, lon_IN, lat_OUT, lon_OUT, N_nearest=10): / np.sum(w, axis = 1)) return var_OUT.reshape(dims_OUT) + def cart_to_azimut_TR(u, v, mode="from"): """ Convert cartesian coordinates or wind vectors to radians using azimuth angle. :param x: the cartesian coordinate - :type x: 1D array - + :type x: 1D array :param y: the cartesian coordinate - :type y: 1D array - + :type y: 1D array :param mode: "to" for the direction that the vector is pointing, "from" for the direction from which the vector is coming - :type mode: str - + :type mode: str :return: ``Theta`` [°] and ``R`` the polar coordinates - """ + if mode == "from": cst = 180 if mode == "to": @@ -998,24 +965,20 @@ def sfc_area_deg(lon1, lon2, lat1, lat2, R=3390000.): \____\lat1 lon1 lon2 :param lon1: longitude from set 1 [°] - :type lon1: float - + :type lon1: float :param lon2: longitude from set 2 [°] - :type lon2: float - + :type lon2: float :param lat1: latitude from set 1 [°] - :type lat1: float - + :type lat1: float :param lat2: longitude from set 2 [°] - :type lat2: float - + :type lat2: float :param R: planetary radius [m] - :type R: int + :type R: int .. note:: qLon and Lat define the corners of the area not the grid cell center. - """ + lat1 *= np.pi/180 lat2 *= np.pi/180 lon1 *= np.pi/180 @@ -1024,6 +987,7 @@ def sfc_area_deg(lon1, lon2, lat1, lat2, R=3390000.): * np.abs(lon1 - lon2) * np.abs(np.sin(lat1) - np.sin(lat2))) + def area_meridional_cells_deg(lat_c, dlon, dlat, normalize=False, R=3390000.): """ Return area of invidual cells for a meridional band of thickness @@ -1040,24 +1004,19 @@ def area_meridional_cells_deg(lat_c, dlon, dlat, normalize=False, R=3390000.): dlon :param lat_c: latitude of cell center [°] - :type lat_c: float - + :type lat_c: float :param dlon: cell angular width [°] - :type dlon: float - + :type dlon: float :param dlat: cell angular height [°] - :type dlat: float - + :type dlat: float :param R: planetary radius [m] - :type R: float - + :type R: float :param normalize: if True, the sum of the output elements = 1 - :type normalize: bool - + :type normalize: bool :return: ``S`` areas of the cells, same size as ``lat_c`` in [m2] or normalized by the total area - """ + # Initialize area_tot = 1. # Compute total area in a longitude band extending from @@ -1071,22 +1030,13 @@ def area_meridional_cells_deg(lat_c, dlon, dlat, normalize=False, R=3390000.): dlat *= np.pi/180 return (2. * R**2 * dlon * np.cos(lat_c) * np.sin(dlat/2.) / area_tot) + def area_weights_deg(var_shape, lat_c, axis = -2): """ Return weights for averaging the variable. - :param var_shape: variable shape - :type var_shape: tuple - - :param lat_c: latitude of cell centers [°] - :type lat_c: float - - :param axis: position of the latitude axis for 2D and higher - dimensional arrays. The default is the SECOND TO LAST dimension - :type axis: int - Expected dimensions are: - + [lat] ``axis`` not needed [lat, lon] ``axis = -2`` or ``axis = 0`` [time, lat, lon] ``axis = -2`` or ``axis = 1`` @@ -1098,13 +1048,20 @@ def area_weights_deg(var_shape, lat_c, axis = -2): may be truncated on either end (e.g., ``lat = [-20 ..., 0... 50]``) but must be continuous. + :param var_shape: variable shape + :type var_shape: tuple + :param lat_c: latitude of cell centers [°] + :type lat_c: float + :param axis: position of the latitude axis for 2D and higher + dimensional arrays. The default is the SECOND TO LAST dimension + :type axis: int :return: ``W`` weights for the variable ready for standard averaging as ``np.mean(var*W)`` [condensed form] or ``np.average(var, weights=W)`` [expanded form] .. note:: Given a variable var: - + ``var = [v1, v2, ...vn]`` The regular average is @@ -1120,13 +1077,13 @@ def area_weights_deg(var_shape, lat_c, axis = -2): ``W = [w1, w2, ... , wn]*N / (w1 + w2 + ...wn)`` Therfore taking a regular average of (``var*W``) with - ``np.mean(var*W)`` or ``np.average(var, weights=W)`` - + ``np.mean(var*W)`` or ``np.average(var, weights=W)`` + returns the weighted average of the variable. Use ``np.average(var, weights=W, axis = X)`` to average over a specific axis. - """ + if (len(np.atleast_1d(lat_c)) == 1 or len(np.atleast_1d(var_shape)) == 1): # If variable or lat is a scalar, do nothing @@ -1155,32 +1112,13 @@ def area_weights_deg(var_shape, lat_c, axis = -2): W = A.reshape(reshape_shape)*len(lat_c) return W*np.ones(var_shape) + def areo_avg(VAR, areo, Ls_target, Ls_angle, symmetric=True): """ Return a value average over a central solar longitude - :param VAR: a variable with ``time`` in the 1st dimension - :type VAR: ND array - - :param areo: solar longitude of the input variable (0-720) - :type areo: 1D array - - :param Ls_target: central solar longitude of interest - :type Ls_target: float - - :param Ls_angle: requested window angle centered at ``Ls_target`` - :type Ls_angle: float - - :param symmetric: If ``True`` and the requested window is out of range, - ``Ls_angle`` is reduced. If False, the time average is performed - on the data available - :type symmetric: bool (defaults to True) - - :return: the variable averaged over solar longitudes - ``Ls_target-Ls_angle/2`` to ``Ls_target+Ls_angle/2`` - EX:: - + ``Ls_target = 90.`` ``Ls_angle = 10.`` @@ -1193,10 +1131,25 @@ def areo_avg(VAR, areo, Ls_target, Ls_angle, symmetric=True): If ``symmetric = False`` and the input data range is Ls = 88-100° then ``88 < Ls_target < 95`` (7°, assymetric) + :param VAR: a variable with ``time`` in the 1st dimension + :type VAR: ND array + :param areo: solar longitude of the input variable (0-720) + :type areo: 1D array + :param Ls_target: central solar longitude of interest + :type Ls_target: float + :param Ls_angle: requested window angle centered at ``Ls_target`` + :type Ls_angle: float + :param symmetric: If ``True`` and the requested window is out of range, + ``Ls_angle`` is reduced. If False, the time average is performed + on the data available + :type symmetric: bool (defaults to True) + :return: the variable averaged over solar longitudes + ``Ls_target-Ls_angle/2`` to ``Ls_target+Ls_angle/2`` + .. note:: The routine can bin data from muliples Mars years - """ + # Take the modulo of solar longitude areo = np.mod(areo, 360) # All dimensions but time @@ -1265,6 +1218,7 @@ def areo_avg(VAR, areo, Ls_target, Ls_angle, symmetric=True): VAR_avg /= count return VAR_avg.reshape(shape_out) + def mass_stream(v_avg, lat, level, type="pstd", psfc=700, H=8000., factor=1.e-8): """ @@ -1279,35 +1233,28 @@ def mass_stream(v_avg, lat, level, type="pstd", psfc=700, H=8000., :param v_avg: zonal wind [m/s] with ``lev`` dimension FIRST and ``lat`` dimension SECOND (e.g., ``[pstd, lat]``, ``[pstd, lat, lon]`` or ``[pstd, lat, lon, time]``) - :type v_avg: ND array - + :type v_avg: ND array :param lat: latitudes [°] - :type lat: 1D array - + :type lat: 1D array :param level: interpolated layers [Pa] or [m] - :type level: 1D array - + :type level: 1D array :param type: interpolation type (``pstd``, ``zstd`` or ``zagl``) - :type type: str - + :type type: str :param psfc: reference surface pressure [Pa] - :type psfc: float - + :type psfc: float :param H: reference scale height [m] when pressures are used - :type H: float - + :type H: float :param factor: normalize the mass stream function by a factor, use ``factor = 1`` for [kg/s] - :type factor: int - + :type factor: int :return: ``MSF`` the meridional mass stream function (in ``factor * [kg/s]``) - .. note:: + .. note:: This routine allows the time and zonal averages to be computed before OR after the MSF calculation. - .. note:: + .. note:: The expressions for MSF use log(pressure) Z coordinates, which integrate better numerically. @@ -1327,8 +1274,8 @@ def mass_stream(v_avg, lat, level, type="pstd", psfc=700, H=8000., ⌠ .g. ⌡ f(z)dz = (Zn-Zn-1){f(Zn) + f(Zn-1)}/2 n-1 - """ + g = 3.72 # m/s2 a = 3400*1000 # m nlev = len(level) @@ -1390,47 +1337,41 @@ def mass_stream(v_avg, lat, level, type="pstd", psfc=700, H=8000., MSF = np.ma.array(MSF, mask = mask) return MSF.reshape(shape_out) + def vw_from_MSF(msf, lat, lev, ztype="pstd", norm=True, psfc=700, H=8000.): """ Return the V and W components of the circulation from the mass stream function. - :param msf: the mass stream function with ``lev`` SECOND TO - LAST and the ``lat`` dimension LAST (e.g., ``[lev, lat]``, - ``[time, lev, lat]``, ``[time, lon, lev, lat]``) - :type msf: ND array - - :param lat: latitude [°] - :type lat: 1D array - - :param lev: level [Pa] or [m] (``pstd``, ``zagl``, ``zstd``) - :type lev: 1D array - - :param ztype: Use ``pstd`` for pressure so vertical - differentation is done in log space. - :type ztype: str - - :param norm: if True, normalize ``lat`` and ``lev`` before - differentiating to avoid having to rescale manually the - vectors in quiver plots - :type norm: bool - - :param psfc: surface pressure for pseudo-height when - ``ztype = pstd`` - :type psfc: float - - :param H: scale height for pseudo-height when ``ztype = pstd`` - :type H: float - + :param msf: the mass stream function with ``lev`` SECOND TO + LAST and the ``lat`` dimension LAST (e.g., ``[lev, lat]``, + ``[time, lev, lat]``, ``[time, lon, lev, lat]``) + :type msf: ND array + :param lat: latitude [°] + :type lat: 1D array + :param lev: level [Pa] or [m] (``pstd``, ``zagl``, ``zstd``) + :type lev: 1D array + :param ztype: Use ``pstd`` for pressure so vertical + differentation is done in log space. + :type ztype: str + :param norm: if True, normalize ``lat`` and ``lev`` before + differentiating to avoid having to rescale manually the + vectors in quiver plots + :type norm: bool + :param psfc: surface pressure for pseudo-height when + ``ztype = pstd`` + :type psfc: float + :param H: scale height for pseudo-height when ``ztype = pstd`` + :type H: float :return: the meditional and altitude components of the mass stream function for plotting as a quiver or streamlines. - .. note:: + .. note:: The components are: ``[v]= g/(2 pi cos(lat)) dphi/dz`` ``[w]= -g/(2 pi cos(lat)) dphi/dlat`` - """ + g = 3.72 # m/s2 lat = lat * np.pi/180 @@ -1468,66 +1409,62 @@ def vw_from_MSF(msf, lat, lev, ztype="pstd", norm=True, psfc=700, H=8000.): * dvar_dh(msf.transpose(T_latIN), xx).transpose(T_latOUT)) return V, W + def alt_KM(press, scale_height_KM=8., reference_press=610.): """ Gives the approximate altitude [km] for a given pressure :param press: the pressure [Pa] - :type press: 1D array - + :type press: 1D array :param scale_height_KM: scale height [km] (default is 8 km, an isothermal at 155K) - :type scale_height_KM: float - + :type scale_height_KM: float :param reference_press: reference surface pressure [Pa] (default is 610 Pa) - :type reference_press: float - + :type reference_press: float :return: ``z_KM`` the equivalent altitude for that pressure [km] - .. note:: + .. note:: Scale height is ``H = rT/g`` - """ + # Pressure -> altitude [km] return (-scale_height_KM * np.log(press/reference_press)) + def press_pa(alt_KM, scale_height_KM=8., reference_press=610.): """ Gives the approximate altitude [km] for a given pressure :param alt_KM: the altitude [km] - :type alt_KM: 1D array - + :type alt_KM: 1D array :param scale_height_KM: scale height [km] (default is 8 km, an isothermal at 155K) - :type scale_height_KM: float - + :type scale_height_KM: float :param reference_press: reference surface pressure [Pa] (default is 610 Pa) - :type reference_press: float - + :type reference_press: float :return: ``press_pa`` the equivalent pressure at that altitude [Pa] - .. note:: + .. note:: Scale height is ``H = rT/g`` - """ + return (reference_press * np.exp(-alt_KM/scale_height_KM)) + def lon180_to_360(lon): """ Transform a float or an array from the -180/180 coordinate system to 0-360 :param lon: longitudes in the -180/180 coordinate system - :type lon: float, 1D array, or 2D array - + :type lon: float, 1D array, or 2D array :return: the equivalent longitudes in the 0-360 coordinate system - """ + lon = np.array(lon) - + if len(np.atleast_1d(lon)) == 1: # ``lon`` is a float if lon < 0: @@ -1539,18 +1476,18 @@ def lon180_to_360(lon): lon = np.append(lon[lon <= 180], lon[lon > 180]) return lon + def lon360_to_180(lon): - lon = np.array(lon) """ Transform a float or an array from the 0-360 coordinate system to -180/180. :param lon: longitudes in the 0-360 coordinate system - :type lon: float, 1D array, or 2D array - + :type lon: float, 1D array, or 2D array :return: the equivalent longitudes in the -180/180 coordinate system - """ + + lon = np.array(lon) if len(np.atleast_1d(lon)) == 1: # ``lon`` is a float if lon > 180: @@ -1562,23 +1499,22 @@ def lon360_to_180(lon): lon = np.append(lon[lon < 0], lon[lon >= 0]) return lon + def shiftgrid_360_to_180(lon, data): """ This function shifts ND data from a 0-360 to a -180/180 grid. :param lon: longitudes in the 0-360 coordinate system - :type lon: 1D array - + :type lon: 1D array :param data: variable with ``lon`` in the last dimension - :type data: ND array - + :type data: ND array :return: shifted data - .. note:: + .. note:: Use ``np.ma.hstack`` instead of ``np.hstack`` to keep the masked array properties - """ + lon = np.array(lon) # convert to +/- 180 lon[lon > 180] -= 360. @@ -1586,19 +1522,18 @@ def shiftgrid_360_to_180(lon, data): data = np.concatenate((data[..., lon < 0], data[..., lon >= 0]), axis = -1) return data + def shiftgrid_180_to_360(lon, data): """ This function shifts ND data from a -180/180 to a 0-360 grid. :param lon: longitudes in the 0-360 coordinate system - :type lon: 1D array - + :type lon: 1D array :param data: variable with ``lon`` in the last dimension - :type data: ND array - + :type data: ND array :return: shifted data - """ + lon = np.array(lon) # convert to 0-360 lon[lon < 0] += 360. @@ -1607,20 +1542,19 @@ def shiftgrid_180_to_360(lon, data): axis = -1) return data + def second_hhmmss(seconds, lon_180=0.): """ Given the time [sec], return local true solar time at a certain longitude. :param seconds: the time [sec] - :type seconds: float - + :type seconds: float :param lon_180: the longitude in -180/180 coordinate - :type lon_180: float - + :type lon_180: float :return: the local time [float] or a tuple (hours, minutes, seconds) - """ + hours = seconds // (60*60) seconds %= (60 * 60) minutes = seconds // 60 @@ -1629,41 +1563,39 @@ def second_hhmmss(seconds, lon_180=0.): hours = np.mod(hours + lon_180/15., 24) return (np.int32(hours), np.int32(minutes), np.int32(seconds)) + def sol_hhmmss(time_sol, lon_180=0.): """ Given the time in days, return return local true solar time at a certain longitude. :param time_sol: the time in sols - :type seconds: float - + :type seconds: float :param lon_180: the longitude in -180/180 coordinate - :type lon_180: float - + :type lon_180: float :return: the local time [float] or a tuple (hours, minutes, seconds) - """ + return second_hhmmss(time_sol * 86400., lon_180) + def UT_LTtxt(UT_sol, lon_180=0., roundmin=None): """ Returns the time in HH:MM:SS at a certain longitude. :param time_sol: the time in sols - :type time_sol: float - + :type time_sol: float :param lon_180: the center longitude in -180/180 coordinates. Increments by 1hr every 15° - :type lon_180: float - + :type lon_180: float :param roundmin: round to the nearest X minute. Typical values are ``roundmin = 1, 15, 60`` - :type roundmin: int + :type roundmin: int - .. note:: + .. note:: If ``roundmin`` is requested, seconds are not shown - """ + def round2num(number, interval): # Internal Function to round a number to the closest range. # e.g., ``round2num(26, 5) = 25``, ``round2num(28, 5) = 30`` @@ -1680,25 +1612,12 @@ def round2num(number, interval): else: return (f"{hh:02}:{mm:02}:{ss:02}") + def dvar_dh(arr, h=None): """ Differentiate an array ``A[dim1, dim2, dim3...]`` w.r.t ``h``. The differentiated dimension must be the first dimension. - If ``h`` is 1D, then ``h``and ``dim1`` must have the same length - - If ``h`` is 2D, 3D or 4D, then ``arr`` and ``h`` must have the - same shape - - :param arr: variable - :type arr: ND array - - :param h: the dimension (``Z``, ``P``, ``lat``, ``lon``) - :type h: str - - Returns: - d_arr: the array differentiated w.r.t ``h``, e.g., d(array)/dh - EX: Compute ``dT/dz`` where ``T[time, lev, lat, lon]`` is the temperature and ``Zkm`` is the array of level heights [km]. @@ -1709,7 +1628,19 @@ def dvar_dh(arr, h=None): dTdz = dvar_dh(t.transpose([1, 0, 2, 3]), Zkm).transpose([1, 0, 2, 3]) + + If ``h`` is 1D, then ``h``and ``dim1`` must have the same length + + If ``h`` is 2D, 3D or 4D, then ``arr`` and ``h`` must have the + same shape + + :param arr: variable + :type arr: ND array + :param h: the dimension (``Z``, ``P``, ``lat``, ``lon``) + :type h: str + :return: d_arr: the array differentiated w.r.t ``h``, e.g., d(array)/dh """ + h = np.array(h) d_arr = np.zeros_like(arr) if h.any(): @@ -1741,26 +1672,27 @@ def dvar_dh(arr, h=None): d_arr[1:-1, ...] = 0.5*(arr[2:, ...] - arr[0:-2, ...]) return d_arr + def zonal_detrend(VAR): """ Substract the zonal average mean value from a field. :param VAR: variable with detrending dimension last - :type VAR: ND array - + :type VAR: ND array :return: detrented field (same size as input) - .. note:: + .. note:: ``RuntimeWarnings`` are expected if the slice contains only NaNs which is the case below the surface and above the model top in the interpolated files. This routine disables such warnings temporarily. - """ + with warnings.catch_warnings(): warnings.simplefilter("ignore", category = RuntimeWarning) return (VAR - np.nanmean(VAR, axis = -1)[..., np.newaxis]) + def get_trend_2D(VAR, LON, LAT, type_trend="wmean"): """ Extract spatial trends from the data. The output can be directly @@ -1769,24 +1701,20 @@ def get_trend_2D(VAR, LON, LAT, type_trend="wmean"): :param VAR: Variable for decomposition. ``lat`` is SECOND TO LAST and ``lon`` is LAST (e.g., ``[time, lat, lon]`` or ``[time, lev, lat, lon]``) - :type VAR: ND array - + :type VAR: ND array :param LON: longitude coordinates - :type LON: 2D array - + :type LON: 2D array :param LAT: latitude coordinates - :type LAT: 2D array - + :type LAT: 2D array :param type_trend: type of averaging to perform: "mean" - use a constant average over all lat/lon "wmean" - use a area-weighted average over all lat/lon "zonal" - detrend over the zonal axis only "2D" - use a 2D planar regression (not area-weighted) - :type type_trend: str - + :type type_trend: str :return: the trend, same size as ``VAR`` - """ + var_shape = np.array(VAR.shape) # Type "zonal" is the easiest as averaging is performed over 1 @@ -1818,23 +1746,21 @@ def get_trend_2D(VAR, LON, LAT, type_trend="wmean"): return None return TREND.reshape(var_shape) + def regression_2D(X, Y, VAR, order=1): """ Linear and quadratic regression on the plane. :param X: first coordinate - :type X: 2D array - + :type X: 2D array :param Y: second coordinate - :type Y: 2D array - + :type Y: 2D array :param VAR: variable of the same size as X - :type VAR: 2D array - + :type VAR: 2D array :param order: 1 (linear) or 2 (quadratic) - :type order: int + :type order: int - .. note:: + .. note:: When ``order = 1``, the equation is: ``aX + bY + C = Z``. When ``order = 2``, the equation is: ``aX^2 + 2bX*Y + cY^2 + 2dX + 2eY + f = Z`` @@ -1851,8 +1777,8 @@ def regression_2D(X, Y, VAR, order=1): The least-squares regression provides the solution that that minimizes ``||b – A x||^2`` - """ + if order == 1: A = np.array([X.flatten(), Y.flatten(), np.ones_like(X.flatten())]).T # An Equivalent notation is: @@ -1885,6 +1811,7 @@ def regression_2D(X, Y, VAR, order=1): XX, YY, XX * YY, XX**2, YY**2], P).reshape(X.shape) return Z + def daily_to_average(varIN, dt_in, nday=5, trim=True): """ Bin a variable from an ``atmos_daily`` file format to the @@ -1892,22 +1819,18 @@ def daily_to_average(varIN, dt_in, nday=5, trim=True): :param varIN: variable with ``time`` dimension first (e.g., ``ts[time, lat, lon]``) - :type varIN: ND array - + :type varIN: ND array :param dt_in: delta of time betwen timesteps in sols (e.g., ``dt_in = time[1] - time[0]``) - :type dt_in: float - + :type dt_in: float :param nday: bining period in sols, default is 5 sols - :type nday: int - + :type nday: int :param trim: whether to discard any leftover data at the end of file before binning - :type trim: bool - + :type trim: bool :return: the variable bin over ``nday`` - .. note:: + .. note:: If ``varIN[time, lat, lon]`` from ``atmos_daily`` is ``[160, 48, 96]`` and has 4 timesteps per day (every 6 hours), then the resulting variable for ``nday = 5`` is @@ -1919,17 +1842,17 @@ def daily_to_average(varIN, dt_in, nday=5, trim=True): time is 133 and last 3 sols the are discarded. If ``trim = False``, the time is 134 and last bin contains only 3 sols of data. - """ + vshape_in = varIN.shape # 0 is the time dimension Nin = vshape_in[0] - + # Add safety check for dt_in if np.isclose(dt_in, 0.0): print("Error: Time difference dt_in is zero or very close to zero.") return None - + iperday = int(np.round(1 / dt_in)) combinedN = int(iperday * nday) N_even = Nin // combinedN @@ -1955,6 +1878,7 @@ def daily_to_average(varIN, dt_in, nday=5, trim=True): axis = 1) return varOUT + def daily_to_diurn(varIN, time_in): """ Bin a variable from an ``atmos_daily`` file into the @@ -1962,12 +1886,10 @@ def daily_to_diurn(varIN, time_in): :param varIN: variable with time dimension first (e.g., ``[time, lat, lon]``) - :type varIN: ND array - + :type varIN: ND array :param time_in: time array in sols. Only the first N elements are actually required if saving memory is important - :type time_in: ND array - + :type time_in: ND array :return: the variable binned in the ``atmos_diurn`` format (``[time, time_of_day, lat, lon]``) and the time of day array [hr] @@ -1983,15 +1905,15 @@ def daily_to_diurn(varIN, time_in): Since the time dimension is first, the output variables may be passed to the ``daily_to_average()`` function for further binning. - """ + dt_in = time_in[1] - time_in[0] - + # Add safety check for dt_in if np.isclose(dt_in, 0.0): print("Error: Time difference dt_in is zero or very close to zero.") return None - + iperday = int(np.round(1/dt_in)) vshape_in = varIN.shape vreshape = np.append([-1, iperday], vshape_in[1:]).astype(int) @@ -2011,19 +1933,22 @@ def daily_to_diurn(varIN, time_in): varOUT = varOUT[:, i_sort, ...] return varOUT + # ====================================================================== # Vertical Grid Utilities # ====================================================================== + def gauss_profile(x, alpha, x0=0.): """ Return Gaussian line shape at x. This can be used to generate a bell-shaped mountain. - """ + return (np.sqrt(np.log(2) / np.pi) / alpha * np.exp(-((x-x0) / alpha)**2 * np.log(2))) + def compute_uneven_sigma(num_levels, N_scale_heights, surf_res, exponent, zero_top): """ @@ -2031,26 +1956,21 @@ def compute_uneven_sigma(num_levels, N_scale_heights, surf_res, and an exponent :param num_levels: the number of levels - :type num_levels: float - + :type num_levels: float :param N_scale_heights: the number of scale heights to the top of the model (e.g., ``N_scale_heights`` = 12.5 ~102 km assuming an 8 km scale height) - :type N_scale_heights: float - + :type N_scale_heights: float :param surf_res: the resolution at the surface - :type surf_res: float - + :type surf_res: float :param exponent: an exponent to increase the thickness of the levels - :type exponent: float - + :type exponent: float :param zero_top: if True, force the top pressure boundary (in N = 0) to 0 Pa - :type zero_top: bool - + :type zero_top: bool :return: an array of sigma layers - """ + b = np.zeros(int(num_levels)+1) for k in range(0, num_levels): # zeta decreases with k @@ -2063,31 +1983,29 @@ def compute_uneven_sigma(num_levels, N_scale_heights, surf_res, b[0] = 0.0 return b + def transition(pfull, p_sigma=0.1, p_press=0.05): """ Return the transition factor to construct ``ak`` and ``bk`` + In the MGCM code, the full pressures are computed from:: + + del(phalf) + pfull = ----------------------------- + log(phalf(k+1/2)/phalf(k-1/2)) + :param pfull: the pressure [Pa] - :type pfull: 1D array - + :type pfull: 1D array :param p_sigma: the pressure level where the vertical grid starts transitioning from sigma to pressure - :type p_sigma: float - + :type p_sigma: float :param p_press: the pressure level above which the vertical grid is pure (constant) pressure - :type p_press: float - + :type p_press: float :return: the transition factor. = 1 for pure sigma, = 0 for pure pressure and =0-1 for the transition - - In the MGCM code, the full pressures are computed from:: - - del(phalf) - pfull = ----------------------------- - log(phalf(k+1/2)/phalf(k-1/2)) - """ + t = np.zeros_like(pfull) for k in range(0, len(pfull)): if(pfull[k] <= p_press): @@ -2100,22 +2018,20 @@ def transition(pfull, p_sigma=0.1, p_press=0.05): t[k] = (np.sin(0.5 * np.pi * x/xx))**2 return t + def swinbank(plev, psfc, ptrans=1.): """ Compute ``ak`` and ``bk`` values with a transition based on Swinbank :param plev: the pressure levels [Pa] - :type plev: 1D array - + :type plev: 1D array :param psfc: the surface pressure [Pa] - :type psfc: 1D array - + :type psfc: 1D array :param ptrans: the transition pressure [Pa] - :type ptrans: 1D array - + :type ptrans: 1D array :return: the coefficients for the new layers - """ + # ``ks`` = number of pure pressure levels ktrans = np.argmin(np.abs(plev - ptrans)) km = len(plev)-1 @@ -2156,6 +2072,7 @@ def swinbank(plev, psfc, ptrans=1.): # ``ks`` would be used for fortran indexing in ``fv_eta.f90`` return aknew, bknew, ks + def polar_warming(T, lat, outside_range=np.nan): """ Return the polar warming, following McDunn et al. 2013: @@ -2164,23 +2081,20 @@ def polar_warming(T, lat, outside_range=np.nan): :param T: temperature with the lat dimension FIRST (transpose as needed) - :type T: ND array - + :type T: ND array :param lat: latitude array - :type lat: 1D array - + :type lat: 1D array :param outside_range: values to set the polar warming to when outside pf the range. Default = NaN but 0 may be desirable. - :type outside_range: float - + :type outside_range: float :return: The polar warming [K] .. note:: ``polar_warming()`` concatenates the results from both hemispheres obtained from the nested function ``PW_half_hemisphere()`` - """ + def PW_half_hemisphere(T_half, lat_half, outside_range=np.nan): # Easy case, T is a 1D on the latitude direction only @@ -2250,34 +2164,31 @@ def PW_half_hemisphere(T_half, lat_half, outside_range=np.nan): PW_half_hemisphere(T_NH, lat_NH, outside_range)), axis = 0)) + def time_shift_calc(var_in, lon, tod, target_times=None): """ Conversion to uniform local time. - + Mars rotates approx. 14.6° lon per Mars-hour (360° ÷ 24.6 hr) Each 14.6° shift in lon represents a 1-hour shift in local time This code uses the more precise calculation: lon_shift = 24.0 * lon / 360. - + :param var_in: variable to be shifted. Assume ``lon`` is the first dimension and ``tod`` is the last dimension - :type var_in: ND array - + :type var_in: ND array :param lon: longitude - :type lon: 1D array - + :type lon: 1D array :param tod: ``time_of_day`` index from the input file - :type tod: 1D array - + :type tod: 1D array :param target_times: local time(s) [hr] to shift to (e.g., ``"3. 15."``) - :type target_times: float (optional) - + :type target_times: float (optional) :return: the array shifted to uniform local time .. note:: If ``target_times`` is not specified, the file is interpolated on the same ``tod`` as the input - """ + if np.shape(var_in) == len(var_in): print('Need longitude and time dimensions') return @@ -2285,19 +2196,19 @@ def time_shift_calc(var_in, lon, tod, target_times=None): # Get dimensions of var_in dims_in = np.shape(var_in) n_dims_in = len(dims_in) - 1 - + # Number of longitudes in file n_lon = dims_in[0] - + # Number of timesteps per day in input n_tod_in = len(tod) if n_tod_in == 0: print('No time steps in input (time_shift_calc in FV3_utils.py)') exit() - + # Store as float for calculations but keep integer version for reshaping n_tod_in_float = float(n_tod_in) - + tod = np.squeeze(tod) # Array dimensions for output @@ -2351,7 +2262,7 @@ def time_shift_calc(var_in, lon, tod, target_times=None): # Time increment of output if target_times is None: - # Preserve original time sampling pattern (in hours) but shift + # Preserve original time sampling pattern (in hours) but shift # it for each lon so # timesteps in output = # timesteps in input dt_out = dt_in else: @@ -2371,7 +2282,7 @@ def time_shift_calc(var_in, lon, tod, target_times=None): # Core calculation # target_times[n]: The target local Mars time we want (e.g., 15:00 local Mars time) # lon_shift: The offset in Mars-hours due to Martian longitude - # The result dtt (delta time transform) tells us which time indices + # The result dtt (delta time transform) tells us which time indices # in the original Mars data to interpolate between for n in range(n_tod_out): # dtt = n*dt_out - lon_shift - target_times[0] + dt_in @@ -2385,7 +2296,7 @@ def time_shift_calc(var_in, lon, tod, target_times=None): # Ensure that local time is bounded by [0, 24] hours kk = np.where(dtt < 0.) # dtt: time in OG data corresponding to time we want - dtt[kk] = dtt[kk] + 24. + dtt[kk] = dtt[kk] + 24. # This is the index into the data aray lower_idx = np.floor(dtt/dt_in) # time step before target local time @@ -2414,7 +2325,7 @@ def time_shift_calc(var_in, lon, tod, target_times=None): frac = fraction[i, n] # Interpolate between the two time levels var_out[i, :, n] = ( - (1.-frac) * var_in[i, :, lower_idx] + (1.-frac) * var_in[i, :, lower_idx] + frac * var_in[i, :, upper_idx] ) @@ -2426,22 +2337,20 @@ def time_shift_calc(var_in, lon, tod, target_times=None): var_out = np.reshape(var_out, dims_out) return var_out + def lin_interp(X_in, X_ref, Y_ref): """ Simple linear interpolation with no dependance on scipy :param X_in: input values - :type X_in: float or array - + :type X_in: float or array :param X_ref x values - :type X_ref: array - + :type X_ref: array :param Y_ref y values - :type Y_ref: array - + :type Y_ref: array :return: y value linearly interpolated at ``X_in`` - """ + X_ref = np.array(X_ref) Y_ref = np.array(Y_ref) @@ -2468,22 +2377,21 @@ def lin_oneElement(x, X_ref, Y_ref): Y_out[i] = lin_oneElement(x_in, X_ref, Y_ref) return Y_out + def add_cyclic(data, lon): """ Add a cyclic (overlapping) point to a 2D array. Useful for azimuth and orthographic projections. :param data: variable of size ``[nlat, nlon]`` - :type data: array - + :type data: array :param lon: longitudes - :type lon: array - + :type lon: array :return: a 2D array of size ``[nlat, nlon+1]`` with last column identical to the 1st; and a 1D array of longitudes size [nlon+1] where the last element is ``lon[-1] + dlon`` - """ + # Compute increment dlon = lon[1]-lon[0] # Create new array, size ``[nlon + 1]`` @@ -2492,6 +2400,7 @@ def add_cyclic(data, lon): data_c[:, -1] = data[:, 0] return data_c, np.append(lon, lon[-1] + dlon) + def spherical_div(U, V, lon_deg, lat_deg, R=3400*1000., spacing="varying"): """ Compute the divergence of the wind fields using finite difference:: @@ -2502,32 +2411,26 @@ def spherical_div(U, V, lon_deg, lat_deg, R=3400*1000., spacing="varying"): :param U: wind field with ``lat`` SECOND TO LAST and ``lon`` as last dimensions (e.g., ``[lat, lon]`` or ``[time, lev, lat, lon``] etc.) - :type U: array - + :type U: array :param V: wind field with ``lat`` SECOND TO LAST and ``lon`` as last dimensions (e.g., ``[lat, lon]`` or ``[time, lev, lat, lon``] etc.) - :type V: array - + :type V: array :param lon_deg: longitude [°] (2D if irregularly-spaced) - :type lon_deg: 1D array - + :type lon_deg: 1D array :param lat_deg: latitude [°] (2D if irregularly-spaced) - :type lat_deg: 1D array - + :type lat_deg: 1D array :param R: planetary radius [m] - :type R: float - + :type R: float :param spacing: when ``lon`` and ``lat`` are 1D arrays, using spacing = "varying" differentiates latitude and longitude. When spacing = "regular", ``dx = lon[1]-lon[0]``, `` dy=lat[1]-lat[0]`` and the ``numpy.gradient()`` method are used - :type spacing: str (defaults to "varying") - + :type spacing: str (defaults to "varying") :return: the horizonal divergence of the wind field [m-1] - """ + lon = lon_deg * np.pi/180 lat = lat_deg * np.pi/180 var_shape = U.shape @@ -2578,6 +2481,7 @@ def spherical_div(U, V, lon_deg, lat_deg, R=3400*1000., spacing="varying"): lat.transpose(T_latIN)).transpose(T_latOUT))) return out + def spherical_curl(U, V, lon_deg, lat_deg, R=3400*1000., spacing="varying"): """ Compute the vertical component of the relative vorticity using @@ -2589,32 +2493,26 @@ def spherical_curl(U, V, lon_deg, lat_deg, R=3400*1000., spacing="varying"): :param U: wind field with ``lat`` SECOND TO LAST and ``lon`` as last dimensions (e.g., ``[lat, lon]`` or ``[time, lev, lat, lon``] etc.) - :type U: array - + :type U: array :param V: wind field with ``lat`` SECOND TO LAST and ``lon`` as last dimensions (e.g., ``[lat, lon]`` or ``[time, lev, lat, lon``] etc.) - :type V: array - + :type V: array :param lon_deg: longitude [°] (2D if irregularly-spaced) - :type lon_deg: 1D array - + :type lon_deg: 1D array :param lat_deg: latitude [°] (2D if irregularly-spaced) - :type lat_deg: 1D array - + :type lat_deg: 1D array :param R: planetary radius [m] - :type R: float - + :type R: float :param spacing: when ``lon`` and ``lat`` are 1D arrays, using spacing = "varying" differentiates latitude and longitude. When spacing = "regular", ``dx = lon[1]-lon[0]``, `` dy=lat[1]-lat[0]`` and the ``numpy.gradient()`` method are used - :type spacing: str (defaults to "varying") - + :type spacing: str (defaults to "varying") :return: the vorticity of the wind field [m-1] - """ + lon = lon_deg * np.pi/180 lat = lat_deg * np.pi/180 @@ -2665,6 +2563,7 @@ def spherical_curl(U, V, lon_deg, lat_deg, R=3400*1000., spacing="varying"): lat.transpose(T_latIN)).transpose(T_latOUT))) return out + def frontogenesis(U, V, theta, lon_deg, lat_deg, R=3400*1000., spacing="varying"): """ @@ -2678,35 +2577,28 @@ def frontogenesis(U, V, theta, lon_deg, lat_deg, R=3400*1000., :param U: wind field with ``lat`` SECOND TO LAST and ``lon`` as last dimensions (e.g., ``[lat, lon]`` or ``[time, lev, lat, lon``] etc.) - :type U: array - + :type U: array :param V: wind field with ``lat`` SECOND TO LAST and ``lon`` as last dimensions (e.g., ``[lat, lon]`` or ``[time, lev, lat, lon``] etc.) - :type V: array - + :type V: array :param theta: potential temperature [K] - :type theta: array - + :type theta: array :param lon_deg: longitude [°] (2D if irregularly-spaced) - :type lon_deg: 1D array - + :type lon_deg: 1D array :param lat_deg: latitude [°] (2D if irregularly-spaced) - :type lat_deg: 1D array - + :type lat_deg: 1D array :param R: planetary radius [m] - :type R: float - + :type R: float :param spacing: when ``lon`` and ``lat`` are 1D arrays, using spacing = "varying" differentiates latitude and longitude. When spacing = "regular", ``dx = lon[1]-lon[0]``, `` dy=lat[1]-lat[0]`` and the ``numpy.gradient()`` method are used - :type spacing: str (defaults to "varying") - + :type spacing: str (defaults to "varying") :return: the frontogenesis field [m-1] - """ + lon = lon_deg*np.pi/180 lat = lat_deg*np.pi/180 @@ -2782,6 +2674,7 @@ def frontogenesis(U, V, theta, lon_deg, lat_deg, R=3400*1000., + U*np.tan(lat_b)/R)) return out + def MGSzmax_ls_lat(ls, lat): """ Return the max altitude for the dust from "MGS scenario" from @@ -2789,14 +2682,12 @@ def MGSzmax_ls_lat(ls, lat): the Martian water cycle as inferred from a general circulation model :param ls: solar longitude [°] - :type ls: array - + :type ls: array :param lat : latitude [°] - :type lat: array - + :type lat: array :return: top altitude for the dust [km] - """ + lat = np.array(lat) * np.pi/180 ls_p = (np.array(ls)-158) * np.pi/180 @@ -2813,14 +2704,12 @@ def MGStau_ls_lat(ls, lat): the Martian water cycle as inferred from a general circulation model :param ls: solar longitude [°] - :type ls: array - + :type ls: array :param lat : latitude [°] - :type lat: array - + :type lat: array :return: top altitude for the dust [km] - """ + lat = np.array(lat) ls_p = (np.array(ls)-250) * np.pi/180 @@ -2844,21 +2733,19 @@ def MGStau_ls_lat(ls, lat): tau[lat > 0] = t_north[lat > 0] return tau + def broadcast(var_1D, shape_out, axis): """ Broadcast a 1D array based on a variable's dimensions :param var_1D: variable (e.g., ``lat`` size = 36, or ``time`` size = 133) - :type var_1D: 1D array - + :type var_1D: 1D array :param shape_out: broadcasting shape (e.g., ``temp.shape = [133, lev, 36, lon]``) - :type shape_out: list - + :type shape_out: list :return: (ND array) broadcasted variables (e.g., size = ``[1,36,1,1]`` for ``lat`` or ``[133,1,1,1]`` for ``time``) - """ # Special case where var_1D has only one element var_1D = np.atleast_1d(var_1D) @@ -2867,14 +2754,14 @@ def broadcast(var_1D, shape_out, axis): reshape_shape[axis] = len(var_1D) return var_1D.reshape(reshape_shape) + def ref_atmosphere_Mars_PTD(Zi): """ Analytical atmospheric model for Martian pressure, temperature, and density. Alex Kling, June 2021 :param Zi: input altitude [m] (must be >= 0) - :type Zi: float or 1D array - + :type Zi: float or 1D array :return: tuple of corresponding pressure [Pa], temperature [K], and density [kg/m3] floats or arrays @@ -2897,9 +2784,9 @@ def ref_atmosphere_Mars_PTD(Zi): altitudes, we provide a fit in the form of ``P = P0 exp(-az-bz^2)`` based on diurnal average of the MCD database at lat = 0, Ls = 150. - """ + # Internal Functions def alt_to_temp_quad(Z, Z0, T0, gam, a): """ @@ -2907,25 +2794,21 @@ def alt_to_temp_quad(Z, Z0, T0, gam, a): temperature in the form ``T(z) = a(z-z0)^2 + b(z-z0) + T0`` :param Z: altitude [m] - :type Z: float - + :type Z: float :param Z0: starting altitude [m] - :type Z0: float - + :type Z0: float :param T0: quadratic coefficient - :type T0: float - + :type T0: float :param gam: quadratic coefficient - :type gam: float - + :type gam: float :param a: quadratic coefficient - :type a: float - + :type a: float :return: temperature at altitude Z [K] - """ + return (T0 + gam*(Z-Z0) + a*(Z-Z0)**2) + def alt_to_press_quad(Z, Z0, P0, T0, gam, a, rgas, g): """ Return the pressure [Pa] in the troposphere as a function of @@ -2934,32 +2817,24 @@ def alt_to_press_quad(Z, Z0, P0, T0, gam, a, rgas, g): ``T(z) = T0 + gam(z-z0) + a*(z-z0)^2`` :param Z: altitude [m] - :type Z: float - + :type Z: float :param Z0: starting altitude [m] - :type Z0: float - + :type Z0: float :param P0: reference pressure at Z0[Pa] - :type P0: float - + :type P0: float :param T0: reference temperature at Z0[Pa] - :type T0: float - + :type T0: float :param gam: linear and quadratic coeff for the temperature - :type gam: float - + :type gam: float :param a: linear and quadratic coeff for the temperature - :type a: float - + :type a: float :param rgas: specific gas constant [J/kg/K] - :type rgas: float - + :type rgas: float :param rg: gravity [m/s2] - :type rg: float - + :type rg: float :return: pressure at alitude Z [Pa] - """ + delta = (4*a*T0 - gam**2) if delta >= 0: @@ -2976,6 +2851,7 @@ def alt_to_press_quad(Z, Z0, P0, T0, gam, a, rgas, g): / (((2*a*(Z-Z0)+gam) + sq_delta) * (gam-sq_delta))) **(-g / (rgas*sq_delta))) + def P_mars_120_300(Zi, Z0=120000., P0=0.00012043158397922564, p1=1.09019694e-04, p2=-3.37385416e-10): """ @@ -2990,16 +2866,19 @@ def P_mars_120_300(Zi, Z0=120000., P0=0.00012043158397922564, ``P = P0 exp(-az - bz^2 - cz^c-d^4 ...)`` :param Z: altitude [m] - :type Z: float - + :type Z: float :return: the equivalent pressure at altitude [Pa] - """ + return (P0 * np.exp(-p1*(Zi-Z0) - p2*(Zi-Z0)**2)) + def T_analytic_scalar(Zi): - # A best fit of globally averaged temperature profiles from - # various sources including: Legacy MGCM, MCS, & MCD + """ + A best fit of globally averaged temperature profiles from + various sources including: Legacy MGCM, MCS, & MCD + """ + if Zi <= 57000: return alt_to_temp_quad(Zi, Z0 = 0, T0 = 225.9, gam = -0.00213479, a = 1.44823e-08) @@ -3012,12 +2891,13 @@ def T_analytic_scalar(Zi): elif 170000 <= Zi: return 174.6 + def P_analytic_scalar(Zi): """ Analytic solution for a pressure derived from a temperature profile - """ + if Zi <= 57000: return alt_to_press_quad(Zi, Z0 = 0, P0 = 610, T0 = 225.9, gam = -0.00213479, a = 1.44823e-08, @@ -3035,6 +2915,7 @@ def P_analytic_scalar(Zi): elif 120000 <= Zi: return P_mars_120_300(Zi) + if len(np.atleast_1d(Zi)) > 1: # Vectorize array P_analytic_scalar = np.vectorize(P_analytic_scalar) @@ -3049,11 +2930,10 @@ def press_to_alt_atmosphere_Mars(Pi): analytical calculation above. :param Pi: input pressure [Pa] (<= 610 Pa) - :type Pi: float or 1D array - + :type Pi: float or 1D array :return: the corresponding altitude [m] (float or 1D array) - """ + # Internal Functions def press_to_alt_quad(P, Z0, P0, T0, gam, a, rgas, g): """ @@ -3063,32 +2943,25 @@ def press_to_alt_quad(P, Z0, P0, T0, gam, a, rgas, g): T(z) = T0 + gam(z-z0) + a*(z-z0)^2 :param P: pressure [Pa] - :type P: float - + :type P: float + :param Z0: referecence altitude [m] - :type Z0: float - + :type Z0: float :param P0: reference pressure at Z0[Pa] - :type P0: float - + :type P0: float :param T0: reference temperature at Z0[Pa] - :type T0: float - + :type T0: float :param gam: linear and quadratic coefficients for temperature - :type gam: float - + :type gam: float :param a: linear and quadratic coefficients for temperature - :type a: float - + :type a: float :param rgas: specific gas constant [J/kg/K] - :type rgas: float - + :type rgas: float :param g: gravity [m/s2] - :type g: - + :type g: :return: the corresponding altitude [m] (float or 1D array) - """ + delta = (4*a*T0 - gam**2) if delta >= 0: @@ -3108,26 +2981,27 @@ def press_to_alt_quad(P, Z0, P0, T0, gam, a, rgas, g): / (gam + sq_delta) * (P/P0)**(-rgas * sq_delta/g) - 1)) + def press_to_alt_mars_120_300(P, Z0=120000., P0=0.00012043158397922564, p1=1.09019694e-04, p2=-3.37385416e-10): """ Martian altitude as a function of pressure from 120-300 km. :param P: pressure [m] - :type P: float - + :type P: float :return: altitude [m] - """ + # ``delta > 0`` on this pressure interval delta = (p1**2 - 4*p2*np.log(P/P0)) return ((-p1 + np.sqrt(delta)) / (2*p2) + Z0) + def alt_analytic_scalar(Pi): """ Analytic solution for the altitude as a function of pressure. - """ + if Pi >= 610: return 0. elif 610 > Pi >= 1.2415639872674782: @@ -3169,32 +3043,29 @@ def alt_analytic_scalar(Pi): alt_analytic_scalar = np.vectorize(alt_analytic_scalar) return alt_analytic_scalar(Pi) -# ============================ Projections ============================= +# ============================ Projections ============================= # The projections below were implemented by Alex Kling following "An # Album of Map Projections," USGS Professional Paper 1453, (1994) # https://pubs.usgs.gov/pp/1453/report.pdf + def azimuth2cart(LAT, LON, lat0, lon0=0): """ Azimuthal equidistant projection. Converts from latitude-longitude to cartesian coordinates. :param LAT: latitudes[°] size [nlat] - :type LAT: 1D or 2D array - + :type LAT: 1D or 2D array :param LON: longitudes [°] size [nlon] - :type LON: 1D or 2D array - + :type LON: 1D or 2D array :param lat0: latitude coordinate of the pole - :type lat0: float - + :type lat0: float :param lon0: longitude coordinate of the pole - :type lon0: float - + :type lon0: float :return: the cartesian coordinates for the latitudes and longitudes - """ + # Convert to radians LAT = LAT * np.pi/180 lat0 = lat0 * np.pi/180 @@ -3209,27 +3080,24 @@ def azimuth2cart(LAT, LON, lat0, lon0=0): - np.sin(lat0) * np.cos(LAT) * np.cos(LON - lon0)) return X, Y + def ortho2cart(LAT, LON, lat0, lon0=0): """ Orthographic projection. Converts from latitude-longitude to cartesian coordinates. :param LAT: latitudes[°] size [nlat] - :type LAT: 1D or 2D array - + :type LAT: 1D or 2D array :param LON: longitudes [°] size [nlon] - :type LON: 1D or 2D array - + :type LON: 1D or 2D array :param lat0: latitude coordinate of the pole - :type lat0: float - + :type lat0: float :param lon0: longitude coordinate of the pole - :type lon0: float - + :type lon0: float :return: the cartesian coordinates for the latitudes and longitudes; and a mask (NaN array) that hides the back side of the planet - """ + # Convert to radians LAT = LAT * np.pi/180 lat0 = lat0 * np.pi/180 @@ -3247,26 +3115,23 @@ def ortho2cart(LAT, LON, lat0, lon0=0): MASK[cosc < 0] = np.nan return X, Y, MASK + def mollweide2cart(LAT, LON): """ Mollweide projection. Converts from latitude-longitude to cartesian coordinates. :param LAT: latitudes[°] size [nlat] - :type LAT: 1D or 2D array - + :type LAT: 1D or 2D array :param LON: longitudes [°] size [nlon] - :type LON: 1D or 2D array - + :type LON: 1D or 2D array :param lat0: latitude coordinate of the pole - :type lat0: float - + :type lat0: float :param lon0: longitude coordinate of the pole - :type lon0: float - + :type lon0: float :return: the cartesian coordinates for the latitudes and longitudes - """ + # Convert to radians LAT = np.array(LAT) * np.pi/180 LON = np.array(LON) * np.pi/180 @@ -3275,8 +3140,8 @@ def mollweide2cart(LAT, LON): def compute_theta(lat): """ Internal Function to compute theta. Latitude is in radians here. - """ + theta0 = lat sum = 0 running = True @@ -3295,6 +3160,7 @@ def compute_theta(lat): "number of iterations") return theta1 + if len(np.atleast_1d(LAT).shape) == 1: # Float or 1D array nlat = len(np.atleast_1d(LAT)) @@ -3316,26 +3182,23 @@ def compute_theta(lat): Y = np.sqrt(2) * np.sin(THETA) return np.squeeze(X), np.squeeze(Y) + def robin2cart(LAT, LON): """ Robinson projection. Converts from latitude-longitude to cartesian coordinates. :param LAT: latitudes[°] size [nlat] - :type LAT: 1D or 2D array - + :type LAT: 1D or 2D array :param LON: longitudes [°] size [nlon] - :type LON: 1D or 2D array - + :type LON: 1D or 2D array :param lat0: latitude coordinate of the pole - :type lat0: float - + :type lat0: float :param lon0: longitude coordinate of the pole - :type lon0: float - + :type lon0: float :return: the cartesian coordinates for the latitudes and longitudes - """ + # Convert to radians lon0 = 0. LAT = np.array(LAT) * np.pi/180 @@ -3369,23 +3232,23 @@ def robin2cart(LAT, LON): Y = 1.3523 * Y1 return X, Y + # ======================= End Projections Section ====================== + def sol2ls(jld, cumulative=False): """ Return the solar longitude (Ls) as a function of the sol number. Sol=0 is the spring equinox. :param jld: sol number after perihelion - :type jld: float or 1D array - + :type jld: float or 1D array :param cumulative: if True, result is cumulative (Ls=0-360, 360-720 etc..) - :type cumulative: bool - + :type cumulative: bool :return: the corresponding solar longitude - """ + # Constants # Year in sols year = 668.0 @@ -3400,6 +3263,7 @@ def sol2ls(jld, cumulative=False): # If ``jld`` is a scalar, reshape to a 1-element array jld = np.array(jld).astype(float).reshape(len(np.atleast_1d(jld))) + # ================================================================== # Internal Function 1: Calculate Ls 0-360 using a numerical solver # ================================================================== @@ -3407,8 +3271,8 @@ def sol2ls_mod(jld): """ Based on Tanguy's ``aerols.py``. Useful link: http://www.jgiesen.de/kepler/kepler.html - """ + # Specify orbit eccentricity ec = .093 # Specify angle of planet inclination @@ -3449,9 +3313,12 @@ def sol2ls_mod(jld): areols[areols < 0.] += 360. return areols + # ================================================================== # Internal Function 2: Calculate cumulative Ls 0-720 # ================================================================== + + def sol2ls_cumu(jld): """ Calculate cumulative Ls. Continuous solar longitudes @@ -3464,8 +3331,8 @@ def sol2ls_cumu(jld): For those edge cases where Ls is close to 359.9, the routine recalculates the Ls at a later time (say 1 sols) to check for outlier points. - """ + # Calculate cumulative Ls using ``sol2ls`` and adding 360 for # every Mars year areols, MY = [np.zeros_like(jld) for _ in range(0, 2)] @@ -3502,20 +3369,20 @@ def sol2ls_cumu(jld): else: return sol2ls_mod(jld) + def ls2sol(Ls_in): """ Ls to sol converter. :param Ls_in: solar longitudes (0-360...720) - :type Ls_in: float or 1D array - + :type Ls_in: float or 1D array :return: the corresponding sol number .. note:: This function simply uses a numerical solver on the ``sol2ls()`` function. - """ + def internal_func(Ls_in): func_int = lambda x: sol2ls(x, cumulative = True) # Distance to the target function diff --git a/amescap/Ncdf_wrapper.py b/amescap/Ncdf_wrapper.py index cac98ad..cc188d5 100644 --- a/amescap/Ncdf_wrapper.py +++ b/amescap/Ncdf_wrapper.py @@ -11,7 +11,6 @@ * ``netCDF4`` * ``os`` * ``datetime`` - """ # Load generic Python modules @@ -56,9 +55,7 @@ class Ncdf(object): :param object: _description_ :type object: _type_ - :return: netCDF file - """ def __init__(self, filename=None, description_txt="", action="w", @@ -89,20 +86,25 @@ def __init__(self, filename=None, description_txt="", action="w", self.var_dict = dict() # print(f"{filename} was created") + def close(self): self.f_Ncdf.close() print(f"{self.filename} was created") + def add_dimension(self, dimension_name, length): self.dim_dict[dimension_name] = ( self.f_Ncdf.createDimension(dimension_name, length)) + def print_dimensions(self): print(self.dim_dict.items()) + def print_variables(self): print(self.var_dict.keys()) + def add_constant(self, variable_name, value, longname_txt="", units_txt=""): if "constant" not in self.dim_dict.keys(): @@ -114,6 +116,7 @@ def add_constant(self, variable_name, value, longname_txt="", units_txt) self.var_dict[variable_name][:] = value + # Private Definitions def _def_variable(self, variable_name, dim_array, longname_txt="", units_txt=""): @@ -124,6 +127,7 @@ def _def_variable(self, variable_name, dim_array, longname_txt="", self.var_dict[variable_name].long_name = longname_txt self.var_dict[variable_name].dim_name = str(dim_array) + def _def_axis1D(self, variable_name, dim_array, longname_txt="", units_txt="", cart_txt=""): self.var_dict[variable_name] = self.f_Ncdf.createVariable(variable_name, @@ -133,6 +137,7 @@ def _def_axis1D(self, variable_name, dim_array, longname_txt="", self.var_dict[variable_name].long_name = longname_txt self.var_dict[variable_name].cartesian_axis = cart_txt + def _test_var_dimensions(self, Ncvar): all_dim_OK = True for s in Ncvar.dimensions: @@ -142,6 +147,7 @@ def _test_var_dimensions(self, Ncvar): all_dim_OK = False return all_dim_OK + def _is_cart_axis(self, Ncvar): """ The cartesian axis attribute is replicated if ``cartesian_axis`` @@ -149,8 +155,8 @@ def _is_cart_axis(self, Ncvar): exclude FV3 T-cell latitudes ``grid_xt_bnds`` and ``grid_yt_bnds``, which are of size ``(lon, bnds)`` & ``(lat, bnds)`` with dimension = 2. - """ + cart_axis = False tmp_cart = getattr(Ncvar, 'cartesian_axis', False) tmp_size = getattr(Ncvar, 'dimensions') @@ -167,8 +173,8 @@ def log_variable(self, variable_name, DATAin, dim_array, longname_txt="", Log.log_variable("sfcT", sfcT, ("time", "Nx"), "soil temperature", "K") - """ + if variable_name not in self.var_dict.keys(): self._def_variable(variable_name, dim_array, longname_txt, units_txt) @@ -177,13 +183,13 @@ def log_variable(self, variable_name, DATAin, dim_array, longname_txt="", self.var_dict[variable_name].units = units_txt self.var_dict[variable_name][:] = DATAin + def log_axis1D(self, variable_name, DATAin, dim_name, longname_txt="", units_txt="", cart_txt=""): """ EX:: Log.log_axis1D("areo", areo, "time", "degree", "T") - """ if variable_name not in self.var_dict.keys(): self._def_axis1D(variable_name, dim_name, longname_txt, units_txt, @@ -193,6 +199,7 @@ def log_axis1D(self, variable_name, DATAin, dim_name, longname_txt="", self.var_dict[variable_name].cartesian_axis = cart_txt self.var_dict[variable_name][:] = DATAin + def add_dim_with_content(self, dimension_name, DATAin, longname_txt="", units_txt="", cart_txt=''): """ @@ -206,8 +213,8 @@ def add_dim_with_content(self, dimension_name, DATAin, longname_txt="", Log.add_dim_with_content("lon", lon_array, "longitudes", "degree", "X") - """ + if dimension_name not in self.dim_dict.keys(): self.add_dimension(dimension_name, len(DATAin)) @@ -228,27 +235,29 @@ def add_dim_with_content(self, dimension_name, DATAin, longname_txt="", # When using ``f=MFDataset(fname,"r")``, ``f.variables[var]`` does # not have a ``name`` attribute but does have ``_name`` + def copy_Ncaxis_with_content(self, Ncdim_var): """ Copy a netCDF DIMENSION variable (e.g., ``Ncdim = f.variables["lon"]``). If the dimension does not exist yet, it will be created - """ + longname_txt = getattr(Ncdim_var, "long_name", Ncdim_var._name) units_txt = getattr(Ncdim_var, "units", "") cart_txt = getattr(Ncdim_var, "cartesian_axis", "") self.add_dim_with_content(Ncdim_var._name, Ncdim_var[:], longname_txt, units_txt, cart_txt) + def copy_Ncvar(self, Ncvar, swap_array=None): """ Copy a netCDF variable from another file (e.g., ``Ncvar = f.variables["ucomp"]``). All dimensions must already exist. If ``swap_array`` is provided, the original values are swapped with this array. - """ + if Ncvar._name not in self.var_dict.keys(): dim_array = Ncvar.dimensions longname_txt = getattr(Ncvar, "long_name", Ncvar._name) @@ -265,13 +274,14 @@ def copy_Ncvar(self, Ncvar, swap_array=None): print(f"***Warning***, '{Ncvar._name}' is already defined, " f"skipping it") + def copy_all_dims_from_Ncfile(self, Ncfile_in, exclude_dim=[], time_unlimited=True): """ Copy all variables, dimensions, and attributes from another netCDF file - """ + # First include dimensions all_dims = Ncfile_in.dimensions.keys() for idim in all_dims: @@ -282,6 +292,7 @@ def copy_all_dims_from_Ncfile(self, Ncfile_in, exclude_dim=[], self.add_dimension(Ncfile_in.dimensions[idim]._name, Ncfile_in.dimensions[idim].size) + def copy_all_vars_from_Ncfile(self, Ncfile_in, exclude_var=[]): # First include variables all_vars = Ncfile_in.variables.keys() @@ -295,16 +306,19 @@ def copy_all_vars_from_Ncfile(self, Ncfile_in, exclude_var=[]): else: self.copy_Ncvar(Ncfile_in.variables[ivar]) + def merge_files_from_list(self, Ncfilename_list, exclude_var=[]): Mf_IN = MFDataset(Ncfilename_list, "r") self.copy_all_dims_from_Ncfile(Mf_IN) self.copy_all_vars_from_Ncfile(Mf_IN, exclude_var = exclude_var) Mf_IN.close() + # ====================================================================== # Wrapper for creating netCDF-like objects from Legacy GCM # Fortran binary files # ====================================================================== + class Fort(object): """ A class that generates an object from a fort.11 file. The new file @@ -330,10 +344,8 @@ class Fort(object): :param object: _description_ :type object: _type_ - :return: _description_ :rtype: _type_ - """ class Fort_var(np.ndarray): @@ -353,15 +365,15 @@ class Fort_var(np.ndarray): :param np.ndarray: _description_ :type np.ndarray: _type_ - :return: _description_ :rtype: _type_ - """ + def __new__(cls, input_vals, *args, **kwargs): return np.asarray(input_vals).view(cls) + def __init__(self, input_vals, name_txt, long_name_txt, units_txt, dimensions_tuple): self.name = name_txt @@ -370,11 +382,13 @@ def __init__(self, input_vals, name_txt, long_name_txt, units_txt, self.dimensions = dimensions_tuple # End inner class + def __init__(self, filename=None, description_txt=""): self.filename = filename self.path, self.name = os.path.split(filename) print(f"Reading {filename}...") self.f = FortranFile(filename) + if len(self.name)==12: # Get output number (e.g. 11 for fort.11_0070) self.fort_type = filename[-7:-5] @@ -407,12 +421,13 @@ def __init__(self, filename=None, description_txt=""): # -1 round to nearest 10 self.fdate = ("%05i"%np.round(self.variables["time"][0], -1)) + # Public methods def write_to_fixed(self): """ Create ``fixed`` file (all static variables) - """ + Log = Ncdf(f"{self.path}/{self.fdate}.fixed.nc") # Define dimensions @@ -440,11 +455,12 @@ def write_to_fixed(self): units_txt = fort_var.units) Log.close() + def write_to_daily(self): """ Create daily file (continuous time series) - """ + Log = Ncdf(f"{self.path}/{self.fdate}.atmos_daily.nc") # Define dimensions @@ -496,11 +512,12 @@ def write_to_daily(self): units_txt = fort_var.units) Log.close() + def write_to_average(self, day_average=5): """ Create average file (e.g., N-day averages [N=5 usually]) - """ + Log = Ncdf(f"{self.path}/{self.fdate}.atmos_average.nc") # Define dimensions @@ -512,7 +529,7 @@ def write_to_average(self, day_average=5): if ivar in ["pfull", "phalf", "zgrid"]: cart_ax = "Z" fort_var = self.variables[ivar] - Log.add_dim_with_content(dimension_name = ivar, + Log.add_dim_with_content(dimension_name = ivar, DATAin = fort_var, longname_txt = fort_var.long_name, units_txt = fort_var.units, @@ -528,19 +545,19 @@ def write_to_average(self, day_average=5): time_in = self.variables["time"] time_out = daily_to_average(varIN = fort_var, dt_in = (time_in[1]-time_in[0]), - nday = day_average, + nday = day_average, trim = True) - Log.log_axis1D(variable_name = "time", + Log.log_axis1D(variable_name = "time", DATAin = time_out, - dim_name = "time", + dim_name = "time", longname_txt = time_in.long_name, - units_txt = time_in.units, + units_txt = time_in.units, cart_txt = "T") # Log static variables for ivar in ["pk", "bk"]: fort_var = self.variables[ivar] - Log.log_variable(variable_name = ivar, + Log.log_variable(variable_name = ivar, DATAin = fort_var, dim_array = fort_var.dimensions, longname_txt = fort_var.long_name, @@ -552,21 +569,22 @@ def write_to_average(self, day_average=5): fort_var = self.variables[ivar] var_out = daily_to_average(varIN = fort_var, dt_in = (time_in[1]-time_in[0]), - nday = day_average, + nday = day_average, trim = True) - Log.log_variable(variable_name = ivar, + Log.log_variable(variable_name = ivar, DATAin = var_out, dim_array = fort_var.dimensions, longname_txt = fort_var.long_name, units_txt = fort_var.units) Log.close() + def write_to_diurn(self, day_average=5): """ Create diurn file (variables organized by time of day & binned (typically 5-day bins) - """ + Log = Ncdf(f"{self.path}/{self.fdate}.atmos_diurn.nc") # Define dimensions @@ -578,7 +596,7 @@ def write_to_diurn(self, day_average=5): if ivar in ["pfull" , "phalf", "zgrid"]: cart_ax="Z" fort_var=self.variables[ivar] - Log.add_dim_with_content(dimension_name = ivar, + Log.add_dim_with_content(dimension_name = ivar, DATAin = fort_var, longname_txt = fort_var.long_name, units_txt = fort_var.units, @@ -640,11 +658,13 @@ def write_to_diurn(self, day_average=5): fort_var.long_name, fort_var.units) Log.close() + # Public method def close(self): self.f.close() print(f"{self.filename} was closed") + # Private methods def _read_Fort11_header(self): """ @@ -662,8 +682,8 @@ def _read_Fort11_header(self): These are saved as attributes (e.g., uses ``f.LAYERS`` to access the number of layers). - """ + Rec = self.f.read_record("f4", "(1, 5)i4", "S7") self.RUNNUM = Rec[0][0] self.JM = Rec[1][0, 0] @@ -690,8 +710,8 @@ def _read_Fort11_constants(self): These are saved as attributes (e.g., uses ``f.rgas`` to access the gas constant for the simulation. - """ + Rec = self.f.read_record(f"(1, {self.LM})f4", f"(1, {self.JM})f4", "f4", "f4", "f4", "f4", "f4", "f4", "f4", @@ -713,11 +733,11 @@ def _read_Fort11_constants(self): self.vinc = Rec[12][0] self.sdepth = np.array(Rec[13][0, :]) self.alicen = Rec[14][0] - self.alices = Rec[15][0] self.egoco2n = Rec[16][0] self.egoco2s = Rec[17][0] + def _read_Fort11_static(self): """ Return values from ``fort.11`` header. @@ -730,8 +750,8 @@ def _read_Fort11_static(self): These are saved as variables (e.g., uses ``f.variables["zsurf"]`` to access the topography. - """ + Rec = self.f.read_record(f"({self.IM},{self.JM})f4", f"({self.IM},{self.JM})f4", f"({self.NL},{self.IM},{self.JM})f4", @@ -751,14 +771,15 @@ def _read_Fort11_static(self): "npcflag", "Polar ice flag", "none", ("lat", "lon")) + def _create_dims(self): """ Create dimension axis from ``IM``, ``JM`` after reading the header. Also compute a vertical grid structure that includes sigma values at the layer boundaries AND midpoints for the radiation code. Total size is ``2*LM+2`` - """ + JM = self.JM # JM = 36 IM = self.IM # IM = 60 LM = self.LM @@ -791,12 +812,13 @@ def _create_dims(self): # starting with the 2nd self.zgrid = self.sdepth[1::2] # TODO check + def _ra_1D(self, new_array, name_txt): """ ``_ra`` stands for "Return array": Append single timesteps along the first (``time``) dimensions - """ + if type(new_array) != np.ndarray: new_array = np.array([new_array]) @@ -811,6 +833,7 @@ def _ra_1D(self, new_array, name_txt): # that ``np.concatenate((x,y))`` takes a tuple as argument. return np.append(self.variables[name_txt], new_array) + def _log_var(self,name_txt, long_name, unit_txt, dimensions, Rec=None, scaling=None): if Rec is None: @@ -852,6 +875,7 @@ def _log_var(self,name_txt, long_name, unit_txt, dimensions, Rec=None, self.variables[name_txt] = self.Fort_var(Rec, name_txt, long_name, unit_txt, dimensions) + def _read_Fort11_dynamic(self): """ Read variables from ``fort.11`` files that change with each @@ -882,8 +906,8 @@ def _read_Fort11_dynamic(self): write(11) surfalb write(11) dheat write(11) geot - """ + # Typically ``nsteps = 16 x 10 = 160`` nsteps = self.nperday * self.nsolfile append = False @@ -904,33 +928,42 @@ def _read_Fort11_dynamic(self): self._ra_1D(Rec[0]/24, "time"), "time", "elapsed time from the start of the run", "days since 0000-00-00 00:00:00", ("time")) + self.variables["areo"] = self.Fort_var( self._ra_1D(Rec[1].reshape([1, 1]), "areo"), "areo", "solar longitude", "degree", ("time", "scalar_axis")) # TODO "areo" monotically increasing? + self.variables["rdist"] = self.Fort_var( self._ra_1D(Rec[2], "rdist"), "rdist", "square of the Sun-Mars distance", "(AU)**2", ("time")) + self.variables["tofday"]=self.Fort_var( self._ra_1D(Rec[3], "tofday"), "npcflag", "time of day", "hours since 0000-00-00 00:00:00", ("time")) # TODO "tofday" edge or center? + self.variables["psf"] = self.Fort_var( self._ra_1D(Rec[4]*100, "psf"), "psf", "Initial global surface pressure", "Pa", ("time")) + self.variables["ptrop"] = self.Fort_var( self._ra_1D(Rec[5], "ptrop"), "ptrop", "pressure at the tropopause", "Pa", ("time")) + self.variables["tautot"]=self.Fort_var( self._ra_1D(Rec[6], "tautot"), "tautot", "Input (global) dust optical depth at the reference pressure", "none", ("time")) + self.variables["rptau"] = self.Fort_var( self._ra_1D(Rec[7]*100, "rptau"), "rptau", "reference pressure for dust optical depth", "Pa", ("time")) + self.variables["sind"] = self.Fort_var( self._ra_1D(Rec[8], "sind"), "sind", "sine of the sub-solar latitude", "none", ("time")) + self.variables["gasp"] = self.Fort_var( self._ra_1D(Rec[9]*100, "gasp"), "gasp", "global average surface pressure", "Pa", ("time")) @@ -942,31 +975,42 @@ def _read_Fort11_dynamic(self): self.variables["nc3"] = self.Fort_var( self._ra_1D(Rec[0], "nc3"), "nc3", "full COMP3 is done every nc3 time steps.", "None", ("time")) + self.variables["ncycle"] = self.Fort_var( self._ra_1D(Rec[1], "ncycle"), "ncycle", "ncycle", "none", ("time")) self._log_var("ps", "surface pressure", "Pa", ("time", "lat", "lon"), scaling = 100) + self._log_var("temp", "temperature", "K", ("time", "pfull", "lat", "lon")) + self._log_var("ucomp", "zonal wind", "m/sec", ("time", "pfull", "lat", "lon")) + self._log_var("vcomp", "meridional wind", "m/s", ("time", "pfull", "lat", "lon")) + self._log_var("ts", "surface temperature", "K", ("time", "lat", "lon")) + self._log_var("snow", "surface amount of CO2 ice on the ground", "kg/m2", ("time", "lat", "lon")) + self._log_var("stressx", "zonal component of surface stress", "kg/m2", ("time", "lat", "lon")) + self._log_var("stressy", "merdional component of surface stress", "kg/m2", ("time", "lat", "lon")) + self._log_var("tstrat", "stratosphere temperature", "K", ("time", "lat", "lon")) + self._log_var("tausurf", "visible dust optical depth at the surface.", "none", ("time", "lat", "lon")) + self._log_var("ssun", "solar energy absorbed by the atmosphere", "W/m2", ("time", "lat", "lon")) @@ -979,16 +1023,21 @@ def _read_Fort11_dynamic(self): self._log_var("dst_mass", "dust aerosol mass mixing ratio", "kg/kg", ("time", "pfull", "lat", "lon"), Rec = Rec[..., 0]) + self._log_var("dst_num", "dust aerosol number", "number/kg", ("time", "pfull", "lat", "lon"), Rec = Rec[..., 1]) + self._log_var("ice_mass", "water ice aerosol mass mixing ratio", "kg/kg", ("time", "pfull", "lat", "lon"), Rec = Rec[..., 2]) + self._log_var("ice_num", "water ice aerosol number", "number/kg", ("time", "pfull", "lat", "lon"), Rec = Rec[..., 3]) + self._log_var("cor_mass", "dust core mass mixing ratio for water ice", "kg/kg", ("time", "pfull", "lat", "lon"), Rec = Rec[..., 4]) + self._log_var("vap_mass", "water vapor mass mixing ratio", "kg/kg", ("time", "pfull", "lat", "lon"), Rec = Rec[..., 5]) @@ -1000,25 +1049,31 @@ def _read_Fort11_dynamic(self): self._log_var("dst_mass_sfc", "dust aerosol mass on the surface", "kg/m2", ("time", "lat", "lon"), Rec = Rec[..., 0]) + self._log_var("dst_num_sfc", "dust aerosol number on the surface", "number/m2", ("time", "lat", "lon"), Rec = Rec[..., 1]) + self._log_var("ice_mass_sfc", "water ice aerosol mass on the surface", "kg/m2", ("time", "lat", "lon"), Rec = Rec[..., 2]) + self._log_var("ice_num_sfc", "water ice aerosol number on the surface", "number/m2", ("time", "lat", "lon"), Rec = Rec[..., 3]) + self._log_var("cor_mass_sfc", "dust core mass for water ice on the surface", "kg/m2", ("time", "lat", "lon"), Rec = Rec[..., 4]) + self._log_var("vap_mass_sfc", "water vapor mass on the surface", "kg/m2", ("time", "lat", "lon"), Rec = Rec[..., 5]) # write(11) stemp Rec = self.f.read_reals("f4").reshape(self.JM, self.IM, self.NL, order = "F") + self._log_var("soil_temp", "sub-surface soil temperature", "K", ("time", "zgrid", "lat", "lon"), Rec = Rec) @@ -1035,11 +1090,14 @@ def _read_Fort11_dynamic(self): self._log_var("fuptopv", "upward visible flux at the top of the atmosphere", "W/m2", ("time", "lat", "lon"), Rec = Rec[0].T) + self._log_var("fdntopv", "downward visible flux at the top of the atmosphere", "W/m2", ("time", "lat", "lon"), Rec = Rec[1].T) + self._log_var("fupsurfv", "upward visible flux at the surface", "W/m2", ("time", "lat", "lon"), Rec = Rec[2].T) + self._log_var("fdnsurfv", "downward visible flux at the surface", "W/m2", ("time", "lat", "lon"), Rec = Rec[3].T) @@ -1056,9 +1114,11 @@ def _read_Fort11_dynamic(self): self._log_var("fuptopir", "upward IR flux at the top of the atmosphere", "W/m2", ("time", "lat", "lon"), Rec = Rec[0].T) + self._log_var("fupsurfir", "upward IR flux at the surface", "W/m2", ("time", "lat", "lon"), Rec = Rec[1].T) + self._log_var("fdnsurfir", "downward IR flux at the surface", "W/m2", ("time", "lat", "lon"), Rec = Rec[2].T) @@ -1072,14 +1132,16 @@ def _read_Fort11_dynamic(self): # write(11) geot self._log_var("dheat", "diabatic heating rate", "K/sol", ("time", "pfull", "lat", "lon")) + self._log_var("geot", "geopotential", "m2/s2", ("time", "pfull", "lat", "lon")) + def _add_axis_as_variables(self): """ Add dimensions to the file as variables - """ + self.variables["lat"] = self.Fort_var(self.lat, "lat", "latitude", "degrees_N", ("lat")) self.variables["lon"] = self.Fort_var(self.lon, "lon", "longitude", @@ -1123,7 +1185,7 @@ def _add_axis_as_variables(self): self.variables["zgrid"] = self.Fort_var( self.zgrid, "zgrid", "depth at the mid-point of each soil layer", "m", ("zgrid")) - + def _linInterpLs(self, Ls, stride=16): """ @@ -1131,18 +1193,17 @@ def _linInterpLs(self, Ls, stride=16): :param Ls: input solar longitude :type Ls: float - :param stride: default stride :type stride: int - :return Ls_out: solar longitude (Ls) [float] - ..NOTE:: In the Legacy GCM fortran binaries, the solar - longitude is only updated once per day, implying that 16 - successive timesteps would have the same ls value. This routine - linearly interpolates the Ls between those successive values. - + ..note:: + In the Legacy GCM fortran binaries, the solar + longitude is only updated once per day, implying that 16 + successive timesteps would have the same ls value. This routine + linearly interpolates the Ls between those successive values. """ + Ls = np.array(Ls) Ls_out = np.zeros_like(Ls) Lsdi = Ls[::stride] diff --git a/amescap/Script_utils.py b/amescap/Script_utils.py index d1e955f..8419736 100644 --- a/amescap/Script_utils.py +++ b/amescap/Script_utils.py @@ -13,7 +13,9 @@ * ``re`` * ``os`` * ``subprocess`` - + * ``sys`` + * ``math`` + * ``matplotlib.colors`` """ # Load generic Python modules @@ -23,6 +25,7 @@ import numpy as np from netCDF4 import Dataset, MFDataset import re +from matplotlib.colors import ListedColormap, hex2color # ====================================================================== # DEFINITIONS @@ -47,19 +50,20 @@ def prYellow(skk): print("\033[93m{}\033[00m".format(skk)) def prPurple(skk): print("\033[95m{}\033[00m".format(skk)) def prLightPurple(skk): print("\033[94m{}\033[00m".format(skk)) + def MY_func(Ls_cont): """ Returns the Mars Year :param Ls_cont: solar longitude (continuous) :type Ls_cont: array - :return: the Mars year :rtype: int - """ + return (Ls_cont)//(360.) + 1 + def find_tod_in_diurn(fNcdf): """ Returns the variable for the local time axis in diurn files @@ -68,18 +72,16 @@ def find_tod_in_diurn(fNcdf): :param fNcdf: a netCDF file :type fNcdf: netCDF file object - :return: the name of the time of day dimension :rtype: str - """ + regex = re.compile("time_of_day.") varset = fNcdf.variables.keys() # Extract the 1st element of the list return [string for string in varset if re.match(regex, string)][0] - def print_fileContent(fileNcdf): """ Prints the contents of a netCDF file to the screen. Variables sorted @@ -87,10 +89,9 @@ def print_fileContent(fileNcdf): :param fileNcdf: full path to the netCDF file :type fileNcdf: str - :return: None - """ + if not os.path.isfile(fileNcdf.name): print(f"{fileNcdf.name} not found") else: @@ -99,10 +100,10 @@ def print_fileContent(fileNcdf): print(list(f.dimensions.keys())) print(str(f.dimensions)) print("===================== CONTENT =====================") - # Get all variables - all_var = f.variables.keys() - # Initialize empty list - all_dims = list() + + all_var = f.variables.keys() # Get all variables + all_dims = list() # Initialize empty list + for ivar in all_var: # Get all the variable dimensions all_dims.append(f.variables[ivar].dimensions) @@ -110,6 +111,7 @@ def print_fileContent(fileNcdf): # Filter duplicates. An object of type set() is an unordered # collection of distinct objects all_dims = set(all_dims) + # Sort dimensions by length (e.g., (lat) will come before # (lat, lon)) all_dims = sorted(all_dims, key = len) @@ -126,7 +128,7 @@ def print_fileContent(fileNcdf): f"{Cyan}{txt_shape}, {Yellow}{txt_long_name} " f"[{txt_units}]{Nclr}" ) - + try: # Skipped if the netCDF file does not contain time t_ini = f.variables["time"][0] @@ -138,13 +140,13 @@ def print_fileContent(fileNcdf): print(f"\nLs ranging from {(np.mod(Ls_ini, 360.)):.2f} to " f"{(np.mod(Ls_end, 360.)):.2f}: {(t_end-t_ini):.2f} days" f" (MY {MY_ini:02}) (MY {MY_end:02})") - except: pass - + f.close() print("=====================================================") + def print_varContent(fileNcdf, list_varfull, print_stat=False): """ Print variable contents from a variable in a netCDF file. Requires @@ -152,28 +154,26 @@ def print_varContent(fileNcdf, list_varfull, print_stat=False): :param fileNcdf: full path to a netcdf file :type fileNcdf: str - :param list_varfull: list of variable names and optional slices (e.g., ``["lon", "ps[:, 10, 20]"]``) :type list_varfull: list - :param print_stat: If True, print min, mean, and max. If False, print values. Defaults to False :type print_stat: bool, optional - :return: None - """ + if not os.path.isfile(fileNcdf.name): print(f"{fileNcdf.name} not found") else: if print_stat: - print(f"{Cyan}____________________________________________________" - f"______________________\n" - f" VAR | MIN | MEAN " - f"| MAX |\n" - f"__________________________|_______________|_______________" - f"|_______________|{Nclr}") + print( + f"{Cyan}" + f"__________________________________________________________________________\n" + f" VAR | MIN | MEAN | MAX |\n" + f"__________________________|_______________|_______________|_______________|" + f"{Nclr}") + for varfull in list_varfull: try: slice = "[:]" @@ -182,6 +182,7 @@ def print_varContent(fileNcdf, list_varfull, print_stat=False): slice = f"[{slice}" else: varname = varfull.strip() + cmd_txt = f"f.variables['{varname}']{slice}" f = Dataset(fileNcdf.name, "r") var = eval(cmd_txt) @@ -192,12 +193,14 @@ def print_varContent(fileNcdf, list_varfull, print_stat=False): Max = np.nanmax(var) print(f"{Cyan}{varfull:>26s}|{Min:>15g}|{Mean:>15g}|" f"{Max:>15g}|{Nclr}") + if varname == "areo": # If variable is areo then print modulo print(f"{Cyan}{varfull:>17s}(mod 360)|" f"({(np.nanmin(var%360)):>13g})|" f"({(np.nanmean(var%360)):>13g})|" f"({(np.nanmax(var%360)):>13g})|{Nclr}") + else: if varname != "areo": print(f"{Cyan}{varfull}= {Nclr}") @@ -211,8 +214,7 @@ def print_varContent(fileNcdf, list_varfull, print_stat=False): else: print(ii, ii%360) - print(f"{Cyan}____________________________________________" - f"__________________________{Nclr}") + print(f"{Cyan}______________________________________________________________________{Nclr}") except: if print_stat: print( @@ -222,65 +224,69 @@ def print_varContent(fileNcdf, list_varfull, print_stat=False): print(f"{Red}{varfull}") if print_stat: # Last line for the table - print(f"{Cyan}__________________________|_______________|_________" - f"______|_______________|{Nclr}") + print(f"{Cyan}__________________________|_______________|_______________|_______________|{Nclr}") f.close() + def give_permission(filename): """ Sets group file permissions for the NAS system - """ + try: # catch error and standard output - subprocess.check_call(["setfacl -v"], - shell = True, - stdout = open(os.devnull, "w"), - stderr = open(os.devnull, "w")) + subprocess.check_call( + ["setfacl -v"], + shell=True, + stdout=open(os.devnull, "w"), + stderr=open(os.devnull, "w") + ) cmd_txt = f"setfacl -R -m g:s0846:r {filename}" - subprocess.call(cmd_txt, shell = True) + subprocess.call(cmd_txt, shell=True) except subprocess.CalledProcessError: pass + def check_file_tape(fileNcdf): """ Checks whether a file exists on the disk. If on a NAS system, also checks if the file needs to be migrated from tape. - + :param fileNcdf: full path to a netcdf file or a file object with a name attribute :type fileNcdf: str or file object """ + # Get the filename, whether passed as string or as file object filename = fileNcdf if isinstance(fileNcdf, str) else fileNcdf.name - + # If filename is not a netCDF file, exit program if not re.search(".nc", filename): print(f"{Red}{filename} is not a netCDF file{Nclr}") exit() - + # First check if the file exists at all if not os.path.isfile(filename): print(f"{Red}File {filename} does not exist{Nclr}") exit() - + # Check if we're on a NAS system by looking for specific environment variables # or filesystem paths that are unique to NAS is_nas = False - + # Method 1: Check for NAS-specific environment variables nas_env_vars = ['PBS_JOBID', 'SGE_ROOT', 'NOBACKUP', 'NASA_ROOT'] for var in nas_env_vars: if var in os.environ: is_nas = True break - + # Method 2: Check for NAS-specific directories nas_paths = ['/nobackup', '/nobackupp', '/u/scicon'] for path in nas_paths: if os.path.exists(path): is_nas = True break - + # Only perform NAS-specific operations if we're on a NAS system if is_nas: try: @@ -288,25 +294,25 @@ def check_file_tape(fileNcdf): subprocess.check_call(["dmls"], shell=True, stdout=open(os.devnull, "w"), stderr=open(os.devnull, "w")) - + # Get the last columns of the ls command (filename and status) cmd_txt = f"dmls -l {filename}| awk '{{print $8,$9}}'" - + # Get identifier from dmls -l command dmls_out = subprocess.check_output(cmd_txt, shell=True).decode("utf-8") - + if dmls_out[1:4] not in ["DUL", "REG", "MIG"]: # If file is OFFLINE, UNMIGRATING etc... - print(f"{Yellow}*** Warning ***\n{dmls_out[6:-1]} " - f"is not loaded on disk (Status: {dmls_out[0:5]}). " - f"Please migrate it to disk with:\n" - f"``dmget {dmls_out[6:-1]}``\n then try again." - f"\n{Nclr}") + print(f"{Yellow}*** Warning ***\n{dmls_out[6:-1]} is not " + f"loaded on disk (Status: {dmls_out[0:5]}). Please " + f"migrate it to disk with: ``dmget {dmls_out[6:-1]}``\n" + f" then try again.\n{Nclr}") exit() except (subprocess.CalledProcessError, FileNotFoundError, IndexError): # If there's any issue with the dmls command, log it but continue if "--debug" in sys.argv: - print(f"{Yellow}Warning: Could not check tape status for {filename}{Nclr}") + print(f"{Yellow}Warning: Could not check tape status for " + f"{filename}{Nclr}") pass # Otherwise, we're not on a NAS system, so just continue @@ -317,16 +323,15 @@ def get_Ncdf_path(fNcdf): .. note:: ``Dataset`` and multi-file dataset (``MFDataset``) have - different attributes for the path, hence the need for this + different attributes for the path, hence the need for this function. :param fNcdf: Dataset or MFDataset object :type fNcdf: netCDF file object - :return: string list for the Dataset (MFDataset) :rtype: str(list) - """ + # Only MFDataset has the_files attribute fname_out = getattr(fNcdf, "_files", False) # Regular Dataset @@ -334,6 +339,7 @@ def get_Ncdf_path(fNcdf): fname_out = getattr(fNcdf, "filepath")() return fname_out + def extract_path_basename(filename): """ Returns the path and basename of a file. If only the filename is @@ -341,14 +347,13 @@ def extract_path_basename(filename): :param filename: name of the netCDF file (may include full path) :type filename: str - :return: full file path & name of file .. note:: This routine does not confirm that the file exists. It operates on the provided input string. - """ + # Get the filename without the path if ("/" in filename or "\\" in filename): @@ -362,6 +367,7 @@ def extract_path_basename(filename): filepath = os.path.expanduser(filepath) return filepath, basename + def FV3_file_type(fNcdf): """ Return the type of the netCDF file (i.e., ``fixed``, ``diurn``, @@ -370,12 +376,11 @@ def FV3_file_type(fNcdf): :param fNcdf: an open Netcdf file :type fNcdf: Netcdf file object - :return: The Ls array type (string, ``fixed``, ``continuous``, or ``diurn``) and the netCDF file type (string ``fixed``, ``diurn``, ``average``, or ``daily``) - """ + # Get the full path from the file fullpath = get_Ncdf_path(fNcdf) @@ -416,6 +421,7 @@ def FV3_file_type(fNcdf): return f_type, interp_type + def alt_FV3path(fullpaths, alt, test_exist=True): """ Returns the original or fixed file given an interpolated daily, @@ -424,21 +430,17 @@ def alt_FV3path(fullpaths, alt, test_exist=True): :param fullpaths: full path to a file or a list of full paths to more than one file :type fullpaths: str - :param alt: type of file to return (i.e., original or fixed) :type alt: str - :param test_exist: Whether file exists on the disk, defaults to True :type test_exist: bool, optional - :raises ValueError: _description_ - :return: path to original or fixed file (e.g., "/u/path/00010.atmos_average.nc" or "/u/path/00010.fixed.nc") :rtype: str - """ + out_list = [] one_element = False @@ -467,9 +469,10 @@ def alt_FV3path(fullpaths, alt, test_exist=True): new_full_path = f"{path}/{file_raw}" if alt == "fixed": new_full_path = f"{path}/{DDDDD}.fixed.nc" - if (test_exist and not os.path.exists(new_full_path)): - raise ValueError(f"In alt_FV3path(), {new_full_path} does not " - f"exist") + if test_exist and not os.path.exists(new_full_path): + raise ValueError( + f"In alt_FV3path(), {new_full_path} does not exist" + ) out_list.append(new_full_path) @@ -478,6 +481,7 @@ def alt_FV3path(fullpaths, alt, test_exist=True): return out_list + def smart_reader(fNcdf, var_list, suppress_warning=False): """ Alternative to ``var = fNcdf.variables["var"][:]`` for handling @@ -486,16 +490,13 @@ def smart_reader(fNcdf, var_list, suppress_warning=False): :param fNcdf: an open netCDF file :type fNcdf: netCDF file object - :param var_list: a variable or list of variables (e.g., ``areo`` or [``pk``, ``bk``, ``areo``]) :type var_list: _type_ - :param suppress_warning: suppress debug statement. Useful if a variable is not expected to be in the file anyway. Defaults to False :type suppress_warning: bool, optional - :return: variable content (single or values to unpack) :rtype: list @@ -517,8 +518,8 @@ def smart_reader(fNcdf, var_list, suppress_warning=False): .. note:: Only the variable content is returned, not attributes - """ + # This out_list is for the variable out_list = [] one_element = False @@ -540,7 +541,8 @@ def smart_reader(fNcdf, var_list, suppress_warning=False): # Try to read in the original file out_list.append(fNcdf.variables[ivar][:]) else: - full_path_try = alt_FV3path(Ncdf_path, alt = "raw", + full_path_try = alt_FV3path(Ncdf_path, + alt = "raw", test_exist = True) if file_is_MF: @@ -582,6 +584,7 @@ def smart_reader(fNcdf, var_list, suppress_warning=False): out_list = out_list[0] return out_list + def regrid_Ncfile(VAR_Ncdf, file_Nc_in, file_Nc_target): """ Regrid a netCDF variable from one file structure to another. @@ -591,28 +594,25 @@ def regrid_Ncfile(VAR_Ncdf, file_Nc_in, file_Nc_target): :param VAR_Ncdf: a netCDF variable object to regrid (e.g., ``f_in.variable["temp"]``) :type VAR_Ncdf: netCDF file variable - :param file_Nc_in: an open netCDF file to source for the variable (e.g., ``f_in = Dataset("filename", "r")``) :type file_Nc_in: netCDF file object - :param file_Nc_target: an open netCDF file with the desired file structure (e.g., ``f_out = Dataset("filename", "r")``) :type file_Nc_target: netCDF file object - :return: the values of the variable interpolated to the target file grid. :rtype: array .. note:: While the KDTree interpolation can handle a 3D dataset - (lon/lat/lev instead of just 2D lon/lat), the grid points in - the vertical are just a few (10--100s) meters in the PBL vs a - few (10-100s) kilometers in the horizontal. This results in - excessive weighting in the vertical, which is why the vertical + (lon/lat/lev instead of just 2D lon/lat), the grid points in + the vertical are just a few (10--100s) meters in the PBL vs a + few (10-100s) kilometers in the horizontal. This results in + excessive weighting in the vertical, which is why the vertical dimension is handled separately. - """ + from amescap.FV3_utils import interp_KDTree, axis_interp ftype_in, zaxis_in = FV3_file_type(file_Nc_in) ftype_t, zaxis_t = FV3_file_type(file_Nc_target) @@ -647,20 +647,31 @@ def regrid_Ncfile(VAR_Ncdf, file_Nc_in, file_Nc_target): # Get array elements var_OUT = VAR_Ncdf[:] - if not (np.array_equal(lat_in,lat_t) and np.array_equal(lon_in,lon_t)): # STEP 1: lat/lon interpolation are always performed unless # target lon and lat are identical if len(np.atleast_1d(lon_in)) == 1: # Special case if input len(lon)=1 (slice or zonal avg), # only interpolate on the latitude axis - var_OUT = axis_interp(var_OUT, lat_in, lat_t, axis = -2, - reverse_input = False, type_int = "lin") + var_OUT = axis_interp( + var_OUT, + lat_in, + lat_t, + axis = -2, + reverse_input = False, + type_int = "lin" + ) elif len(np.atleast_1d(lat_in)) == 1: # Special case if input len(lat)=1 (slice or medidional avg), # only interpolate on the longitude axis - var_OUT = axis_interp(var_OUT, lon_in, lon_t, axis = -1, - reverse_input = False, type_int = "lin") + var_OUT = axis_interp( + var_OUT, + lon_in, + lon_t, + axis = -1, + reverse_input = False, + type_int = "lin" + ) else: # Bi-directional interpolation var_OUT = interp_KDTree(var_OUT, lat_in, lon_in, lat_t, lon_t) @@ -683,16 +694,26 @@ def regrid_Ncfile(VAR_Ncdf, file_Nc_in, file_Nc_target): intType = "lin" elif zaxis_in == "pstd": intType = "log" - var_OUT = axis_interp(var_OUT, lev_in,lev_t, pos_axis, - reverse_input = reverse_input, - type_int = intType) + var_OUT = axis_interp( + var_OUT, + lev_in, + lev_t, + pos_axis, + reverse_input = reverse_input, + type_int = intType + ) if "time" in VAR_Ncdf.dimensions: # STEP 3: Linear interpolation in Ls pos_axis = 0 - var_OUT = axis_interp(var_OUT, np.squeeze(areo_in)%360, - np.squeeze(areo_t)%360, pos_axis, - reverse_input = False, type_int = "lin") + var_OUT = axis_interp( + var_OUT, + np.squeeze(areo_in)%360, + np.squeeze(areo_t)%360, + pos_axis, + reverse_input = False, + type_int = "lin" + ) if ftype_in == "diurn": # STEP 4: Linear interpolation in time of day @@ -704,25 +725,31 @@ def regrid_Ncfile(VAR_Ncdf, file_Nc_in, file_Nc_target): tod_name_t = find_tod_in_diurn(file_Nc_target) tod_in = file_Nc_in.variables[tod_name_in][:] tod_t = file_Nc_target.variables[tod_name_t][:] - var_OUT = axis_interp(var_OUT, tod_in, tod_t, pos_axis, - reverse_input = False, - type_int = "lin") + var_OUT = axis_interp( + var_OUT, + tod_in, + tod_t, + pos_axis, + reverse_input = False, + type_int = "lin" + ) return var_OUT + def progress(k, Nmax): """ Displays a progress bar to monitor heavy calculations. :param k: current iteration of the outer loop :type k: int - :param Nmax: max iteration of the outer loop :type Nmax: int - """ + # For rounding to the 2nd digit from math import ceil progress = float(k)/Nmax + # Modify barLength to change length of progress bar barLength = 10 status = "" @@ -743,6 +770,7 @@ def progress(k, Nmax): sys.stdout.write(text) sys.stdout.flush() + def section_content_amescap_profile(section_ID): """ Executes first code section in ``~/.amescap_profile`` to read in @@ -751,10 +779,9 @@ def section_content_amescap_profile(section_ID): :param section_ID: the section to load (e.g., Pressure definitions for pstd) :type section_ID: str - :return: the relevant line with Python syntax - """ + import os import numpy as np input_file = os.environ["HOME"]+"/.amescap_profile" @@ -786,11 +813,13 @@ def section_content_amescap_profile(section_ID): f"{Cyan} ``cp AmesCAP/mars_templates/amescap_profile " f"~/.amescap_profile``") exit() + except Exception as exception: # Return the error print(f"{Red}Error") print(exception) + def filter_vars(fNcdf, include_list=None, giveExclude=False): """ Filters the variable names in a netCDF file for processing. Returns @@ -800,18 +829,15 @@ def filter_vars(fNcdf, include_list=None, giveExclude=False): :param fNcdf: an open netCDF object for a diurn, daily, or average file :type fNcdf: netCDF file object - :param include_list:list of variables to include (e.g., [``ucomp``, ``vcomp``], defaults to None :type include_list: list or None, optional - :param giveExclude: if True, returns variables to be excluded from the file, defaults to False :type giveExclude: bool, optional - :return: list of variable names to include in the processed file - """ + var_list = fNcdf.variables.keys() if include_list is None: # If no list is provided, return all variables: @@ -844,6 +870,7 @@ def filter_vars(fNcdf, include_list=None, giveExclude=False): out_list = exclude_list return out_list + def find_fixedfile(filename): """ Finds the relevant fixed file for a given average, daily, or diurn @@ -852,7 +879,6 @@ def find_fixedfile(filename): :param filename: an average, daily, or diurn netCDF file :type filename: str - :return: full path to the correspnding fixed file :rtype: str @@ -865,9 +891,10 @@ def find_fixedfile(filename): atmos_average.tileX_plevs.nc -> fixed.tileX.nc atmos_average.tileX_plevs_custom.nc -> fixed.tileX.nc atmos_average_custom.tileX_plevs.nc -> fixed.tileX.nc - """ + filepath, fname = extract_path_basename(filename) + if "tile" in fname: # Try the tile or standard version of the fixed files name_fixed = f"{filepath}/fixed.tile{fname.split('tile')[1][0]}.nc" @@ -879,6 +906,7 @@ def find_fixedfile(filename): name_fixed = "FixedFileNotFound" return name_fixed + def get_longname_unit(fNcdf, varname): """ Returns the longname and unit attributes of a variable in a netCDF @@ -887,30 +915,28 @@ def get_longname_unit(fNcdf, varname): :param fNcdf: an open netCDF file :type fNcdf: netCDF file object - :param varname: variable to extract attribute from :type varname: str - :return: longname and unit attributes :rtype: str .. note:: Some functions in MarsVars edit the units - (e.g., [kg] -> [kg/m]), therefore the empty string is 4 - characters in length (" " instead of "") to allow for + (e.g., [kg] -> [kg/m]), therefore the empty string is 4 + characters in length (" " instead of "") to allow for editing by ``editing units_txt[:-2]``, for example. - """ + return (getattr(fNcdf.variables[varname], "long_name", " "), getattr(fNcdf.variables[varname], "units", " ")) + def wbr_cmap(): """ Returns a color map that goes from white -> blue -> green -> yellow -> red - """ - from matplotlib.colors import ListedColormap + tmp_cmap = np.zeros((254, 4)) tmp_cmap[:, 3] = 1. tmp_cmap[:, 0:3] = np.array([ @@ -980,13 +1006,13 @@ def wbr_cmap(): [150, 22, 26], [146, 21, 25]])/255. return ListedColormap(tmp_cmap) + def rjw_cmap(): """ Returns John Wilson's preferred color map (red -> jade -> wisteria) - """ - from matplotlib.colors import ListedColormap + tmp_cmap = np.zeros((55, 4)) tmp_cmap[:, 3] = 1. tmp_cmap[:, 0:3] = np.array([ @@ -1006,13 +1032,13 @@ def rjw_cmap(): [255, 20, 11], [255, 0, 0], [237, 17, 0]])/255. return ListedColormap(tmp_cmap) + def hot_cold_cmap(): """ Returns Dark blue > light blue>white>yellow>red colormap Based on Matlab's bipolar colormap - """ - from matplotlib.colors import ListedColormap + tmp_cmap = np.zeros((128,4)) tmp_cmap [:,3]=1. #set alpha tmp_cmap[:,0:3]=np.array([ @@ -1051,14 +1077,14 @@ def hot_cold_cmap(): return ListedColormap(tmp_cmap) + def dkass_dust_cmap(): """ Returns a color map useful for dust cross-sections. (yellow -> orange -> red -> purple) Provided by Courtney Batterson. - """ - from matplotlib.colors import ListedColormap, hex2color + tmp_cmap = np.zeros((256, 4)) tmp_cmap[:, 3] = 1. dkass_cmap = [ @@ -1109,14 +1135,14 @@ def dkass_dust_cmap(): tmp_cmap[:, 0:3] = RGB_T return ListedColormap(tmp_cmap) + def dkass_temp_cmap(): """ Returns a color map that highlights the 200K temperatures. (black -> purple -> blue -> green -> yellow -> orange -> red) Provided by Courtney Batterson. - """ - from matplotlib.colors import ListedColormap, hex2color + tmp_cmap = np.zeros((256, 4)) tmp_cmap[:, 3] = 1. dkass_cmap = [ @@ -1167,6 +1193,7 @@ def dkass_temp_cmap(): tmp_cmap[:, 0:3] = RGB_T return ListedColormap(tmp_cmap) + def pretty_print_to_fv_eta(var, varname, nperline=6): """ Print the ``ak`` or ``bk`` coefficients for copying to @@ -1174,16 +1201,13 @@ def pretty_print_to_fv_eta(var, varname, nperline=6): :param var: ak or bk data :type var: array - :param varname: the variable name ("a" or "b") :type varname: str - :param nperline: the number of elements per line, defaults to 6 :type nperline: int, optional - :return: a print statement for copying into ``fv_eta.f90`` - """ + NLAY = len(var) - 1 ni = 0 # Print the piece of code to copy/paste in ``fv_eta.f90`` @@ -1231,23 +1255,21 @@ def replace_dims(Ncvar_dim, vert_dim_name=None): :param Ncvar_dim: netCDF variable dimensions (e.g., ``f_Ncdf.variables["temp"].dimensions``) :type Ncvar_dim: str - :param vert_dim_name: the vertical dimension if it is ambiguous (``pstd``, ``zstd``, or ``zagl``). Defaults to None :type vert_dim_name: str, optional - :return: updated dimensions :rtype: str - """ + # Set input dictionary options recognizable as MGCM variables lat_dic = ["lat", "lats", "latitudes", "latitude"] lon_dic = ["lon", "lon", "longitude", "longitudes"] lev_dic = ["pressure", "altitude"] - areo_dic = ["ls"] # Set the desired output names dims_out = list(Ncvar_dim).copy() + for ii, idim in enumerate(Ncvar_dim): # Rename axes if idim in lat_dic: dims_out[ii] = "lat" @@ -1262,6 +1284,7 @@ def replace_dims(Ncvar_dim, vert_dim_name=None): dims_out[ii] = vert_dim_name return tuple(dims_out) + def ak_bk_loader(fNcdf): """ Return ``ak`` and ``bk`` arrays from the current netCDF file. If @@ -1270,19 +1293,18 @@ def ak_bk_loader(fNcdf): :param fNcdf: an open netCDF file :type fNcdf: a netCDF file object - :return: the ``ak`` and ``bk`` arrays .. note:: This routine will look for both ``ak`` and ``bk``. There - are cases when it is convenient to load the ``ak``, ``bk`` once - when the files are first opened in ``MarsVars``, but the ``ak`` - and ``bk`` arrays may not be necessary for in the calculation - as is the case for ``MarsVars XXXXX.atmos_average_psd.nc - --add msf``, which operates on a pressure interpolated + are cases when it is convenient to load the ``ak``, ``bk`` once + when the files are first opened in ``MarsVars``, but the ``ak`` + and ``bk`` arrays may not be necessary for in the calculation + as is the case for ``MarsVars XXXXX.atmos_average_psd.nc + --add msf``, which operates on a pressure interpolated (``_pstd.nc``) file. - """ + # First try to read ak and bk in the current netCDF file: allvars = fNcdf.variables.keys() @@ -1298,9 +1320,8 @@ def ak_bk_loader(fNcdf): else: ak = np.array(fNcdf.variables["pk"]) bk = np.array(fNcdf.variables["bk"]) - - else: + try: filepath, fname = extract_path_basename(fullpath_name) if not 'average' in os.listdir(filepath): @@ -1308,7 +1329,7 @@ def ak_bk_loader(fNcdf): name_average = "FixedFileNotFound" else: name_average = list(filter(lambda s: 'average' in s, os.listdir()))[0] - + f_average = Dataset(name_average, "r", format = "NETCDF4_CLASSIC") allvars = f_average.variables.keys() @@ -1319,6 +1340,7 @@ def ak_bk_loader(fNcdf): ak = np.array(f_average.variables["pk"]) bk = np.array(f_average.variables["bk"]) f_average.close() + except: try: name_fixed = find_fixedfile(fullpath_name) @@ -1333,139 +1355,155 @@ def ak_bk_loader(fNcdf): ak = np.array(f_fixed.variables["pk"]) bk = np.array(f_fixed.variables["bk"]) f_fixed.close() - + except: print(f"{Red}Fixed file does not exist in {filepath}. " f"Make sure the fixed file you are referencing " f"matches the FV3 filetype (e.g., ``fixed.tileX.nc`` " f"for operations on tile X)") exit() + return ak, bk + def read_variable_dict_amescap_profile(f_Ncdf=None): """ - Inspect a Netcdf file and return the name of the variables and + Inspect a Netcdf file and return the name of the variables and dimensions based on the content of ~/.amescap_profile. - + Calling this function allows to remove hard-coded calls in CAP. - For example, to f.variables['ucomp'] is replaced by + For example, to f.variables['ucomp'] is replaced by f.variables["ucomp"], with "ucomp" taking the values of'ucomp', 'U' - + :param f_Ncdf: An opened Netcdf file object :type f_Ncdf: File object - - :return: Model, a dictionary with the dimensions and variables, + :return: Model, a dictionary with the dimensions and variables, e.g. "ucomp"='U' or "dim_lat"='latitudes' .. note:: - The defaut names for variables are defined in () + The defaut names for variables are defined in () parenthesis in ~/.amescap_profile:: - + 'X direction wind [m/s] (ucomp)>' - The defaut names for dimensions are defined in {} parenthesis in + The defaut names for dimensions are defined in {} parenthesis in ~/.amescap_profile:: - + Ncdf Y latitude dimension [integer] {lat}>lats - The dimensions (lon, lat, pfull, pstd) are loaded in the dictionary + The dimensions (lon, lat, pfull, pstd) are loaded in the dictionary as "dim_lon", "dim_lat" - """ if f_Ncdf is not None: - var_list_Ncdf=list(f_Ncdf.variables.keys()) - dim_list_Ncdf=list(f_Ncdf.dimensions.keys()) + var_list_Ncdf = list(f_Ncdf.variables.keys()) + dim_list_Ncdf = list(f_Ncdf.dimensions.keys()) else: - var_list_Ncdf=[] - dim_list_Ncdf=[] + var_list_Ncdf = [] + dim_list_Ncdf = [] - all_lines=section_content_amescap_profile('Variable dictionary') - lines=all_lines.split('\n') - #Remove empty lines: + all_lines = section_content_amescap_profile('Variable dictionary') + lines = all_lines.split('\n') + # Remove empty lines: while("" in lines):lines.remove("") - #Initialize model + # Initialize model class model(object): pass - MOD=model() + MOD = model() - #Read through all lines in the Variable dictionary section of amesgcm_profile: + # Read through all lines in the Variable dictionary section of amesgcm_profile: for il in lines: - var_list=[] - #e.g. 'X direction wind [m/s] (ucomp)>U,u' + var_list = [] + # e.g. 'X direction wind [m/s] (ucomp)>U,u' left,right=il.split('>') #Split on either side of '>' - #If using {var}, current entry is a dimension. If using (var), it is a variable + # If using {var}, current entry is a dimension. If using (var), it is a variable if '{' in left: - sep1='{';sep2='}';type_input='dimension' + sep1 = '{';sep2='}';type_input='dimension' elif '(' in left: - sep1='(';sep2=')';type_input='variable' + sep1 = '(';sep2=')';type_input='variable' # First get 'ucomp' from 'X direction wind [m/s] (ucomp) - _,tmp=FV3_var=left.split(sep1) - FV3_var=tmp.replace(sep2,'').strip() #THIS IS THE FV3 NAME OF THE CURRENT VARIABLE - #Then, get the list of variable on the righ-hand side, e.g. 'U,u' - all_vars=right.split(',') + _,tmp = FV3_var = left.split(sep1) + FV3_var = tmp.replace(sep2,'').strip() # FV3 NAME OF CURRENT VAR + + # Get the list of vars on the right-hand side, e.g., 'U,u' + all_vars = right.split(',') for ii in all_vars: - var_list.append(ii.strip()) #var_list IS A LIST OF POTENTIAL CORRESPONDING VARIABLES - #Set the attribute to the dictionary - #If the list is empty, e.g just [''], use the default FV3 variable presents in () or {} - if len(var_list)==1 and var_list[0]=='':var_list[0]=FV3_var #var_list IS A LIST OF POTENTIAL CORRESPONDING VARIABLES - found_list=[] - - #Place the input in the appropriate varialbe () or dimension {} dictionary - #print('var_list>>>',var_list) - if type_input=='variable': + var_list.append(ii.strip()) + # var_list = list of potential corresponding variables + + # Set the attribute to the dictionary + # If the list is empty, e.g just [''], use the default FV3 + # variable presents in () or {} + if len(var_list) == 1 and var_list[0] == '': + var_list[0] = FV3_var + # var_list = list of potential corresponding variables + + found_list = [] + + # Place the input in the appropriate variable () or dimension + # {} dictionary + if type_input == 'variable': for ivar in var_list: if ivar in var_list_Ncdf:found_list.append(ivar) - if len(found_list)==0: + if len(found_list) == 0: setattr(MOD,FV3_var,FV3_var) - elif len(found_list)==1: - setattr(MOD,FV3_var,found_list[0]) + elif len(found_list) == 1: + setattr(MOD,FV3_var, found_list[0]) else: - setattr(MOD,FV3_var,found_list[0]) - prYellow("***Warning*** more than one possible variable '%s' found in file: %s"%(FV3_var,found_list)) - if type_input=='dimension': + setattr(MOD,FV3_var, found_list[0]) + prYellow( + f"***Warning*** more than one possible variable " + f"'{FV3_var}' found in file: {found_list}" + ) + + if type_input == 'dimension': for ivar in var_list: if ivar in dim_list_Ncdf:found_list.append(ivar) - if len(found_list)==0: - setattr(MOD,'dim_'+FV3_var,FV3_var) - elif len(found_list)==1: - setattr(MOD,'dim_'+FV3_var,found_list[0]) + if len(found_list) == 0: + setattr(MOD, f"dim_{FV3_var}", FV3_var) + elif len(found_list) == 1: + setattr(MOD, f"dim_{FV3_var}", found_list[0]) else: - setattr(MOD,'dim_'+FV3_var,found_list[0]) - prYellow("***Warning*** more than one possible dimension '%s' found in file: %s"%(FV3_var,found_list)) - + setattr(MOD, f"dim_{FV3_var}", found_list[0]) + prYellow( + f"***Warning*** more than one possible dimension " + f"'{FV3_var}' found in file: {found_list}" + ) return MOD + def reset_FV3_names(MOD): """ - This function reset the model dictionary to the native FV3's + This function reset the model dictionary to the native FV3's variables, e.g.:: - + model.dim_lat = 'latitude' > model.dim_lat = 'lat' model.ucomp = 'U' > model.ucomp = 'ucomp' :param MOD: Generated with read_variable_dict_amescap_profile() :type MOD: class object - - :return: same object with updated names for the dimensions and + :return: same object with updated names for the dimensions and variables - """ - atts_list=dir(MOD) #Get all attributes - vars_list=[k for k in atts_list if '__' not in k] #do not touch all the __init__ etc.. + + atts_list = dir(MOD) # Get all attributes + vars_list = [k for k in atts_list if '__' not in k] # do not touch all the __init__ etc.. for ivar in vars_list: - name=ivar #get the native name, e.g ucomp + # Get the native name, e.g., ucomp + name = ivar if 'dim_' in ivar: - name=ivar[4:] #if attribute is dim_lat, just get the 'lat' part - setattr(MOD,ivar,name) #reset the original names + # If attribute is dim_lat, just get the 'lat' part + name = ivar[4:] + setattr(MOD,ivar,name) # Reset the original names return MOD -def except_message(debug,exception,varname,ifile,pre="",ext=""): + +def except_message(debug, exception, varname, ifile, pre="", ext=""): """ This function prints an error message in the case of an exception. It also contains a special error in the case of an already existing @@ -1473,77 +1511,74 @@ def except_message(debug,exception,varname,ifile,pre="",ext=""): :param debug: Flag for debug mode :type debug: logical - :param exception: Exception from try statement :type exception: class object - :param varname: Name of variable causing exception :type varname: string - :param ifile: Name of input file :type ifile: string - :param pre: Prefix to new variable :type pre: string - :param ext: Extension to new variable :type ext: string - """ + if debug: raise - if str(exception)[0:35] == ( - "NetCDF: String match to name in use"): - print(f"{Yellow}***Error*** Variable already " - f"exists in file.\nDelete the existing " - f"variable with " - f"``MarsVars {ifile} -rm " + + if str(exception)[0:35] == ("NetCDF: String match to name in use"): + print(f"{Yellow}***Error*** Variable already exists in file.\n" + f"Delete the existing variable with ``MarsVars {ifile} -rm " f"{pre}{varname}{ext}``{Nclr}") else: print(f"{Red}***Error*** {str(exception)}") + def check_bounds(values, min_val, max_val, dx): """ Check if all values in an array are within specified bounds. Exits program if any value is out of bounds. - + Parameters: - values : array-like - Single value or array of values to check - min_val : float - Minimum allowed value - max_val : float - Maximum allowed value - - Returns: - values : array or float - The validated value(s) + :param values: Single value or array of values to check + :type values: array-like + :param min_val: Minimum allowed value + :type min_val: float + :param max_val: Maximum allowed value + :type max_val: float + :return values: The validated value(s) + :rtype: array or float """ + try: # Handle both single values and arrays is_scalar = np.isscalar(values) values_array = np.array([values]) if is_scalar else np.array(values) - + # Check for non-numeric values if not np.issubdtype(values_array.dtype, np.number): print(f"Error: Input contains non-numeric values") sys.exit(1) - + # Find any out-of-bounds values - mask_invalid = (values_array < (min_val-dx)) | (values_array > (max_val+dx)) - + mask_invalid = ( + values_array < (min_val-dx)) | (values_array > (max_val+dx) + ) + if np.any(mask_invalid): # Get the invalid values invalid_values = values_array[mask_invalid] invalid_indices = np.where(mask_invalid)[0] - - print(f"Error: Values out of allowed range [{min_val}, {max_val}]") - print(f"Invalid values at indices {invalid_indices}: {invalid_values}") + + print( + f"Error: Values out of allowed range [{min_val}, {max_val}]\n" + f"Invalid values at indices {invalid_indices}: {invalid_values}" + ) exit() - + # Return original format (scalar or array) return values if is_scalar else values_array - + except Exception as e: print(f"Error: {str(e)}") exit() diff --git a/amescap/Spectral_utils.py b/amescap/Spectral_utils.py index 4801728..ad270ce 100644 --- a/amescap/Spectral_utils.py +++ b/amescap/Spectral_utils.py @@ -10,12 +10,11 @@ * ``scipy.signal`` * ``ImportError`` * ``Exception`` - + * ``pyshtools`` (optional) """ # Load generic Python modules import numpy as np - from amescap.Script_utils import Yellow, Cyan, Nclr, progress try: @@ -31,6 +30,7 @@ # ====================================================================== # DEFINITIONS # ====================================================================== + # Try to import pyshtools with proper error handling try: import pyshtools @@ -53,6 +53,7 @@ f"__________________{Nclr}\n\n" ) + def diurn_extract(VAR, N, tod, lon): """ Extract the diurnal component of a field. Original code by John @@ -61,23 +62,19 @@ def diurn_extract(VAR, N, tod, lon): :param VAR: field to process. Time of day dimension must be first (e.g., ``[tod, time, lat, lon]`` or ``[tod]`` :type VAR: 1D or ND array - :param N: number of harmonics to extract (``N=1`` for diurnal, ``N=2`` for diurnal AND semi diurnal, etc.) :type N: int - :param tod: universal time of day in sols (``0-1.``) If provided in hours (``0-24``), it will be normalized. :type tod: 1D array - :param lon: longitudes ``0-360`` :type lon: 1D array or float - :return: the amplitudes & phases for the Nth first harmonics, (e.g., size ``[Nh, time, lat, lon]``) :rtype: ND arrays - """ + dimsIN = VAR.shape nsteps = len(tod) period = 24 @@ -130,10 +127,10 @@ def diurn_extract(VAR, N, tod, lon): for nn in range(1,N+1): # Python does zero-indexing, start at nn-1 - cosser = np.dot(VAR[:, ...].T, np.cos(nn * arg)).squeeze() - sinser = np.dot(VAR[:, ...].T, np.sin(nn * arg)).squeeze() + cosser = np.dot(VAR[:, ...].T, np.cos(nn*arg)).squeeze() + sinser = np.dot(VAR[:, ...].T, np.sin(nn*arg)).squeeze() - amp[nn-1, :] = 2 * rnorm * np.sqrt(cosser**2 + sinser**2) + amp[nn-1, :] = 2*rnorm * np.sqrt(cosser**2 + sinser**2) phas[nn-1, :] = (180/np.pi) * np.arctan2(sinser, cosser) # Apply local time correction to the phase @@ -150,27 +147,22 @@ def reconstruct_diurn(amp, phas, tod, lon, sumList=[]): :param amp: amplitude of the signal. Harmonics dimension FIRST (e.g., ``[N, time, lat, lon]``) :type amp: array - :param phas: phase of the signal [hr]. Harmonics dimension FIRST :type phas: array - :param tod: time of day in universal time [hr] :type tod: 1D array - :param lon: longitude for converting universal -> local time :type lon: 1D array or float - :param sumList: the harmonics to include when reconstructing the wave (e.g., ``sumN=[1, 2, 4]``), defaults to ``[]`` :type sumList: list, optional - :return: a variable with reconstructed harmonics with N dimension FIRST and time of day SECOND (``[N, tod, time, lat, lon]``). If sumList is provided, the wave output harmonics will be aggregated (i.e., size = ``[tod, time, lat, lon]``) :rtype: _type_ - """ + dimsIN = amp.shape N = dimsIN[0] dimsSUM = np.append([len(tod)], dimsIN[1:]) @@ -188,7 +180,7 @@ def reconstruct_diurn(amp, phas, tod, lon, sumList=[]): dimAXIS[:] = 1 dimAXIS[-1] = len(lon) lon = lon.reshape(dimAXIS) - + # Reshape tod array dimAXIS = np.arange(len(dimsSUM)) dimAXIS[:] = 1 @@ -201,7 +193,7 @@ def reconstruct_diurn(amp, phas, tod, lon, sumList=[]): varSUM = np.zeros(dimsSUM) for nn in range(1, N+1): # Compute each harmonic varOUT[nn-1, ...] = ( - amp[nn-1, ...] + amp[nn-1, ...] * np.cos(nn * (tod-phas[nn-1, ...] + DT) / 24*2*np.pi) ) # If a sum of harmonics is requested, sum it @@ -223,31 +215,26 @@ def space_time(lon, timex, varIN, kmx, tmx): :param lon: longitude [°] (0-360) :type lon: 1D array - :param timex: time [sol] (e.g., 1.5 days sampled every hour is ``[0/24, 1/24, 2/24,.. 1,.. 1.5]``) :type timex: 1D array - :param varIN: variable for the Fourier analysis. First axis must be ``lon`` and last axis must be ``time`` (e.g., ``varIN[lon, time]``, ``varIN[lon, lat, time]``, or ``varIN[lon, lev, lat, time]``) :type varIN: array - :param kmx: the number of longitudinal wavenumbers to extract (max = ``nlon/2``) :type kmx: int - :param tmx: the number of tidal harmonics to extract (max = ``nsamples/2``) :type tmx: int - :return: (ampe) East propagating wave amplitude [same unit as varIN]; (ampw) West propagating wave amplitude [same unit as varIN]; (phasee) East propagating phase [°]; (phasew) West propagating phase [°] - .. NOTE:: 1. ``ampe``, ``ampw``, ``phasee``, and ``phasew`` have + .. note:: 1. ``ampe``, ``ampw``, ``phasee``, and ``phasew`` have dimensions ``[kmx, tmx]`` or ``[kmx, tmx, lat]`` or ``[kmx, tmx, lev, lat]`` etc.\n 2. The x and y axes may be constructed as follows, which will @@ -258,8 +245,8 @@ def space_time(lon, timex, varIN, kmx, tmx): KTIME, KLON = np.meshgrid(ktime, klon) amplitude = np.concatenate((ampw[:, ::-1], ampe), axis = 1) phase = np.concatenate((phasew[:, ::-1], phasee), axis = 1) - """ + # Get input variable dimensions dims = varIN.shape lon_id = dims[0] @@ -275,8 +262,9 @@ def space_time(lon, timex, varIN, kmx, tmx): varIN = np.reshape(varIN, (lon_id, jd, time_id)) # Initialize 4 empty arrays - ampw, ampe, phasew, phasee = ([np.zeros((kmx, tmx, jd)) for _x - in range(0, 4)]) + ampw, ampe, phasew, phasee = ( + [np.zeros((kmx, tmx, jd)) for _x in range(0, 4)] + ) #TODO not implemented yet: # zamp, zphas = [np.zeros((jd, tmx)) for _x in range(0, 2)] @@ -286,7 +274,7 @@ def space_time(lon, timex, varIN, kmx, tmx): argx = lon * 2*np.pi / 360 rnorm = 2. / len(argx) # If timex = [0/24, 1/24, 2/24,.. 1] arg cycles for m [0, 2Pi] - arg = timex * 2*np.pi + arg = timex * 2*np.pi # Nyquist cut off rnormt = 2. / len(arg) @@ -328,7 +316,6 @@ def space_time(lon, timex, varIN, kmx, tmx): ampe[kk, nn, :] = ae.T phasew[kk, nn, :] = pw.T phasee[kk, nn, :] = pe.T - #End loop ampw = np.reshape(ampw, (kmx, tmx) + dim_sup_id) ampe = np.reshape(ampe, (kmx, tmx) + dim_sup_id) @@ -378,9 +365,9 @@ def space_time(lon, timex, varIN, kmx, tmx): # if len(dims) > 2: # stamp = reshape(stamp, [kmx dims(2:end-1)]) # stphs = reshape(stphs, [kmx dims(2:end-1)]) - return ampe, ampw, phasee, phasew + def zeroPhi_filter(VAR, btype, low_highcut, fs, axis=0, order=4, add_trend=False): """ @@ -390,34 +377,27 @@ def zeroPhi_filter(VAR, btype, low_highcut, fs, axis=0, order=4, :param VAR: values for filtering 1D or ND array. Filtered dimension must be FIRST. Adjusts axis as necessary. :type VAR: array - :param btype: filter type (i.e., "low", "high" or "band") :type btype: str - :param low_high_cut: low, high, or [low, high] cutoff frequency depending on the filter [Hz or m-1] :type low_high_cut: int - :param fs: sampling frequency [Hz or m-1] :type fs: int - :param axis: if data is an ND array, this identifies the filtering dimension :type axis: int - :param order: order for the filter :type order: int - :param add_trend: if True, return the filtered output. If false, return the trend and filtered output. :type add_trend: bool - :return: the filtered data - .. NOTE:: ``Wn=[low, high]`` are expressed as a function of the + .. note:: ``Wn=[low, high]`` are expressed as a function of the Nyquist frequency - """ + # Create the filter low_highcut = np.array(low_highcut) nyq = 0.5*fs @@ -437,6 +417,7 @@ def zeroPhi_filter(VAR, btype, low_highcut, fs, axis=0, order=4, else: return VAR_f + def zonal_decomposition(VAR): """ Decomposition into spherical harmonics. [A. Kling, 2020] @@ -444,17 +425,16 @@ def zonal_decomposition(VAR): :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 + .. 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" @@ -470,8 +450,9 @@ def zonal_decomposition(VAR): 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]) + 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) @@ -484,6 +465,7 @@ def zonal_decomposition(VAR): return COEFFS, psd.reshape(psd_out_shape) + def zonal_construct(COEFFS_flat, VAR_shape, btype=None, low_highcut=None): """ Recomposition into spherical harmonics @@ -492,21 +474,17 @@ def zonal_construct(COEFFS_flat, VAR_shape, btype=None, low_highcut=None): 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 - + :type low_high_cut: int or list :return: detrended variable reconstructed to original size (e.g., [time, lev, lat, lon]) - .. NOTE:: The minimum and maximum wavelenghts in [km] are computed:: + .. 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 @@ -514,8 +492,8 @@ def zonal_construct(COEFFS_flat, VAR_shape, btype=None, low_highcut=None): 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" @@ -550,8 +528,9 @@ def zonal_construct(COEFFS_flat, VAR_shape, btype=None, low_highcut=None): 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 \ No newline at end of file + + return PYSHTOOLS_AVAILABLE diff --git a/amescap/cli.py b/amescap/cli.py index c4baff2..5988fe5 100644 --- a/amescap/cli.py +++ b/amescap/cli.py @@ -6,7 +6,7 @@ from amescap.Script_utils import Yellow, Nclr, Green, Cyan def get_install_info(): - # Get the location and timestamp of cli + # Get the location and timestamp of cli cli_path = os.path.abspath(__file__) install_time = time.ctime(os.path.getctime(cli_path)) return f""" @@ -31,7 +31,7 @@ def format_help(self): {Cyan}How do I use CAP? -----------------{Nclr} -Below is a list of the executables in CAP. Use this list to find the executable that performs the operation you desire. +Below is a list of the executables in CAP. Use this list to find the executable that performs the operation you desire. To see the arguments for each executable, use: {Green} -h Example: MarsVars -h{Nclr} @@ -73,7 +73,7 @@ def format_help(self): formatter_class=CustomHelpFormatter, add_help=False # This prevents the default help message ) - + parser.add_argument('-h', '--help', action='help', default=argparse.SUPPRESS, help=argparse.SUPPRESS) parser.add_argument('command', nargs='?', default='help', @@ -85,7 +85,7 @@ def format_help(self): if args.command in ['version', 'info']: print(get_install_info()) return 0 - + # Print help for all cases parser.print_help() return 0 diff --git a/amescap/pdf2image.py b/amescap/pdf2image.py index 080e42b..08fc8d4 100644 --- a/amescap/pdf2image.py +++ b/amescap/pdf2image.py @@ -13,7 +13,8 @@ * ``os`` * ``subprocess`` * ``PIL`` - + * ``uuid`` + * ``Pillow`` """ # Load generic Python modules @@ -25,6 +26,7 @@ from subprocess import Popen, PIPE from PIL import Image + def convert_from_path(pdf_path, dpi=200, output_folder=None, first_page=None, last_page=None, fmt="ppm", thread_count=1, userpw=None, use_cropbox=False): @@ -34,33 +36,25 @@ def convert_from_path(pdf_path, dpi=200, output_folder=None, first_page=None, :param pdf_path: path to the PDF that you want to convert :type pdf_path: str - :param dpi: image quality in DPI (default 200) :type dpi: int - :param output_folder: folder to write the images to (instead of directly in memory) :type output_folder: str - :param first_page: first page to process :type first_page: int - :param last_page: last page to process before stopping :type last_page: int - :param fmt: output image format :type fmt: str - :param thread_count: how many threads to spawn for processing :type thread_count: int - :param userpw: PDF password :type userpw: str - :param use_cropbox: use cropbox instead of mediabox :type use_cropbox: bool - """ + page_count = __page_count(pdf_path, userpw) if thread_count < 1: @@ -86,17 +80,23 @@ def convert_from_path(pdf_path, dpi=200, output_folder=None, first_page=None, # A unique identifier for our files if the directory is not empty uid = str(uuid.uuid4()) # Get the number of pages the thread will be processing - thread_page_count = page_count // thread_count + int(reminder > 0) + thread_page_count = page_count//thread_count + int(reminder > 0) # Build the command accordingly args, parse_buffer_func = __build_command( - ["pdftoppm", "-r", str(dpi), pdf_path], output_folder, - current_page, (current_page + thread_page_count - 1), fmt, uid, - userpw, use_cropbox) + ["pdftoppm", "-r", str(dpi), pdf_path], + output_folder, + current_page, + (current_page + thread_page_count - 1), + fmt, + uid, + userpw, + use_cropbox + ) # Update page values current_page = current_page + thread_page_count reminder -= int(reminder > 0) # Spawn the process and save its uuid - processes.append((uid, Popen(args, stdout = PIPE, stderr = PIPE))) + processes.append((uid, Popen(args, stdout=PIPE, stderr=PIPE))) images = [] for uid, proc in processes: @@ -107,53 +107,50 @@ def convert_from_path(pdf_path, dpi=200, output_folder=None, first_page=None, images += parse_buffer_func(data) return images + def convert_from_bytes(pdf_file, dpi=200, output_folder=None, first_page=None, last_page=None, fmt='ppm', thread_count=1, userpw=None, use_cropbox=False): """ - Convert PDF to Image will throw an error whenever one of the condition is reached + + Convert PDF to Image will throw an error whenever one of the + condition is reached :param pdf_file: Bytes representing the PDF file :type pdf_file: float - :param dpi: image quality in DPI (default 200) :type dpi: int - :param output_folder: folder to write the images to (instead of directly in memory) :type output_folder: str - :param first_page: first page to process :type first_page: int - :param last_page: last page to process before stopping :type last_page: int - :param fmt: output image format :type fmt: str - :param thread_count: how many threads to spawn for processing :type thread_count: int - :param userpw: PDF password :type userpw: str - :param use_cropbox: use cropbox instead of mediabox :type use_cropbox: bool - """ + with tempfile.NamedTemporaryFile('wb') as f: f.write(pdf_file) f.flush() - return convert_from_path(f.name, - dpi = dpi, - output_folder = output_folder, - first_page = first_page, - last_page = last_page, - fmt = fmt, - thread_count = thread_count, - userpw = userpw, - use_cropbox = use_cropbox) + return convert_from_path( + f.name, + dpi = dpi, + output_folder = output_folder, + first_page = first_page, + last_page = last_page, + fmt = fmt, + thread_count = thread_count, + userpw = userpw, + use_cropbox = use_cropbox + ) def __build_command(args, output_folder, first_page, last_page, fmt, uid, userpw, use_cropbox): @@ -163,7 +160,9 @@ def __build_command(args, output_folder, first_page, last_page, fmt, uid, args.extend(["-f", str(first_page)]) if last_page is not None: args.extend(["-l", str(last_page)]) + parsed_format, parse_buffer_func = __parse_format(fmt) + if parsed_format != "ppm": args.append("-" + parsed_format) if output_folder is not None: @@ -172,6 +171,7 @@ def __build_command(args, output_folder, first_page, last_page, fmt, uid, args.extend(["-upw", userpw]) return args, parse_buffer_func + def __parse_format(fmt): if fmt[0] == ".": fmt = fmt[1:] @@ -180,44 +180,51 @@ def __parse_format(fmt): return "jpeg", __parse_buffer_to_jpeg if fmt == "png": return "png", __parse_buffer_to_png - # Unable to parse the format so use the default return "ppm", __parse_buffer_to_ppm + def __parse_buffer_to_ppm(data): images = [] index = 0 while index < len(data): code, size, rgb = tuple(data[index:index+40].split(b"\n")[0:3]) size_x, size_y = tuple(size.split(b" ")) - file_size = (len(code) + len(size) + len(rgb) + 3 - + int(size_x) * int(size_y) * 3) + file_size = ( + len(code) + len(size) + len(rgb) + 3 + int(size_x) * int(size_y)*3 + ) images.append(Image.open(BytesIO(data[index:index + file_size]))) index += file_size return images + def __parse_buffer_to_jpeg(data): # Last element is obviously empty return [Image.open(BytesIO(image_data + b"\xff\xd9")) for image_data in data.split(b"\xff\xd9")[:-1]] + def __parse_buffer_to_png(data): images = [] index = 0 while index < len(data): # 4 bytes for IEND + 4 bytes for CRC file_size = data[index:].index(b"IEND") + 8 - images.append(Image.open(BytesIO(data[index:index+file_size]))) + images.append(Image.open(BytesIO(data[index:index + file_size]))) index += file_size return images + def __page_count(pdf_path, userpw=None): try: if userpw is not None: - proc = Popen(["pdfinfo", pdf_path, "-upw", userpw], stdout = PIPE, + proc = Popen(["pdfinfo", pdf_path, "-upw", userpw], + stdout = PIPE, stderr = PIPE) else: - proc = Popen(["pdfinfo", pdf_path], stdout = PIPE, stderr = PIPE) + proc = Popen(["pdfinfo", pdf_path], + stdout = PIPE, + stderr = PIPE) out, err = proc.communicate() except: raise Exception("Unable to get page count. If not on a Linux system, " @@ -225,11 +232,14 @@ def __page_count(pdf_path, userpw=None): try: # This will throw an error if we are unable to get page count - return int(re.search(r"Pages:\s+(\d+)", - out.decode("utf8", "ignore")).group(1)) + return int( + re.search(r"Pages:\s+(\d+)", out.decode("utf8", "ignore")).group(1) + ) except: - raise Exception(f"Unable to get page count. " - f"{err.decode('utf8', 'ignore')}") + raise Exception( + f"Unable to get page count. {err.decode('utf8', 'ignore')}" + ) + def __load_from_output_folder(output_folder, uid): return [Image.open(os.path.join(output_folder, f)) for f diff --git a/bin/MarsCalendar.py b/bin/MarsCalendar.py index 84e88b7..f971b31 100755 --- a/bin/MarsCalendar.py +++ b/bin/MarsCalendar.py @@ -17,7 +17,10 @@ * ``numpy`` * ``argparse`` - + * ``functools`` + * ``traceback`` + * ``sys`` + * ``amescap`` """ # Make print statements appear in color @@ -35,11 +38,13 @@ # Load amesCAP modules from amescap.FV3_utils import (sol2ls, ls2sol) + def debug_wrapper(func): """ - A decorator that wraps a function with error handling based on the + A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -57,6 +62,7 @@ def wrapper(*args, **kwargs): return 1 # Error exit code return wrapper + # ====================================================================== # ARGUMENT PARSER # ====================================================================== @@ -142,10 +148,12 @@ def wrapper(*args, **kwargs): f"help.{Nclr}") exit() + # ====================================================================== # DEFINITIONS # ====================================================================== + def parse_array(len_input): """ Formats the input array for conversion. @@ -156,17 +164,15 @@ def parse_array(len_input): :param len_input: The input Ls or sol to convert. Can either be one number (e.g., 300) or start stop step (e.g., 300 310 2). - :type len_input: float - + :type len_input: float or list of floats :raises: Error if neither ``[-ls --ls]`` or ``[-sol --sol]`` are provided. - :return: ``input_as_arr`` An array formatted for input into ``ls2sol`` or ``sol2ls``. If ``len_input = 300``, then ``input_as_arr=[300]``. If ``len_input = 300 310 2``, then ``input_as_arr = [300, 302, 304, 306, 308]``.\n - """ + if len(len_input) == 1: input_as_arr = len_input @@ -183,10 +189,12 @@ def parse_array(len_input): exit() return(input_as_arr) + # ====================================================================== # MAIN PROGRAM # ====================================================================== + @debug_wrapper def main(): # Load in user-specified Mars year, if any. Default = 0 @@ -229,10 +237,12 @@ def main(): print("\n") + # ====================================================================== # END OF PROGRAM # ====================================================================== + if __name__ == "__main__": exit_code = main() sys.exit(exit_code) diff --git a/bin/MarsFiles.py b/bin/MarsFiles.py index 55fdf0c..b8db80c 100755 --- a/bin/MarsFiles.py +++ b/bin/MarsFiles.py @@ -73,9 +73,11 @@ get_longname_unit, check_bounds ) + # ------------------------------------------------------ # EXTENSION FUNCTION # ------------------------------------------------------ + class ExtAction(argparse.Action): def __init__(self, *args, ext_content=None, parser=None, **kwargs): self.parser = parser @@ -135,6 +137,7 @@ def debug_wrapper(func): A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -152,6 +155,7 @@ def wrapper(*args, **kwargs): return 1 # Error exit code return wrapper + # ------------------------------------------------------ # ARGUMENT PARSER # ------------------------------------------------------ @@ -609,6 +613,7 @@ def wrapper(*args, **kwargs): f"argument{Nclr}") exit() + # ------------------------------------------------------ # EXTENSION FUNCTION # ------------------------------------------------------ @@ -639,9 +644,12 @@ def wrapper(*args, **kwargs): # Append extension, if any: out_ext = (f"{out_ext}_{args.extension}") + # ------------------------------------------------------ # DEFINITIONS # ------------------------------------------------------ + + def concatenate_files(file_list, full_file_list): """ Concatenates sequential output files in chronological order. @@ -651,6 +659,7 @@ def concatenate_files(file_list, full_file_list): :param full_file_list: list of file names and full paths :type full_file_list: list """ + print(f"{Yellow}Using internal method for concatenation{Nclr}") # For fixed files, deleting all but the first file has the same @@ -717,6 +726,7 @@ def concatenate_files(file_list, full_file_list): return + def split_files(file_list, split_dim): """ Extracts variables in the file along the time dimension, unless @@ -726,6 +736,7 @@ def split_files(file_list, split_dim): :type split_dim: dimension along which to perform extraction :returns: new file with sliced dimensions """ + if split_dim not in ['time','areo', 'lev', 'lat', 'lon']: print(f"{Red}Split dimension must be one of the following:" f" time, areo, lev, lat, lon{Nclr}") @@ -1016,6 +1027,7 @@ def split_files(file_list, split_dim): fNcdf.close() return + # ------------------------------------------------------ # Time-Shifting Implementation # Victoria H. @@ -1028,6 +1040,7 @@ def process_time_shift(file_list): :param file_list: list of file names :type file_list: list """ + if args.time_shift == 999: # Target local times requested by user target_list = None @@ -1169,9 +1182,11 @@ def process_time_shift(file_list): fdiurn.close() return + # ------------------------------------------------------ # MAIN PROGRAM # ------------------------------------------------------ + @debug_wrapper def main(): global data_dir @@ -1193,6 +1208,7 @@ def main(): ) exit() + # ------------------------------------------------------------------ # Conversion Legacy -> MGCM Format # Richard U. and Alex. K. @@ -1254,6 +1270,7 @@ def main(): # Time-shift files process_time_shift(file_list) + # ------------------------------------------------------------------ # Bin a daily file as an average file # Alex K. @@ -1344,6 +1361,7 @@ def main(): fnew.copy_Ncvar(fdaily.variables[ivar]) fnew.close() + # ------------------------------------------------------------------ # Bin a daily file as a diurn file # Alex K. @@ -1434,6 +1452,7 @@ def main(): fnew.copy_Ncvar(fdaily.variables[ivar]) fnew.close() + # ------------------------------------------------------------------ # Temporal Filtering # Alex K. & R. J. Wilson @@ -1564,6 +1583,7 @@ def main(): fnew.copy_Ncvar(fdaily.variables[ivar]) fnew.close() + # ------------------------------------------------------------------ # Zonal Decomposition Analysis # Alex K. @@ -1720,6 +1740,7 @@ def main(): fnew.copy_Ncvar(fname.variables[ivar]) fnew.close() + # ------------------------------------------------------------------ # Tidal Analysis # Alex K. & R. J. Wilson @@ -1875,6 +1896,7 @@ def main(): ) fnew.close() + # ------------------------------------------------------------------ # Regridding Routine # Alex K. @@ -1935,6 +1957,7 @@ def main(): fnew.close() fNcdf_t.close() + # ------------------------------------------------------------------ # Zonal Averaging # Alex K. @@ -2008,6 +2031,7 @@ def main(): f"--bin_files --concatenate, --time_shift, --bin_average, " f"--bin_diurn etc ...``{Nclr}") + # ---------------------------------------------------------------------- # DEFINITIONS # ---------------------------------------------------------------------- @@ -2028,6 +2052,7 @@ def make_FV3_files(fpath, typelistfv3, renameFV3=True): :return: The MGCM-like files: ``XXXXX.atmos_average.nc``, ``XXXXX.atmos_daily.nc``, ``XXXXX.atmos_diurn.nc``. """ + historyDir = os.getcwd() histfile = Dataset(fpath, "r", format="NETCDF4_CLASSIC") histdims = histfile.dimensions.keys() @@ -2036,6 +2061,7 @@ def make_FV3_files(fpath, typelistfv3, renameFV3=True): # Convert the first Ls in file to a sol number fdate = f"{(ls2sol_1year(histfile.variables['ls'][0])):05}" + def proccess_file(newf, typefv3): """ Creates required variables and inputs them into the new @@ -2050,6 +2076,7 @@ def proccess_file(newf, typefv3): :type typefv3: str :return: netCDF file with minimum required variables """ + for dname in histdims: if dname == "nlon": var = histfile.variables["longitude"] @@ -2186,6 +2213,7 @@ def proccess_file(newf, typefv3): except OSError as e: print(f"Warning: Could not copy fixed file: {e}") + def do_avg_vars(histfile, newf, avgtime, avgtod, bin_period=5): """ Performs a time average over all fields in a file. @@ -2205,6 +2233,7 @@ def do_avg_vars(histfile, newf, avgtime, avgtod, bin_period=5): :type bin_period: int, optional :return: a time-averaged file """ + histvars = histfile.variables.keys() for vname in histvars: var = histfile.variables[vname] @@ -2364,6 +2393,7 @@ def do_avg_vars(histfile, newf, avgtime, avgtod, bin_period=5): ) return 0 + def change_vname_longname_unit(vname, longname_txt, units_txt): """ Update variable ``name``, ``longname``, and ``units``. This is @@ -2377,6 +2407,7 @@ def change_vname_longname_unit(vname, longname_txt, units_txt): :type units_txt: str :return: variable name and corresponding description and unit """ + if vname == "psurf": vname = "ps" longname_txt = "surface pressure" @@ -2435,6 +2466,7 @@ def replace_dims(dims, todflag): :type todflag: bool :return: new dimension names for the variable """ + newdims = dims if "nlat" in dims: newdims = replace_at_index(newdims, newdims.index("nlat"), "lat") @@ -2451,6 +2483,7 @@ def replace_dims(dims, todflag): ) return newdims + def replace_at_index(tuple_dims, idx, new_name): """ Updates variable dimensions. @@ -2465,11 +2498,13 @@ def replace_at_index(tuple_dims, idx, new_name): :type new_name: str :return: updated dimensions """ + if new_name is None: return tuple_dims[:idx] + tuple_dims[idx+1:] else: return tuple_dims[:idx] + (new_name,) + tuple_dims[idx+1:] + def ls2sol_1year(Ls_deg, offset=True, round10=True): """ Returns a sol number from the solar longitude. @@ -2486,6 +2521,7 @@ def ls2sol_1year(Ls_deg, offset=True, round10=True): For the moment, this is consistent with 0 <= Ls <= 359.99, but not for monotically increasing Ls. """ + Ls_perihelion = 250.99 # Ls at perihelion tperi = 485.35 # Time (in sols) at perihelion Ns = 668.6 # Number of sols in 1 MY @@ -2513,9 +2549,11 @@ def ls2sol_1year(Ls_deg, offset=True, round10=True): Ds = np.round(Ds, -1) return Ds + # ------------------------------------------------------ # END OF PROGRAM # ------------------------------------------------------ + if __name__ == "__main__": exit_code = main() - sys.exit(exit_code) \ No newline at end of file + sys.exit(exit_code) diff --git a/bin/MarsFormat.py b/bin/MarsFormat.py index 6e8fe40..3ba4f74 100755 --- a/bin/MarsFormat.py +++ b/bin/MarsFormat.py @@ -1,9 +1,9 @@ #!/usr/bin/env python3 """ The MarsFormat executable is for converting non-MGCM data, such as that -from EMARS, OpenMARS, PCM, and MarsWRF, into MGCM-like netCDF data +from EMARS, OpenMARS, PCM, and MarsWRF, into MGCM-like netCDF data products. The MGCM is the NASA Ames Mars Global Climate Model developed -and maintained by the Mars Climate Modeling Center (MCMC). The MGCM +and maintained by the Mars Climate Modeling Center (MCMC). The MGCM data repository is available at data.nas.nasa.gov/mcmc. The executable requires two arguments: @@ -24,11 +24,15 @@ * ``sys`` * ``argparse`` * ``os`` + * ``re`` + * ``functools`` + * ``traceback`` + * ``xarray`` + * ``amescap`` List of Functions: * download - Queries the requested file from the NAS Data Portal. - """ # Make print statements appear in color @@ -55,11 +59,13 @@ xr.set_options(keep_attrs=True) + def debug_wrapper(func): """ - A decorator that wraps a function with error handling based on the + A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -77,9 +83,11 @@ def wrapper(*args, **kwargs): return 1 # Error exit code return wrapper + # ====================================================== # ARGUMENT PARSER # ====================================================== + parser=argparse.ArgumentParser( prog=('MarsFormat'), description=( @@ -171,22 +179,24 @@ def wrapper(*args, **kwargs): path2data = os.getcwd() ref_press = 725 # TODO hard-coded reference pressure + def get_time_dimension_name(DS, model): """ Find the time dimension name in the dataset. Updates the model object with the correct dimension name. - - Args: - DS: The xarray Dataset - model: Model object with dimension information - - Returns: - str: The actual time dimension name found + + :param DS: The xarray Dataset + :type DS: xarray.Dataset + :param model: Model object with dimension information + :type model: object + :return: The actual time dimension name found + :rtype: str """ + # First try the expected dimension name if model.dim_time in DS.dims: return model.dim_time - + # Check alternative names possible_names = ['Time', 'time', 'ALSO_Time'] for name in possible_names: @@ -195,7 +205,7 @@ def get_time_dimension_name(DS, model): f"instead of '{model.dim_time}'{Nclr}") model.dim_time = name return name - + # If no time dimension is found, raise an error raise KeyError(f"No time dimension found in dataset. Expected one " f"of: {model.dim_time}, {', '.join(possible_names)}") @@ -203,7 +213,7 @@ def get_time_dimension_name(DS, model): @debug_wrapper def main(): ext = '' # Initialize empty extension - + if args.gcm_name not in ['marswrf', 'openmars', 'pcm', 'emars']: print(f"{Yellow}***Notice*** No operation requested. Use " f"'-gcm' and specify openmars, marswrf, pcm, emars") @@ -238,7 +248,7 @@ def main(): # Open dataset with xarray DS = xr.open_dataset(fullnameIN, decode_times=False) - + # Store the original time values and units before any modifications original_time_vals = DS[model.time].values.copy() # This will always exist original_time_units = DS[model.time].attrs.get('units', '') @@ -248,17 +258,17 @@ def main(): f"'{original_time_units}' and description " f"'{original_time_desc}'" ) - + # Find and update time dimension name time_dim = get_time_dimension_name(DS, model) - + # -------------------------------------------------------------- # MarsWRF Processing # -------------------------------------------------------------- if model_type == 'marswrf': # print(f"{Cyan}Current variables at top of marswrf " # f"processing: \n{list(DS.variables)}{Nclr}\n") - + # First, save all variable descriptions in attrs longname for var_name in DS.data_vars: var = DS[var_name] @@ -280,60 +290,60 @@ def main(): # Time conversion (minutes to days) time = (DS[model.time] / 60 / 24) if 'time' in DS else None - # Handle latitude and longitude + # Handle latitude and longitude if len(DS[model.lat].shape) > 1: lat = DS[model.lat][0, :, 0] lon = DS[model.lon][0, 0, :] else: lat = DS[model.lat] lon = DS[model.lon] - + # Convert longitudes to 0-360 lon360 = (lon + 360)%360 - # Update coordinates + # Update coordinates if time is not None: DS[model.time] = time DS[model.lon] = lon360 DS[model.lat] = lat - + # Derive phalf # This employs ZNU (half, mass levels and ZNW (full, w) levels phalf = DS.P_TOP.values[0] + DS.ZNW.values[0,:]*DS.P0 pfull = DS.P_TOP.values[0] + DS.ZNU.values[0,:]*DS.P0 - + DS = DS.assign_coords(pfull=(model.dim_pfull, pfull)) DS = DS.assign_coords(phalf=(model.dim_phalf, phalf)) N_phalf=len(DS.bottom_top)+1 ak = np.zeros(N_phalf) bk = np.zeros(N_phalf) - + ak[-1] = DS.P_TOP[0] # MarsWRF pressure increases w/N bk[:] = np.array(DS.ZNW[0,:], copy=True) - + # Fill ak, bk, pfull, phalf arrays DS = DS.assign(ak=(model.dim_phalf, ak)) DS = DS.assign(bk=(model.dim_phalf, bk)) - + DS.phalf.attrs['description'] = ( '(ADDED POST-PROCESSING) pressure at layer interfaces') DS.phalf.attrs['units'] = ('Pa') - + DS['ak'].attrs['description'] = ( '(ADDED POST-PROCESSING) pressure part of the hybrid coordinate') DS['bk'].attrs['description'] = ( '(ADDED POST-PROCESSING) vertical coordinate sigma value') DS['ak'].attrs['units']='Pa' DS['bk'].attrs['units']='None' - - zagl_lvl = ((DS.PH[:, :-1, :, :] + DS.PHB[0, :-1, :, :]) + + zagl_lvl = ((DS.PH[:, :-1, :, :] + DS.PHB[0, :-1, :, :]) /DS.G - DS.HGT[0, :, :]) - + zfull3D = ( 0.5*(zagl_lvl[:, :-1, :, :] + zagl_lvl[:, 1:, :, :]) ) - + # Derive atmospheric temperature [K] # ---------------------------------- gamma = DS.CP / (DS.CP - DS.R_D) @@ -343,24 +353,24 @@ def main(): DS['temp'].attrs['description'] = ('(ADDED POST-PROCESSING) Temperature') DS['temp'].attrs['long_name'] = ('(ADDED POST-PROCESSING) Temperature') DS['temp'].attrs['units'] = 'K' - + # Unstagger U, V, W, Zfull onto Regular Grid # ------------------------------------------ - # For variables staggered x (lon) + # For variables staggered x (lon) # [t,z,y,x'] -> regular mass grid [t,z,y,x]: - # Step 1: Identify variables with the dimension - # 'west_east_stag' and _U not in the variable name (these + # Step 1: Identify variables with the dimension + # 'west_east_stag' and _U not in the variable name (these # are staggered grid identifiers) variables_with_west_east_stag = [var for var in DS.variables if 'west_east_stag' in DS[var].dims and '_U' not in var] - + print( f"{Cyan}Interpolating Staggered Variables to Standard Grid{Nclr}" ) - # dims_list finds the dims of the variable and replaces + # dims_list finds the dims of the variable and replaces # west_east_stag with west_east print('From west_east_stag to west_east: ' + ', '.join(variables_with_west_east_stag)) for var_name in variables_with_west_east_stag: - # Inspiration: pyhton-wrf destag.py + # Inspiration: pyhton-wrf destag.py # https://github.com/NCAR/wrf-python/blob/57116836593b7d7833e11cf11927453c6388487b/src/wrf/destag.py#L9 var = getattr(DS, var_name) dims = var.dims @@ -372,13 +382,13 @@ def main(): new_dims = tuple(dims_list) # Note that XLONG_U is cyclic # LON[x,0] = LON[x,-1] = 0 - + transformed_var = ( - 0.5 * (var.isel(west_east_stag=slice(None, -1)) + 0.5 * (var.isel(west_east_stag=slice(None, -1)) + var.isel(west_east_stag=slice(1, None))) ) - DS[var_name] = xr.DataArray(transformed_var, - dims=new_dims, + DS[var_name] = xr.DataArray(transformed_var, + dims=new_dims, coords={'XLAT':DS['XLAT']}) print(f"\n{DS[var_name].attrs['description']}") @@ -391,13 +401,13 @@ def main(): DS[var_name].attrs['stagger'] = ( ('USTAGGERED IN POST-PROCESSING') ) - + # For variables staggered y (lat) # [t,z,y',x] -> [t,z,y,x] variables_with_south_north_stag = [var for var in DS.variables if 'south_north_stag' in DS[var].dims and '_V' not in var] - + print('From south_north_stag to south_north: ' + ', '.join(variables_with_south_north_stag)) - + for var_name in variables_with_south_north_stag: var = getattr(DS, var_name) dims = var.dims @@ -407,15 +417,15 @@ def main(): dims_list[i]='south_north' break # Stop the loop once the replacement is made new_dims = tuple(dims_list) - + transformed_var = ( - 0.5 * (var.isel(south_north_stag=slice(None, -1)) + 0.5 * (var.isel(south_north_stag=slice(None, -1)) + var.isel(south_north_stag=slice(1, None))) ) - DS[var_name] = xr.DataArray(transformed_var, - dims=new_dims, + DS[var_name] = xr.DataArray(transformed_var, + dims=new_dims, coords={'XLONG':DS['XLONG']}) - + DS[var_name].attrs['description'] = ( '(UNSTAGGERED IN POST-PROCESSING) ' + DS[var_name].attrs['description'] ) @@ -425,13 +435,13 @@ def main(): DS[var_name].attrs['stagger'] = ( 'USTAGGERED IN POST-PROCESSING' ) - + # For variables staggered p/z (height) # [t,z',y,x] -> [t,z,y,x] variables_with_bottom_top_stag = [var for var in DS.variables if 'bottom_top_stag' in DS[var].dims and 'ZNW' not in var and 'phalf' not in var] - + print('From bottom_top_stag to bottom_top: ' + ', '.join(variables_with_bottom_top_stag)) - + for var_name in variables_with_bottom_top_stag: var = getattr(DS, var_name) dims = var.dims @@ -442,9 +452,9 @@ def main(): break # Stop the loop once the replacement is made new_dims = tuple(dims_list) transformed_var = ( - 0.5 * (var.sel(bottom_top_stag=slice(None, -1)) + 0.5 * (var.sel(bottom_top_stag=slice(None, -1)) + var.sel(bottom_top_stag=slice(1, None)))) - + DS[var_name] = xr.DataArray(transformed_var, dims=new_dims) DS[var_name].attrs['description'] = ( '(UNSTAGGERED IN POST-PROCESSING) ' + DS[var_name].attrs['description'] @@ -454,14 +464,14 @@ def main(): ) DS[var_name].attrs['stagger'] = ( 'USTAGGERED IN POST-PROCESSING' - ) - + ) + # Find layer heights above topography; m zfull3D = 0.5 * (zagl_lvl[:,:-1,:,:] + zagl_lvl[:,1:,:,:]) print(f"{Red} Dropping 'Times' variable with non-numerical values") DS = DS.drop_vars("Times") - + # -------------------------------------------------------------- # OpenMars Processing # -------------------------------------------------------------- @@ -490,10 +500,10 @@ def main(): # add ak,bk as variables # add p_half dimensions as vertical grid coordinate - # Compute sigma values. Swap the sigma array upside down - # twice with [::-1] because layers_mid_point_to_boundary() + # Compute sigma values. Swap the sigma array upside down + # twice with [::-1] because layers_mid_point_to_boundary() # needs (sigma[0] = 0, sigma[-1] = 1). - # Then reorganize in the original openMars format with + # Then reorganize in the original openMars format with # (sigma[0] = 1, sigma[-1] = 0) bk = layers_mid_point_to_boundary(DS[model.dim_pfull][::-1], 1.)[::-1] ak = np.zeros(len(DS[model.dim_pfull]) + 1) @@ -507,11 +517,11 @@ def main(): ) DS.phalf.attrs['units'] = ('Pa') - DS = DS.assign(bk=(model.dim_phalf, + DS = DS.assign(bk=(model.dim_phalf, np.array(bk))) - DS = DS.assign(ak=(model.dim_phalf, + DS = DS.assign(ak=(model.dim_phalf, np.zeros(len(DS[model.dim_pfull]) + 1))) - + # Update Variable Description & Longname DS['ak'].attrs['long_name'] = ( '(ADDED POST-PROCESSING) pressure part of the hybrid coordinate' @@ -533,17 +543,17 @@ def main(): variables_with_latu = [var for var in DS.variables if 'latu' in DS[var].dims] variables_with_lonv = [var for var in DS.variables if 'lonv' in DS[var].dims] - + print(f"{Cyan}Changing time units from hours to sols]") DS[model.time] = DS[model.time].values/24. - DS[model.time].attrs['long_name'] = 'time' + DS[model.time].attrs['long_name'] = 'time' DS[model.time].attrs['units'] = 'days since 0000-00-00 00:00:00' - + print(f"{Cyan}Converting reference pressure to [Pa]") DS[model.pfull] = DS[model.pfull].values*100 DS[model.pfull].attrs['units'] = 'Pa' - - # dims_list process finds dims of the variable and replaces + + # dims_list process finds dims of the variable and replaces # west_east_stag with west_east print('From latu to lat: ' + ', '.join(variables_with_latu )) for var_name in variables_with_latu: @@ -569,7 +579,7 @@ def main(): #newlat_val=np.append(newlat_val,0) #newlat_val ==lat transformed_var = ( - 0.5 * (var.isel(latu=slice(None, -1)).values + 0.5 * (var.isel(latu=slice(None, -1)).values + var.isel(latu=slice(1, None)).values) ) @@ -579,25 +589,25 @@ def main(): for i, dim in enumerate(new_dims): if dim == 'lat': list_pads.append((0,1)) - # (begining, end)=(0, 1), meaning 1 padding at + # (begining, end)=(0, 1), meaning 1 padding at # the end of the array else: - list_pads.append((0,0)) + list_pads.append((0,0)) # (begining, end)=(0, 0), no pad on that axis - - transformed_var = np.pad(transformed_var, - list_pads, - mode='constant', + + transformed_var = np.pad(transformed_var, + list_pads, + mode='constant', constant_values=0) - - DS[var_name] = xr.DataArray(transformed_var, - dims=new_dims, + + DS[var_name] = xr.DataArray(transformed_var, + dims=new_dims, coords={'lat':DS['lat']}) DS[var_name].attrs['long_name'] = ( f'(UNSTAGGERED IN POST-PROCESSING) {longname_txt}' ) DS[var_name].attrs['units'] = units_txt - + for var_name in variables_with_lonv: var = getattr(DS, var_name) dims = var.dims @@ -613,34 +623,34 @@ def main(): new_dims = tuple(dims_list) transformed_var = ( - 0.5 * (var.isel(lonv=slice(None, -1)).values + 0.5 * (var.isel(lonv=slice(None, -1)).values + var.isel(lonv=slice(1, None)).values) ) # This is equal to lon[0:-1] - # Add padding + # Add padding list_pads = [] for i, dim in enumerate(new_dims): if dim == 'lon': - list_pads.append((0, 1)) - # (begining, end)=(0, 1), meaning 1 padding at + list_pads.append((0, 1)) + # (begining, end)=(0, 1), meaning 1 padding at # the end of the array else: - list_pads.append((0, 0)) + list_pads.append((0, 0)) # (begining, end)=(0, 0), no pad on that axis - transformed_var = np.pad(transformed_var, list_pads, + transformed_var = np.pad(transformed_var, list_pads, mode='wrap') #TODO with this method V[0] =V[-1]: can we add cyclic point before destaggering? - DS[var_name] = xr.DataArray(transformed_var, - dims=new_dims, + DS[var_name] = xr.DataArray(transformed_var, + dims=new_dims, coords={'lon':DS['lon']}) - + DS[var_name].attrs['long_name'] = ( f'(UNSTAGGERED IN POST-PROCESSING) {longname_txt}' ) DS[var_name].attrs['units'] = units_txt DS.drop_vars(['latu','lonv']) - + # -------------------------------------------------------------- # PCM Processing # -------------------------------------------------------------- @@ -675,7 +685,7 @@ def main(): if 'ap' in DS and 'bp' in DS: # Calculate phalf values from ap and bp phalf_values = (DS.ap.values + DS.bp.values*ref_press) - + # Check the order - for vertical pressure coordinates, we want: # - Lowest pressure (top of atmosphere) at index 0 # - Highest pressure (surface) at index -1 @@ -688,19 +698,19 @@ def main(): # Already in the correct orientation print(f"{Green}PCM phalf values already in correct orientation (lowest at index 0)") DS.attrs['vertical_dimension_flipped'] = False - + # Store phalf values in the dataset DS['phalf'] = (['interlayer'], phalf_values) DS['phalf'].attrs['long_name'] = '(ADDED POST-PROCESSING) pressure at layer interfaces' DS['phalf'].attrs['units'] = 'Pa' - + # Also need to fix pfull to match phalf orientation pfull = DS['pfull'].values if DS.attrs['vertical_dimension_flipped'] and pfull[0] > pfull[-1]: # If we flipped phalf, also ensure pfull has lowest pressure at index 0 if DS['pfull'].values[0] > DS['pfull'].values[-1]: DS['pfull'] = (['altitude'], DS['pfull'].values[::-1]) - print(f"{Yellow}Also flipped pfull values to match phalf orientation") + print(f"{Yellow}Also flipped pfull values to match phalf orientation") # -------------------------------------------------------------- # START PROCESSING FOR ALL MODELS @@ -716,7 +726,7 @@ def main(): DS = DS.isel(**{model.dim_phalf: slice(None, None, -1)}) print(f"{Red}NOTE: all variables flipped along vertical dimension. " f"Top of the atmosphere is now index = 0") - + # Reorder dimensions print(f"{Cyan} Transposing variable dimensions to match order " f"expected in CAP") @@ -747,7 +757,7 @@ def main(): DS[model.areo].attrs['long_name'] = ( '(SCALAR AXIS ADDED POST-PROCESSING)' ) - + print(f"{Red}NOTE: scalar axis added to aerocentric longitude") # Standardize variables names if requested @@ -762,43 +772,43 @@ def main(): # Model has {'ucomp':'U','temp':'T', 'dim_lon'='XLON'...} # Create a reversed dictionary, e.g. {'U':'ucomp','T':'temp'...} # to revert to the original variable names before archiving - # Note that model.() is constructed only from + # Note that model.() is constructed only from # .amescap_profile and may include names not present in file model_dims = dict() model_vars = dict() - + # Loop over optential dims and vars in model for key_i in model.__dict__.keys(): val = getattr(model,key_i) if key_i[0:4] != 'dim_': # Potential variables if val in (list(DS.keys()) + list(DS.coords)): - # Check if the key is a variable (e.g. temp) + # Check if the key is a variable (e.g. temp) # or coordinate (e.g. lat) model_vars[val] = key_i else: # Potential dimensions if val in list(DS.dims): model_dims[val] = key_i[4:] - + # Sort the dictionaries to remove duplicates - # Remove key/val duplicates: e.g if {'lat':'lat'}, there is + # Remove key/val duplicates: e.g if {'lat':'lat'}, there is # no need to rename that variable model_dims = {key: val for key, val in model_dims.items() if key != val} model_vars = {key: val for key, val in model_vars.items() if key != val} - + # Avoiding conflict with derived 'temp' variable in MarsWRF # T is perturb T in MarsWRF, and temp was derived earlier - if (model_type == 'marswrf' and - 'T' in model_vars and + if (model_type == 'marswrf' and + 'T' in model_vars and model_vars['T'] == 'temp'): print(f"{Yellow}Note: Removing 'T' from variable mapping for " f"MarSWRF to avoid conflict with derived 'temp'") del model_vars['T'] # Remove the T -> temp mapping - + print(f"DEBUG: Model dimensions: {model_dims}") print(f"DEBUG: Model variables: {model_vars}") - + # Special handling for PCM to avoid dimension swap errors dimension_swap_failed = False if model_type == 'pcm': @@ -816,7 +826,7 @@ def main(): else: # Normal processing for other GCM types DS = DS.swap_dims(dims_dict = model_dims) - + # Continue with variable renaming regardless of dimension swap status if not dimension_swap_failed: DS = DS.rename_vars(name_dict = model_vars) @@ -830,18 +840,18 @@ def main(): # Add the _nat suffix as if -rn was used, but we still renamed variables if '_nat' not in ext: ext = f'{ext}_nat' - + # -------------------------------------------------------------- # CREATE ATMOS_DAILY, ATMOS_AVERAGE, & ATMOS_DIURN FILES # -------------------------------------------------------------- if args.bin_average and not args.bin_diurn: ext = f'{ext}_average' nday = args.bin_average - + # Calculate time step from original unmodified values dt_in = float(original_time_vals[1] - original_time_vals[0]) print(f"DEBUG: Using original time values with dt_in = {dt_in}") - + # Convert time step to days based on original units dt_days = dt_in if 'minute' in original_time_units.lower() or 'minute' in original_time_desc.lower(): @@ -852,27 +862,27 @@ def main(): print(f"DEBUG: Converting {dt_in} hours to {dt_days} days") else: print(f"DEBUG: No time unit found in original attributes, assuming 'days'") - + # Check if bin size is appropriate if dt_days >= nday: print(f"{Red}***Error***: Requested bin size ({nday} days) is smaller than or equal to " f"the time step in the data ({dt_days:.2f} days)") continue # Skip to next file - + # Calculate samples per day and samples per bin samples_per_day = 1.0 / dt_days samples_per_bin = nday * samples_per_day - + # Need at least one sample per bin if samples_per_bin < 1: print(f"{Red}***Error***: Time sampling in file ({1.0/samples_per_day:.2f} days " f"between samples) is too coarse for {nday}-day bins") continue # Skip to next file - + # Round to nearest integer for coarsen function combinedN = max(1, int(round(samples_per_bin))) print(f"DEBUG: Using {combinedN} time steps per {nday}-day bin") - + # Coarsen and average DS_average = DS.coarsen(**{model.dim_time:combinedN}, boundary='trim').mean() @@ -890,11 +900,11 @@ def main(): if phalf_vals[0] > phalf_vals[-1]: print(f"{Yellow}Warning: phalf orientation incorrect after binning, fixing...") DS_average['phalf'] = (['interlayer'], phalf_vals[::-1]) - + # Create New File, set time dimension as unlimitted base_name = os.path.splitext(fullnameIN)[0] fullnameOUT = f"{base_name}{ext}.nc" - DS_average.to_netcdf(fullnameOUT, unlimited_dims=model.dim_time, + DS_average.to_netcdf(fullnameOUT, unlimited_dims=model.dim_time, format='NETCDF4_CLASSIC') elif args.bin_diurn: @@ -906,13 +916,13 @@ def main(): nday = args.bin_average else: nday = 5 - + print(f"Using {nday}-day bins then diurnal binning") - + # Calculate time step from original unmodified values dt_in = float(original_time_vals[1] - original_time_vals[0]) print(f"DEBUG: Using original time values with dt_in = {dt_in}") - + # Convert time step to days based on original units dt_days = dt_in if 'minute' in original_time_units.lower() or 'minute' in original_time_desc.lower(): @@ -923,19 +933,19 @@ def main(): print(f"DEBUG: Converting {dt_in} hours to {dt_days} days") else: print(f"DEBUG: No time unit found in original attributes, assuming 'days'") - + # Calculate samples per day and check if valid samples_per_day = 1.0 / dt_days if samples_per_day < 1: print(f"{Red}***Error***: Operation not permitted because " f"time sampling in file < one time step per day") continue # Skip to next file - + # Calculate number of steps to bin iperday = int(round(samples_per_day)) combinedN = int(iperday * nday) print(f"DEBUG: Using {combinedN} time steps per {nday}-day bin with {iperday} samples per day") - + # Output Binned Data to New **atmos_diurn.nc file # create a new time of day dimension tod_name = f"time_of_day_{iperday:02d}" @@ -947,11 +957,11 @@ def main(): for i in range(0, int(days/nday)): # Slice original dataset in 5 sol increments downselect = ( - DS.isel(**{model.dim_time:slice(i*combinedN, + DS.isel(**{model.dim_time:slice(i*combinedN, i*combinedN + combinedN)}) ) - # Rename time dimension to time of day and find the + # Rename time dimension to time of day and find the # local time equivalent downselect = downselect.rename({model.dim_time: tod_name}) downselect[tod_name] = ( @@ -964,7 +974,7 @@ def main(): # Add back in the time dimensionn idx = idx.expand_dims({model.dim_time: [i]}) - # Concatenate into new diurn array with a local time + # Concatenate into new diurn array with a local time # and time dimension (stack along time) if DS_diurn is None: DS_diurn = idx @@ -972,27 +982,27 @@ def main(): DS_diurn = xr.concat([DS_diurn, idx], dim=model.dim_time) #TODO # ==== Overwrite the ak, bk arrays=== [AK] - #For some reason I can't track down, the ak('phalf') and bk('phalf') + #For some reason I can't track down, the ak('phalf') and bk('phalf') # turn into ak ('time', 'time_of_day_12','phalf'), in PCM which messes # the pressure interpolation. # Safe approach to fix the dimensions for ak/bk arrays # First check if these variables exist in the diurn dataset ak_dims = None bk_dims = None - + if model.ak in DS_diurn and model.bk in DS_diurn: # Get the dimensions and print for debugging ak_dims = DS_diurn[model.ak].dims bk_dims = DS_diurn[model.bk].dims print(f"DEBUG: {ak_dims} and {bk_dims}") - + # Use explicit dimension checks (safer than len(dims) > 1) if any(dim != model.dim_phalf for dim in ak_dims): print(f"DEBUG: Fixing dimensions for {model.ak} and {model.bk} in diurn file") # Ensure we're assigning the correct structure DS_diurn[model.ak] = DS[model.ak].copy() DS_diurn[model.bk] = DS[model.bk].copy() - + # replace the time dimension with the time dimension from DS_average time_DS = DS[model.dim_time] time_avg_DS = time_DS.coarsen(**{model.dim_time:combinedN},boundary='trim').mean() @@ -1002,7 +1012,7 @@ def main(): DS_diurn[model.dim_time].attrs['long_name'] = ( f'time averaged over {nday} sols' ) - + # Safe phalf check for PCM files if model_type == 'pcm' and DS_diurn is not None and 'phalf' in DS_diurn: try: @@ -1017,19 +1027,20 @@ def main(): DS_diurn['phalf'] = (DS_diurn['phalf'].dims, phalf_vals[::-1]) except Exception as e: print(f"{Yellow}Note: Could not check phalf orientation: {str(e)}") - + # Create New File, set time dimension as unlimitted fullnameOUT = f'{fullnameIN[:-3]}{ext}.nc' - DS_diurn.to_netcdf(fullnameOUT, unlimited_dims=model.dim_time, + DS_diurn.to_netcdf(fullnameOUT, unlimited_dims=model.dim_time, format='NETCDF4_CLASSIC') else: ext = f'{ext}_daily' fullnameOUT = f'{fullnameIN[:-3]}{ext}.nc' - DS.to_netcdf(fullnameOUT, unlimited_dims=model.dim_time, + DS.to_netcdf(fullnameOUT, unlimited_dims=model.dim_time, format='NETCDF4_CLASSIC') print(f"{Cyan}{fullnameOUT} was created") + if __name__ == "__main__": exit_code = main() sys.exit(exit_code) diff --git a/bin/MarsInterp.py b/bin/MarsInterp.py index e7d18c0..644c814 100755 --- a/bin/MarsInterp.py +++ b/bin/MarsInterp.py @@ -26,7 +26,11 @@ * ``os`` * ``time`` * ``matplotlib`` - + * ``re`` + * ``functools`` + * ``traceback`` + * ``sys`` + * ``amescap`` """ # Make print statements appear in color @@ -60,11 +64,13 @@ ) from amescap.Ncdf_wrapper import Ncdf + def debug_wrapper(func): """ - A decorator that wraps a function with error handling based on the + A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -82,6 +88,7 @@ def wrapper(*args, **kwargs): return 1 # Error exit code return wrapper + # ====================================================================== # ARGUMENT PARSER # ====================================================================== @@ -184,6 +191,7 @@ def wrapper(*args, **kwargs): parser.error(f"{Red}{file.name} is not a netCDF file{Nclr}") exit() + # ====================================================================== # DEFINITIONS # ====================================================================== @@ -208,6 +216,7 @@ def wrapper(*args, **kwargs): # MAIN PROGRAM # ====================================================================== + @debug_wrapper def main(): start_time = time.time() @@ -308,6 +317,7 @@ def main(): else: newname = (f"{filepath}/{ifile[:-3]}_{interp_type}.nc") + # ============================================================== # Interpolation # ============================================================== @@ -453,19 +463,12 @@ def main(): else: fnew.copy_Ncvar(fNcdf.variables[ivar]) - - - - - - - - print("\r ", end="") fNcdf.close() fnew.close() print(f"Completed in {(time.time() - start_time):3f} sec") + # ====================================================================== # END OF PROGRAM # ====================================================================== diff --git a/bin/MarsPlot.py b/bin/MarsPlot.py index 8e416c5..9437006 100755 --- a/bin/MarsPlot.py +++ b/bin/MarsPlot.py @@ -19,7 +19,7 @@ * ``warnings`` * ``subprocess`` * ``matplotlib`` - + * ``pypdf`` """ # Make print statements appear in color @@ -76,11 +76,13 @@ global current_version current_version = 3.5 + def debug_wrapper(func): """ - A decorator that wraps a function with error handling based on the + A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -98,6 +100,7 @@ def wrapper(*args, **kwargs): return 1 # Error exit code return wrapper + # ====================================================== # ARGUMENT PARSER # ====================================================== @@ -338,6 +341,7 @@ def wrapper(*args, **kwargs): f"MarsPlot -template -trim{Nclr}") exit() + # ====================================================================== # MAIN PROGRAM # ====================================================================== @@ -487,7 +491,7 @@ def main(): progress(100, 100, "Done") # ============ For Multipage PDF ============ - # Make multipage PDF out of figures in /plots. Remove individual + # Make multipage PDF out of figures in /plots. Remove individual # plots. Debug files when complete. if out_format == "pdf" and len(fig_list) > 0: print("Merging figures...") @@ -549,8 +553,8 @@ def main(): except subprocess.CalledProcessError: # If gs command fails again, prompt user to try # generating PNG instead - print("ERROR with merging PDF, please " - "try a different format, such as PNG.") + print("ERROR with merging PDF, please try a different format, " + "such as PNG.") if debug: raise else: @@ -559,6 +563,8 @@ def main(): f"to use MarsPlot. Type 'MarsPlot -h' if you need " f"more assistance.{Nclr}") exit() + + # ====================================================================== # DATA OPERATION UTILITIES # ====================================================================== @@ -593,13 +599,11 @@ def mean_func(arr, axis): :param arr: the array to be averaged :type arr: array - :param axis: the axis over which to average the array :type axis: int - :return: the mean over the time axis - """ + with warnings.catch_warnings(): warnings.simplefilter("ignore", category = RuntimeWarning) if include_NaNs: @@ -614,23 +618,19 @@ def shift_data(lon, data): :param lon: 1D array of longitude :type lon: array [lon] - :param data: 2D array with last dimension = longitude :type data: array [1,lon] - :raises ValueError: Longitude coordinate type is not recognized. - :return: longitude (-180/180) :rtype: array [lon] - :return: shifted data :rtype: array [1,lon] - + .. note:: Use ``np.ma.hstack`` instead of ``np.hstack`` to keep the masked array properties - """ + nlon = len(lon) # If 1D plot, reshape array if len(data.shape) <= 1: @@ -657,11 +657,10 @@ def MY_func(Ls_cont): :param Ls_cont: solar longitude (``areo``; continuous) :type Ls_cont: array [areo] - :return: the Mars year :rtype: int - """ + return (Ls_cont)//(360.)+1 @@ -673,20 +672,17 @@ def get_lon_index(lon_query_180, lons): :param lon_query_180: longitudes in -180/180: value, ``[min, max]``, or `None` :type lon_query_180: list - :param lons: longitude in 0-360 :type lons: array [lon] - :return: 1D array of file indices :rtype: array - :return: text descriptor for the extracted longitudes :rtype: str - + .. note:: The keyword ``all`` passed as ``-99999`` by the rT() functions - """ + Nlon = len(lons) lon_query_180 = np.array(lon_query_180) @@ -758,10 +754,13 @@ def get_lon_index(lon_query_180, lons): np.arange(0, loni_bounds[1]+1)) txt_lon = (f", lon=avg[{lons[loni_bounds[0]]:.1f}" f"<->{lons[loni_bounds[1]]:.1f}]") - #.. note:: is lon dimension is degenerate, e.g. (time,lev,lat,1) - #loni must be a scalar, otherwise f.variables['var'][time,lev,lat,loni] returns an error + + #.. note:: if lon dimension is degenerate, e.g. (time,lev,lat,1) + # loni must be a scalar, otherwise + # f.variables['var'][time,lev,lat,loni] returns an error if len(np.atleast_1d(loni))==1 and not np.isscalar(loni): - loni=loni[0] + loni = loni[0] + return loni, txt_lon @@ -772,24 +771,23 @@ def get_lat_index(lat_query, lats): :param lat_query: requested latitudes (-90/+90) :type lat_query: list - :param lats: latitude :type lats: array [lat] - :return: 1d array of file indices :rtype: text descriptor for the extracted longitudes - + .. note::T The keyword ``all`` passed as ``-99999`` by the ``rt()`` function - """ + Nlat = len(lats) lat_query = np.array(lat_query) if None in lat_query: # Default to equator lat_query = np.array(0.) + if lat_query.size == 1: # If one latitude provided if lat_query == -99999: @@ -800,6 +798,7 @@ def get_lat_index(lat_query, lats): # Get closest value lati = np.argmin(abs(lat_query-lats)) txt_lat = f", lat={lats[lati]:g}" + elif lat_query.size == 2: # If range of latitudes provided lat_bounds = np.array([np.argmin(abs(lat_query[0] - lats)), @@ -809,6 +808,7 @@ def get_lat_index(lat_query, lats): lat_bounds = np.flipud(lat_bounds) lati = np.arange(lat_bounds[0], lat_bounds[1]+1) txt_lat = f", lat=avg[{lats[lati[0]]:g}<->{lats[lati[-1]]:g}]" + return lati, txt_lat @@ -819,21 +819,18 @@ def get_tod_index(tod_query, tods): :param tod_query: requested time of day (0-24) :type tod_query: list - :param tods: times of day :type tods: array [tod] - :return: file indices :rtype: array [tod] - :return: descriptor for the extracted time of day :rtype: str - + .. note:: The keyword ``all`` is passed as ``-99999`` by the ``rT()`` function - """ + Ntod = len(tods) tod_query = np.array(tod_query) @@ -852,6 +849,7 @@ def get_tod_index(tod_query, tods): todi = np.argmin(abs(tod_query-tods)) txt_tmp = UT_LTtxt(tods[todi]/24., lon_180 = 0., roundmin = 1) txt_tod = f", tod= {txt_tmp}" + elif tod_query.size == 2: # If range of times of day provided tod_bounds = np.array([np.argmin(abs(tod_query[0] - tods)), @@ -866,6 +864,7 @@ def get_tod_index(tod_query, tods): txt_tmp = UT_LTtxt(tods[todi[0]]/24., lon_180 = 0., roundmin = 1) txt_tmp2 = UT_LTtxt(tods[todi[-1]]/24., lon_180 = 0., roundmin = 1) txt_tod = f", tod=avg[{txt_tmp}<->{txt_tmp2}]" + return todi, txt_tod @@ -876,21 +875,18 @@ def get_level_index(level_query, levs): :param level_query: requested pressure [Pa] (depth [m]) :type level_query: float - :param levs: levels (in the native coordinates) :type levs: array [lev] - :return: file indices :rtype: array - :return: descriptor for the extracted pressure (depth) :rtype: str - + .. note:: The keyword ``all`` is passed as ``-99999`` by the ``rT()`` functions - """ + level_query = np.array(level_query) Nz = len(levs) @@ -898,6 +894,7 @@ def get_level_index(level_query, levs): # Default to surface # If level_query >>> Psfc (even for a 10-bar Early Mars sim) level_query = np.array(2*10**7) + if level_query.size == 1: # If one level provided if level_query == -99999: @@ -913,6 +910,7 @@ def get_level_index(level_query, levs): txt_level = ", at sfc" else: txt_level = f", lev={levs[levi]:1.2e} Pa/m" + elif level_query.size == 2: # Bounds are provided levi_bounds = np.array([np.argmin(abs(level_query[0] - levs)), @@ -927,6 +925,7 @@ def get_level_index(level_query, levs): lev_bounds = np.flipud(lev_bounds) txt_level = (f", lev=avg[{lev_bounds[0]:1.2e}" f"<->{lev_bounds[1]:1.2e}] Pa/m") + return levi, txt_level @@ -941,30 +940,29 @@ def get_time_index(Ls_query_360, LsDay): :param Ls_query_360: requested solar longitudes :type Ls_query_360: list - :param LsDay: continuous solar longitudes :type LsDay: array [areo] - :return: file indices :rtype: array - :return: descriptor for the extracted solar longitudes :rtype: str - + .. note:: The keyword ``all`` is passed as ``-99999`` by the ``rT()`` function - """ + if len(np.atleast_1d(LsDay)) == 1: # Special case: file has 1 timestep, transform LsDay -> array LsDay = np.array([LsDay]) Nt = len(LsDay) Ls_query_360 = np.array(Ls_query_360) + if None in Ls_query_360: # Defaultto last timestep Ls_query_360 = np.mod(LsDay[-1], 360.) + if Ls_query_360.size == 1: # If one time provided if Ls_query_360 == -99999: @@ -974,6 +972,7 @@ def get_time_index(Ls_query_360, LsDay): else: # Get Mars Year (MY) of last timestep in file MY_end = MY_func(LsDay[-1]) + if MY_end >= 1: # Check if desired Ls available in this MY Ls_query = Ls_query_360 + (MY_end - 1)*360. @@ -981,12 +980,15 @@ def get_time_index(Ls_query_360, LsDay): else: Ls_query = Ls_query_360 + if Ls_query > LsDay[-1] and MY_end > 1: # Lok one year back MY_end -= 1 Ls_query = Ls_query_360 + (MY_end - 1)*360. + ti = np.argmin(abs(Ls_query - LsDay)) txt_time = f", Ls= (MY{MY_end:02}) {np.mod(LsDay[ti], 360.):.2f}" + elif Ls_query_360.size == 2: # If a range of times provided MY_last = MY_func(LsDay[-1]) @@ -995,10 +997,12 @@ def get_time_index(Ls_query_360, LsDay): Ls_query_last = Ls_query_360[1] + (MY_last-1)*360. else: Ls_query_last = Ls_query_360[1] + if Ls_query_last > LsDay[-1] and MY_last > 1: # Look one MY back MY_last -= 1 Ls_query_last = Ls_query_360[1] + (MY_last-1)*360. + ti_last = np.argmin(abs(Ls_query_last - LsDay)) # Then get first value for that MY MY_beg = MY_last.copy() @@ -1018,8 +1022,10 @@ def get_time_index(Ls_query_360, LsDay): txt_time = (f", Ls= avg [(MY{MY_beg:02}) " f"{np.mod(Ls_bounds[0], 360.):.2f} <-> (MY{MY_last:02}) " f"{np.mod(Ls_bounds[1], 360.):.2f}]") + return ti, txt_time + # ====================================================================== # TEMPLATE UTILITIES # ====================================================================== @@ -1030,15 +1036,13 @@ def filter_input(txt, typeIn="char"): :param txt: text input into ``Custom.in`` to the right of an equal sign :type txt: str - :param typeIn: type of data expected: ``char``, ``float``, ``int``, ``bool``, defaults to ``char`` :type typeIn: str, optional - :return: text input reformatted to ``[val1, val2]`` :rtype: float or array - """ + if txt == "None" or not txt: # If None or empty string return None @@ -1057,6 +1061,7 @@ def filter_input(txt, typeIn="char"): if typeIn == "bool": answ.append(txt.split(",")[i].strip() == "True") return answ + else: if typeIn == "char": answ = txt @@ -1088,11 +1093,10 @@ def rT(typeIn="char"): :param typeIn: type of data expected: ``char``, ``float``, ``int``, ``bool``, defaults to ``char`` :type typeIn: str, optional - :return: text input reformatted to ``[val1, val2]`` :rtype: float or array - """ + global customFileIN raw_input = customFileIN.readline() @@ -1112,6 +1116,7 @@ def rT(typeIn="char"): if raw_input[i] == "=": record = True txt = current_varfull.strip() + return filter_input(txt, typeIn) @@ -1122,24 +1127,19 @@ def read_axis_options(axis_options_txt): :param axis_options_txt: a copy of the last line ``Axis Options`` in ``Custom.in`` templates :type axis_options_txt: str - :return: X-axis bounds as a numpy array or ``None`` if undedefined :rtype: array or None - :return: Y-axis bounds as a numpy array or ``None`` if undedefined :rtype: array or None - :return: colormap (e.g., ``jet``, ``nipy_spectral``) or line options (e.g., ``--r`` for dashed red) :rtype: str - :return: linear (``lin``) or logarithmic (``log``) color scale :rtype: str - :return: projection (e.g., ``ortho -125,45``) :rtype: str - """ + list_txt = axis_options_txt.split(":")[1].split("|") # Xaxis: get bounds @@ -1189,20 +1189,16 @@ def split_varfull(varfull): :param varfull: a ``varfull`` object (e.g, ``atmos_average@2.zsurf``, ``02400.atmos_average@2.zsurf``) :type varfull: str - :return: (sol_array) a sol number or ``None`` (if none provided) :rtype: int or None - :return: (filetype) file type (e.g, ``atmos_average``) :rtype: str - :return: (var) variable of interest (e.g, ``zsurf``) :rtype: str - :return: (``simuID``) simulation ID (Python indexing starts at 0) :rtype: int - """ + if varfull.count(".") == 1: # Default: no sol number provided (e.g., atmos_average2.zsurf). # Extract variables and file from varfull @@ -1243,16 +1239,15 @@ def remove_whitespace(raw_input): :param raw_input: user input for variable, (e.g., ``[atmos_average.temp] + 2)`` :type raw_input: str - :return: raw_input without whitespaces (e.g., ``[atmos_average.temp]+2)`` :rtype: str - """ processed_input = "" for i in range(0, len(raw_input)): if raw_input[i] != " ": processed_input += raw_input[i] + return processed_input @@ -1263,12 +1258,11 @@ def clean_comma_whitespace(raw_input): :param raw_input: dimensions specified by user input to Variable (e.g., ``lat=3. , lon=2 , lev = 10.``) :type raw_input: str - :return: raw_input without whitespaces (e.g., ``lat=3.,lon=2,lev=10.``) :rtype: str - """ + processed_input = "" for i in range(0, len(raw_input)): if raw_input[i] != ",": @@ -1283,15 +1277,15 @@ def get_list_varfull(raw_input): :param raw_input: complex user input to Variable (e.g., ``2*[atmos_average.temp]+[atmos_average2.ucomp]*1000``) :type raw_input: str - :return: list required variables (e.g., [``atmos_average.temp``, ``atmos_average2.ucomp``]) :rtype: str - """ + var_list = [] record = False current_name = "" + for i in range(0, len(raw_input)): if raw_input[i] == "]": record = False @@ -1301,6 +1295,7 @@ def get_list_varfull(raw_input): current_name += raw_input[i] if raw_input[i] == "[": record = True + return var_list @@ -1319,23 +1314,19 @@ def get_overwrite_dim_2D(varfull_bracket, plot_type, fdim1, fdim2, ftod): :param varfull_bracket: a ``varfull`` object with ``{}`` (e.g., ``atmos_average.temp{lev=10;ls=350;lon=155;lat=25}``) :type varfull_bracket: str - :param plot_type: the type of the plot template :type plot_type: str - :param fdim1: X axis dimension for plot :type fdim1: str - :param fdim2: Y axis dimension for plot :type fdim2: str - :return: (varfull) required file and variable (e.g., ``atmos_average.temp``); (fdim_out1) X axis dimension for plot; (fdim_out2) Y axis dimension for plot; (ftod_out) if X or Y axis dimension is time of day - """ + # Initialization: use dimension provided in template fdim_out1 = fdim1 fdim_out2 = fdim2 @@ -1363,31 +1354,37 @@ def get_overwrite_dim_2D(varfull_bracket, plot_type, fdim1, fdim2, ftod): print(f"{Yellow}*** Warning*** Ignoring dimension: " f"{split_dim[i].split('=')[0]} because it is not recognized." f"Valid dimensions = ls, lev, lon, lat, or tod{Nclr}") + if plot_type == "2D_lon_lat": if split_dim[i].split("=")[0] == "ls": fdim_out1 = filter_input(split_dim[i].split("=")[1], "float") if split_dim[i].split("=")[0] == "lev": fdim_out2 = filter_input(split_dim[i].split("=")[1], "float") + if plot_type == "2D_lat_lev": if split_dim[i].split("=")[0] == "ls": fdim_out1 = filter_input(split_dim[i].split("=")[1], "float") if split_dim[i].split("=")[0] == "lon": fdim_out2 = filter_input(split_dim[i].split("=")[1], "float") + if plot_type == "2D_time_lat": if split_dim[i].split("=")[0] == "lon": fdim_out1 = filter_input(split_dim[i].split("=")[1], "float") if split_dim[i].split("=")[0] == "lev": fdim_out2 = filter_input(split_dim[i].split("=")[1], "float") + if plot_type == "2D_lon_lev": if split_dim[i].split("=")[0] == "ls": fdim_out1 = filter_input(split_dim[i].split("=")[1], "float") if split_dim[i].split("=")[0] == "lat": fdim_out2 = filter_input(split_dim[i].split("=")[1], "float") + if plot_type == "2D_time_lev": if split_dim[i].split("=")[0] == "lat": fdim_out1 = filter_input(split_dim[i].split("=")[1], "float") if split_dim[i].split("=")[0] == "lon": fdim_out2 = filter_input(split_dim[i].split("=")[1], "float") + if plot_type == "2D_lon_time": if split_dim[i].split("=")[0] == "lat": fdim_out1 = filter_input(split_dim[i].split("=")[1], "float") @@ -1395,11 +1392,13 @@ def get_overwrite_dim_2D(varfull_bracket, plot_type, fdim1, fdim2, ftod): fdim_out2 = filter_input(split_dim[i].split("=")[1], "float") ftod_out = None + if split_dim[i].split("=")[0] == "tod": # Always get time of day ftod_out = filter_input(split_dim[i].split("=")[1], "float") # .. note:: filter_input() converts text (3 or 4, 5) to variable: # (e.g., numpy.array([3.]) or numpy.array([4., 5.])) + return varfull_no_bracket, fdim_out1, fdim_out2, ftod_out @@ -1412,22 +1411,16 @@ def get_overwrite_dim_1D(varfull_bracket, t_in, lat_in, lon_in, lev_in, :param varfull_bracket: a ``varfull`` object with ``{}`` (e.g., ``atmos_average.temp{lev=10;ls=350;lon=155;lat=25}``) :type varfull_bracket: str - :param t_in: self.t variable :type t_in: array [time] - :param lat_in: self.lat variable :type lat_in: array [lat] - :param lon_in: self.lon variable :type lon_in: array [lon] - :param lev_in: self.lev variable :type lev_in: array [lev] - :param ftod_in: self.ftod variable :type ftod_in: array [tod] - :return: ``varfull`` object without brackets (e.g., ``atmos_average.temp``); :return: (t_out) dimension to update; @@ -1435,8 +1428,8 @@ def get_overwrite_dim_1D(varfull_bracket, t_in, lat_in, lon_in, lev_in, :return: (lon_out) dimension to update; :return: (lev_out) dimension to update; :return: (ftod_out) dimension to update; - """ + # Initialization: Use dimension provided in template t_out = t_in lat_out = lat_in @@ -1476,11 +1469,13 @@ def get_overwrite_dim_1D(varfull_bracket, t_in, lat_in, lon_in, lev_in, ftod_out = filter_input(split_dim[i].split("=")[1], "float") # .. note:: filter_input() converts text ("3" or "4,5") to variable: # (e.g., numpy.array([3.]) or numpy.array([4.,5.])) + return varfull_no_bracket, t_out, lat_out, lon_out, lev_out, ftod_out def create_exec(raw_input, varfull_list): expression_exec = raw_input + for i in range(0, len(varfull_list)): swap_txt = f"[{varfull_list[i]}]" expression_exec = expression_exec.replace(swap_txt, f"VAR[{i:0}]") @@ -1493,19 +1488,16 @@ def fig_layout(subID, nPan, vertical_page=False): :param subID: current subplot number :type subID: int - :param nPan: number of panels desired on page (max = 64, 8x8) :type nPan: int - :param vertical_page: reverse the tuple for portrait format if ``True`` :type vertical_page: bool - :return: plot layout (e.g., ``plt.subplot(nrows = out[0], ncols = out[1], plot_number = out[2])``) :rtype: tuple - """ + out = list((0, 0, 0)) if nPan == 1: @@ -1554,10 +1546,10 @@ def fig_layout(subID, nPan, vertical_page=False): def make_template(): """ Generate the ``Custom.in`` template file. - - :return: Custom.in blank template + :return: Custom.in blank template """ + global customFileIN # Will be modified global current_version newname = os.path.join(output_path,"Custom.in") @@ -1681,8 +1673,8 @@ def give_permission(filename): :param filename: name of the file :type filename: str - """ + # NAS system only: set group permissions to file try: # Catch error and standard output @@ -1701,10 +1693,9 @@ def namelist_parser(Custom_file): :param Custom_file: full path to ``Custom.in`` file :type Custom_file: str - :return: updated global variables, ``FigLayout``, ``objectList`` - """ + global objectList global customFileIN global input_paths @@ -1884,15 +1875,13 @@ def get_figure_header(line_txt): :param line_txt: template header from Custom.in (e.g., ``<<<<<<<<<| Plot 2D lon X lat = True |>>>>>>>>``) :type line_txt: str - :return: (figtype) figure type (e.g., ``Plot 2D lon X lat``) :rtype: str - :return: (boolPlot) whether to plot (``True``) or skip (``False``) figure :rtype: bool - """ + # Plot 2D lon X lat = True line_cmd = line_txt.split("|")[1].strip() # Plot 2D lon X lat @@ -1909,14 +1898,12 @@ def format_lon_lat(lon_lat, type): :param lon_lat: latitude or longitude (+180/-180) :type lon_lat: float - :param type: ``lat`` or ``lon`` :type type: str - :return: formatted label :rtype: str - """ + letter = "" if type == "lon": if lon_lat < 0: @@ -1943,8 +1930,8 @@ def get_Ncdf_num(): :return: a sorted array of sols :rtype: array - """ + # e.g., 00350.fixed.nc list_dir = os.listdir(input_paths[0]) avail_fixed = [k for k in list_dir if ".fixed.nc" in k] @@ -1964,14 +1951,12 @@ def select_range(Ncdf_num, bound): :param Ncdf_num: a sorted array of sols :type Ncdf_num: array - :param bound: a sol (e.g., 0350) or range of sols ``[min max]`` :type bound: int or array - :return: a sorted array of sols within the bounds :rtype: array - """ + bound = np.array(bound) if bound.size == 1: Ncdf_num = Ncdf_num[Ncdf_num == bound] @@ -1997,12 +1982,11 @@ def create_name(root_name): :param root_name: path + default name for the file type (e.g., ``/path/custom.in`` or ``/path/figure.png``) :type root_name: str - :return: the modified name if the file already exists (e.g., ``/path/custom_01.in`` or ``/path/figure_01.png``) :rtype: str - """ + n = 1 # Get extension length (e.g., 2 for *.nc, 3 for *.png) len_ext = len(root_name.split(".")[-1]) @@ -2025,13 +2009,11 @@ def progress(k, Nmax, txt="", success=True): :param k: current iteration of the outer loop :type k: float - :param Nmax: max iteration of the outer loop :type Nmax: float - :return: progress bar (EX: ``Running... [#---------] 10.64 %``) - """ + progress = float(k)/Nmax barLength = 10 block = int(round(barLength*progress)) @@ -2058,22 +2040,18 @@ def prep_file(var_name, file_type, simuID, sol_array): :param var_name: variable to extract (e.g., ``ucomp``) :type var_name: str - :param file_type: MGCM output file type (e.g., ``average``) :type file_name: str - :param simuID: simulation ID number (e.g., 2 for 2nd simulation) :type simuID: int - :param sol_array: date in file name (e.g., [3340,4008]) :type sol_array: list - :return: Dataset or MFDataset object; (var_info) longname and units; (dim_info) dimensions e.g., (``time``, ``lat``,``lon``); (dims) shape of the array e.g., [133,48,96] - """ + global input_paths # Holds sol numbers (e.g., [1500,2400]) global Ncdf_num @@ -2133,6 +2111,7 @@ def __call__(self, x, pos=None): else: return "{x:g}".format(x = x) + # ====================================================================== # FIGURE DEFINITIONS # ====================================================================== @@ -2702,11 +2681,8 @@ def solid_contour(self, xdata, ydata, var, contours): plt.clabel(CS, inline = 1, fontsize = 14, fmt = "%g") -# =============================== - class Fig_2D_lon_lat(Fig_2D): - # Make_template calls method from parent class def make_template(self): super(Fig_2D_lon_lat, self).make_template( @@ -2727,13 +2703,10 @@ def get_topo_2D(self, varfull, plot_type): :param varfull: variable input to main_variable in Custom.in (e.g., ``03340.atmos_average.ucomp``) :type varfull: str - :param plot_type: plot type (e.g., ``Plot 2D lon X time``) :type plot_type: str - :return: topography or ``None`` if no matching ``fixed`` file - """ if not "[" in varfull: @@ -3194,7 +3167,6 @@ def do_plot(self): class Fig_2D_lat_lev(Fig_2D): - def make_template(self): # Calls method from parent class super(Fig_2D_lat_lev, self).make_template("Plot 2D lat X lev", @@ -3253,12 +3225,11 @@ def do_plot(self): class Fig_2D_lon_lev(Fig_2D): - def make_template(self): """ Calls method from parent class - """ + super(Fig_2D_lon_lev, self).make_template("Plot 2D lon X lev", "Ls 0-360 ", "Latitude", "Lon +/-180", "Level[Pa/m]") @@ -3267,8 +3238,8 @@ def make_template(self): def do_plot(self): """ Create figure - """ + ax = super(Fig_2D_lon_lev, self).fig_init() try: # Try to create figure, else return error @@ -3321,14 +3292,12 @@ def do_plot(self): class Fig_2D_time_lev(Fig_2D): - def make_template(self): # Calls method from parent class super(Fig_2D_time_lev, self).make_template("Plot 2D time X lev", "Latitude", "Lon +/-180", "Ls", "Level[Pa/m]") - def do_plot(self): # Create figure ax = super(Fig_2D_time_lev, self).fig_init() @@ -3402,7 +3371,6 @@ def do_plot(self): class Fig_2D_lon_time(Fig_2D): - def make_template(self): # Calls method from parent class super(Fig_2D_lon_time, self).make_template("Plot 2D lon X time", @@ -3556,8 +3524,8 @@ def get_plot_type(self): not passed a template. :return: type of 1D plot to create (1D_time, 1D_lat, etc.) - """ + ncheck = 0 graph_type = "Error" if self.t == -88888 or self.t == "AXIS": @@ -3681,40 +3649,29 @@ def read_NCDF_1D(self, var_name, file_type, simuID, sol_array, :param var_name: variable name (e.g., ``temp``) :type var_name: str - :param file_type: MGCM output file type. Must be ``fixed`` or ``average`` :type file_type: str - :param simuID: number identifier for netCDF file directory :type simuID: str - :param sol_array: sol if different from default (e.g., ``02400``) :type sol_array: str - :param plot_type: ``1D_lon``, ``1D_lat``, ``1D_lev``, or ``1D_time`` :type plot_type: str - :param t_req: Ls requested :type t_req: str - :param lat_req: lat requested :type lat_req: str - :param lon_req: lon requested :type lon_req: str - :param lev_req: level [Pa/m] requested :type lev_req: str - :param ftod_req: time of day requested :type ftod_req: str - :return: (dim_array) the axis (e.g., an array of longitudes), (var_array) the variable extracted - """ f, var_info, dim_info, dims = prep_file(var_name, file_type, diff --git a/bin/MarsPull.py b/bin/MarsPull.py index ba2f98a..4f5d4f1 100755 --- a/bin/MarsPull.py +++ b/bin/MarsPull.py @@ -24,7 +24,6 @@ List of Functions: * download - Queries the requested file from the NAS Data Portal. - """ # make print statements appear in color @@ -42,11 +41,13 @@ import traceback # For printing stack traces import requests # Download data from website + def debug_wrapper(func): """ - A decorator that wraps a function with error handling based on the + A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -64,6 +65,7 @@ def wrapper(*args, **kwargs): return 1 # Error exit code return wrapper + # ------------------------------------------------------ # ARGUMENT PARSER # ------------------------------------------------------ @@ -146,6 +148,7 @@ def wrapper(*args, **kwargs): args = parser.parse_args() debug = args.debug + # ------------------------------------------------------ # DEFINITIONS # ------------------------------------------------------ @@ -168,24 +171,21 @@ def wrapper(*args, **kwargs): 310, 316, 321, 327, 333, 338, 344, 349, 354, 0 ]) + def download(url, file_name): """ Downloads a file from the NAS Data Portal (data.nas.nasa.gov). - + :param url: The url to download from, e.g 'https://data.nas.nasa.\ gov/legacygcm/download_data.php?file=/legacygcmdata/LegacyGCM_\ Ls000_Ls004.nc' :type url: str - :param file_name: The local file_name e.g '/lou/la4/akling/Data/L\ egacyGCM_Ls000_Ls004.nc' :type file_name: str - :return: The requested file(s), downloaded and saved to the \ current directory. - :raises FileNotFoundError: A file-not-found error. - """ _, fname = os.path.split(file_name) @@ -217,11 +217,13 @@ def download(url, file_name): f.write(response.content) print(f"{fname} Done") + def print_file_list(list_of_files): print("Available files:") for file in list_of_files: print(file) + # ------------------------------------------------------ # MAIN FUNCTION # ------------------------------------------------------ @@ -232,7 +234,7 @@ def main(): if not args.list_files and not args.directory_name: print("Error: You must specify either -list or a directory.") sys.exit(1) - + if args.list_files: # Send an HTTP GET request to the URL and store the response. legacy_data = requests.get( @@ -242,7 +244,7 @@ def main(): 'https://data.nas.nasa.gov/mcmcref/fv3betaout1/' ) - # Access the text content of the response, which contains the + # Access the text content of the response, which contains the # webpage's HTML. legacy_dir_text = legacy_data.text fv3_dir_text = fv3_data.text @@ -253,7 +255,7 @@ def main(): ) legacy_urls = re.findall( - fr"{legacy_dir_search}[a-zA-Z0-9_\-\.~:/?#\[\]@!$&'()*+,;=]+", + fr"{legacy_dir_search}[a-zA-Z0-9_\-\.~:/?#\[\]@!$&'()*+,;=]+", legacy_dir_text ) @@ -265,9 +267,9 @@ def main(): legacy_data = requests.get(legacy_urls[0]) legacy_dir_text = legacy_data.text - legacy_files_available = re.findall(r'download="(fort\.11_[0-9]+)"', + legacy_files_available = re.findall(r'download="(fort\.11_[0-9]+)"', legacy_dir_text) - fv3_files_available = re.findall(r'href="[^"]*\/([^"\/]+\.nc)"', + fv3_files_available = re.findall(r'href="[^"]*\/([^"\/]+\.nc)"', fv3_dir_text) print_file_list(legacy_files_available) @@ -280,7 +282,7 @@ def main(): ]: requested_url = ( "https://data.nas.nasa.gov/legacygcm/legacygcmdata/" - + portal_dir + + portal_dir + "/" ) elif portal_dir in ['FV3BETAOUT1']: @@ -292,7 +294,7 @@ def main(): prYellow("ERROR No file requested. Use [-ls --ls] or " "[-f --filename] to specify a file to download.") sys.exit(1) # Return a non-zero exit code - + if args.ls: data_input = np.asarray(args.ls) if len(data_input) == 1: @@ -313,9 +315,9 @@ def main(): i_end += 1 file_list = np.arange(i_start, i_end + 1) - + prCyan(f"Saving {len(file_list)} file(s) to {save_dir}") - + for ii in file_list: if portal_dir == 'ACTIVECLDS_NCDF': # Legacy .nc files @@ -338,12 +340,12 @@ def main(): file_name = save_dir + f print(f"Downloading {url}...") download(url, file_name) - + elif not args.list_files: # If no directory is provided and its not a -list request prYellow("ERROR: A directory must be specified unless using -list.") sys.exit(1) - + # ------------------------------------------------------ # END OF PROGRAM # ------------------------------------------------------ diff --git a/bin/MarsVars.py b/bin/MarsVars.py index 2932b5c..08428d8 100755 --- a/bin/MarsVars.py +++ b/bin/MarsVars.py @@ -40,7 +40,7 @@ * ``time`` * ``io`` * ``locale`` - + * ``amescap`` """ # Make print statements appear in color @@ -78,15 +78,17 @@ ) from amescap.Ncdf_wrapper import Ncdf + # ====================================================== # DEFINITIONS # ====================================================== def debug_wrapper(func): """ - A decorator that wraps a function with error handling based on the + A decorator that wraps a function with error handling based on the --debug flag. """ + @functools.wraps(func) def wrapper(*args, **kwargs): global debug @@ -290,6 +292,7 @@ def wrapper(*args, **kwargs): ], } + def add_help(var_list): help_text = (f"{'VARIABLE':9s} {'DESCRIPTION':33s} {'UNIT':11s} " f"{'REQUIRED VARIABLES':24s} {'SUPPORTED FILETYPES'}" @@ -304,6 +307,7 @@ def add_help(var_list): ) return(help_text) + # ====================================================================== # ARGUMENT PARSER # ====================================================================== @@ -581,6 +585,7 @@ def add_help(var_list): C_dst = (4/3) * (rho_dst/Qext_dst) * Reff_dst # = 12114.286 [m-2] C_ice = (4/3) * (rho_ice/Qext_ice) * Reff_ice # = 2188.874 [m-2] + # ====================================================================== # Helper functions for cross-platform file operations # ====================================================================== @@ -589,17 +594,18 @@ def ensure_file_closed(filepath, delay=0.5): """ Try to ensure a file is not being accessed by the system. This is especially helpful for Windows environments. - + :param filepath: Path to the file :param delay: Delay in seconds to wait for handles to release """ + if not os.path.exists(filepath): return - + # Force garbage collection to release file handles import gc gc.collect() - + # For Windows systems, try to explicitly close open handles if os.name == 'nt': try: @@ -610,64 +616,68 @@ def ensure_file_closed(filepath, delay=0.5): # If we can't open it, wait a bit for any handles to be released print(f"{Yellow}File {filepath} appears to be locked, waiting...{Nclr}") time.sleep(delay) - + # Give the system time to release any file locks time.sleep(delay) + def safe_remove_file(filepath, max_attempts=5, delay=1): """ Safely remove a file with retries for Windows file locking issues - + :param filepath: Path to the file to remove :param max_attempts: Number of attempts to make :param delay: Delay between attempts in seconds :return: True if successful, False otherwise """ + if not os.path.exists(filepath): return True - + print(f"Removing file: {filepath}") - + for attempt in range(max_attempts): try: # Try to ensure file is not locked ensure_file_closed(filepath) - + # Try to remove the file os.remove(filepath) - + # Verify removal if not os.path.exists(filepath): print(f"{Green}File removal successful on attempt {attempt+1}{Nclr}") return True - + except Exception as e: print(f"{Yellow}File removal attempt {attempt+1} failed: {e}{Nclr}") if attempt < max_attempts - 1: print(f"Retrying in {delay} seconds...") time.sleep(delay) - + print(f"{Red}Failed to remove file after {max_attempts} attempts{Nclr}") return False + def safe_move_file(src_file, dst_file, max_attempts=5, delay=1): """ Safely move a file with retries for Windows file locking issues - + :param src_file: Source file path :param dst_file: Destination file path :param max_attempts: Number of attempts to make :param delay: Delay between attempts in seconds :return: True if successful, False otherwise """ + print(f"Moving file: {src_file} -> {dst_file}") - + for attempt in range(max_attempts): try: # Ensure both files have all handles closed ensure_file_closed(src_file) ensure_file_closed(dst_file) - + # On Windows, try to remove the destination first if it exists if os.path.exists(dst_file): if not safe_remove_file(dst_file): @@ -688,20 +698,20 @@ def safe_move_file(src_file, dst_file, max_attempts=5, delay=1): else: # No existing destination, just do a normal move shutil.move(src_file, dst_file) - + # Verify the move was successful if os.path.exists(dst_file) and not os.path.exists(src_file): print(f"{Green}File move successful on attempt {attempt+1}{Nclr}") return True else: raise Exception("File move verification failed") - + except Exception as e: print(f"{Yellow}File move attempt {attempt+1} failed: {e}{Nclr}") if attempt < max_attempts - 1: print(f"Retrying in {delay} seconds...") time.sleep(delay * (attempt + 1)) # Increasing delay for subsequent attempts - + # Last resort: try copy and then remove if move fails after all attempts try: print(f"{Yellow}Trying final fallback: copy + remove{Nclr}") @@ -712,7 +722,7 @@ def safe_move_file(src_file, dst_file, max_attempts=5, delay=1): return True except Exception as e: print(f"{Red}Fallback also failed: {e}{Nclr}") - + print(f"{Red}Failed to move file after {max_attempts} attempts{Nclr}") return False @@ -720,7 +730,10 @@ def safe_move_file(src_file, dst_file, max_attempts=5, delay=1): # ==== FIX FOR UNICODE ENCODING ERROR IN HELP MESSAGE ==== # Helper function to handle Unicode output properly on Windows def safe_print(text): - """Print text safely, handling encoding issues on Windows""" + """ + Print text safely, handling encoding issues on Windows + """ + try: # Try to print directly print(text) @@ -736,7 +749,7 @@ def patched_print_message(self, message, file=None): """Patched version of _print_message that handles Unicode encoding errors""" if file is None: file = sys.stdout - + try: # Try the original method first original_print_message(self, message, file) @@ -754,60 +767,63 @@ def patched_print_message(self, message, file=None): def force_close_netcdf_files(file_or_dir, delay=1.0): """ Aggressively try to ensure netCDF files are closed on Windows systems - + :param file_or_dir: Path to the file or directory to process :param delay: Delay in seconds after forcing closure """ + import gc - + # Only needed on Windows if os.name != 'nt': return - + # Force Python's garbage collection multiple times for _ in range(3): gc.collect() - + # On Windows, add delay to allow file handles to be fully released time.sleep(delay) + def safe_copy_replace(src_file, dst_file, max_attempts=5, delay=1.0): """ Windows-specific approach to copy file contents and replace destination This avoids move operations which are more likely to fail with locking - + :param src_file: Source file path :param dst_file: Destination file path :param max_attempts: Maximum number of retry attempts :param delay: Base delay between attempts (increases with retries) :return: True if successful, False otherwise """ + import gc - + print(f"Performing copy-replace: {src_file} -> {dst_file}") - + # Force garbage collection to release file handles force_close_netcdf_files(os.path.dirname(dst_file), delay=delay) - + # Check if source file exists if not os.path.exists(src_file): print(f"{Red}Source file does not exist: {src_file}{Nclr}") return False - + for attempt in range(max_attempts): try: # Rather than moving, copy the contents with open(src_file, 'rb') as src: src_data = src.read() - + # Close the source file and force GC gc.collect() time.sleep(delay) - + # Write to destination with open(dst_file, 'wb') as dst: dst.write(src_data) - + # Verify file sizes match if os.path.getsize(src_file) == os.path.getsize(dst_file): # Now remove the source file @@ -815,25 +831,23 @@ def safe_copy_replace(src_file, dst_file, max_attempts=5, delay=1.0): os.remove(src_file) except: print(f"{Yellow}Warning: Source file {src_file} could not be removed, but destination was updated{Nclr}") - + print(f"{Green}File successfully replaced on attempt {attempt+1}{Nclr}") return True else: raise Exception("File sizes don't match after copy") - + except Exception as e: print(f"{Yellow}File replace attempt {attempt+1} failed: {e}{Nclr}") # Wait longer with each attempt time.sleep(delay * (attempt + 1)) # Force GC again force_close_netcdf_files(os.path.dirname(dst_file), delay=delay*(attempt+1)) - + print(f"{Red}Failed to replace file after {max_attempts} attempts{Nclr}") return False -# =========================== - def compute_p_3D(ps, ak, bk, shape_out): """ Compute the 3D pressure at layer midpoints. @@ -851,11 +865,13 @@ def compute_p_3D(ps, ak, bk, shape_out): :return: ``p_3D`` The full 3D pressure array (Pa) :rtype: array [time, lev, lat, lon] """ + p_3D = fms_press_calc(ps, ak, bk, lev_type="full") # Swap dimensions 0 and 1 (time and lev) p_3D = p_3D.transpose(lev_T) return p_3D.reshape(shape_out) + # ===================================================================== def compute_rho(p_3D, temp): """ @@ -868,8 +884,10 @@ def compute_rho(p_3D, temp): :return: Density (kg/m^3) :rtype: array [time, lev, lat, lon] """ + return p_3D / (rgas*temp) + # ===================================================================== def compute_xzTau(q, temp, lev, const, f_type): """ @@ -889,12 +907,13 @@ def compute_xzTau(q, temp, lev, const, f_type): :return: ``xzTau`` Dust or ice extinction rate (km-1) :rtype: array [time, lev, lat, lon] """ + if f_type == "diurn": - # Handle diurn files + # Handle diurn files PT = np.repeat(lev, (q.shape[0] * q.shape[1] * q.shape[3] * q.shape[4])) PT = np.reshape(PT, - (q.shape[2], q.shape[0], q.shape[1], q.shape[3], + (q.shape[2], q.shape[0], q.shape[1], q.shape[3], q.shape[4]) ) # (lev, tim, tod, lat, lon) -> (tim, tod, lev, lat, lon) @@ -903,7 +922,7 @@ def compute_xzTau(q, temp, lev, const, f_type): # For average and daily files, ensure proper broadcasting across all times # Create a properly sized pressure field with correct time dimension P = np.zeros_like(q) - + # Fill P with the appropriate pressure level for each vertical index for z in range(len(lev)): if len(q.shape) == 4: # Standard [time, lev, lat, lon] format @@ -918,6 +937,7 @@ def compute_xzTau(q, temp, lev, const, f_type): xzTau = (rho_z * (q*1.e6)/const) * 1000 return xzTau + # ===================================================================== def compute_mmr(xTau, temp, lev, const, f_type): """ @@ -937,13 +957,14 @@ def compute_mmr(xTau, temp, lev, const, f_type): :return: ``q``, Dust or ice mass mixing ratio (ppm) :rtype: array [time, lev, lat, lon] """ + if f_type == "diurn": # Handle diurnal files PT = np.repeat( lev, (xTau.shape[0] * xTau.shape[1] * xTau.shape[3] * xTau.shape[4]) ) PT = np.reshape( - PT, (xTau.shape[2], xTau.shape[0], xTau.shape[1], xTau.shape[3], + PT, (xTau.shape[2], xTau.shape[0], xTau.shape[1], xTau.shape[3], xTau.shape[4]) ) # (lev, tim, tod, lat, lon) -> (tim, tod, lev, lat, lon) @@ -951,7 +972,7 @@ def compute_mmr(xTau, temp, lev, const, f_type): else: # For average and daily files, create properly broadcast pressure array P = np.zeros_like(xTau) - + # Fill P with the appropriate pressure level for each vertical index for z in range(len(lev)): if len(xTau.shape) == 4: # Standard [time, lev, lat, lon] format @@ -959,13 +980,14 @@ def compute_mmr(xTau, temp, lev, const, f_type): else: # Handle other shapes appropriately P[..., z, :, :] = lev[z] - + rho_z = P / (Rd*temp) # Convert extinction (xzTau) from km-1 -> m-1 # Convert mass mixing ratio (q) from ppm (kg/kg) -> mg/kg q = (const * (xTau/1000) / rho_z) / 1.e6 return q + # ===================================================================== def compute_Vg_sed(xTau, nTau, T): """ @@ -980,6 +1002,7 @@ def compute_Vg_sed(xTau, nTau, T): :return: ``Vg`` Dust sedimentation rate (m/s) :rtype: array [time, lev, lat, lon] """ + r0 = ( ((3.*xTau) / (4.*np.pi*rho_dst*nTau)) ** (1/3) * np.exp(-3 * sigma**2 / 2) @@ -994,6 +1017,7 @@ def compute_Vg_sed(xTau, nTau, T): Vg = c * (1 + alpha*Kn)/eta return Vg + # ===================================================================== def compute_w_net(Vg, wvar): """ @@ -1009,9 +1033,11 @@ def compute_w_net(Vg, wvar): :return: `w_net` Net vertical wind speed (m/s) :rtype: array [time, lev, lat, lon] """ + w_net = np.subtract(wvar, Vg) return w_net + # ===================================================================== def compute_theta(p_3D, ps, T, f_type): """ @@ -1028,6 +1054,7 @@ def compute_theta(p_3D, ps, T, f_type): :return: Potential temperature (K) :rtype: array [time, lev, lat, lon] """ + theta_exp = R / (M_co2*Cp) # Broadcast dimensions if f_type == "diurn": @@ -1039,6 +1066,7 @@ def compute_theta(p_3D, ps, T, f_type): return T * (np.reshape(ps, ps_shape)/p_3D) ** theta_exp + # ===================================================================== def compute_w(rho, omega): """ @@ -1064,8 +1092,10 @@ def compute_w(rho, omega): :return: vertical wind (m/s) :rtype: array [time, lev, lat, lon] """ + return -omega / (rho*g) + # ===================================================================== def compute_zfull(ps, ak, bk, T): """ @@ -1082,6 +1112,7 @@ def compute_zfull(ps, ak, bk, T): :return: ``zfull`` (m) :rtype: array [time, lev, lat, lon] """ + zfull = fms_Z_calc( ps, ak, bk, T.transpose(lev_T), topo=0., lev_type="full" ) @@ -1092,6 +1123,7 @@ def compute_zfull(ps, ak, bk, T): zfull = zfull.transpose(lev_T_out) return zfull + # ===================================================================== def compute_zhalf(ps, ak, bk, T): """ @@ -1108,6 +1140,7 @@ def compute_zhalf(ps, ak, bk, T): :return: ``zhalf`` (m) :rtype: array [time, lev, lat, lon] """ + zhalf = fms_Z_calc( ps, ak, bk, T.transpose(lev_T), topo=0., lev_type="half" ) @@ -1118,6 +1151,7 @@ def compute_zhalf(ps, ak, bk, T): zhalf = zhalf.transpose(lev_T_out) return zhalf + # ===================================================================== def compute_DZ_full_pstd(pstd, T, ftype="average"): """ @@ -1145,6 +1179,7 @@ def compute_DZ_full_pstd(pstd, T, ftype="average"): :return: DZ_full_pstd, Layer thicknesses (Pa) :rtype: array [time, lev, lat, lon] """ + # Determine whether the lev dimension is located at i = 1 or i = 2 if ftype == "diurn": axis = 2 @@ -1163,17 +1198,17 @@ def compute_DZ_full_pstd(pstd, T, ftype="average"): # Reshape pstd according to new_shape pstd_reshaped = pstd.reshape(new_shape) - # Ensure pstd is broadcast to match the shape of T + # Ensure pstd is broadcast to match the shape of T # (along non-level dimensions) broadcast_shape = list(T.shape) broadcast_shape[0] = len(pstd) # Keep level dimension the same pstd_broadcast = np.broadcast_to(pstd_reshaped, broadcast_shape) - + # Compute thicknesses using avg. temperature of both layers DZ_full_pstd = np.zeros_like(T) DZ_full_pstd[0:-1, ...] = ( - -rgas * 0.5 - * (T[1:, ...] + T[0:-1, ...]) + -rgas * 0.5 + * (T[1:, ...] + T[0:-1, ...]) / g * np.log(pstd_broadcast[1:, ...] / pstd_broadcast[0:-1, ...]) ) @@ -1188,6 +1223,7 @@ def compute_DZ_full_pstd(pstd, T, ftype="average"): return np.swapaxes(DZ_full_pstd, 0, axis) + # ===================================================================== def compute_N(theta, zfull): """ @@ -1200,6 +1236,7 @@ def compute_N(theta, zfull): :return: ``N``, Brunt Vaisala freqency [rad/s] :rtype: array [time, lev, lat, lon] """ + # Differentiate theta w.r.t. zfull to obdain d(theta)/dz dtheta_dz = dvar_dh(theta.transpose(lev_T), zfull.transpose(lev_T)).transpose(lev_T) @@ -1212,6 +1249,7 @@ def compute_N(theta, zfull): return N + # ===================================================================== def compute_Tco2(P_3D): """ @@ -1224,6 +1262,7 @@ def compute_Tco2(P_3D): :return: CO2 frost point [K] :rtype: array [time, lev, lat, lon] """ + # Set some constants B = -3167.8 # K CO2_triple_pt_P = 518000 # Pa @@ -1240,6 +1279,7 @@ def compute_Tco2(P_3D): return np.where(condition, temp_where_true, temp_where_false) + # ===================================================================== def compute_scorer(N, ucomp, zfull): """ @@ -1254,6 +1294,7 @@ def compute_scorer(N, ucomp, zfull): :return: ``scorer_wl`` Scorer horizontal wavelength (m) :rtype: array [time, lev, lat, lon] """ + # Differentiate U w.r.t. zfull TWICE to obdain d^2U/dz^2 dUdz = dvar_dh(ucomp.transpose(lev_T), zfull.transpose(lev_T)).transpose(lev_T) @@ -1273,6 +1314,7 @@ def compute_scorer(N, ucomp, zfull): return scorer_wl + # ===================================================================== def compute_DP_3D(ps, ak, bk, shape_out): """ @@ -1291,6 +1333,7 @@ def compute_DP_3D(ps, ak, bk, shape_out): :return: ``DP`` Layer thickness in pressure units (Pa) :rtype: array [time, lev, lat, lon] """ + # Get the 3D pressure field from fms_press_calc p_half3D = fms_press_calc(ps, ak, bk, lev_type="half") # fms_press_calc will swap dimensions 0 and 1 so p_half3D has @@ -1304,6 +1347,7 @@ def compute_DP_3D(ps, ak, bk, shape_out): DP = DP_3D.reshape(shape_out) return DP + # ===================================================================== def compute_DZ_3D(ps, ak, bk, temp, shape_out): """ @@ -1341,6 +1385,7 @@ def compute_DZ_3D(ps, ak, bk, temp, shape_out): return DZ + # ===================================================================== def compute_Ep(temp): """ @@ -1353,8 +1398,10 @@ def compute_Ep(temp): :return: ``Ep`` Wave potential energy (J/kg) :rtype: array [time, lev, lat, lon] """ + return 0.5 * g**2 * (zonal_detrend(temp) / (temp*N))**2 + # ===================================================================== def compute_Ek(ucomp, vcomp): """ @@ -1369,8 +1416,10 @@ def compute_Ek(ucomp, vcomp): :return: ``Ek`` Wave kinetic energy (J/kg) :rtype: array [time, lev, lat, lon] """ + return 0.5 * (zonal_detrend(ucomp)**2 + zonal_detrend(vcomp)**2) + # ===================================================================== def compute_MF(UVcomp, w): """ @@ -1383,8 +1432,10 @@ def compute_MF(UVcomp, w): :return: ``u'w'`` or ``v'w'``, Zonal/meridional momentum flux (J/kg) :rtype: array [time, lev, lat, lon] """ + return zonal_detrend(UVcomp) * zonal_detrend(w) + # ===================================================================== def compute_WMFF(MF, rho, lev, interp_type): """ @@ -1414,6 +1465,7 @@ def compute_WMFF(MF, rho, lev, interp_type): :return: The zonal or meridional wave-mean flow forcing (m/s2) :rtype: array [time, lev, lat, lon] """ + # Differentiate the momentum flux (MF) darr_dz = dvar_dh((rho*MF).transpose(lev_T), lev).transpose(lev_T) # Manually swap dimensions 0 and 1 so lev_T has lev for first @@ -1429,12 +1481,12 @@ def compute_WMFF(MF, rho, lev, interp_type): # ===================================================================== -def check_dependencies(f, var, master_list, add_missing=True, +def check_dependencies(f, var, master_list, add_missing=True, dependency_chain=None): """ Check if all dependencies (deps.) for a variable are present in the file, and optionally try to add missing deps.. - + :param f: NetCDF file object :param var: Variable to check deps. for :param master_list: Dict of supported vars and their deps. @@ -1442,56 +1494,57 @@ def check_dependencies(f, var, master_list, add_missing=True, :param dependency_chain: List of vars in the current dep. chain (for detecting cycles) :return: True if all deps. are present or successfully added, False otherwise """ + # Initialize dependency chain if None if dependency_chain is None: dependency_chain = [] - + # Check if we're in a circular dependency if var in dependency_chain: print(f"{Red}Circular dependency detected: " f"{' -> '.join(dependency_chain + [var])}{Nclr}") return False - + # Add current variable to dependency chain dependency_chain = dependency_chain + [var] - + if var not in master_list: print(f"{Red}Variable `{var}` is not in the master list of supported " f"variables.{Nclr}") return False - + # Get the list of required variables for this variable required_vars = master_list[var][2] - + # Check each required variable missing_vars = [] for req_var in required_vars: if req_var not in f.variables: missing_vars.append(req_var) - + if not missing_vars: # All dependencies are present return True - + if not add_missing: # Dependencies are missing but we're not adding them dependency_list = ", ".join(missing_vars) print(f"{Red}Missing dependencies for {var}: {dependency_list}{Nclr}") return False - + # Try to add missing dependencies successfully_added = [] failed_to_add = [] - + for missing_var in missing_vars: # Check if we can add this dependency (must be in master_list) if missing_var in master_list: - # Recursively check dependencies for this variable, passing + # Recursively check dependencies for this variable, passing # the current dependency chain - if check_dependencies(f, - missing_var, - master_list, - add_missing=True, + if check_dependencies(f, + missing_var, + master_list, + add_missing=True, dependency_chain=dependency_chain): # If dependencies are satisfied, try to add the variable try: @@ -1499,17 +1552,17 @@ def check_dependencies(f, var, master_list, add_missing=True, f"be added{Nclr}") # Get the file type and interpolation type f_type, interp_type = FV3_file_type(f) - + # Check if the interpolation type is compatible with this variable if interp_type not in master_list[missing_var][3]: print(f"{Red}Cannot add {missing_var}: incompatible " f"file type {interp_type}{Nclr}") failed_to_add.append(missing_var) continue - + # Mark it as successfully checked successfully_added.append(missing_var) - + except Exception as e: print(f"{Red}Error checking dependency {missing_var}: " f"{str(e)}{Nclr}") @@ -1522,7 +1575,7 @@ def check_dependencies(f, var, master_list, add_missing=True, print(f"{Red}Dependency {missing_var} for {var} is not in the " f"master list and cannot be added automatically{Nclr}") failed_to_add.append(missing_var) - + # Check if all dependencies were added if not failed_to_add: return True @@ -1537,12 +1590,13 @@ def check_dependencies(f, var, master_list, add_missing=True, # ===================================================================== def check_variable_exists(var_name, file_vars): """ - Check if a variable exists in the file, considering alternative + Check if a variable exists in the file, considering alternative naming conventions """ + if var_name in file_vars: return True - + # Handle _micro/_mom naming variations if var_name.endswith('_micro'): alt_name = var_name.replace('_micro', '_mom') @@ -1556,16 +1610,19 @@ def check_variable_exists(var_name, file_vars): # # Special case for dst_num_mom/dst_num_micro # if 'dst_num_micro' in file_vars: # return True - + return False # ===================================================================== def get_existing_var_name(var_name, file_vars): - """Get the actual variable name that exists in the file""" + """ + Get the actual variable name that exists in the file + """ + if var_name in file_vars: return var_name - + # Check alternative names if var_name.endswith('_micro'): alt_name = var_name.replace('_micro', '_mom') @@ -1579,7 +1636,7 @@ def get_existing_var_name(var_name, file_vars): # # Special case for dst_num_mom/dst_num_micro # if 'dst_num_micro' in file_vars: # return 'dst_num_micro' - + return var_name # Return original if no match found @@ -1587,26 +1644,27 @@ def get_existing_var_name(var_name, file_vars): def process_add_variables(file_name, add_list, master_list, debug=False): """ Process the list of variables to add, handling dependencies appropriately. - + :param file_name: Input file path :param add_list: List of variables to add :param master_list: Dictionary of supported variables and their dependencies :param debug: Whether to show debug information """ + # Create a topologically sorted list of variables to add variables_to_add = [] already_in_file = [] - + # First check if all requested variables already exist in the file with Dataset(file_name, "r", format="NETCDF4_CLASSIC") as f: file_vars = set(f.variables.keys()) - + # Check if all requested variables are already in the file for var in add_list: if check_variable_exists(var, file_vars): existing_name = get_existing_var_name(var, file_vars) already_in_file.append((var, existing_name)) - + # If all requested variables exist, report and exit if len(already_in_file) == len(add_list): if len(add_list) == 1: @@ -1626,40 +1684,41 @@ def process_add_variables(file_name, add_list, master_list, debug=False): else: print(f"{Yellow} - {var} (as '{actual_var}'){Nclr}") return - + + def add_with_dependencies(var): """ Helper function to add a variable and its dependencies to the list """ - + # Skip if already processed if var in variables_to_add: return True - + # Open the file to check dependencies with Dataset(file_name, "a", format="NETCDF4_CLASSIC") as f: file_vars = set(f.variables.keys()) - + # Skip if already in file if check_variable_exists(var, file_vars): return True f_type, interp_type = FV3_file_type(f) - + # Check file compatibility if interp_type not in master_list[var][3]: compat_file_fmt = ", ".join(master_list[var][3]) print(f"{Red}ERROR: Variable '{Yellow}{var}{Red}' can only be " f"added to file type: {Yellow}{compat_file_fmt}{Nclr}") return False - + # Check each dependency all_deps_ok = True for dep in master_list[var][2]: # Skip if already in file (including alternative names) if check_variable_exists(dep, file_vars): continue - + # If dependency can be added, try to add it if dep in master_list: if not add_with_dependencies(dep): @@ -1671,20 +1730,20 @@ def add_with_dependencies(var): all_deps_ok = False print(f"{Red}Cannot add {var}: Required dependency {dep} " f"is not in the list of supported variables{Nclr}") - + if all_deps_ok: variables_to_add.append(var) return True else: return False - + # Check all requested variables for var in add_list: if var not in master_list: print(f"{Red}Variable '{var}' is not supported and cannot be " f"added to the file.{Nclr}") continue - + # Skip if already in file with Dataset(file_name, "r", format="NETCDF4_CLASSIC") as f: if check_variable_exists(var, f.variables.keys()): @@ -1696,25 +1755,25 @@ def add_with_dependencies(var): print(f"{Yellow}Variable '{var}' is already in the file." f"{Nclr}") continue - + # Try to add the variable and its dependencies add_with_dependencies(var) - + # Now add the variables in the correct order for var in variables_to_add: try: f = Dataset(file_name, "a", format="NETCDF4_CLASSIC") - + # Skip if already in file (double-check) if check_variable_exists(var, f.variables.keys()): f.close() continue - + print(f"Processing: {var}...") - + # Define lev_T and lev_T_out for this file f_type, interp_type = FV3_file_type(f) - + if f_type == "diurn": lev_T = [2, 1, 0, 3, 4] lev_T_out = [1, 2, 0, 3, 4] @@ -1723,21 +1782,21 @@ def add_with_dependencies(var): lev_T = [1, 0, 2, 3] lev_T_out = lev_T lev_axis = 1 - + # Make lev_T and lev_T_out available to compute functions globals()['lev_T'] = lev_T globals()['lev_T_out'] = lev_T_out - + # temp and ps are always required. Get dimension dim_out = f.variables["temp"].dimensions temp = f.variables["temp"][:] shape_out = temp.shape - + if interp_type == "pfull": # Load ak and bk for pressure calculation. # Usually required. ak, bk = ak_bk_loader(f) - + # level, ps, and p_3d are often required. lev = f.variables["pfull"][:] ps = f.variables["ps"][:] @@ -1850,8 +1909,8 @@ def add_with_dependencies(var): ucomp = f.variables["ucomp"][:] vcomp = f.variables["vcomp"][:] - - # lev_T swaps dims 0 & 1, ensuring level is the first + + # lev_T swaps dims 0 & 1, ensuring level is the first # dimension du_dz = dvar_dh( ucomp.transpose(lev_T), @@ -1911,13 +1970,13 @@ def add_with_dependencies(var): # [lev, lat, t, tod, lon] # -> [t, tod, lev, lat, lon] # [0 1 2 3 4] -> [2 3 0 1 4] - OUT = mass_stream(vcomp.transpose([2, 3, 0, 1, 4]), - lat, + OUT = mass_stream(vcomp.transpose([2, 3, 0, 1, 4]), + lat, lev, type=interp_type).transpose([2, 3, 0, 1, 4]) else: - OUT = mass_stream(vcomp.transpose([1, 2, 3, 0]), - lat, + OUT = mass_stream(vcomp.transpose([1, 2, 3, 0]), + lat, lev, type=interp_type).transpose([3, 0, 1, 2]) # [t, lev, lat, lon] @@ -1972,10 +2031,11 @@ def add_with_dependencies(var): # After successfully adding print(f"{Green}*** Variable '{var}' added successfully ***{Nclr}") f.close() - + except Exception as e: except_message(debug, e, var, file_name) - + + # ====================================================== # MAIN PROGRAM # ====================================================== @@ -2004,16 +2064,16 @@ def __init__(self, name): file_wrapper = FileWrapper(input_file) check_file_tape(file_wrapper) - + # Before any operations, ensure file is accessible ensure_file_closed(input_file) - + # ============================================================== # Remove Function # ============================================================== if args.remove_variable: remove_list = args.remove_variable - + # Create path for temporary file using os.path for cross-platform ifile_tmp = os.path.splitext(input_file)[0] + "_tmp.nc" @@ -2023,7 +2083,7 @@ def __init__(self, name): os.remove(ifile_tmp) except: print(f"{Yellow}Warning: Could not remove existing temporary file: {ifile_tmp}{Nclr}") - + # Open, copy, and close files try: f_IN = Dataset(input_file, "r", format="NETCDF4_CLASSIC") @@ -2032,7 +2092,7 @@ def __init__(self, name): Log.copy_all_vars_from_Ncfile(f_IN, remove_list) f_IN.close() Log.close() - + # Handle differently based on platform if os.name == 'nt': # On Windows, use our specialized copy-replace method @@ -2044,7 +2104,7 @@ def __init__(self, name): # On Unix systems, use standard move shutil.move(ifile_tmp, input_file) print(f"{Cyan}{input_file} was updated{Nclr}") - + except Exception as e: print(f"{Red}Error in remove_variable: {str(e)}{Nclr}") # Clean up temporary file if it exists @@ -2060,27 +2120,27 @@ def __init__(self, name): if args.extract_copy: # Ensure any existing files are properly closed ensure_file_closed(input_file) - + # Create path for extract file using os.path for cross-platform ifile_extract = os.path.splitext(input_file)[0] + "_extract.nc" - + # Remove any existing extract file if os.path.exists(ifile_extract): safe_remove_file(ifile_extract) - + try: f_IN = Dataset(input_file, "r", format="NETCDF4_CLASSIC") # The variable to exclude exclude_list = filter_vars(f_IN, args.extract_copy, giveExclude = True) - + Log = Ncdf(ifile_extract, "Edited in postprocessing") Log.copy_all_dims_from_Ncfile(f_IN) Log.copy_all_vars_from_Ncfile(f_IN, exclude_list) f_IN.close() Log.close() - + # Verify the extract file was created successfully if os.path.exists(ifile_extract): print(f"{Cyan}Extract file created: {ifile_extract}{Nclr}\n") @@ -2098,9 +2158,9 @@ def __init__(self, name): # If the list is not empty, load ak and bk for the pressure # calculation. ak and bk are always necessary. if args.add_variable: - process_add_variables(input_file, args.add_variable, master_list, + process_add_variables(input_file, args.add_variable, master_list, debug) - + # ============================================================== # Vertical Differentiation # ============================================================== @@ -2114,18 +2174,18 @@ def __init__(self, name): f"in {input_file}{Nclr}") f.close() continue - + if interp_type == "pfull": ak, bk = ak_bk_loader(f) - + print(f"Differentiating: {idiff}...") - + if f_type == "diurn": lev_T = [2, 1, 0, 3, 4] else: # If [t, lat, lon] -> [lev, t, lat, lon] lev_T = [1, 0, 2, 3] - + try: # Get the actual variable name in case of alternative names actual_var_name = get_existing_var_name(idiff, f.variables.keys()) @@ -2136,7 +2196,7 @@ def __init__(self, name): # to [kg/m]) new_unit = f"{unit_text[:-2]}/m]" new_lname = f"vertical gradient of {lname_text}" - + # temp and ps are always required. Get dimension dim_out = f.variables["temp"].dimensions if interp_type == "pfull": @@ -2146,11 +2206,11 @@ def __init__(self, name): temp = f.variables["temp"][:] ps = f.variables["ps"][:] # Z is the first axis - zfull = fms_Z_calc(ps, - ak, - bk, - temp.transpose(lev_T), - topo=0., + zfull = fms_Z_calc(ps, + ak, + bk, + temp.transpose(lev_T), + topo=0., lev_type="full").transpose(lev_T) # Average file: zfull = [lev, t, lat, lon] @@ -2175,12 +2235,12 @@ def __init__(self, name): dzfull_pstd = compute_DZ_full_pstd(lev, temp) darr_dz = (dvar_dh( var.transpose(lev_T) - ).transpose(lev_T) + ).transpose(lev_T) / dzfull_pstd) elif interp_type in ["zagl", "zstd"]: lev = f.variables[interp_type][:] - darr_dz = dvar_dh(var.transpose(lev_T), + darr_dz = dvar_dh(var.transpose(lev_T), lev).transpose(lev_T) # .. note:: lev_T swaps dims 0 & 1, ensuring level is # the first dimension for the differentiation @@ -2190,10 +2250,10 @@ def __init__(self, name): var_Ncdf.long_name = (new_lname + cap_str) var_Ncdf.units = new_unit var_Ncdf[:] = darr_dz - + f.close() print(f"{Green}d_dz_{idiff}: Done{Nclr}") - + except Exception as e: except_message(debug, e, idiff, input_file, pre="d_dz_") @@ -2202,21 +2262,21 @@ def __init__(self, name): # ============================================================== for izdetrend in args.zonal_detrend: f = Dataset(input_file, "a", format="NETCDF4_CLASSIC") - + # Use check_variable_exists instead of direct key lookup if not check_variable_exists(izdetrend, f.variables.keys()): print(f"{Red}zonal detrend error: variable {izdetrend} is " f"not in {input_file}{Nclr}") f.close() continue - + print(f"Detrending: {izdetrend}...") - + try: # Get the actual variable name in case of alternative names actual_var_name = get_existing_var_name(izdetrend, f.variables.keys()) var = f.variables[actual_var_name][:] - + lname_text, unit_text = get_longname_unit(f, actual_var_name) new_lname = f"zonal perturbation of {lname_text}" @@ -2228,10 +2288,10 @@ def __init__(self, name): var_Ncdf.long_name = (new_lname + cap_str) var_Ncdf.units = unit_text var_Ncdf[:] = zonal_detrend(var) - + f.close() print(f"{Green}{izdetrend}_p: Done{Nclr}") - + except Exception as e: except_message(debug, e, izdetrend, input_file, ext="_p") @@ -2241,7 +2301,7 @@ def __init__(self, name): for idp_to_dz in args.dp_to_dz: f = Dataset(input_file, "a", format="NETCDF4_CLASSIC") f_type, interp_type = FV3_file_type(f) - + # Use check_variable_exists instead of direct key lookup if not check_variable_exists(idp_to_dz, f.variables.keys()): print(f"{Red}dp_to_dz error: variable {idp_to_dz} is not " @@ -2253,23 +2313,23 @@ def __init__(self, name): # Get the actual variable name in case of alternative names actual_var_name = get_existing_var_name(idp_to_dz, f.variables.keys()) var = f.variables[actual_var_name][:] - + # Ensure required variables (DP, DZ) exist - if not (check_variable_exists('DP', f.variables.keys()) and + if not (check_variable_exists('DP', f.variables.keys()) and check_variable_exists('DZ', f.variables.keys())): print(f"{Red}Error: DP and DZ variables required for " f"conversion. Add them with CAP:\n{Nclr}MarsVars " f"{input_file} -add DP DZ") f.close() continue - + print(f"Converting: {idp_to_dz}...") - - new_unit = (getattr(f.variables[actual_var_name], "units", "") + + new_unit = (getattr(f.variables[actual_var_name], "units", "") + "/m") - new_lname = (getattr(f.variables[actual_var_name], "long_name", "") + new_lname = (getattr(f.variables[actual_var_name], "long_name", "") + " rescaled to meter-1") - + # Get dimension dim_out = f.variables[actual_var_name].dimensions @@ -2280,7 +2340,7 @@ def __init__(self, name): var_Ncdf[:] = ( var * f.variables["DP"][:] / f.variables["DZ"][:] ) - + f.close() print(f"{Green}{idp_to_dz}_dp_to_dz: Done{Nclr}") @@ -2290,7 +2350,7 @@ def __init__(self, name): for idz_to_dp in args.dz_to_dp: f = Dataset(input_file, "a", format="NETCDF4_CLASSIC") f_type, interp_type = FV3_file_type(f) - + # Use check_variable_exists instead of direct key lookup if not check_variable_exists(idz_to_dp, f.variables.keys()): print(f"{Red}dz_to_dp error: variable {idz_to_dp} is not " @@ -2304,20 +2364,20 @@ def __init__(self, name): # Get the actual variable name in case of alternative names actual_var_name = get_existing_var_name(idz_to_dp, f.variables.keys()) var = f.variables[actual_var_name][:] - + # Ensure required variables (DP, DZ) exist - if not (check_variable_exists('DP', f.variables.keys()) and + if not (check_variable_exists('DP', f.variables.keys()) and check_variable_exists('DZ', f.variables.keys())): print(f"{Red}Error: DP and DZ variables required for " f"conversion{Nclr}") f.close() continue - - new_unit = (getattr(f.variables[actual_var_name], "units", "") + + new_unit = (getattr(f.variables[actual_var_name], "units", "") + "/m") - new_lname = (getattr(f.variables[actual_var_name], "long_name", "") + new_lname = (getattr(f.variables[actual_var_name], "long_name", "") + " rescaled to Pa-1") - + # Get dimension dim_out = f.variables[actual_var_name].dimensions @@ -2328,7 +2388,7 @@ def __init__(self, name): var_Ncdf[:] = ( var * f.variables["DP"][:] / f.variables["DZ"][:] ) - + f.close() print(f"{Green}{idz_to_dp}_dz_to_dp: Done{Nclr}") @@ -2353,39 +2413,40 @@ def __init__(self, name): /__ var (dp/g) p_top """ + for icol in args.column_integrate: f = Dataset(input_file, "a") f_type, interp_type = FV3_file_type(f) if interp_type == "pfull": ak, bk = ak_bk_loader(f) - + # Use check_variable_exists instead of direct key lookup if not check_variable_exists(icol, f.variables.keys()): print(f"{Red}column integration error: variable {icol} is " f"not in {input_file}{Nclr}") f.close() continue - + print(f"Performing column integration: {icol}...") - + try: # Get the actual variable name in case of alternative names actual_var_name = get_existing_var_name(icol, f.variables.keys()) var = f.variables[actual_var_name][:] - + lname_text, unit_text = get_longname_unit(f, actual_var_name) # turn "kg/kg" -> "kg/m2" new_unit = f"{unit_text[:-3]}/m2" new_lname = f"column integration of {lname_text}" - + # temp and ps always required # Get dimension dim_in = f.variables["temp"].dimensions shape_in = f.variables["temp"].shape - + # TODO edge cases where time = 1 - + if f_type == "diurn": # if [t, tod, lat, lon] lev_T = [2, 1, 0, 3, 4] @@ -2412,28 +2473,28 @@ def __init__(self, name): var_Ncdf.long_name = (new_lname + cap_str) var_Ncdf.units = new_unit var_Ncdf[:] = out - + f.close() print(f"{Green}{icol}_col: Done{Nclr}") except Exception as e: except_message(debug, e, icol, input_file, ext="_col") - + if args.edit_variable: # Create path for temporary file using os.path for cross-platform ifile_tmp = os.path.splitext(input_file)[0] + "_tmp.nc" - + # Remove any existing temporary file if os.path.exists(ifile_tmp): try: os.remove(ifile_tmp) except: print(f"{Yellow}Warning: Could not remove existing temporary file: {ifile_tmp}{Nclr}") - + try: # Open input file in read mode f_IN = Dataset(input_file, "r", format="NETCDF4_CLASSIC") - + # Create a new temporary file Log = Ncdf(ifile_tmp, "Edited in postprocessing") Log.copy_all_dims_from_Ncfile(f_IN) @@ -2468,7 +2529,7 @@ def __init__(self, name): Log.log_axis1D( name_text, vals, dim_out, lname_text, unit_text, cart_text ) - + # Close files to release handles f_IN.close() Log.close() @@ -2484,7 +2545,7 @@ def __init__(self, name): # On Unix systems, use standard move shutil.move(ifile_tmp, input_file) print(f"{Cyan}{input_file} was updated{Nclr}") - + except Exception as e: print(f"{Red}Error in edit_variable: {str(e)}{Nclr}") # Clean up temporary file if it exists @@ -2494,6 +2555,7 @@ def __init__(self, name): except: pass + # ====================================================================== # END OF PROGRAM # ======================================================================