diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..22cd098
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,59 @@
+name: CI Workflow
+
+on:
+ push:
+ branches:
+ - main
+ pull_request:
+ branches:
+ - main
+ workflow_dispatch:
+
+permissions:
+ contents: write
+
+jobs:
+ format-check:
+ name: Run Black Formatter
+ runs-on: ubuntu-latest
+ steps:
+ - name: Checkout Code
+ uses: actions/checkout@v4
+
+ - name: Set up Python
+ uses: actions/setup-python@v5
+ with:
+ python-version: "3.10"
+
+ - name: Install Black
+ run: pip install black==26.3.0
+
+ - name: Run Black
+ run: black --check .
+
+ docs:
+ name: Build Documentation
+ runs-on: ubuntu-latest
+ steps:
+ - uses: actions/checkout@v4
+ - uses: ConorMacBride/install-package@v1
+ with:
+ apt: libgeos-dev graphviz
+ - uses: actions/setup-python@v5
+ with:
+ python-version: "3.10.5"
+ - name: Install dependencies
+ run: |
+ pip install -r requirements.txt
+ pip install sphinx furo sphinx-changelog
+ - name: Sphinx build
+ run: |
+ sphinx-build docs/source _build
+ - name: Deploy to GitHub Pages
+ uses: peaceiris/actions-gh-pages@v3
+ if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
+ with:
+ publish_branch: gh-pages
+ github_token: ${{ secrets.GITHUB_TOKEN }}
+ publish_dir: _build/
+ force_orphan: true
diff --git a/.github/workflows/documentation.yml b/.github/workflows/documentation.yml
deleted file mode 100644
index fe5b0a9..0000000
--- a/.github/workflows/documentation.yml
+++ /dev/null
@@ -1,33 +0,0 @@
-name: docs
-
-on: [push, pull_request, workflow_dispatch]
-
-permissions:
- contents: write
-
-jobs:
- docs:
- runs-on: ubuntu-latest
- steps:
- - uses: actions/checkout@v3
- - uses: ConorMacBride/install-package@v1
- with:
- apt: libgeos-dev graphviz
- - uses: actions/setup-python@v3
- with:
- python-version: '3.10.5'
- - name: Install dependencies
- run: |
- pip install -r requirements.txt
- pip install sphinx furo sphinx-changelog
- - name: Sphinx build
- run: |
- sphinx-build docs/source _build
- - name: Deploy to GitHub Pages
- uses: peaceiris/actions-gh-pages@v3
- if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' }}
- with:
- publish_branch: gh-pages
- github_token: ${{ secrets.GITHUB_TOKEN }}
- publish_dir: _build/
- force_orphan: true
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index a1845aa..48e99eb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,6 +7,23 @@
*.pdf
*.png
*.json
+*.bat
+*.log
+*.egg-info
+*.swp
+*.bak
/docs/build/*
.VSCodeCounter/*
+/notebooks/*
+/preprint/*
+/poster/*
+*submission/*
+manuscript/*
+first_revision/*
+outputs/*
+local_archive/*
+
+# Local configuration (never commit!)
+pycsa/local_paths.py
+setup_paths_local.sh
diff --git a/README.md b/README.md
index 0c7dfc4..3c9e4ef 100644
--- a/README.md
+++ b/README.md
@@ -1,15 +1,15 @@
-
- 
+
+ 
-Constrained Spectral Approximation Method
+Constrained Spectral Approximation
-
-
+
+
@@ -20,7 +20,7 @@
-The Constrained Spectral Approximation Method (CSAM) is a physically sound and robust method for approximating the spectrum of subgrid-scale orography. It operates under the following constraints:
+The Constrained Spectral Approximation (CSA) method is a physically sound and robust method for approximating the spectrum of subgrid-scale orography. It operates under the following constraints:
* Utilises a limited number of spectral modes (no more than 100)
* Significantly reduces the complexity of physical terrain by over 500 times
@@ -32,15 +32,15 @@ This method is primarily used to represent terrain for weather forecasting purpo
---
-**[Read the documentation here](https://ray-chew.github.io/pyCSAM/index.html)**
+**[Read the documentation here](https://ray-chew.github.io/pyCSA/index.html)**
---
## Requirements
-See [`requirements.txt`](https://github.com/ray-chew/pyCSAM/blob/main/requirements.txt)
+See [`requirements.txt`](https://github.com/ray-chew/pyCSA/blob/main/requirements.txt)
-> **NOTE:** The Sphinx dependencies can be found in [`docs/conf.py`](https://github.com/ray-chew/pyCSAM/blob/main/docs/source/conf.py).
+> **NOTE:** The Sphinx dependencies can be found in [`docs/conf.py`](https://github.com/ray-chew/pyCSA/blob/main/docs/source/conf.py).
## Usage
@@ -51,17 +51,17 @@ Fork this repository and clone your remote fork.
### Configuration
-The user-defined input parameters are in the [`inputs`](https://github.com/ray-chew/pyCSAM/tree/main/inputs) subpackage. These parameters are imported into the run scripts in [`runs`](https://github.com/ray-chew/pyCSAM/tree/main/runs).
+The user-defined input parameters are in the [`inputs`](https://github.com/ray-chew/pyCSA/tree/main/inputs) subpackage. These parameters are imported into the run scripts in [`runs`](https://github.com/ray-chew/pyCSA/tree/main/runs).
### Execution
-A simple setup can be found in [`runs.idealised_isosceles`](https://github.com/ray-chew/pyCSAM/blob/main/runs/idealised_isosceles.py). To execute this run script:
+A simple setup can be found in [`runs.idealised_isosceles`](https://github.com/ray-chew/pyCSA/blob/main/runs/idealised_isosceles.py). To execute this run script:
```console
python3 ./runs/idealised_isosceles.py
```
-However, the codebase is structured such that the user can easily assemble a run script to define their own experiments. Refer to the documentation for the [available APIs](https://ray-chew.github.io/pyCSAM/api.html).
+However, the codebase is structured such that the user can easily assemble a run script to define their own experiments. Refer to the documentation for the [available APIs](https://ray-chew.github.io/pyCSA/api.html).
## License
diff --git a/docs/source/conf.py b/docs/source/conf.py
index aa97dfb..a7ee091 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -18,9 +18,9 @@
# -- Project information -----------------------------------------------------
-project = "CSAM"
-copyright = "2024, Ray Chew, Stamen Dolaptchiev, Maja-Sophie Wedel, Ulrich Achatz"
-author = "Ray Chew, Stamen Dolaptchiev, Maja-Sophie Wedel, Ulrich Achatz"
+project = "CSA"
+copyright = "2024, Ray Chew"
+author = "Ray Chew"
# The full version, including alpha/beta/rc tags
release = "v0.95.1"
diff --git a/docs/source/index.rst b/docs/source/index.rst
index e44c33e..1723362 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -1,4 +1,4 @@
-CSAM's Home
+CSA's Home
===========
.. toctree::
@@ -19,7 +19,7 @@ CSAM's Home
-This page documents the codebase for the Constrained Spectral Approximation Method (CSAM). CSAM is a physically sound and robust method for approximating the spectrum of subgrid-scale orography. It operates under the following constraints:
+This page documents the codebase for the Constrained Spectral Approximation Method (CSA). CSA is a physically sound and robust method for approximating the spectrum of subgrid-scale orography. It operates under the following constraints:
* Utilises a limited number of spectral modes (no more than 100)
* Significantly reduces the complexity of physical terrain by over 500 times
diff --git a/docs/source/modules/runs.icon_usgs_test.rst b/docs/source/modules/runs.icon_usgs_test.rst
index 1be017b..5682646 100644
--- a/docs/source/modules/runs.icon_usgs_test.rst
+++ b/docs/source/modules/runs.icon_usgs_test.rst
@@ -1,7 +1,7 @@
runs.icon_usgs_test
===================
-Run script for CSAM experiments involving the ICON grid and the USGS GMTED 2010 orographic dataset.
+Run script for CSA experiments involving the ICON grid and the USGS GMTED 2010 orographic dataset.
diff --git a/docs/source/quick_start.rst b/docs/source/quick_start.rst
index fc6f7dd..c1d19a3 100644
--- a/docs/source/quick_start.rst
+++ b/docs/source/quick_start.rst
@@ -1,6 +1,6 @@
Quickstart
==========
-A quick and dirty guide to using the CSAM codebase
+A quick and dirty guide to using the CSA codebase
Requirements
^^^^^^^^^^^^
@@ -13,9 +13,9 @@ To run the code, make sure the following packages are installed, preferably in a
Overview
^^^^^^^^
-The CSAM codebase is structured modularly, see :numref:`structure` for a graphical overview.
+The CSA codebase is structured modularly, see :numref:`structure` for a graphical overview.
-The package :mod:`wrappers` provides interfaces to the core code components in :mod:`src` and :mod:`vis`. For example, it defines the First and Second Approximation steps in the CSAM algorithm and applies the tapering of the physical data. Refer to the :doc:`APIs ` for more details.
+The package :mod:`wrappers` provides interfaces to the core code components in :mod:`src` and :mod:`vis`. For example, it defines the First and Second Approximation steps in the CSA algorithm and applies the tapering of the physical data. Refer to the :doc:`APIs ` for more details.
Helper functions and data structures are provided for the processing of user-defined topographies (:mod:`src.var.topo`), grids (:mod:`src.var.grid`), and input parameters (:mod:`src.var.params`).
@@ -24,8 +24,8 @@ These *building blocks* are the assembled for different kinds of experiments in
.. graphviz::
:align: center
:name: structure
- :alt: CSAM program structure
- :caption: CSAM program structure
+ :alt: CSA program structure
+ :caption: CSA program structure
digraph {
graph [
@@ -209,4 +209,4 @@ Alternatively, the run script could be executed via ``ipython``.
.. note::
- The development of the CSAM codebase frontend is currently ongoing. The current design approach of the program structure aims to simplify debugging and diagnostics using an ``ipython`` environment.
\ No newline at end of file
+ The development of the CSA codebase frontend is currently ongoing. The current design approach of the program structure aims to simplify debugging and diagnostics using an ``ipython`` environment.
\ No newline at end of file
diff --git a/examples/etopo_loader_example.py b/examples/etopo_loader_example.py
new file mode 100644
index 0000000..cd90449
--- /dev/null
+++ b/examples/etopo_loader_example.py
@@ -0,0 +1,96 @@
+"""
+Example script demonstrating how to use the ETOPO 2022 15 arc-second loader
+
+This script shows how to:
+1. Set up parameters for ETOPO data loading
+2. Load a regional topography dataset
+3. Apply coarse-graining for different resolutions
+"""
+
+import numpy as np
+from pycsa.core import io, var
+
+
+class params:
+ """Simple parameter class for ETOPO loading"""
+
+ def __init__(self):
+ # Path to ETOPO data directory (must end with /)
+ self.path_etopo = "/home/ray/git-projects/spec_appx/data/etopo_15s/"
+
+ # Define region of interest [lat_min, lat_max]
+ self.lat_extent = [30.0, 45.0]
+
+ # Define region of interest [lon_min, lon_max]
+ self.lon_extent = [-120.0, -105.0]
+
+ # Coarse-graining factor (1 = no coarse-graining, 2 = 2x2 average, etc.)
+ # ETOPO 15" has ~3600 points per 15 degrees, so coarse-graining is useful
+ # etopo_cg = 2 -> ~30" resolution
+ # etopo_cg = 4 -> ~60" resolution (1 arc-minute)
+ # etopo_cg = 8 -> ~120" resolution (2 arc-minutes)
+ self.etopo_cg = 1 # Default: no coarse-graining
+
+
+# Example 1: Load high-resolution data (15 arc-seconds, no coarse-graining)
+print("Example 1: Loading high-resolution ETOPO data...")
+params1 = params()
+params1.etopo_cg = 1
+cell1 = var.topo_cell()
+
+loader1 = io.ncdata.read_etopo_topo(cell1, params1, verbose=True)
+print(f"Loaded: {len(cell1.lat)} x {len(cell1.lon)} = {cell1.topo.shape}")
+print(f"Lat range: {cell1.lat.min():.4f} to {cell1.lat.max():.4f}")
+print(f"Lon range: {cell1.lon.min():.4f} to {cell1.lon.max():.4f}")
+print(f"Elevation range: {cell1.topo.min():.1f} to {cell1.topo.max():.1f} meters")
+print()
+
+
+# Example 2: Load with 4x coarse-graining (~60" resolution)
+print("Example 2: Loading with 4x coarse-graining...")
+params2 = params()
+params2.etopo_cg = 4
+cell2 = var.topo_cell()
+
+loader2 = io.ncdata.read_etopo_topo(cell2, params2)
+print(f"Loaded: {len(cell2.lat)} x {len(cell2.lon)} = {cell2.topo.shape}")
+print(f"Data reduction factor: {cell1.topo.size / cell2.topo.size:.1f}x")
+print()
+
+
+# Example 3: Load a small region
+print("Example 3: Loading a small region (35-37°N, -115 to -110°W)...")
+params3 = params()
+params3.lat_extent = [35.0, 37.0]
+params3.lon_extent = [-115.0, -110.0]
+params3.etopo_cg = 1
+cell3 = var.topo_cell()
+
+loader3 = io.ncdata.read_etopo_topo(cell3, params3)
+print(f"Loaded: {len(cell3.lat)} x {len(cell3.lon)} = {cell3.topo.shape}")
+print(f"Elevation range: {cell3.topo.min():.1f} to {cell3.topo.max():.1f} meters")
+print()
+
+
+# Example 4: Cross-dateline region (if needed)
+print("Example 4: Region spanning across dateline...")
+params4 = params()
+params4.lat_extent = [40.0, 50.0]
+params4.lon_extent = [170.0, -170.0] # Crosses dateline
+params4.etopo_cg = 8
+cell4 = var.topo_cell()
+
+try:
+ loader4 = io.ncdata.read_etopo_topo(cell4, params4)
+ print(f"Loaded: {len(cell4.lat)} x {len(cell4.lon)} = {cell4.topo.shape}")
+except Exception as e:
+ print(f"Note: Dateline crossing may need verification: {e}")
+print()
+
+
+print("Done! All loaders completed successfully.")
+print("\nUsage tips:")
+print('- Set etopo_cg = 1 for full 15" resolution (very high-res!)')
+print('- Set etopo_cg = 4 for ~60" (~1.8 km at equator)')
+print('- Set etopo_cg = 8 for ~120" (~3.6 km at equator)')
+print("- Coarse-graining reduces memory and speeds up processing")
diff --git a/inputs/archive/debug_run.py b/inputs/archive/debug_run.py
index f6c4fa0..39fc6e3 100644
--- a/inputs/archive/debug_run.py
+++ b/inputs/archive/debug_run.py
@@ -1,5 +1,4 @@
-"""User-defined parameters used in the debugger
-"""
+"""User-defined parameters used in the debugger"""
import numpy as np
from src import var
diff --git a/inputs/archive/lam_alaska_pmf_selector.py b/inputs/archive/lam_alaska_pmf_selector.py
index cc73547..6fd33d3 100644
--- a/inputs/archive/lam_alaska_pmf_selector.py
+++ b/inputs/archive/lam_alaska_pmf_selector.py
@@ -3,7 +3,6 @@
import matplotlib.pyplot as plt
import pandas as pd
-
# %%
pmf_diffs = [
-0.0652774741607357,
diff --git a/inputs/icon_global_run.py b/inputs/icon_global_run.py
new file mode 100644
index 0000000..039bc73
--- /dev/null
+++ b/inputs/icon_global_run.py
@@ -0,0 +1,42 @@
+from pycsa.core import var, utils
+from pycsa import local_paths
+
+params = var.params()
+
+params.fn_output = "icon_merit_global"
+utils.transfer_attributes(params, local_paths.paths, prefix="path")
+
+### alaska
+params.lat_extent = [48.0, 64.0, 64.0]
+params.lon_extent = [-148.0, -148.0, -112.0]
+
+### Tierra del Fuego
+params.lat_extent = [-38.0, -56.0, -56.0]
+params.lon_extent = [-76.0, -76.0, -53.0]
+
+### South Pole
+params.lat_extent = [-75.0, -61.0, -61.0]
+params.lon_extent = [-77.0, -50.0, -50.0]
+
+params.tri_set = [13, 104, 105, 106]
+
+params.merit_cg = 100
+
+# Setup the Fourier parameters and object.
+params.nhi = 32
+params.nhj = 64
+
+params.n_modes = 100
+params.padding = 10
+
+params.U, params.V = 10.0, 0.0
+
+params.rect = True
+
+params.debug = False
+params.dfft_first_guess = False
+params.refine = False
+params.verbose = False
+
+params.plot = False
+params.plot_output = True
diff --git a/inputs/icon_regional_run.py b/inputs/icon_regional_run.py
new file mode 100644
index 0000000..0c54552
--- /dev/null
+++ b/inputs/icon_regional_run.py
@@ -0,0 +1,43 @@
+import numpy as np
+from pycsa.core import var, utils
+from pycsa import local_paths
+
+params = var.params()
+
+params.fn_output = "icon_merit_reg"
+utils.transfer_attributes(params, local_paths.paths, prefix="path")
+
+### alaska
+params.lat_extent = [48.0, 64.0, 64.0]
+params.lon_extent = [-148.0, -148.0, -112.0]
+
+### Tierra del Fuego
+params.lat_extent = [-38.0, -56.0, -56.0]
+params.lon_extent = [-76.0, -76.0, -53.0]
+
+### South Pole
+params.lat_extent = [-75.0, -61.0, -61.0]
+params.lon_extent = [-77.0, -50.0, -50.0]
+
+params.tri_set = [13, 104, 105, 106]
+
+params.merit_cg = 100
+
+# Setup the Fourier parameters and object.
+params.nhi = 24
+params.nhj = 48
+
+params.n_modes = 50
+params.padding = 10
+
+params.U, params.V = 10.0, 0.0
+
+params.rect = True
+
+params.debug = False
+params.dfft_first_guess = False
+params.refine = False
+params.verbose = False
+
+params.plot = False
+params.plot_output = True
diff --git a/inputs/lam_run.py b/inputs/lam_run.py
index 396ecd2..5e8aeeb 100644
--- a/inputs/lam_run.py
+++ b/inputs/lam_run.py
@@ -7,21 +7,24 @@
"""
import numpy as np
-from src import var
+from pycsa.core import var, utils
+from pycsa import local_paths
params = var.params()
+utils.transfer_attributes(params, local_paths.paths, prefix="path")
+
run_case = "R2B4"
# run_case = "R2B5"
# run_case = "R2B4_STRW"
-run_case = "R2B4_NN"
-run_case = "R2B4_NE"
-run_case = "R2B4_SE"
-run_case = "R2B4_SS"
-run_case = "R2B4_SW"
-run_case = "R2B4_WW"
-run_case = "R2B4_NW"
+# run_case = "R2B4_NN"
+# run_case = "R2B4_NE"
+# run_case = "R2B4_SE"
+# run_case = "R2B4_SS"
+# run_case = "R2B4_SW"
+# run_case = "R2B4_WW"
+# run_case = "R2B4_NW"
if run_case == "R2B4":
coarse = True
diff --git a/inputs/local_paths_example.py b/inputs/local_paths_example.py
new file mode 100644
index 0000000..4a1b4fb
--- /dev/null
+++ b/inputs/local_paths_example.py
@@ -0,0 +1,12 @@
+from pycsa.core import var
+
+paths = var.obj()
+
+paths.compact_grid = "..."
+paths.compact_topo = "..."
+
+paths.icon_grid = "..."
+paths.output = "..."
+
+paths.merit = "..."
+paths.rema = "..."
diff --git a/inputs/selected_run.py b/inputs/selected_run.py
index 4928a24..ddd78d9 100644
--- a/inputs/selected_run.py
+++ b/inputs/selected_run.py
@@ -3,13 +3,15 @@
* Potential Biases (``POT_BIAS``)
* Iterative refinement (``ITER_REF``)
* FFT vs LSFF in the First Approximation step (``DFFT_FA`` and ``LSFF_FA``)
- * Complementary study on the flux computation; does not appear in the manuscript (``FLUX_SDY``)
+ * Complementary study on the flux computation; does not appear in the manuscript (``FLUX_SDY``)
"""
import numpy as np
-from src import var
+from pycsa import var, utils
+from pycsa import local_paths
params = var.params()
+utils.transfer_attributes(params, local_paths.paths, prefix="path")
# potential biases study
# run_case = "POT_BIAS"
@@ -66,7 +68,7 @@
params.dfft_first_guess = False
params.nhi = 32
params.nhj = 64
- params.rect_set = np.sort([158])
+ params.rect_set = np.sort([210])
params.recompute_rhs = True
params.plot = True
@@ -81,6 +83,8 @@
dfft_tag = "dfft" if params.dfft_first_guess else "lsff"
params.run_case = run_case
params.fn_tag = "selected_alaska%s_%s_fa" % (suffix_tag, dfft_tag)
+params.path_etopo = "./data/etopo_15s/"
+params.etopo_cg = 1 # Coarse-graining factor for ETOPO 15" data
params.lat_extent = [48.0, 64.0, 64.0]
params.lon_extent = [-148.0, -148.0, -112.0]
diff --git a/notebooks/nc_compactifier.ipynb b/notebooks/nc_compactifier.ipynb
deleted file mode 100644
index a2e4425..0000000
--- a/notebooks/nc_compactifier.ipynb
+++ /dev/null
@@ -1,221 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "bce0cd9c-34d5-4d71-9c7a-d61702d9fb09",
- "metadata": {},
- "outputs": [],
- "source": [
- "import numpy as np\n",
- "import netCDF4 as nc\n",
- "import matplotlib.pyplot as plt\n",
- "from topoPy import *\n",
- "\n",
- "df = nc.Dataset('../data/icon_grid_0010_R02B04_G_linked.nc')\n",
- "\n",
- "clat = df.variables['clat'][:]\n",
- "clon = df.variables['clon'][:]\n",
- "clat_verts = df.variables['clat_vertices'][:]\n",
- "clon_verts = df.variables['clon_vertices'][:]\n",
- "links = df.variables['links'][:]\n",
- "\n",
- "# clat = clat*(180/np.pi)\n",
- "# clon = clon*(180/np.pi)\n",
- "# clat_verts = clat_verts*(180/np.pi)\n",
- "# clon_verts = clon_verts*(180/np.pi)\n",
- "\n",
- "datfile = '../data/GMTED2010_topoGlobal_SGS_30ArcSec.nc'\n",
- "var = {'name':'topo','units':'m'}\n",
- "\n",
- "np.random.seed(555)\n",
- "# icon_cell_indexes = np.sort([440, 19442, 5595, 5026, 4793, 4631])\n",
- "# icon_cell_indexes = np.random.randint(0,np.size(clat)-1,36)\n",
- "icon_cell_indexes = np.array([ 343, 1021, 1367, 2045, 2391, 3069, 3415, 4093, 4439,\n",
- " 5117, 5588, 5603, 5985, 6012, 6612, 6627, 7009, 7036,\n",
- " 7636, 7651, 8033, 8060, 8660, 8675, 9057, 9084, 9684,\n",
- " 9699, 10081, 10108]) # cells that are not being found on the grid...\n",
- "\n",
- "# icon_cell_indexes = [3027,3028,3029]\n",
- "# Mount Ebrus; Firehorn; Taunus; Pirin; Langtang; ???\n",
- "print(icon_cell_indexes)\n",
- "print(icon_cell_indexes.shape)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "b1606672-cac9-4388-8664-322f8ff53fcd",
- "metadata": {},
- "outputs": [],
- "source": [
- "comp_clat = clat[icon_cell_indexes]\n",
- "comp_clon = clon[icon_cell_indexes]\n",
- "comp_clat_verts = clat_verts[icon_cell_indexes]\n",
- "comp_clon_verts = clon_verts[icon_cell_indexes]\n",
- "\n",
- "ncfile = Dataset('../data/icon_compact.nc',mode='w') \n",
- "print(ncfile)\n",
- "\n",
- "cell = ncfile.createDimension('cell', np.size(comp_clat)) # latitude axis\n",
- "nv = ncfile.createDimension('nv', 3) # longitude axis\n",
- "for dim in ncfile.dimensions.items():\n",
- " print(dim)\n",
- "\n",
- "ncfile.title='Compact ICON grid for testing and debugging purposes'\n",
- "print(ncfile.title)\n",
- "\n",
- "clat = ncfile.createVariable('clat', np.float32, ('cell',))\n",
- "clat.units = 'radian'\n",
- "clat.long_name = 'center latitude'\n",
- "\n",
- "clon = ncfile.createVariable('clon', np.float32, ('cell',))\n",
- "clon.units = 'radian'\n",
- "clon.long_name = 'center longitude'\n",
- "\n",
- "clat_verts = ncfile.createVariable('clat_vertices', np.float32, ('cell','nv',))\n",
- "clat_verts.units = 'radian'\n",
- "\n",
- "clon_verts = ncfile.createVariable('clon_vertices', np.float32, ('cell','nv',))\n",
- "clon_verts.units = 'radian'\n",
- "\n",
- "clat[:] = comp_clat\n",
- "clon[:] = comp_clon\n",
- "clat_verts[:,:] = comp_clat_verts\n",
- "clon_verts[:,:] = comp_clon_verts\n",
- "\n",
- "print(clon_verts[:,:])\n",
- "\n",
- "ncfile.close(); print('Dataset is closed!')"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "63ac01e1-ac7b-4ab9-86ae-c0720df0965a",
- "metadata": {},
- "outputs": [],
- "source": [
- "links_tmp = links[icon_cell_indexes].flatten()\n",
- "links_tmp = links_tmp[np.where(links_tmp > 0)]\n",
- "links_tmp = np.sort(list(set(links_tmp)))\n",
- "links_tmp -= 1\n",
- "\n",
- "print(links_tmp)\n",
- "\n",
- "lon, lat, z = readnc(datfile, var)\n",
- "nrecords = np.shape(z)[0]; nlon = np.shape(lon)[1]; nlat = np.shape(lat)[1]\n",
- "\n",
- "sz_tmp = np.size(links_tmp)\n",
- "# print(nlat,nlon, np.size(links_tmp))\n",
- "\n",
- "compactified_topo = np.zeros((sz_tmp,nlat, nlon))\n",
- "print(compactified_topo.shape)\n",
- "\n",
- "compactified_lat = np.zeros((sz_tmp,nlat))\n",
- "compactified_lon = np.zeros((sz_tmp,nlon))\n",
- " \n",
- "for i,lnk in enumerate(links_tmp):\n",
- " print(\"i, lnk = \", (i, lnk))\n",
- " compactified_lat[i] = lat[lnk]\n",
- " compactified_lon[i] = lon[lnk]\n",
- " compactified_topo[i] = z[lnk]\n",
- " \n",
- "del lat, lon, z\n",
- "\n",
- "ncfile = Dataset('../data/topo_compact.nc',mode='w',format='NETCDF4_CLASSIC') \n",
- "print(ncfile)\n",
- "\n",
- "nfiles = ncfile.createDimension('nfiles', sz_tmp)\n",
- "lat = ncfile.createDimension('lat', nlat)\n",
- "lon = ncfile.createDimension('lon', nlon)\n",
- "for dim in ncfile.dimensions.items():\n",
- " print(dim)\n",
- "\n",
- "ncfile.title='Compact GMTED2010 USGS Topography grid for testing and debugging purposes'\n",
- "print(ncfile.title)\n",
- "\n",
- "lat = ncfile.createVariable('lat', np.float32, ('nfiles','lat'))\n",
- "lat.units = 'degrees'\n",
- "\n",
- "lon = ncfile.createVariable('lon', np.float32, ('nfiles','lon'))\n",
- "lon.units = 'degrees'\n",
- "\n",
- "topo = ncfile.createVariable('topo', np.float32, ('nfiles','lat','lon'))\n",
- "topo.units = 'meters'\n",
- "\n",
- "lat[:,:] = compactified_lat\n",
- "lon[:,:] = compactified_lon\n",
- "topo[:,:,:] = compactified_topo\n",
- "\n",
- "ncfile.close(); print('Dataset is closed!')"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "298dad73-5d5e-4544-8973-80b51da762a9",
- "metadata": {},
- "outputs": [],
- "source": [
- "lon, lat, z = readnc(datfile, var)\n",
- "nrecords = np.shape(z)[0]; nlon = np.shape(lon)[1]; nlat = np.shape(lat)[1]"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "af54da18-5352-4419-a4d6-a2cf99ecb32c",
- "metadata": {},
- "outputs": [],
- "source": [
- "print(lon[1][:])\n",
- "print(lon[2][:])\n",
- "print(lon[3][:])\n",
- "print(lon[4][:])"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "cf97c8a1-bbf8-4041-bca2-1e838b73c218",
- "metadata": {},
- "outputs": [],
- "source": [
- "print(lat[1][:])\n",
- "print(lat[2][:])\n",
- "print(lat[3][:])\n",
- "print(lat[4][:])"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "id": "9ad749df-3b5d-4e91-b0d4-1036b896c992",
- "metadata": {},
- "outputs": [],
- "source": []
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "Python 3 (ipykernel)",
- "language": "python",
- "name": "python3"
- },
- "language_info": {
- "codemirror_mode": {
- "name": "ipython",
- "version": 3
- },
- "file_extension": ".py",
- "mimetype": "text/x-python",
- "name": "python",
- "nbconvert_exporter": "python",
- "pygments_lexer": "ipython3",
- "version": "3.9.15"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/notebooks/prepare_orog.ipynb b/notebooks/prepare_orog.ipynb
new file mode 100644
index 0000000..d2ccd8c
--- /dev/null
+++ b/notebooks/prepare_orog.ipynb
@@ -0,0 +1,525 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 10,
+ "id": "41815348-c600-4691-a06c-01289a389066",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import sys\n",
+ "# setting path\n",
+ "sys.path.append('..')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 11,
+ "id": "eae8ab31-3641-4ff0-9023-955f97fd6d27",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ ""
+ ]
+ },
+ "execution_count": 11,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "import numpy as np\n",
+ "import netCDF4 as nc\n",
+ "import matplotlib.pyplot as plt\n",
+ "\n",
+ "from src import io, var, utils, fourier, lin_reg, reconstruction\n",
+ "from vis import plotter\n",
+ "\n",
+ "import importlib\n",
+ "importlib.reload(io)\n",
+ "importlib.reload(var)\n",
+ "importlib.reload(utils)\n",
+ "importlib.reload(fourier)\n",
+ "importlib.reload(lin_reg)\n",
+ "importlib.reload(reconstruction)\n",
+ "\n",
+ "importlib.reload(plotter)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "id": "bb9cd4be-ba2f-4921-94ce-448da4ec394b",
+ "metadata": {},
+ "source": [
+ "Prepare orography by generating underlying lat-lon grid of interest."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 12,
+ "id": "7848fa25-c08a-4f87-807e-2b6b05c3b782",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "### tierra del fuego\n",
+ "lat_min = -56.0\n",
+ "lat_max = -38.0\n",
+ "\n",
+ "lon_min = -76.0\n",
+ "lon_max = -53.0\n",
+ "\n",
+ "### alaska\n",
+ "lat_min = 48.0\n",
+ "lat_max = 64.0\n",
+ "\n",
+ "lon_min = -148.0\n",
+ "lon_max = -112.0\n",
+ "\n",
+ "### south pole (REMA)\n",
+ "lat_min = -89.0 \n",
+ "lat_max = -61.0 \n",
+ "\n",
+ "lon_min = -77.0\n",
+ "lon_max = -50.0\n",
+ "\n",
+ "# dlat, dlon in degs\n",
+ "dlat = 0.05\n",
+ "dlon = 0.05\n",
+ "\n",
+ "lat = np.arange(lat_min - dlat, lat_max + dlat, dlat)\n",
+ "lon = np.arange(lon_min - dlon, lon_max + dlon, dlon)\n",
+ "\n",
+ "lat = np.deg2rad(lat)\n",
+ "lon = np.deg2rad(lon)\n",
+ "\n",
+ "lat_mgrid, lon_mgrid = np.meshgrid(lat,lon)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 13,
+ "id": "c81b0521-19d1-4c61-8785-c026c7cd1221",
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "grid = var.grid()\n",
+ " \n",
+ "reader = io.ncdata()\n",
+ "fn = '../data/icon_grid_0012_R02B04_G_linked.nc'\n",
+ "reader.read_dat(fn, grid)\n",
+ "# grid.apply_f(utils.rad2deg)\n",
+ "\n",
+ "vids = []\n",
+ "for lat_ref in lat:\n",
+ " for lon_ref in lon:\n",
+ " vid = utils.pick_cell(lat_ref, lon_ref, grid)\n",
+ " vids.append(vid)\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 14,
+ "id": "87f97fd3-0fa8-4449-8836-74ad08a96d1e",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[ 1.55152168 1.52214143 1.53717002 ... -0.48917738 -0.49075103\n",
+ " -0.47920845]\n",
+ "[[75 79 83 ... 0 0 0]\n",
+ " [75 79 0 ... 0 0 0]\n",
+ " [75 79 83 ... 0 0 0]\n",
+ " ...\n",
+ " [26 50 0 ... 0 0 0]\n",
+ " [26 50 0 ... 0 0 0]\n",
+ " [26 50 0 ... 0 0 0]]\n"
+ ]
+ }
+ ],
+ "source": [
+ "icon_cell_indexes = np.array(list((set(vids))))\n",
+ "\n",
+ "clat = grid.clat\n",
+ "clat_verts = grid.clat_vertices\n",
+ "clon = grid.clon\n",
+ "clon_verts = grid.clon_vertices\n",
+ "links = grid.links\n",
+ "\n",
+ "print(clat)\n",
+ "print(links)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 15,
+ "id": "84e3c0b9-8579-4d72-8c6d-909cbb8650cb",
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "text/plain": [
+ "(120, 108)"
+ ]
+ },
+ "execution_count": 15,
+ "metadata": {},
+ "output_type": "execute_result"
+ }
+ ],
+ "source": [
+ "links[icon_cell_indexes].shape"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 16,
+ "id": "da65b094",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[ 82 86 101 103 105]\n",
+ "[[3 4 0 ... 0 0 0]\n",
+ " [4 5 0 ... 0 0 0]\n",
+ " [3 4 0 ... 0 0 0]\n",
+ " ...\n",
+ " [2 0 0 ... 0 0 0]\n",
+ " [1 2 0 ... 0 0 0]\n",
+ " [2 0 0 ... 0 0 0]]\n"
+ ]
+ }
+ ],
+ "source": [
+ "comp_clat = clat[icon_cell_indexes]\n",
+ "comp_clon = clon[icon_cell_indexes]\n",
+ "comp_clat_verts = clat_verts[icon_cell_indexes]\n",
+ "comp_clon_verts = clon_verts[icon_cell_indexes]\n",
+ "comp_links = links[icon_cell_indexes]\n",
+ "\n",
+ "sorted_unique_links = np.sort(list(set(comp_links[np.where(comp_links > 0)])))\n",
+ "print(sorted_unique_links)\n",
+ "\n",
+ "for new_id, link_id in enumerate(sorted_unique_links):\n",
+ " comp_links[np.where(comp_links == link_id)] = new_id + 1\n",
+ "\n",
+ "print(comp_links)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 17,
+ "id": "0870e3eb-76ca-4443-8612-627a5aba3853",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "\n",
+ "root group (NETCDF4 data model, file format HDF5):\n",
+ " dimensions(sizes): \n",
+ " variables(dimensions): \n",
+ " groups: \n",
+ "('cell', : name = 'cell', size = 120)\n",
+ "('nv', : name = 'nv', size = 3)\n",
+ "('nlinks', : name = 'nlinks', size = 108)\n",
+ "Compact ICON grid for testing and debugging purposes\n",
+ "[[-1.2566371 -0.62831855 -1.2566371 ]\n",
+ " [-1.2566371 -1.2566371 -1.546781 ]\n",
+ " [-0.8575162 -1.2566371 -0.62831855]\n",
+ " [-1.2566371 -0.8575162 -1.2566371 ]\n",
+ " [-0.8575162 -0.62831855 -0.9664932 ]\n",
+ " [-0.9664932 -1.2566371 -0.8575162 ]\n",
+ " [-1.2566371 -0.9664932 -1.2566371 ]\n",
+ " [-1.2566371 -1.2566371 -1.4840158 ]\n",
+ " [-1.2566371 -1.2566371 -1.4434999 ]\n",
+ " [-1.2566371 -1.2566371 -1.4151917 ]\n",
+ " [-1.2566371 -1.2566371 -1.3943659 ]\n",
+ " [-1.3943659 -1.4151917 -1.2566371 ]\n",
+ " [-1.4151917 -1.4434999 -1.2566371 ]\n",
+ " [-0.85713506 -1.0292583 -0.7667262 ]\n",
+ " [-1.0292583 -1.2566371 -0.9664932 ]\n",
+ " [-0.9664932 -0.7667262 -1.0292583 ]\n",
+ " [-1.2566371 -1.0292583 -1.2566371 ]\n",
+ " [-1.0292583 -0.85713506 -1.0697742 ]\n",
+ " [-1.0697742 -1.2566371 -1.0292583 ]\n",
+ " [-1.2566371 -1.0697742 -1.2566371 ]\n",
+ " [-0.85713506 -0.7273876 -0.9199494 ]\n",
+ " [-0.9199494 -0.80066335 -0.9659675 ]\n",
+ " [-0.9659675 -1.0980824 -0.9199494 ]\n",
+ " [-1.0980824 -1.2566371 -1.0697742 ]\n",
+ " [-1.0697742 -0.9199494 -1.0980824 ]\n",
+ " [-0.9199494 -1.0697742 -0.85713506]\n",
+ " [-1.2566371 -1.0980824 -1.2566371 ]\n",
+ " [-1.0980824 -0.9659675 -1.1189082 ]\n",
+ " [-1.1189082 -1.2566371 -1.0980824 ]\n",
+ " [-1.2566371 -1.1189082 -1.2566371 ]\n",
+ " [-1.2566371 -1.2566371 -1.3784156 ]\n",
+ " [-1.2566371 -1.2566371 -1.3658272 ]\n",
+ " [-1.3658272 -1.3784156 -1.2566371 ]\n",
+ " [-1.2566371 -1.2566371 -1.3556495 ]\n",
+ " [-1.2566371 -1.2566371 -1.3472605 ]\n",
+ " [-1.3472605 -1.3556495 -1.2566371 ]\n",
+ " [-1.3556495 -1.3658272 -1.2566371 ]\n",
+ " [-1.2566371 -1.2566371 -1.340239 ]\n",
+ " [-1.2566371 -1.2566371 -1.3342832 ]\n",
+ " [-1.3342832 -1.340239 -1.2566371 ]\n",
+ " [-1.340239 -1.3342832 -1.4168724 ]\n",
+ " [-1.2566371 -1.2566371 -1.329176 ]\n",
+ " [-1.3247527 -1.329176 -1.2566371 ]\n",
+ " [-1.329176 -1.3247527 -1.396579 ]\n",
+ " [-1.3342832 -1.329176 -1.4059856 ]\n",
+ " [-1.329176 -1.3342832 -1.2566371 ]\n",
+ " [-1.0414002 -0.981052 -1.054351 ]\n",
+ " [-0.981052 -1.0414002 -0.96296936]\n",
+ " [-0.96296936 -0.90536904 -0.981052 ]\n",
+ " [-1.3472605 -1.340239 -1.4296209 ]\n",
+ " [-1.340239 -1.3472605 -1.2566371 ]\n",
+ " [-0.90536904 -0.96296936 -0.88184917]\n",
+ " [-0.88184917 -0.82768697 -0.90536904]\n",
+ " [-0.8559369 -0.9359902 -0.81551445]\n",
+ " [-0.9359902 -1.0284123 -0.90066475]\n",
+ " [-0.90066475 -0.81551445 -0.9359902 ]\n",
+ " [-1.0284123 -1.1348586 -1.0009478 ]\n",
+ " [-1.1348586 -1.2566371 -1.1189082 ]\n",
+ " [-1.1189082 -1.0009478 -1.1348586 ]\n",
+ " [-1.0009478 -1.1189082 -0.9659675 ]\n",
+ " [-0.9659675 -0.85670334 -1.0009478 ]\n",
+ " [-0.85670334 -0.76629126 -0.90066475]\n",
+ " [-0.90066475 -1.0009478 -0.85670334]\n",
+ " [-1.0009478 -0.90066475 -1.0284123 ]\n",
+ " [-0.85670334 -0.9659675 -0.80066335]\n",
+ " [-1.3784156 -1.3943659 -1.2566371 ]\n",
+ " [-1.2566371 -1.1348586 -1.2566371 ]\n",
+ " [-1.1348586 -1.0284123 -1.1474469 ]\n",
+ " [-1.1474469 -1.2566371 -1.1348586 ]\n",
+ " [-1.2566371 -1.1474469 -1.2566371 ]\n",
+ " [-1.0284123 -0.9359902 -1.0504717 ]\n",
+ " [-0.9359902 -0.8559369 -0.964892 ]\n",
+ " [-0.964892 -1.0504717 -0.9359902 ]\n",
+ " [-1.0504717 -0.964892 -1.0685759 ]\n",
+ " [-1.0685759 -1.1576246 -1.0504717 ]\n",
+ " [-1.1576246 -1.2566371 -1.1474469 ]\n",
+ " [-1.1474469 -1.0504717 -1.1576246 ]\n",
+ " [-1.0504717 -1.1474469 -1.0284123 ]\n",
+ " [-1.2566371 -1.1576246 -1.2566371 ]\n",
+ " [-1.1576246 -1.0685759 -1.1660136 ]\n",
+ " [-1.1660136 -1.2566371 -1.1576246 ]\n",
+ " [-1.2566371 -1.1660136 -1.2566371 ]\n",
+ " [-0.8559369 -0.7866259 -0.8895696 ]\n",
+ " [-0.8895696 -0.8233815 -0.9179507 ]\n",
+ " [-0.8547899 -0.9179507 -0.8233815 ]\n",
+ " [-0.9179507 -0.8547899 -0.94213307]\n",
+ " [-0.8547899 -0.798587 -0.88184917]\n",
+ " [-0.88184917 -0.94213307 -0.8547899 ]\n",
+ " [-0.94213307 -0.88184917 -0.96296936]\n",
+ " [-0.96296936 -1.0265045 -0.94213307]\n",
+ " [-1.0265045 -1.0964017 -1.0092112 ]\n",
+ " [-1.0092112 -0.94213307 -1.0265045 ]\n",
+ " [-0.94213307 -1.0092112 -0.9179507 ]\n",
+ " [-1.0964017 -1.173035 -1.0836533 ]\n",
+ " [-1.173035 -1.2566371 -1.1660136 ]\n",
+ " [-1.1660136 -1.0836533 -1.173035 ]\n",
+ " [-1.0836533 -1.1660136 -1.0685759 ]\n",
+ " [-1.0685759 -0.9889414 -1.0836533 ]\n",
+ " [-0.9889414 -0.9179507 -1.0092112 ]\n",
+ " [-1.0092112 -1.0836533 -0.9889414 ]\n",
+ " [-1.0836533 -1.0092112 -1.0964017 ]\n",
+ " [-0.9179507 -0.9889414 -0.8895696 ]\n",
+ " [-0.9889414 -1.0685759 -0.964892 ]\n",
+ " [-0.964892 -0.8895696 -0.9889414 ]\n",
+ " [-0.8895696 -0.964892 -0.8559369 ]\n",
+ " [-1.2566371 -1.173035 -1.2566371 ]\n",
+ " [-1.173035 -1.0964017 -1.1789908 ]\n",
+ " [-1.1789908 -1.2566371 -1.173035 ]\n",
+ " [-1.2566371 -1.1789908 -1.2566371 ]\n",
+ " [-1.0964017 -1.0265045 -1.1072886 ]\n",
+ " [-1.0265045 -0.96296936 -1.0414002 ]\n",
+ " [-1.0414002 -1.1072886 -1.0265045 ]\n",
+ " [-1.1072886 -1.0414002 -1.1166952 ]\n",
+ " [-1.1166952 -1.1840981 -1.1072886 ]\n",
+ " [-1.1840981 -1.2566371 -1.1789908 ]\n",
+ " [-1.1789908 -1.1072886 -1.1840981 ]\n",
+ " [-1.1072886 -1.1789908 -1.0964017 ]\n",
+ " [-1.2566371 -1.1840981 -1.2566371 ]\n",
+ " [-1.1840981 -1.1166952 -1.1885214 ]\n",
+ " [-1.1885214 -1.2566371 -1.1840981 ]]\n",
+ "Dataset is closed!\n"
+ ]
+ }
+ ],
+ "source": [
+ "ncfile = nc.Dataset('../data/icon_compact.nc',mode='w') \n",
+ "print(ncfile)\n",
+ "\n",
+ "cell = ncfile.createDimension('cell', np.size(comp_clat)) # latitude axis\n",
+ "nv = ncfile.createDimension('nv', 3) # longitude axis\n",
+ "nlinks = ncfile.createDimension('nlinks', links.shape[1]) # link length\n",
+ "for dim in ncfile.dimensions.items():\n",
+ " print(dim)\n",
+ "\n",
+ "ncfile.title='Compact ICON grid for testing and debugging purposes'\n",
+ "print(ncfile.title)\n",
+ "\n",
+ "clat = ncfile.createVariable('clat', np.float32, ('cell',))\n",
+ "clat.units = 'radian'\n",
+ "clat.long_name = 'center latitude'\n",
+ "\n",
+ "clon = ncfile.createVariable('clon', np.float32, ('cell',))\n",
+ "clon.units = 'radian'\n",
+ "clon.long_name = 'center longitude'\n",
+ "\n",
+ "clat_verts = ncfile.createVariable('clat_vertices', np.float32, ('cell','nv',))\n",
+ "clat_verts.units = 'radian'\n",
+ "\n",
+ "clon_verts = ncfile.createVariable('clon_vertices', np.float32, ('cell','nv',))\n",
+ "clon_verts.units = 'radian'\n",
+ "\n",
+ "clinks = ncfile.createVariable('links', np.int32, ('cell','nlinks',))\n",
+ "clinks.units = ''\n",
+ "\n",
+ "clat[:] = comp_clat\n",
+ "clon[:] = comp_clon\n",
+ "clat_verts[:,:] = comp_clat_verts\n",
+ "clon_verts[:,:] = comp_clon_verts\n",
+ "clinks[:,:] = comp_links\n",
+ "\n",
+ "print(clon_verts[:,:])\n",
+ "\n",
+ "ncfile.close(); print('Dataset is closed!')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 18,
+ "id": "43652c93-3d56-4251-8241-8671f251d2ca",
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "[ 81 85 100 102 104]\n",
+ "(5, 2400, 3600)\n",
+ "i, lnk = (0, 81)\n",
+ "i, lnk = (1, 85)\n",
+ "i, lnk = (2, 100)\n",
+ "i, lnk = (3, 102)\n",
+ "i, lnk = (4, 104)\n",
+ "\n",
+ "root group (NETCDF4_CLASSIC data model, file format HDF5):\n",
+ " dimensions(sizes): \n",
+ " variables(dimensions): \n",
+ " groups: \n",
+ "('nfiles', : name = 'nfiles', size = 5)\n",
+ "('lat', : name = 'lat', size = 2400)\n",
+ "('lon', : name = 'lon', size = 3600)\n",
+ "Compact GMTED2010 USGS Topography grid for testing and debugging purposes\n",
+ "Dataset is closed!\n"
+ ]
+ }
+ ],
+ "source": [
+ "links_tmp = links[icon_cell_indexes].flatten()\n",
+ "links_tmp = links_tmp[np.where(links_tmp > 0)]\n",
+ "links_tmp = np.sort(list(set(links_tmp)))\n",
+ "links_tmp -= 1\n",
+ "\n",
+ "print(links_tmp)\n",
+ "\n",
+ "topo = var.topo()\n",
+ "fn = '../data/GMTED2010_topoGlobal_SGS_30ArcSec.nc'\n",
+ "reader.read_dat(fn, topo)\n",
+ "\n",
+ "lon = topo.lon\n",
+ "lat = topo.lat\n",
+ "\n",
+ "z = topo.topo\n",
+ "\n",
+ "\n",
+ "del topo\n",
+ "\n",
+ "# lon, lat, z = readnc(datfile, var)\n",
+ "nrecords = np.shape(z)[0]; nlon = np.shape(lon)[1]; nlat = np.shape(lat)[1]\n",
+ "\n",
+ "sz_tmp = np.size(links_tmp)\n",
+ "\n",
+ "compactified_topo = np.zeros((sz_tmp,nlat, nlon))\n",
+ "print(compactified_topo.shape)\n",
+ "\n",
+ "compactified_lat = np.zeros((sz_tmp,nlat))\n",
+ "compactified_lon = np.zeros((sz_tmp,nlon))\n",
+ " \n",
+ "for i,lnk in enumerate(links_tmp):\n",
+ " print(\"i, lnk = \", (i, lnk))\n",
+ " compactified_lat[i] = lat[lnk]\n",
+ " compactified_lon[i] = lon[lnk]\n",
+ " compactified_topo[i] = z[lnk]\n",
+ " \n",
+ "del lat, lon, z\n",
+ "\n",
+ "ncfile = nc.Dataset('../data/topo_compact.nc',mode='w',format='NETCDF4_CLASSIC') \n",
+ "print(ncfile)\n",
+ "\n",
+ "nfiles = ncfile.createDimension('nfiles', sz_tmp)\n",
+ "lat = ncfile.createDimension('lat', nlat)\n",
+ "lon = ncfile.createDimension('lon', nlon)\n",
+ "for dim in ncfile.dimensions.items():\n",
+ " print(dim)\n",
+ "\n",
+ "ncfile.title='Compact GMTED2010 USGS Topography grid for testing and debugging purposes'\n",
+ "print(ncfile.title)\n",
+ "\n",
+ "lat = ncfile.createVariable('lat', np.float32, ('nfiles','lat'))\n",
+ "lat.units = 'degrees'\n",
+ "\n",
+ "lon = ncfile.createVariable('lon', np.float32, ('nfiles','lon'))\n",
+ "lon.units = 'degrees'\n",
+ "\n",
+ "topo = ncfile.createVariable('topo', np.float32, ('nfiles','lat','lon'))\n",
+ "topo.units = 'meters'\n",
+ "\n",
+ "lat[:,:] = compactified_lat\n",
+ "lon[:,:] = compactified_lon\n",
+ "topo[:,:,:] = compactified_topo\n",
+ "\n",
+ "ncfile.close(); print('Dataset is closed!')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "id": "3b34ddf0-f9f4-4c14-bd97-50c43cdee9f7",
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.10.5"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/pycsa/__init__.py b/pycsa/__init__.py
new file mode 100644
index 0000000..e23924f
--- /dev/null
+++ b/pycsa/__init__.py
@@ -0,0 +1,43 @@
+"""
+pyCSA: Constrained Spectral Approximation Method
+
+A Python package for spectral approximation methods applied to topographic analysis.
+"""
+
+__version__ = "0.95.1"
+
+# Core modules - commonly used data structures and utilities
+from pycsa.core import (
+ var,
+ utils,
+ io,
+ physics,
+ fourier,
+ delaunay,
+ reconstruction,
+ lin_reg,
+)
+
+# Wrappers - high-level interfaces
+from pycsa.wrappers import interface, diagnostics
+
+# Plotting - visualization tools
+from pycsa.plotting import plotter, cart_plot
+
+__all__ = [
+ # Core
+ "var",
+ "utils",
+ "io",
+ "physics",
+ "fourier",
+ "delaunay",
+ "reconstruction",
+ "lin_reg",
+ # Wrappers
+ "interface",
+ "diagnostics",
+ # Plotting
+ "plotter",
+ "cart_plot",
+]
diff --git a/src/__init__.py b/pycsa/core/__init__.py
similarity index 100%
rename from src/__init__.py
rename to pycsa/core/__init__.py
diff --git a/pycsa/core/buffer_pool.py b/pycsa/core/buffer_pool.py
new file mode 100644
index 0000000..0e32ea5
--- /dev/null
+++ b/pycsa/core/buffer_pool.py
@@ -0,0 +1,145 @@
+"""
+Dynamic buffer pool for reusing NumPy arrays across multiple computations.
+
+This module provides memory-efficient buffer management for spectral approximation
+computations where array sizes may vary between cells (e.g., different amounts of
+topography data per cell).
+"""
+
+import numpy as np
+
+
+class BufferPool:
+ """Dynamic buffer pool that auto-grows to handle variable array sizes.
+
+ Strategy:
+ - Keeps the largest buffer seen for each key
+ - Returns views (slices) for smaller requests → zero-copy!
+ - Auto-grows when larger size requested
+ - Tracks usage statistics for performance analysis
+
+ This is particularly effective for workflows processing many cells with
+ varying data sizes, as it eliminates repeated memory allocations while
+ adapting to size variations.
+
+ Examples
+ --------
+ >>> pool = BufferPool()
+ >>> # First call allocates
+ >>> arr1 = pool.get_or_create('coeff', (1000, 100), np.float64)
+ >>> # Second call with same size reuses buffer
+ >>> arr2 = pool.get_or_create('coeff', (1000, 100), np.float64)
+ >>> # Smaller size returns a view of existing buffer
+ >>> arr3 = pool.get_or_create('coeff', (500, 100), np.float64)
+ >>> # Larger size triggers reallocation
+ >>> arr4 = pool.get_or_create('coeff', (2000, 100), np.float64)
+ """
+
+ def __init__(self):
+ """Initialize empty buffer pool."""
+ self.buffers = {} # key -> (max_shape, array)
+ self.stats = {} # key -> {hits, misses, grows}
+
+ def get_or_create(self, key, shape, dtype=np.float64):
+ """Get buffer from pool, creating or growing as needed.
+
+ Parameters
+ ----------
+ key : str
+ Identifier for this buffer (e.g., 'coeff', 'E_tilda_lm')
+ shape : tuple of int
+ Requested shape for the array
+ dtype : numpy dtype, optional
+ Data type for the array (default: np.float64)
+
+ Returns
+ -------
+ numpy.ndarray
+ Array of requested shape and dtype. May be a view into a larger buffer.
+
+ Notes
+ -----
+ The returned array should be treated as writable. If you need the data
+ to persist beyond the next call to get_or_create with the same key,
+ make a copy.
+ """
+ # Initialize stats for new keys
+ if key not in self.stats:
+ self.stats[key] = {"hits": 0, "misses": 0, "grows": 0}
+
+ if key in self.buffers:
+ current_shape, buf = self.buffers[key]
+
+ # Check if requested size fits in current buffer
+ if all(req <= curr for req, curr in zip(shape, current_shape)):
+ # Cache hit! Return view of existing buffer
+ self.stats[key]["hits"] += 1
+ # Create view with appropriate slice for each dimension
+ slices = tuple(slice(0, s) for s in shape)
+ return buf[slices]
+
+ # Need bigger buffer - reallocate
+ self.stats[key]["grows"] += 1
+ # Keep maximum of current and requested for each dimension
+ new_shape = tuple(max(c, r) for c, r in zip(current_shape, shape))
+ self.buffers[key] = (new_shape, np.empty(new_shape, dtype=dtype))
+
+ # Return view of newly allocated buffer
+ slices = tuple(slice(0, s) for s in shape)
+ return self.buffers[key][1][slices]
+
+ # First allocation for this key
+ self.stats[key]["misses"] += 1
+ self.buffers[key] = (shape, np.empty(shape, dtype=dtype))
+ return self.buffers[key][1]
+
+ def clear(self):
+ """Free all buffers and reset statistics.
+
+ Use this when done processing a batch of cells to release memory.
+ In Dask workflows, buffers are automatically released when the
+ worker process terminates, so calling clear() is optional.
+ """
+ self.buffers.clear()
+ self.stats.clear()
+
+ def get_stats(self):
+ """Get buffer usage statistics for performance analysis.
+
+ Returns
+ -------
+ dict
+ Dictionary mapping buffer keys to statistics:
+ - 'hits': Number of times buffer was reused
+ - 'misses': Number of times buffer was allocated
+ - 'grows': Number of times buffer was grown
+
+ Examples
+ --------
+ >>> pool = BufferPool()
+ >>> # ... use pool ...
+ >>> stats = pool.get_stats()
+ >>> print(f"Coefficient buffer hit rate: {stats['coeff']['hits'] /
+ ... (stats['coeff']['hits'] + stats['coeff']['misses']):.1%}")
+ """
+ return self.stats.copy()
+
+ def get_memory_usage(self):
+ """Get current memory usage of all buffers.
+
+ Returns
+ -------
+ dict
+ Dictionary with:
+ - 'total_mb': Total memory used by all buffers in MB
+ - 'buffers': Dict mapping keys to individual buffer sizes in MB
+ """
+ total_bytes = 0
+ buffer_sizes = {}
+
+ for key, (shape, buf) in self.buffers.items():
+ size_bytes = buf.nbytes
+ total_bytes += size_bytes
+ buffer_sizes[key] = size_bytes / (1024**2) # Convert to MB
+
+ return {"total_mb": total_bytes / (1024**2), "buffers": buffer_sizes}
diff --git a/src/delaunay.py b/pycsa/core/delaunay.py
similarity index 94%
rename from src/delaunay.py
rename to pycsa/core/delaunay.py
index 005059c..a8d5479 100644
--- a/src/delaunay.py
+++ b/pycsa/core/delaunay.py
@@ -1,6 +1,6 @@
import numpy as np
from scipy.spatial import Delaunay
-from src import utils, var
+from pycsa.core import utils, var
def get_decomposition(topo, xnp=11, ynp=6, padding=0):
@@ -68,8 +68,8 @@ def get_land_cells(tri, topo, height_tol=0.5, percent_tol=0.95):
Parameters
----------
- tri : :class:`scipy.spatial.qhull.Delaunay` instance
- scipy Delaunay triangulation instance containing tuples of the three vertice coordinates of a triangle
+ tri : instance containing tuples of the three vertice coordinates of a triangle
+ E.g., :class:`scipy.spatial.qhull.Delaunay`
topo : array-like
2D topographic data
height_tol : float, optional
diff --git a/src/fourier.py b/pycsa/core/fourier.py
similarity index 83%
rename from src/fourier.py
rename to pycsa/core/fourier.py
index 3ce1ffd..66a190a 100644
--- a/src/fourier.py
+++ b/pycsa/core/fourier.py
@@ -1,12 +1,44 @@
import numpy as np
+try:
+ import numba as nb
+
+ NUMBA_AVAILABLE = True
+except ImportError:
+ NUMBA_AVAILABLE = False
+
+
+# Numba-optimized functions for hot computational loops
+if NUMBA_AVAILABLE:
+
+ @nb.njit(parallel=True, fastmath=True, cache=True)
+ def _compute_trig_terms(tt_sum_flat, bcos_out, bsin_out):
+ """Numba-optimized computation of sin and cos terms.
+
+ Computes both sin and cos in a single pass with SIMD vectorization.
+ This is faster than calling np.sin and np.cos separately.
+ """
+ two_pi = 2.0 * np.pi
+ n = tt_sum_flat.shape[0]
+ m = tt_sum_flat.shape[1]
+
+ for i in nb.prange(n):
+ for j in range(m):
+ arg = two_pi * tt_sum_flat[i, j]
+ bcos_out[i, j] = np.cos(arg)
+ bsin_out[i, j] = np.sin(arg)
+
+else:
+ # Fallback if Numba not available
+ _compute_trig_terms = None
+
class f_trans(object):
"""
Fourier transformer class
"""
- def __init__(self, nhar_i, nhar_j):
+ def __init__(self, nhar_i, nhar_j, buffer_pool=None):
"""
Initalises a discrete spectral space with the corresponding Fourier coefficients spanning ``nhar_i`` and ``nhar_j``.
@@ -16,9 +48,12 @@ def __init__(self, nhar_i, nhar_j):
number of spectral modes in the first horizontal direction
nhar_j : int
number of spectral modes in the second horizontal direction
+ buffer_pool : BufferPool, optional
+ Buffer pool for memory-efficient array reuse
"""
self.nhar_i = nhar_i
self.nhar_j = nhar_j
+ self.buffer_pool = buffer_pool
self.m_i = None
self.m_j = None
@@ -139,12 +174,11 @@ def do_full(self, cell, grad=False):
self.__get_IJ(cell)
self.__prepare_terms(cell)
- self.term1 = np.expand_dims(self.term1, -1)
- self.term1 = np.repeat(self.term1, self.nhar_j, -1)
- self.term2 = np.expand_dims(self.term2, 1)
- self.term2 = np.repeat(self.term2, self.nhar_i, 1)
-
- tt_sum = self.term1 + self.term2
+ # Optimized: Use broadcasting instead of expand_dims + repeat
+ # Old approach created large intermediate arrays
+ # New approach: term1[:, :, None] broadcasts with term2[:, None, :]
+ # This is equivalent but avoids memory allocation and copying
+ tt_sum = self.term1[:, :, np.newaxis] + self.term2[:, np.newaxis, :]
del self.term1
del self.term2
@@ -154,8 +188,18 @@ def do_full(self, cell, grad=False):
else:
tt_sum = tt_sum.reshape(tt_sum.shape[0], -1)
- bcos = np.cos(2.0 * np.pi * (tt_sum))
- bsin = np.sin(2.0 * np.pi * (tt_sum))
+ # Compute both sin and cos - use Numba if available for speedup
+ if NUMBA_AVAILABLE and _compute_trig_terms is not None:
+ # Numba-optimized path: pre-allocate and compute in-place
+ bcos = np.empty_like(tt_sum)
+ bsin = np.empty_like(tt_sum)
+ _compute_trig_terms(tt_sum, bcos, bsin)
+ else:
+ # NumPy fallback path
+ two_pi_tt = 2.0 * np.pi * tt_sum
+ bcos = np.cos(two_pi_tt)
+ bsin = np.sin(two_pi_tt)
+ del two_pi_tt
del tt_sum
@@ -272,7 +316,7 @@ def get_freq_grid(self, a_m):
cos_terms = a_m[: len(self.k_idx)]
sin_terms = a_m[len(self.k_idx) :]
- fourier_coeff = np.zeros((nhar_i, nhar_j), dtype=np.complex_)
+ fourier_coeff = np.zeros((nhar_i, nhar_j), dtype=np.complex128)
for cnt, (row, col) in enumerate(zip(self.k_idx, self.l_idx)):
fourier_coeff[row, col] = cos_terms[cnt] + 1.0j * sin_terms[cnt]
diff --git a/pycsa/core/io.py b/pycsa/core/io.py
new file mode 100644
index 0000000..dd424e6
--- /dev/null
+++ b/pycsa/core/io.py
@@ -0,0 +1,1777 @@
+"""
+Input/Output routines
+"""
+
+import netCDF4 as nc
+import numpy as np
+import h5py
+import os
+import threading
+
+from datetime import datetime
+from scipy import interpolate
+from tqdm import tqdm
+
+from pycsa.core import utils
+
+# ============================================================================
+# CRITICAL: Global lock for NetCDF/HDF5 operations
+# HDF5 is NOT thread-safe by default. Even opening different files from
+# different threads can cause crashes if HDF5 wasn't compiled with --enable-threadsafe.
+# This lock serializes ALL NetCDF Dataset operations across all threads.
+# ============================================================================
+_NETCDF_GLOBAL_LOCK = threading.Lock()
+
+
+class ncdata(object):
+ """Helper class to read NetCDF4 topographic data"""
+
+ def __init__(self, read_merit=False, padding=0, padding_tol=50):
+ """
+
+ Parameters
+ ----------
+ read_merit : bool, optional
+ toggles between the `MERIT DEM `_ and `USGS GMTED 2010 `_ data files. By default False, i.e., read USGS GMTED 2010 data files.
+ padding : int, optional
+ number of data points to pad the loaded topography file, by default 0
+ padding_tol : int, optional
+ padding tolerance is added no matter the user-defined ``padding``, by default 50
+ """
+ self.read_merit = read_merit
+ self.padding = padding_tol + padding
+ self.is_open = False
+
+ def read_dat(self, fn, obj):
+ """Reads data by attributes defined in the ``obj`` class.
+
+ Parameters
+ ----------
+ fn : str
+ filename
+ obj : :class:`src.var.grid` or :class:`src.var.topo` or :class:`src.var.topo_cell`
+ any data object in :mod:`src.var` accepting topography attributes
+ """
+ df = nc.Dataset(fn, "r")
+
+ for key, _ in vars(obj).items():
+ if key in df.variables:
+ setattr(obj, key, df.variables[key][:])
+
+ df.close()
+
+ # def open(self, fn):
+ # self.df = nc.Dataset(fn, "r")
+ # self.is_open = True
+
+ # def close(self):
+ # if self.is_open and hasattr(self, "df"):
+ # self.df.close()
+
+ def __get_truths(self, arr, vert_pts, d_pts):
+ """Assembles Boolean array selecting for data points within a given lat-lon range, including padded boundary."""
+ return (arr >= (vert_pts.min() - self.padding * d_pts)) & (
+ arr <= vert_pts.max() + self.padding * d_pts
+ )
+
+ def read_topo(self, topo, cell, lon_vert, lat_vert):
+ """Reads USGS GMTED 2010 dataset
+
+ Parameters
+ ----------
+ topo : :class:`src.var.topo` or :class:`src.var.topo_cell`
+ instance of a topography class containing the full regional or global topography loaded via :func:`src.io.read_dat`.
+ cell : :class:`src.var.topo_cell`
+ instance of a cell object
+ lon_vert : list
+ extent of the longitudinal coordinates encompassing the region to be loaded
+ lat_vert : list
+ extent of the latitudinal coordinates encompassing the region to be loaded
+
+ .. note:: Loading the global topography in the ``topo`` argument may not be memory efficient. The notebook ``nc_compactifier.ipynb`` contains a script to extract a region of interest from the global GMTED 2010 dataset.
+ """
+ lon, lat, z = topo.lon, topo.lat, topo.topo
+
+ nrecords = np.shape(z)[0]
+
+ bool_arr = np.zeros_like(z).astype(bool)
+ lat_arr = np.zeros_like(z)
+ lon_arr = np.zeros_like(z)
+
+ z = z[:, ::-1, :]
+
+ for n in range(nrecords):
+ lat_n = lat[n]
+ lon_n = lon[n]
+
+ dlat, dlon = np.diff(lat_n).mean(), np.diff(lon_n).mean()
+
+ lon_nm, lat_nm = np.meshgrid(lon_n, lat_n)
+
+ bool_arr[n] = self.__get_truths(lon_nm, lon_vert, dlon) & self.__get_truths(
+ lat_nm, lat_vert, dlat
+ )
+
+ lat_arr[n] = lat_nm
+ lon_arr[n] = lon_nm
+
+ lon_res = lon_arr[bool_arr]
+ lat_res = lat_arr[bool_arr]
+ z_res = z[bool_arr].data
+
+ # ---- processing of the lat,lon,topo to get the regular 2D grid for topography
+ lon_uniq, lat_uniq = np.unique(lon_res), np.unique(
+ lat_res
+ ) # get unique values of lon,lat
+ nla = len(lat_uniq)
+ nlo = len(lon_uniq)
+
+ lat_res_sort_idx = np.argsort(lat_res)
+ lon_res_sort_idx = np.argsort(
+ lon_res[lat_res_sort_idx].reshape(nla, nlo), axis=1
+ )
+ z_res = z_res[lat_res_sort_idx]
+ z_res = np.take_along_axis(z_res.reshape(nla, nlo), lon_res_sort_idx, axis=1)
+ topo_2D = z_res.reshape(nla, nlo)
+
+ print("Data fetched...")
+ cell.lon = lon_uniq
+ cell.lat = lat_uniq
+ cell.topo = topo_2D
+
+ class read_merit_topo(object):
+ """Subclass to read MERIT topographic data"""
+
+ def __init__(self, cell, params, verbose=False, is_parallel=False):
+ """Populates ``cell`` object instance with arguments from ``params``
+
+ Parameters
+ ----------
+ cell : :class:`src.var.topo` or :class:`src.var.topo_cell`
+ instance of an object with topograhy attribute
+ params : :class:`src.var.params`
+ user-defined run parameters
+ verbose : bool, optional
+ prints loading progression, by default False
+ """
+ self.dir = params.path_merit
+ self.verbose = verbose
+ self.opened_dfs = []
+ # Thread-local storage: each thread gets its own file handles
+ # This prevents concurrent access to the same NetCDF Dataset object
+ self._thread_local = threading.local()
+
+ self.fn_lon = np.array(
+ [
+ -180.0,
+ -150.0,
+ -120.0,
+ -90.0,
+ -60.0,
+ -30.0,
+ 0.0,
+ 30.0,
+ 60.0,
+ 90.0,
+ 120.0,
+ 150.0,
+ 180.0,
+ ]
+ )
+ self.fn_lat = np.array([90.0, 60.0, 30.0, 0.0, -30.0, -60.0, -90.0])
+
+ self.lat_verts = np.array(params.lat_extent)
+ self.lon_verts = np.array(params.lon_extent)
+
+ self.merit_cg = params.merit_cg
+ self.split_EW = False
+ self.span = False
+ self.interp_lons = []
+
+ if not is_parallel:
+ self.get_topo(cell)
+
+ self.is_parallel = is_parallel
+
+ def _get_cached_file(self, filepath):
+ """
+ Get a thread-local cached NetCDF file handle with global locking.
+
+ Uses global lock because HDF5 is not thread-safe on this system.
+ Even opening different files from different threads causes crashes.
+ """
+ # Get or create thread-local file cache
+ if not hasattr(self._thread_local, "file_cache"):
+ self._thread_local.file_cache = {}
+
+ cache = self._thread_local.file_cache
+
+ if filepath not in cache:
+ if self.verbose:
+ print(
+ f"[Thread {threading.current_thread().name}] Opening: {filepath}"
+ )
+
+ # CRITICAL: Use global lock to serialize HDF5 file opens
+ with _NETCDF_GLOBAL_LOCK:
+ cache[filepath] = nc.Dataset(filepath, "r")
+
+ return cache[filepath]
+
+ def close_cached_files(self):
+ """Close all cached NetCDF files in current thread."""
+ if hasattr(self._thread_local, "file_cache"):
+ for filepath, ds in self._thread_local.file_cache.items():
+ try:
+ ds.close()
+ except Exception as e:
+ print(f"Warning: Error closing {filepath}: {e}")
+ self._thread_local.file_cache.clear()
+
+ def get_topo(self, cell):
+
+ # if lat_verts
+
+ if (self.lon_verts.max() - self.lon_verts.min()) > 180.0:
+ self.split_EW = True
+
+ if self.split_EW:
+ min_lon = (
+ max(
+ np.where(
+ self.lon_verts < 0.0, self.lon_verts + 360.0, self.lon_verts
+ )
+ )
+ - 360.0
+ )
+ max_lon = min(
+ np.where(
+ self.lon_verts < 0.0, self.lon_verts + 360.0, self.lon_verts
+ )
+ )
+ else:
+ min_lon = self.lon_verts.min()
+ max_lon = self.lon_verts.max()
+
+ lat_min_idx = self.__compute_idx(self.lat_verts.min(), "min", "lat")
+ lat_max_idx = self.__compute_idx(self.lat_verts.max(), "max", "lat")
+
+ if not self.split_EW:
+ lon_min_idx = self.__compute_idx(min_lon, "min", "lon")
+ lon_max_idx = self.__compute_idx(max_lon, "max", "lon")
+ else:
+ lon_min_idx = self.__compute_idx(min_lon, "max", "lon")
+ lon_max_idx = self.__compute_idx(max_lon, "min", "lon")
+
+ if (self.lon_verts.max() - self.lon_verts.min()) > 180.0:
+ lon_idx_rng = list(range(lon_max_idx, len(self.fn_lon) - 1)) + list(
+ range(0, lon_min_idx + 1)
+ )
+
+ else:
+ if lon_min_idx == lon_max_idx:
+ lon_max_idx += 1
+ lon_idx_rng = list(range(lon_min_idx, lon_max_idx))
+
+ lat_idx_rng = list(range(lat_max_idx, lat_min_idx))
+
+ fns, dirs, lon_cnt, lat_cnt = self.__get_fns(lat_idx_rng, lon_idx_rng)
+
+ self.__load_topo(
+ cell, fns, dirs, lon_cnt, lat_cnt, lat_idx_rng, lon_idx_rng
+ )
+
+ def __compute_idx(self, vert, typ, direction):
+ """Given a point ``vert``, look up which MERIT NetCDF file contains this point."""
+ if direction == "lon":
+ fn_int = self.fn_lon
+ else:
+ fn_int = self.fn_lat
+
+ where_idx = np.argmin(np.abs(fn_int - vert))
+
+ if self.verbose:
+ print(fn_int, where_idx)
+
+ if typ == "min":
+ if (vert - fn_int[where_idx]) < 0.0:
+ if direction == "lon":
+ # if not self.split_EW:
+ where_idx -= 1
+ else:
+ where_idx += 1
+ elif typ == "max":
+ if (vert - fn_int[where_idx]) > 0.0:
+ if direction == "lon":
+ if not self.split_EW:
+ where_idx += 1
+ else:
+ where_idx -= 1
+
+ if (where_idx == (len(fn_int) - 1)) and self.split_EW:
+ where_idx -= 1
+
+ where_idx = int(where_idx)
+
+ if self.verbose:
+ print("where_idx, vert, fn_int[where_idx] for typ:")
+ print(where_idx, vert, fn_int[where_idx], typ)
+ print("")
+
+ return where_idx
+
+ def __get_fns(self, lat_idx_rng, lon_idx_rng):
+ """Construct the full filenames required for the loading of the topographic data from the indices identified in :func:`src.io.ncdata.read_merit_topo.__compute_idx`"""
+ fns = []
+ dirs = []
+
+ for lat_cnt, lat_idx in enumerate(lat_idx_rng):
+ l_lat_bound, r_lat_bound = (
+ self.fn_lat[lat_idx],
+ self.fn_lat[lat_idx + 1],
+ )
+ l_lat_tag, r_lat_tag = self.__get_NSEW(
+ l_lat_bound, "lat"
+ ), self.__get_NSEW(r_lat_bound, "lat")
+
+ if (l_lat_tag == "S" and r_lat_tag == "S") and (
+ l_lat_bound == -60 and r_lat_bound == -90
+ ):
+ merit_or_rema = "REMA_BKG"
+ self.rema = True
+ self.dir = self.dir.replace("MERIT", "REMA")
+ else:
+ merit_or_rema = "MERIT"
+ self.rema = False
+ self.dir = self.dir.replace("REMA", "MERIT")
+
+ for lon_cnt, lon_idx in enumerate(lon_idx_rng):
+ l_lon_bound, r_lon_bound = (
+ self.fn_lon[lon_idx],
+ self.fn_lon[lon_idx + 1],
+ )
+ l_lon_tag, r_lon_tag = self.__get_NSEW(
+ l_lon_bound, "lon"
+ ), self.__get_NSEW(r_lon_bound, "lon")
+
+ name = "%s_%s%.2d-%s%.2d_%s%.3d-%s%.3d.nc4" % (
+ merit_or_rema,
+ l_lat_tag,
+ np.abs(l_lat_bound),
+ r_lat_tag,
+ np.abs(r_lat_bound),
+ l_lon_tag,
+ np.abs(l_lon_bound),
+ r_lon_tag,
+ np.abs(r_lon_bound),
+ )
+
+ fns.append(name)
+ dirs.append(self.dir)
+
+ return fns, dirs, lon_cnt, lat_cnt
+
+ def __load_topo(
+ self,
+ cell,
+ fns,
+ dirs,
+ lon_cnt,
+ lat_cnt,
+ lat_idx_rng,
+ lon_idx_rng,
+ init=True,
+ populate=True,
+ ):
+ """
+ This method assembles a contiguous array in ``cell.topo`` containing the regional topography to be loaded.
+
+ However, this full regional array is assembled from an array of block arrays. Each block array is loaded from a separated MERIT data file and varies in shape that is not known beforehand.
+
+ Therefore, the ``get_topo`` method is run recursively:
+ 1. The first run determines the shape of each constituting block array and subsequently the shape of the full regional array. An empty array in initialised.
+ 2. The second run populates the empty array with the information of the block arrays obtained in the first run.
+ """
+ if (cell.topo is None) and (init):
+ self.__load_topo(
+ cell,
+ fns,
+ dirs,
+ lon_cnt,
+ lat_cnt,
+ lat_idx_rng,
+ lon_idx_rng,
+ init=False,
+ populate=False,
+ )
+
+ if not populate:
+ n_col = 0
+ n_row = 0
+ nc_lon = 0
+ nc_lat = 0
+ else:
+ n_col = 0
+ n_row = 0
+ lon_sz_old = 0
+ lat_sz_old = 0
+ cell.lat = []
+ cell.lon = []
+
+ ### Handles the case where a cell spans four topographic datasets
+ cnt_lat = 0
+ cnt_lon = 0
+
+ for cnt, fn in enumerate(fns):
+ ############################################
+ #
+ # Open data file (using cache for performance)
+ #
+ ############################################
+ filepath = dirs[cnt] + fn
+ test = self._get_cached_file(filepath)
+ if test not in self.opened_dfs:
+ self.opened_dfs.append(test)
+
+ ############################################
+ #
+ # Load lat data
+ #
+ ############################################
+
+ lat = test["lat"]
+ lat_min_idx = np.argmin(
+ np.abs((lat - np.sign(lat) * 1e-4) - self.lat_verts.min())
+ )
+ lat_max_idx = np.argmin(
+ np.abs((lat + np.sign(lat) * 1e-4) - self.lat_verts.max())
+ )
+
+ lat_high = np.max((lat_min_idx, lat_max_idx))
+ lat_low = np.min((lat_min_idx, lat_max_idx))
+
+ lat = test["lat"]
+
+ ############################################
+ #
+ # Load lon data
+ #
+ ############################################
+
+ # in the case where fns contains both MERIT and REMA dataset, then for the n_row = 0, we do...
+ if (
+ any("REMA" in fn for fn in fns)
+ and any("MERIT" in fn for fn in fns)
+ and (not populate)
+ ):
+ if n_row == 0:
+ # run MERIT and REMA interpolation
+ new_lon = self.__do_interp_lon_1D(
+ dirs, fns, cnt_lon, lon_cnt, n_col, lon_idx_rng
+ )
+ self.interp_lons.append(new_lon)
+
+ # flag stating that we have MERIT+REMA mix
+ self.span = True
+
+ lon = test["lon"]
+
+ lon_low, lon_high = self.__get_lon_idxs(lon, lon_idx_rng, n_col)
+
+ if not populate:
+ if n_row == 0:
+
+ # if (cnt_lon < (lon_cnt + 1)) and lon_nc_change:
+ if not self.span:
+ nc_lon += lon_high - lon_low
+ else:
+ nc_lon += len(new_lon)
+ cnt_lon += 1
+
+ if n_col == 0:
+ # if (cnt_lat < (lat_cnt + 1)) and lat_nc_change:
+ nc_lat += lat_high - lat_low
+ cnt_lat += 1
+
+ n_col += 1
+ if n_col == (lon_cnt + 1):
+ n_col = 0
+ n_row += 1
+
+ else:
+ topo = test["Elevation"][lat_low:lat_high, lon_low:lon_high]
+
+ curr_lon = lon[lon_low:lon_high].tolist()
+
+ if n_col == 0:
+ curr_lat = lat[lat_low:lat_high].tolist()
+ cell.lat += curr_lat
+ if not self.span:
+ if n_row == 0:
+ cell.lon += curr_lon
+ else: # interpolate topo data to new lon grid
+ new_lon = self.interp_lons[n_col]
+ topo = self.__interp_topo_2D(topo, curr_lat, curr_lon, new_lon)
+
+ if n_row == 0:
+ cell.lon += new_lon.tolist()
+
+ # # current dataset at n_row = 0 is a MERIT dataset
+ # if "MERIT" in fn:
+ # self.merit = True
+
+ # # topographic data is read over MERIT and REMA interface:
+ # if n_row > 0:
+ # if ("REMA" in fn) and (self.prev_merit):
+
+ if not self.span:
+ lon_sz = lon_high - lon_low
+ else:
+ lon_sz = len(self.interp_lons[n_col])
+ lat_sz = lat_high - lat_low
+
+ cell.topo[
+ lat_sz_old : lat_sz_old + lat_sz,
+ lon_sz_old : lon_sz_old + lon_sz,
+ ] = topo
+
+ n_col += 1
+ lon_sz_old += np.copy(lon_sz)
+
+ if n_col == (lon_cnt + 1):
+ n_col = 0
+ lon_sz_old = 0
+
+ n_row += 1
+ lat_sz_old = np.copy(lat_sz)
+
+ # Note: Files are kept open in cache for reuse (closed via close_cached_files())
+
+ if not populate:
+ cell.topo = np.zeros((nc_lat, nc_lon))
+ else:
+
+ if self.split_EW:
+ cell.lon = np.array(cell.lon)
+ cell.lon[cell.lon < 0.0] += 360.0
+
+ iint = self.merit_cg
+
+ if max(cell.lat) < -85.0:
+ iint *= 5
+
+ cell.lat = utils.sliding_window_view(
+ np.sort(cell.lat), (iint,), (iint,)
+ ).mean(axis=-1)
+ cell.lon = utils.sliding_window_view(
+ np.sort(cell.lon), (iint,), (iint,)
+ ).mean(axis=-1)
+
+ cell.topo = utils.sliding_window_view(
+ cell.topo, (iint, iint), (iint, iint)
+ ).mean(axis=(-1, -2))[::-1, :]
+
+ def __do_interp_lon_1D(self, dirs, fns, cnt_lon, lon_cnt, n_col, lon_idx_rng):
+ # Note: MERIT is always on n_row = 0 and REMA on n_row = 1
+
+ merit_path = dirs[cnt_lon] + fns[cnt_lon]
+ merit_dat = self._get_cached_file(merit_path)
+ merit_lon = merit_dat["lon"]
+
+ rema_path = dirs[cnt_lon + lon_cnt + 1] + fns[cnt_lon + lon_cnt + 1]
+ rema_dat = self._get_cached_file(rema_path)
+ rema_lon = rema_dat["lon"]
+
+ merit_lon_low, merit_lon_high = self.__get_lon_idxs(
+ merit_lon, lon_idx_rng, n_col
+ )
+ rema_lon_low, rema_lon_high = self.__get_lon_idxs(
+ rema_lon, lon_idx_rng, n_col
+ )
+
+ merit_lon = merit_lon[merit_lon_low:merit_lon_high].tolist()
+ rema_lon = rema_lon[rema_lon_low:rema_lon_high].tolist()
+
+ new_max = min(max(merit_lon), max(rema_lon))
+ new_min = max(min(merit_lon), min(rema_lon))
+ # we always use the number of data points in the merit lon grid:
+ new_sz = min(len(merit_lon), len(rema_lon))
+
+ new_lon = np.linspace(new_min, new_max, new_sz)
+
+ # Files kept open in cache (no close needed)
+
+ return new_lon
+
+ @staticmethod
+ def __interp_topo_2D(topo, curr_lat, curr_lon, new_lon):
+ interp = interpolate.RegularGridInterpolator((curr_lat, curr_lon), topo)
+ XX, YY = np.meshgrid(new_lon, curr_lat)
+ return interp((YY, XX))
+
+ def __get_lon_idxs(
+ self,
+ lon,
+ lon_idx_rng,
+ n_col,
+ ):
+ l_lon_bound, r_lon_bound = (
+ self.fn_lon[lon_idx_rng[n_col]],
+ self.fn_lon[lon_idx_rng[n_col] + 1],
+ )
+
+ lon_rng = r_lon_bound - l_lon_bound
+
+ lon_in_file = self.lon_verts[
+ ((self.lon_verts - l_lon_bound) > 0)
+ & ((self.lon_verts - l_lon_bound) <= lon_rng)
+ ]
+
+ if len(lon_in_file) == 0:
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+
+ else:
+ if not self.split_EW:
+ if lon_in_file.max() == self.lon_verts.max():
+ lon_high = np.argmin(np.abs(lon - lon_in_file.max()))
+ else:
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+
+ if lon_in_file.min() == self.lon_verts.min():
+ lon_low = np.argmin(np.abs(lon - lon_in_file.min()))
+ else:
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+
+ else:
+ # Handle dateline crossing cases
+ negative_lons = self.lon_verts[self.lon_verts < 0.0]
+
+ # Check if we have negative longitudes before using min/max
+ if len(negative_lons) > 0 and lon_in_file.max() == min(
+ np.where(
+ self.lon_verts < 0.0, self.lon_verts + 360.0, self.lon_verts
+ )
+ ):
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+ lon_low = np.argmin(np.abs(lon - lon_in_file.min()))
+ else:
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+
+ # Check if we have negative longitudes before using max
+ if len(negative_lons) > 0 and lon_in_file.min() == (
+ max(negative_lons + 360.0) - 360.0
+ ):
+ lon_high = np.argmin(np.abs(lon - lon_in_file.max()))
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+ else:
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+
+ return lon_low, lon_high
+
+ def close_all(self):
+ for df in self.opened_dfs:
+ df.close()
+
+ @staticmethod
+ def __get_NSEW(vert, typ):
+ """Method to determine `NSEW` in MERIT filename"""
+ if typ == "lat":
+ if vert >= 0.0:
+ dir_tag = "N"
+ else:
+ dir_tag = "S"
+ if typ == "lon":
+ if vert >= 0.0:
+ dir_tag = "E"
+ else:
+ dir_tag = "W"
+
+ return dir_tag
+
+ class read_etopo_topo(object):
+ """Subclass to read ETOPO 2022 15 arc-second topographic data"""
+
+ def __init__(self, cell, params, verbose=False, is_parallel=False):
+ """Populates ``cell`` object instance with arguments from ``params``
+
+ Parameters
+ ----------
+ cell : :class:`src.var.topo` or :class:`src.var.topo_cell`
+ instance of an object with topography attribute
+ params : :class:`src.var.params`
+ user-defined run parameters
+ verbose : bool, optional
+ prints loading progression, by default False
+ is_parallel : bool, optional
+ flag for parallel processing, by default False
+ """
+ self.dir = params.path_etopo
+ self.verbose = verbose
+ self.opened_dfs = []
+ # Thread-local storage: each thread gets its own file handles
+ # This prevents concurrent access to the same NetCDF Dataset object
+ self._thread_local = threading.local()
+
+ # ETOPO 2022 tiles are at 15 degree intervals
+ self.fn_lon = np.array(
+ [
+ -180,
+ -165,
+ -150,
+ -135,
+ -120,
+ -105,
+ -90,
+ -75,
+ -60,
+ -45,
+ -30,
+ -15,
+ 0,
+ 15,
+ 30,
+ 45,
+ 60,
+ 75,
+ 90,
+ 105,
+ 120,
+ 135,
+ 150,
+ 165,
+ 180,
+ ]
+ )
+ self.fn_lat = np.array(
+ [90, 75, 60, 45, 30, 15, 0, -15, -30, -45, -60, -75, -90]
+ )
+
+ self.lat_verts = np.array(params.lat_extent)
+ self.lon_verts = np.array(params.lon_extent)
+
+ self.etopo_cg = params.etopo_cg if hasattr(params, "etopo_cg") else 1
+ self.split_EW = False
+
+ if not is_parallel:
+ self.get_topo(cell)
+
+ self.is_parallel = is_parallel
+
+ def _get_cached_file(self, filepath):
+ """
+ Get a thread-local cached NetCDF file handle with global locking.
+
+ Uses global lock because HDF5 is not thread-safe on this system.
+ Even opening different files from different threads causes crashes.
+ """
+ # Get or create thread-local file cache
+ if not hasattr(self._thread_local, "file_cache"):
+ self._thread_local.file_cache = {}
+
+ cache = self._thread_local.file_cache
+
+ if filepath not in cache:
+ if self.verbose:
+ print(
+ f"[Thread {threading.current_thread().name}] Opening: {filepath}"
+ )
+
+ import time
+
+ max_retries = 3
+ retry_delay = 0.5
+
+ for attempt in range(max_retries):
+ try:
+ # CRITICAL: Use global lock to serialize HDF5 file opens
+ with _NETCDF_GLOBAL_LOCK:
+ cache[filepath] = nc.Dataset(filepath, "r")
+ break
+ except (OSError, RuntimeError, TypeError) as e:
+ if attempt < max_retries - 1:
+ # Retry with exponential backoff
+ if self.verbose:
+ print(
+ f"Warning: Attempt {attempt+1} failed for {filepath}, retrying: {e}"
+ )
+ time.sleep(retry_delay * (2**attempt))
+ else:
+ raise RuntimeError(
+ f"Failed to open {filepath} after {max_retries} attempts: {e}"
+ )
+
+ return cache[filepath]
+
+ def close_cached_files(self):
+ """Close all cached NetCDF files in current thread."""
+ if hasattr(self._thread_local, "file_cache"):
+ for filepath, ds in self._thread_local.file_cache.items():
+ try:
+ ds.close()
+ except Exception as e:
+ print(f"Warning: Error closing {filepath}: {e}")
+ self._thread_local.file_cache.clear()
+
+ def get_topo(self, cell):
+ """Main method to load ETOPO topography data"""
+
+ # Compute longitude span
+ lon_span = self.lon_verts.max() - self.lon_verts.min()
+
+ # A true dateline crossing occurs when:
+ # 1. We have longitudes on both sides of ±180° (some positive, some negative)
+ # 2. AND the span wraps around (e.g., 170° to -170° = 340° wrap, not 20°)
+ # The key is to check if converting all to [0, 360) would reduce the span
+ lon_verts_360 = np.where(
+ self.lon_verts < 0.0, self.lon_verts + 360.0, self.lon_verts
+ )
+ span_360 = lon_verts_360.max() - lon_verts_360.min()
+
+ # If converting to [0, 360) reduces the span, it's a true dateline crossing
+ crosses_dateline = (span_360 < lon_span) and (lon_span > 180.0)
+
+ if self.verbose:
+ print(f"DEBUG get_topo: lon_verts = {self.lon_verts}")
+ print(f"DEBUG get_topo: lon_span = {lon_span}, span_360 = {span_360}")
+ print(f"DEBUG get_topo: crosses_dateline = {crosses_dateline}")
+
+ # Determine loading strategy
+ if lon_span >= 360.0:
+ # Full global extent: load all tiles
+ self.split_EW = False
+ lon_idx_rng = list(range(0, len(self.fn_lon) - 1))
+ if self.verbose:
+ print(f"Full global extent detected (span={lon_span}°)")
+ print(f"Loading all {len(lon_idx_rng)} longitude tiles")
+
+ elif crosses_dateline:
+ # True dateline crossing (e.g., [170, -170])
+ # Work in [0, 360) representation to compute tile indices
+ self.split_EW = True
+
+ # Use [0, 360) representation for proper wraparound
+ min_lon_360 = lon_verts_360.min()
+ max_lon_360 = lon_verts_360.max()
+
+ # Find tile indices in [0, 360) space, then convert back
+ # Western tiles: from max_lon (e.g., ~170°) to 180°
+ # Eastern tiles: from -180° to min_lon (e.g., ~-170° = 190° in [0,360))
+
+ # Convert back to [-180, 180) for tile index lookup
+ # since fn_lon is in [-180, 180) space
+ min_lon = min_lon_360 if min_lon_360 <= 180 else min_lon_360 - 360
+ max_lon = max_lon_360 if max_lon_360 <= 180 else max_lon_360 - 360
+
+ # Compute indices using the [-180, 180) values
+ lon_min_idx = self.__compute_idx(min_lon, "min", "lon")
+ lon_max_idx = self.__compute_idx(max_lon, "max", "lon")
+
+ if self.verbose:
+ print(f"DEBUG dateline: min_lon={min_lon}, max_lon={max_lon}")
+ print(
+ f"DEBUG dateline: lon_min_idx={lon_min_idx}, lon_max_idx={lon_max_idx}"
+ )
+
+ # For dateline crossing, we need tiles covering the span from min_lon to max_lon
+ # Since we're crossing the dateline, the span wraps around ±180°
+ # In [-180, 180) representation:
+ # - min_lon is the easternmost extent (e.g., 144°)
+ # - max_lon is the westernmost extent (e.g., -144°)
+ # We need tiles from min_lon eastward to 180°, then from -180° eastward to max_lon
+ # In tile index space: from lon_min_idx to end (index 24), plus from start (index 0) to lon_max_idx
+
+ # Special case: if both indices are the same, we only need that tile and possibly neighbors
+ if lon_min_idx == lon_max_idx:
+ # Both edges are in the same tile - check if we need neighbors
+ lon_idx_rng = [lon_min_idx]
+ if lon_min_idx >= len(self.fn_lon) - 2: # Near the end of the array
+ # Also include the dateline tile(s)
+ lon_idx_rng.append(0) # Add first tile for wraparound
+ else:
+ # Normal dateline crossing: go from min_idx to end (excluding the duplicate at 180°),
+ # then from start to max_idx
+ # Note: fn_lon[-1] = 180° maps to same tile as fn_lon[0] = -180°, so exclude index len-1
+ lon_idx_rng = list(range(lon_min_idx, len(self.fn_lon) - 1)) + list(
+ range(0, lon_max_idx + 1)
+ )
+
+ if self.verbose:
+ print(f"DEBUG dateline: lon_idx_rng={lon_idx_rng}")
+
+ if self.verbose:
+ print(
+ f"Dateline crossing detected: [{self.lon_verts.min():.2f}, {self.lon_verts.max():.2f}]"
+ )
+ print(f" In [0,360): [{min_lon:.2f}, {max_lon:.2f}]")
+ print(f" lon_min_idx={lon_min_idx}, lon_max_idx={lon_max_idx}")
+ print(f" Loading tiles: {lon_idx_rng}")
+
+ else:
+ # Normal case: straightforward longitude range (including large spans like [-90, 180])
+ self.split_EW = False
+ min_lon = self.lon_verts.min()
+ max_lon = self.lon_verts.max()
+
+ lon_min_idx = self.__compute_idx(min_lon, "min", "lon")
+ lon_max_idx = self.__compute_idx(max_lon, "max", "lon")
+
+ if lon_min_idx == lon_max_idx:
+ lon_max_idx += 1
+ lon_idx_rng = list(range(lon_min_idx, lon_max_idx))
+
+ # Latitude indices (same for all cases)
+ lat_min_idx = self.__compute_idx(self.lat_verts.min(), "min", "lat")
+ lat_max_idx = self.__compute_idx(self.lat_verts.max(), "max", "lat")
+ lat_idx_rng = list(range(lat_max_idx, lat_min_idx))
+
+ # Get filenames and load data
+ fns, lon_cnt, lat_cnt = self.__get_fns(lat_idx_rng, lon_idx_rng)
+
+ if self.verbose:
+ print(
+ f"DEBUG: Generated {len(fns)} files, lon_cnt={lon_cnt}, lat_cnt={lat_cnt}"
+ )
+ print(f"DEBUG: First few files: {fns[:min(5, len(fns))]}")
+ print(f"DEBUG: Last few files: {fns[-min(5, len(fns)):]}")
+
+ self.__load_topo(cell, fns, lon_cnt, lat_cnt, lat_idx_rng, lon_idx_rng)
+
+ def __compute_idx(self, vert, typ, direction):
+ """Given a point ``vert``, look up which ETOPO NetCDF file contains this point."""
+ if direction == "lon":
+ fn_int = self.fn_lon
+ else:
+ fn_int = self.fn_lat
+
+ where_idx = np.argmin(np.abs(fn_int - vert))
+
+ if self.verbose:
+ print(fn_int, where_idx)
+
+ if typ == "min":
+ if (vert - fn_int[where_idx]) < 0.0:
+ if direction == "lon":
+ where_idx -= 1
+ else:
+ where_idx += 1
+ elif typ == "max":
+ if (vert - fn_int[where_idx]) > 0.0:
+ if direction == "lon":
+ if not self.split_EW:
+ where_idx += 1
+ else:
+ where_idx -= 1
+
+ if (where_idx == (len(fn_int) - 1)) and self.split_EW:
+ where_idx -= 1
+
+ where_idx = int(where_idx)
+
+ if self.verbose:
+ print("where_idx, vert, fn_int[where_idx] for typ:")
+ print(where_idx, vert, fn_int[where_idx], typ)
+ print("")
+
+ return where_idx
+
+ def __get_fns(self, lat_idx_rng, lon_idx_rng):
+ """Construct the full filenames required for loading topographic data"""
+ fns = []
+
+ # Initialize to avoid UnboundLocalError if ranges are empty
+ lon_cnt = 0
+ lat_cnt = 0
+
+ for lat_cnt, lat_idx in enumerate(lat_idx_rng):
+ l_lat_bound = self.fn_lat[lat_idx]
+ l_lat_tag = self.__get_NSEW(l_lat_bound, "lat")
+
+ for lon_cnt, lon_idx in enumerate(lon_idx_rng):
+ l_lon_bound = self.fn_lon[lon_idx]
+ l_lon_tag = self.__get_NSEW(l_lon_bound, "lon")
+
+ # ETOPO filename format: ETOPO_2022_v1_15s_N00E000_surface.nc
+ name = "ETOPO_2022_v1_15s_%s%.2d%s%.3d_surface.nc" % (
+ l_lat_tag,
+ np.abs(l_lat_bound),
+ l_lon_tag,
+ np.abs(l_lon_bound),
+ )
+
+ fns.append(name)
+
+ return fns, lon_cnt, lat_cnt
+
+ def __load_topo(
+ self,
+ cell,
+ fns,
+ lon_cnt,
+ lat_cnt,
+ lat_idx_rng,
+ lon_idx_rng,
+ init=True,
+ populate=True,
+ ):
+ """
+ Assembles a contiguous array in ``cell.topo`` containing the regional topography.
+
+ This method runs recursively:
+ 1. First run determines the shape of each block array and initializes the full regional array.
+ 2. Second run populates the array with the actual topography data.
+ """
+ if (cell.topo is None) and (init):
+ self.__load_topo(
+ cell,
+ fns,
+ lon_cnt,
+ lat_cnt,
+ lat_idx_rng,
+ lon_idx_rng,
+ init=False,
+ populate=False,
+ )
+
+ if not populate:
+ n_col = 0
+ n_row = 0
+ nc_lon = 0
+ nc_lat = 0
+ else:
+ n_col = 0
+ n_row = 0
+ lon_sz_old = 0
+ lat_sz_old = 0
+ cell.lat = []
+ cell.lon = []
+
+ cnt_lat = 0
+ cnt_lon = 0
+
+ for cnt, fn in enumerate(fns):
+ ############################################
+ # Open data file (using cache for performance)
+ ############################################
+ filepath = self.dir + fn
+ test = self._get_cached_file(filepath)
+ if test not in self.opened_dfs:
+ self.opened_dfs.append(test)
+
+ ############################################
+ # Load lat data
+ ############################################
+ lat = test["lat"]
+
+ # Extract latitude data based on requested extent
+ # Always use the precise extraction based on lat_verts, don't try to be clever
+ lat_min_idx = np.argmin(
+ np.abs((lat - np.sign(lat) * 1e-4) - self.lat_verts.min())
+ )
+ lat_max_idx = np.argmin(
+ np.abs((lat + np.sign(lat) * 1e-4) - self.lat_verts.max())
+ )
+
+ lat_high = np.max((lat_min_idx, lat_max_idx))
+ lat_low = np.min((lat_min_idx, lat_max_idx))
+
+ ############################################
+ # Load lon data
+ ############################################
+ lon = test["lon"]
+ lon_low, lon_high = self.__get_lon_idxs(lon, lon_idx_rng, n_col)
+
+ if not populate:
+ if n_row == 0:
+ nc_lon += lon_high - lon_low
+ cnt_lon += 1
+
+ if n_col == 0:
+ nc_lat += lat_high - lat_low
+ cnt_lat += 1
+
+ n_col += 1
+ if n_col == (lon_cnt + 1):
+ n_col = 0
+ n_row += 1
+
+ else:
+ # ETOPO uses 'z' for elevation, map to 'topo'
+ # Convert masked array to regular array to avoid issues
+ topo = test["z"][lat_low:lat_high, lon_low:lon_high].data
+
+ curr_lon = lon[lon_low:lon_high].data.tolist()
+
+ if n_col == 0:
+ curr_lat = lat[lat_low:lat_high].data.tolist()
+ cell.lat += curr_lat
+
+ if n_row == 0:
+ cell.lon += curr_lon
+
+ lon_sz = lon_high - lon_low
+ lat_sz = lat_high - lat_low
+
+ cell.topo[
+ lat_sz_old : lat_sz_old + lat_sz,
+ lon_sz_old : lon_sz_old + lon_sz,
+ ] = topo
+
+ n_col += 1
+ lon_sz_old += np.copy(lon_sz)
+
+ if n_col == (lon_cnt + 1):
+ n_col = 0
+ lon_sz_old = 0
+
+ n_row += 1
+ lat_sz_old += np.copy(
+ lat_sz
+ ) # FIX: Add to offset, don't replace!
+
+ # Note: Files are kept open in cache for reuse (closed via close_cached_files())
+
+ if not populate:
+ cell.topo = np.zeros((nc_lat, nc_lon))
+ else:
+ if self.split_EW:
+ cell.lon = np.array(cell.lon)
+ cell.lon[cell.lon < 0.0] += 360.0
+
+ # Apply coarse-graining if specified
+ iint = self.etopo_cg
+
+ # Convert lists to numpy arrays
+ lat_arr = np.array(cell.lat)
+ lon_arr = np.array(cell.lon)
+
+ # Sort latitude and longitude indices to reorder topo array
+ lat_sort_idx = np.argsort(lat_arr)
+ lon_sort_idx = np.argsort(lon_arr)
+
+ lat_sorted = lat_arr[lat_sort_idx]
+ lon_sorted = lon_arr[lon_sort_idx]
+
+ # Reorder topo array rows and columns to match sorted lat/lon
+ # Use np.ix_ for proper 2D indexing
+ topo_sorted = cell.topo[np.ix_(lat_sort_idx, lon_sort_idx)]
+
+ if iint > 1:
+ # Apply coarse-graining using sliding window
+ try:
+ cell.lat = utils.sliding_window_view(
+ lat_sorted, (iint,), (iint,)
+ ).mean(axis=-1)
+ cell.lon = utils.sliding_window_view(
+ lon_sorted, (iint,), (iint,)
+ ).mean(axis=-1)
+
+ cell.topo = utils.sliding_window_view(
+ topo_sorted, (iint, iint), (iint, iint)
+ ).mean(axis=(-1, -2))
+ except (ValueError, MemoryError) as e:
+ # If coarse-graining fails, fall back to no coarse-graining
+ print(
+ f"Warning: Coarse-graining failed ({e}), using full resolution"
+ )
+ cell.lat = lat_sorted
+ cell.lon = lon_sorted
+ cell.topo = topo_sorted
+ else:
+ cell.lat = lat_sorted
+ cell.lon = lon_sorted
+ cell.topo = topo_sorted
+
+ def __get_lon_idxs(self, lon, lon_idx_rng, n_col):
+ """Get longitude indices for data extraction"""
+ l_lon_bound = self.fn_lon[lon_idx_rng[n_col]]
+
+ # Handle wraparound at dateline: index 24 (180°) wraps to index 0 (-180°)
+ # since both map to the same W180 tile
+ r_idx = lon_idx_rng[n_col] + 1
+ if r_idx >= len(self.fn_lon):
+ r_idx = (
+ 1 # Skip index 0 (-180°), go to index 1 (-165°) for proper bounds
+ )
+ r_lon_bound = self.fn_lon[r_idx]
+
+ lon_rng = r_lon_bound - l_lon_bound
+
+ lon_in_file = self.lon_verts[
+ ((self.lon_verts - l_lon_bound) >= 0)
+ & ((self.lon_verts - l_lon_bound) <= lon_rng)
+ ]
+
+ if len(lon_in_file) == 0:
+ # No user-requested extent falls within this tile's bounds
+ # Extract entire tile (this handles full global and wraparound cases)
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+ else:
+ if not self.split_EW:
+ if lon_in_file.max() == self.lon_verts.max():
+ lon_high = np.argmin(np.abs(lon - lon_in_file.max()))
+ else:
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+
+ if lon_in_file.min() == self.lon_verts.min():
+ lon_low = np.argmin(np.abs(lon - lon_in_file.min()))
+ else:
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+ else:
+ # Handle dateline crossing cases
+ negative_lons = self.lon_verts[self.lon_verts < 0.0]
+
+ # Check if we have negative longitudes before using min/max
+ if len(negative_lons) > 0 and lon_in_file.max() == min(
+ np.where(
+ self.lon_verts < 0.0, self.lon_verts + 360.0, self.lon_verts
+ )
+ ):
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+ lon_low = np.argmin(np.abs(lon - lon_in_file.min()))
+ else:
+ lon_high = np.argmin(np.abs(lon - r_lon_bound))
+
+ # Check if we have negative longitudes before using max
+ if len(negative_lons) > 0 and lon_in_file.min() == (
+ max(negative_lons + 360.0) - 360.0
+ ):
+ lon_high = np.argmin(np.abs(lon - lon_in_file.max()))
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+ else:
+ lon_low = np.argmin(np.abs(lon - l_lon_bound))
+
+ return lon_low, lon_high
+
+ def close_all(self):
+ """Close all opened NetCDF files"""
+ for df in self.opened_dfs:
+ df.close()
+
+ @staticmethod
+ def __get_NSEW(vert, typ):
+ """Method to determine `NSEW` in ETOPO filename"""
+ if typ == "lat":
+ if vert >= 0.0:
+ dir_tag = "N"
+ else:
+ dir_tag = "S"
+ if typ == "lon":
+ # Special case: 180° uses W180 in ETOPO naming convention
+ # (since 180°E and 180°W are the same meridian, ETOPO uses W)
+ if vert == 180.0:
+ dir_tag = "W"
+ elif vert >= 0.0:
+ dir_tag = "E"
+ else:
+ dir_tag = "W"
+
+ return dir_tag
+
+
+class writer(object):
+ """
+ HDF5 writer class
+
+ Contains methods to create HDF5 file, create data sets and populate them with output variables.
+
+ .. note:: This class was taken from an I/O routine originally written for the numerical flow solver used in `Chew et al. (2022) `_ and `Chew et al. (2023) `_.
+ """
+
+ def __init__(self, fn, idxs, sfx="", debug=False):
+ """
+ Creates an empty HDF5 file with filename ``fn`` and a group for each index in ``idxs``
+
+ Parameters
+ ----------
+ fn : str
+ filename
+ idxs : list
+ list of cell indices
+ sfx : str, optional
+ suffixes to the filename, by default ''
+ debug : bool, optional
+ debug flag, by default False
+ """
+
+ self.FORMAT = ".h5"
+ self.OUTPUT_FOLDER = "../outputs/"
+ self.OUTPUT_FILENAME = fn
+ self.OUTPUT_FULLPATH = self.OUTPUT_FOLDER + self.OUTPUT_FILENAME
+ self.SUFFIX = sfx
+ self.DEBUG = debug
+
+ self.IDXS = idxs
+ self.PATHS = [
+ # vars from the 'tri' object
+ "tri_lat_verts",
+ "tri_lon_verts",
+ "tri_clats",
+ "tri_clons",
+ "points",
+ "simplices",
+ # vars from the 'cell' object
+ "lon",
+ "lat",
+ "lon_grid",
+ "lat_grid",
+ # vars from the 'analysis' object
+ "ampls",
+ "kks",
+ "lls",
+ "recon",
+ ]
+
+ self.ATTRS = [
+ # vars from the 'analysis' object
+ "wlat",
+ "wlon",
+ ]
+
+ if debug:
+ self.PATHS = np.append(
+ self.PATHS,
+ [
+ "mask",
+ "topo_ref",
+ "pmf_ref",
+ "spectrum_ref",
+ "spectrum_fg",
+ "recon_fg",
+ "pmf_fg",
+ ],
+ )
+
+ self.io_create_file(self.IDXS)
+
+ def io_create_file(self, paths):
+ """
+ Helper function to create file.
+
+ Parameters
+ ----------
+ paths : list
+ List of strings containing the name of the groups.
+
+ Notes
+ -----
+ Currently, if the filename of the HDF5 file already exists, this function will append the existing filename with '_old' and create an empty HDF5 file with the same filename in its place.
+
+ """
+ # If directory does not exist, create it.
+ if not os.path.exists(self.OUTPUT_FOLDER):
+ os.mkdir(self.OUTPUT_FOLDER)
+
+ # If file exists, rename it with old.
+ if os.path.exists(self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT):
+ os.rename(
+ self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT,
+ self.OUTPUT_FULLPATH + self.SUFFIX + "_old" + self.FORMAT,
+ )
+
+ file = h5py.File(self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT, "a")
+ for path in paths:
+ path = str(path)
+ # check if groups have been created
+ # if not created, create empty groups
+ if not (path in file):
+ file.create_group(path, track_order=True)
+
+ file.close()
+
+ def write_all(self, idx, *args):
+ """Write all attributes and datasets of a given class instance to the group ``idx``.
+
+ Parameters
+ ----------
+ idx : str or int
+ group name to write the attributes or datasets
+ """
+ for arg in args:
+ for attr in self.PATHS:
+ if hasattr(arg, attr):
+ self.populate(idx, attr, getattr(arg, attr))
+
+ for attr in self.ATTRS:
+ if hasattr(arg, attr):
+ self.write_attr(idx, attr, getattr(arg, attr))
+
+ def write_attr(self, idx, key, value):
+ """Write HDF5 attributes for a group
+
+ Parameters
+ ----------
+ idx : str or int
+ group name to write the attributes
+ key : str
+ attribute name
+ value : any
+ attribute value that is accepted by HDF5
+ """
+ file = h5py.File(self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT, "r+")
+
+ try:
+ file[str(idx)].attrs.create(str(key), value)
+ except:
+ file[str(idx)].attrs.create(
+ str(key), repr(value), dtype=" 0)
+
+ H_spec_var = grp.createVariable("H_spec", "f8", ("nspec",))
+ H_spec_var[:] = self.__pad_zeros(analysis.ampls[pick_idx], self.n_modes)
+
+ kks_var = grp.createVariable("kks", "f8", ("nspec",))
+ kks_var[:] = self.__pad_zeros(analysis.kks[pick_idx], self.n_modes)
+
+ lls_var = grp.createVariable("lls", "f8", ("nspec",))
+ lls_var[:] = self.__pad_zeros(analysis.lls[pick_idx], self.n_modes)
+
+ rootgrp.close()
+
+ def duplicate(self, id, struct):
+
+ rootgrp = nc.Dataset(self.path + self.fn, "a", format="NETCDF4")
+
+ grp = rootgrp.createGroup(str(id))
+
+ is_land_var = grp.createVariable("is_land", "i4")
+ is_land_var[:] = struct.is_land
+
+ clat_var = grp.createVariable("clat", "f8")
+ clat_var[:] = struct.clat
+ clon_var = grp.createVariable("clon", "f8")
+ clon_var[:] = struct.clon
+
+ # Add cell_area if available
+ if struct.cell_area is not None:
+ cell_area_var = grp.createVariable("cell_area", "f8")
+ cell_area_var[:] = struct.cell_area
+ cell_area_var.units = "m^2"
+ cell_area_var.long_name = "Area of ICON grid cell"
+
+ if struct.is_land:
+ dk_var = grp.createVariable("dk", "f8")
+ dk_var[:] = struct.dk
+ dl_var = grp.createVariable("dl", "f8")
+ dl_var[:] = struct.dl
+
+ pick_idx = np.where(struct.ampls > 0)
+
+ H_spec_var = grp.createVariable("H_spec", "f8", ("nspec",))
+ H_spec_var[:] = self.__pad_zeros(struct.ampls[pick_idx], self.n_modes)
+
+ kks_var = grp.createVariable("kks", "f8", ("nspec",))
+ kks_var[:] = self.__pad_zeros(struct.kks[pick_idx], self.n_modes)
+
+ lls_var = grp.createVariable("lls", "f8", ("nspec",))
+ lls_var[:] = self.__pad_zeros(struct.lls[pick_idx], self.n_modes)
+
+ rootgrp.close()
+
+ def duplicate_all(self, data):
+
+ rootgrp = nc.Dataset(self.path + self.fn, "a", format="NETCDF4")
+
+ for id, struct in enumerate(tqdm(data)):
+ grp = rootgrp.createGroup(str(id))
+
+ is_land_var = grp.createVariable("is_land", "i4")
+ is_land_var[:] = struct.is_land
+
+ clat_var = grp.createVariable("clat", "f8")
+ clat_var[:] = struct.clat
+ clon_var = grp.createVariable("clon", "f8")
+ clon_var[:] = struct.clon
+
+ if struct.is_land:
+ dk_var = grp.createVariable("dk", "f8")
+ dk_var[:] = struct.dk
+ dl_var = grp.createVariable("dl", "f8")
+ dl_var[:] = struct.dl
+
+ pick_idx = np.where(struct.ampls > 0)
+
+ H_spec_var = grp.createVariable("H_spec", "f8", ("nspec",))
+ H_spec_var[:] = self.__pad_zeros(struct.ampls[pick_idx], self.n_modes)
+
+ kks_var = grp.createVariable("kks", "f8", ("nspec",))
+ kks_var[:] = self.__pad_zeros(struct.kks[pick_idx], self.n_modes)
+
+ lls_var = grp.createVariable("lls", "f8", ("nspec",))
+ lls_var[:] = self.__pad_zeros(struct.lls[pick_idx], self.n_modes)
+
+ rootgrp.close()
+
+ @staticmethod
+ def read_dat(path, fn, id, struct):
+ try:
+ rootgrp = nc.Dataset(path + fn, "a", format="NETCDF4")
+ except:
+ return False
+
+ grp = rootgrp[str(id)]
+
+ struct.is_land = grp["is_land"][:]
+ struct.clat = grp["clat"][:]
+ struct.clon = grp["clon"][:]
+
+ if struct.is_land:
+ struct.dk = grp["dk"][:]
+ struct.dl = grp["dl"][:]
+
+ struct.ampls = grp["H_spec"][:]
+ struct.kks = grp["kks"][:]
+ struct.lls = grp["lls"][:]
+
+ rootgrp.close()
+
+ return True
+
+ class grp_struct(object):
+ def __init__(self, c_idx, clat, clon, is_land, analysis=None, cell_area=None):
+ self.c_idx = c_idx
+ self.clat = clat
+ self.clon = clon
+ self.is_land = is_land
+ self.cell_area = cell_area
+
+ self.dk = None
+ self.dl = None
+
+ self.ampls = None
+ self.kks = None
+ self.lls = None
+
+ if analysis is not None:
+ for key, value in vars(analysis).items():
+ setattr(self, key, value)
+
+ @staticmethod
+ def __pad_zeros(lst, n_modes):
+
+ if lst.size < n_modes:
+ pad_len = n_modes - lst.size
+ else:
+ pad_len = 0
+
+ return np.concatenate((lst, np.zeros((pad_len))))
+
+
+class reader(object):
+ """Simple reader class to read HDF5 output written by :class:`src.io.writer`"""
+
+ def __init__(self, fn):
+ """
+ Parameters
+ ----------
+ fn : str
+ filename of the file to be read
+ """
+ self.fn = fn
+
+ self.names = {
+ "lat": "lat",
+ "lon": "lon",
+ "recon": "data",
+ "ampls": "spec",
+ "pmf_sg": "pmf",
+ }
+
+ def get_params(self, params):
+ """Get the user-defined parameters from the HDF5 file attributes
+
+ Parameters
+ ----------
+ params : :class:`src.var.params`
+ empty instance of the user-defined parameters class to be populated
+ """
+ file = h5py.File(self.fn)
+
+ for key in file.attrs.keys():
+ setattr(params, key, file.attrs[key])
+
+ file.close()
+
+ def read_data(self, idx, name):
+ """Read a particular dataset ``name`` from a group ``idx``
+
+ Parameters
+ ----------
+ idx : str or int
+ the group name
+ name : str
+ the dataset name
+
+ Returns
+ -------
+ array-like
+ the dataset
+ """
+ file = h5py.File(self.fn)
+ dat = file[str(idx)][name][:]
+ file.close()
+
+ return np.array(dat)
+
+ def read_all(self, idx, cell):
+ """Populate ``cell`` with all datasets in a group ``idx``
+
+ Parameters
+ ----------
+ idx : int or str
+ the group name
+ cell : :class:`src.var.topo_cell`
+ empty instance of a cell object to be populated
+ """
+ file = h5py.File(self.fn)
+
+ idx = str(idx)
+ for key, value in self.names.items():
+ setattr(cell, value, file[idx][key][:])
+
+ file.close()
+
+
+def fn_gen(params):
+ """Automatically generates HDF5 output filename from :class:`src.var.params`.
+
+ Parameters
+ ----------
+ params : :class:`src.var.params`
+ instance of the user parameter class
+
+ Returns
+ -------
+ str
+ automatically generated filename
+ """
+
+ if hasattr(params, "fn_tag"):
+ tag = params.fn_tag
+ else:
+ tag = "unnamed"
+
+ if params.enable_merit:
+ topo_dat = "merit"
+ else:
+ topo_dat = "usgs"
+
+ now = datetime.now()
+
+ date = now.strftime("%d%m%y")
+ time = now.strftime("%H%M%S")
+
+ ord = ["tag", "topo_dat", "date", "time"]
+
+ fn = ""
+ for item in ord:
+ fn += locals()[item]
+ fn += "_"
+
+ return fn[:-1]
diff --git a/pycsa/core/lin_reg.py b/pycsa/core/lin_reg.py
new file mode 100644
index 0000000..bcd7996
--- /dev/null
+++ b/pycsa/core/lin_reg.py
@@ -0,0 +1,204 @@
+"""
+Linear regression module with buffer pool and sparse solver support
+"""
+
+import numpy as np
+import scipy.linalg as la
+from scipy.sparse.linalg import gmres
+from scipy.linalg import blas
+from scipy.sparse import csr_matrix, eye
+from scipy.sparse.linalg import spsolve
+
+
+def get_coeffs(fobj, buffer_pool=None):
+ """Assembles the Fourier coefficients from the sine and cosine terms generated in the :class:`Fourier transformer class `.
+
+ Parameters
+ ----------
+ fobj : :class:`src.fourier.f_trans` instance
+ instance of the Fourier transformer class.
+ buffer_pool : BufferPool, optional
+ Buffer pool for memory-efficient array reuse
+
+ Returns
+ -------
+ array-like
+ 2D array corresponding to the ``M`` matrix.
+ """
+ Ncos = fobj.bf_cos
+ Nsin = fobj.bf_sin
+
+ n_points = Ncos.shape[0]
+ n_modes = Ncos.shape[1] + Nsin.shape[1]
+
+ if buffer_pool:
+ # Use buffer pool - handles variable sizes dynamically
+ coeff = buffer_pool.get_or_create("coeff", (n_points, n_modes), Ncos.dtype)
+ coeff[:, : Ncos.shape[1]] = Ncos
+ coeff[:, Ncos.shape[1] :] = Nsin
+ else:
+ # Fallback for backward compatibility
+ coeff = np.hstack([Ncos, Nsin])
+
+ del fobj.bf_cos
+ del fobj.bf_sin
+
+ if fobj.grad:
+ if buffer_pool:
+ # Allocate larger buffer for gradient stacking
+ coeff_grad = buffer_pool.get_or_create(
+ "coeff_grad", (2 * n_points, n_modes), Ncos.dtype
+ )
+ coeff_grad[:n_points] = coeff
+ coeff_grad[n_points:] = coeff
+ return coeff_grad
+ else:
+ coeff = np.vstack([coeff, coeff])
+
+ return coeff
+
+
+def do(
+ fobj,
+ cell,
+ lmbda=0.0,
+ iter_solve=True,
+ save_coeffs=False,
+ buffer_pool=None,
+ use_sparse=False,
+):
+ """
+ Does the linear regression with optional buffer pool and sparse solver
+
+ Parameters
+ ----------
+ fobj : :class:`src.fourier.f_trans` instance
+ instance of the Fourier transformer class.
+ cell : :class:`src.var.topo_cell` instance
+ cell object instance
+ lmbda : float, optional
+ regularisation parameter, by default 0.0
+ iter_solve : bool, optional
+ toggles between using direct or iterative solver, by default True
+ save_coeffs : bool, optional
+ skips the linear regression and just saves the generated ``M`` matrix for diagnostics and debugging, by default False
+ buffer_pool : BufferPool, optional
+ Buffer pool for memory-efficient array reuse
+ use_sparse : bool, optional
+ Use sparse matrix solver (automatic for few modes), by default False
+
+ Returns
+ -------
+ a_m : list
+ list of Fourier amplitudes corresponding to the unknown vector in the linear problem
+ data_recons : like
+ vector-like topography reconstructed from ``a_m``
+ """
+ if fobj.grad:
+ cell.get_grad()
+ data = cell.grad_topo_m
+ else:
+ data = cell.topo_m
+
+ coeff = get_coeffs(fobj, buffer_pool)
+
+ if save_coeffs:
+ fobj.coeff = coeff
+ return None, None
+
+ # Determine if sparse solver should be used
+ # Criteria: pick_kls enabled AND <10% of total modes selected
+ use_sparse_solver = use_sparse or (
+ getattr(fobj, "pick_kls", False)
+ and hasattr(fobj, "k_idx")
+ and len(fobj.k_idx) < 0.1 * (fobj.nhar_i * fobj.nhar_j)
+ )
+
+ if use_sparse_solver:
+ # ============================================================
+ # SPARSE PATH: For Second Approximation with few modes
+ # ============================================================
+ # Convert to sparse matrix (CSR format is efficient for matrix ops)
+ coeff_sparse = csr_matrix(coeff)
+ coeff_T_sparse = coeff_sparse.T
+
+ # Compute sparse normal equations
+ h_tilda_l_sparse = coeff_T_sparse @ data.reshape(-1, 1)
+ E_tilda_lm_sparse = coeff_T_sparse @ coeff_sparse
+
+ # Add regularization to sparse matrix
+ if lmbda > 0:
+ trace = E_tilda_lm_sparse.diagonal().mean() * lmbda
+ E_tilda_lm_sparse = E_tilda_lm_sparse + trace * eye(
+ E_tilda_lm_sparse.shape[0]
+ )
+
+ # Solve with sparse solver (direct solver for sparse SPD matrices)
+ # Convert RHS to dense array if it's sparse, otherwise use as-is
+ if hasattr(h_tilda_l_sparse, "toarray"):
+ rhs = h_tilda_l_sparse.toarray().flatten()
+ else:
+ rhs = np.asarray(h_tilda_l_sparse).flatten()
+ a_m = spsolve(E_tilda_lm_sparse, rhs)
+
+ # Reconstruct (sparse @ dense is efficient)
+ recons_result = coeff_sparse @ a_m
+ if hasattr(recons_result, "toarray"):
+ data_recons = recons_result.toarray().flatten()
+ else:
+ data_recons = np.asarray(recons_result).flatten()
+
+ else:
+ # ============================================================
+ # DENSE PATH: Standard approach with optional buffer reuse
+ # ============================================================
+ # Compute RHS
+ h_tilda_l = np.dot(coeff.T, data.reshape(-1, 1)).flatten()
+
+ # Compute LHS with optional buffer reuse
+ if buffer_pool:
+ n_modes = coeff.shape[1]
+ E_tilda_lm = buffer_pool.get_or_create(
+ "E_tilda_lm", (n_modes, n_modes), np.float64
+ )
+ # Compute and store in buffer
+ E_tilda_lm[:] = np.dot(coeff.T, coeff)
+ else:
+ E_tilda_lm = np.dot(coeff.T, coeff)
+
+ # Add regularization to diagonal (vectorized for speed)
+ if lmbda > 0:
+ trace = np.trace(E_tilda_lm) / E_tilda_lm.shape[0] * lmbda
+ np.fill_diagonal(E_tilda_lm, np.diag(E_tilda_lm) + trace)
+
+ # E_tilda_lm is symmetric positive definite (M^T M form with regularization)
+ # Use Cholesky decomposition for 2-5x speedup vs GMRES
+ if iter_solve:
+ try:
+ # Attempt Cholesky factorization (fastest for SPD matrices)
+ c, lower = la.cho_factor(E_tilda_lm, lower=True, check_finite=False)
+ a_m = la.cho_solve((c, lower), h_tilda_l, check_finite=False)
+ except la.LinAlgError:
+ # Fallback to GMRES if matrix is not positive definite
+ szc = E_tilda_lm.shape[0]
+ a_m, info = gmres(
+ E_tilda_lm,
+ h_tilda_l,
+ tol=1e-8, # Convergence tolerance
+ atol=1e-10, # Absolute tolerance
+ maxiter=min(szc, 100),
+ ) # Limit iterations
+ if info != 0:
+ # GMRES didn't converge, warn user
+ import warnings
+
+ warnings.warn(
+ f"GMRES did not converge (info={info}), solution may be inaccurate"
+ )
+ else:
+ # Direct inversion (slower, but kept for compatibility)
+ a_m = la.inv(E_tilda_lm).dot(h_tilda_l)
+
+ data_recons = coeff.dot(a_m)
+
+ return a_m, data_recons
diff --git a/src/physics.py b/pycsa/core/physics.py
similarity index 58%
rename from src/physics.py
rename to pycsa/core/physics.py
index 17d343d..9bbd84e 100644
--- a/src/physics.py
+++ b/pycsa/core/physics.py
@@ -45,12 +45,6 @@ def compute_uw_pmf(self, analysis, summed=True):
U = self.U
V = self.V
- wlat = analysis.wlat
- wlon = analysis.wlon
-
- kks = analysis.kks * 2.0 * np.pi
- lls = analysis.lls * 2.0 * np.pi
-
# if ((kks.ndim == 1) and (lls.ndim == 1)):
# print(True)
# ampls = analysis.ampls[np.nonzero(analysis.ampls)]
@@ -58,34 +52,37 @@ def compute_uw_pmf(self, analysis, summed=True):
# ampls = analysis.ampls
ampls = np.copy(analysis.ampls)
- wla = wlat # * self.AE
- wlo = wlon # * self.AE
-
- kks = kks / wlo
- lls = lls / wla
+ kks = analysis.kks
+ lls = analysis.lls
om = -kks * U - lls * V
omsq = om**2
- mms = (N**2 * (kks**2 + lls**2) / omsq) - (kks**2 + lls**2)
- # ampls[np.where(mms <= 0.0)] = 0.0
- mms[np.isnan(mms)] = 0.0
- mms = np.sqrt(mms)
-
- # wave-action density
- Ag = -0.5 * ((ampls) ** 2 * N**2 / om)
- Ag[np.isinf(Ag)] = 0.0
- Ag[np.isnan(Ag)] = 0.0
-
- # group velocity in z-direction
- cgz = (
- self.N
- * (kks**2 + lls**2) ** 0.5
- * mms
- / (kks**2 + lls**2 + mms**2) ** (3 / 2)
- )
-
- cgz[np.isnan(cgz)] = 0.0
+ # Compute mms safely: avoid divide-by-zero and sqrt of negatives.
+ # We intentionally silence expected divide/invalid warnings and map singularities to 0.
+ base = kks**2 + lls**2
+ with np.errstate(divide="ignore", invalid="ignore"):
+ frac = np.divide(N**2 * base, omsq, out=np.zeros_like(omsq), where=omsq > 0)
+ mms = frac - base
+ # Clip negatives to zero before sqrt to avoid invalid warnings
+ mms = np.sqrt(np.clip(mms, 0.0, None))
+
+ # wave-action density (Ag): safe division with zeros where om == 0
+ with np.errstate(divide="ignore", invalid="ignore"):
+ Ag = -0.5 * np.divide(
+ (ampls**2) * N**2, om, out=np.zeros_like(om), where=om != 0
+ )
+ Ag = np.nan_to_num(Ag, nan=0.0, posinf=0.0, neginf=0.0)
+
+ # group velocity in z-direction, computed safely
+ denom = (base + mms**2) ** 1.5
+ with np.errstate(divide="ignore", invalid="ignore"):
+ cgz = (
+ self.N
+ * np.sqrt(base)
+ * np.divide(mms, denom, out=np.zeros_like(denom), where=denom > 0)
+ )
+ cgz = np.nan_to_num(cgz, nan=0.0, posinf=0.0, neginf=0.0)
uw_pmf = Ag * kks * cgz
diff --git a/src/reconstruction.py b/pycsa/core/reconstruction.py
similarity index 72%
rename from src/reconstruction.py
rename to pycsa/core/reconstruction.py
index b857c50..c664a52 100644
--- a/src/reconstruction.py
+++ b/pycsa/core/reconstruction.py
@@ -17,14 +17,8 @@ def recon_2D(recons_z, cell):
array-like
2D reconstructed topography, values outside the mask are set to zero.
"""
- lon, lat = cell.lon, cell.lat
-
+ # Vectorized implementation - replaces nested Python loops with NumPy indexing
recons_z_2D = np.zeros(np.shape(cell.topo))
- c = 0
- for i in range(len(lat)):
- for j in range(len(lon)):
- if cell.mask[i, j] == 1:
- recons_z_2D[i, j] = recons_z[c]
- c = c + 1
+ recons_z_2D[cell.mask] = recons_z
return recons_z_2D
diff --git a/pycsa/core/tile_cache.py b/pycsa/core/tile_cache.py
new file mode 100644
index 0000000..07f22c2
--- /dev/null
+++ b/pycsa/core/tile_cache.py
@@ -0,0 +1,933 @@
+"""
+Topography tile caching system for efficient parallel processing.
+
+This module provides a caching layer for MERIT/ETOPO topography tiles to avoid
+repeatedly opening/closing NetCDF files during parallel cell processing.
+"""
+
+import netCDF4 as nc
+import numpy as np
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+import logging
+
+from pycsa.core.io import _NETCDF_GLOBAL_LOCK
+from pycsa.core import utils
+
+logger = logging.getLogger(__name__)
+
+
+# ETOPO 2022 15 arc-second tile grid (15° spacing in both lat and lon)
+_ETOPO_FN_LON = np.array(
+ [
+ -180,
+ -165,
+ -150,
+ -135,
+ -120,
+ -105,
+ -90,
+ -75,
+ -60,
+ -45,
+ -30,
+ -15,
+ 0,
+ 15,
+ 30,
+ 45,
+ 60,
+ 75,
+ 90,
+ 105,
+ 120,
+ 135,
+ 150,
+ 165,
+ 180,
+ ]
+)
+_ETOPO_FN_LAT = np.array([90, 75, 60, 45, 30, 15, 0, -15, -30, -45, -60, -75, -90])
+
+
+def compute_split_EW(lon_verts: np.ndarray) -> bool:
+ """Determine whether a cell's longitude extent truly crosses the dateline.
+
+ Uses the robust span-comparison formula: a true crossing occurs only when
+ converting to the [0, 360) representation reduces the span AND the original
+ span exceeds 180°. This avoids the false positives that plagued cells in
+ the western hemisphere near the dateline (e.g. Aleutian cells).
+ """
+ lon_verts = np.asarray(lon_verts)
+ lon_span = lon_verts.max() - lon_verts.min()
+ lon_verts_360 = np.where(lon_verts < 0.0, lon_verts + 360.0, lon_verts)
+ span_360 = lon_verts_360.max() - lon_verts_360.min()
+ return bool((span_360 < lon_span) and (lon_span > 180.0))
+
+
+def _etopo_NSEW(vert: float, typ: str) -> str:
+ """N/S for latitude, E/W for longitude with the +180° → 'W' convention."""
+ if typ == "lat":
+ return "N" if vert >= 0.0 else "S"
+ # longitude — note ETOPO's quirk: 180° always uses 'W' (since 180°E ≡ 180°W)
+ if vert == 180.0:
+ return "W"
+ return "E" if vert >= 0.0 else "W"
+
+
+def _etopo_tile_filename(lat_bound: float, lon_bound: float) -> str:
+ """ETOPO 2022 15s tile filename for the (lat, lon) tile origin."""
+ return "ETOPO_2022_v1_15s_%s%.2d%s%.3d_surface.nc" % (
+ _etopo_NSEW(lat_bound, "lat"),
+ np.abs(int(lat_bound)),
+ _etopo_NSEW(lon_bound, "lon"),
+ np.abs(int(lon_bound)),
+ )
+
+
+class TopographyTileCache:
+ """
+ Cache for topography data tiles.
+
+ Pre-loads all required MERIT/ETOPO/REMA tiles into memory and provides
+ fast access to subsets for individual grid cells.
+
+ This dramatically speeds up parallel processing by avoiding repeated
+ file I/O operations.
+
+ Parameters
+ ----------
+ data_dir : str or Path
+ Base directory containing topography data tiles
+ tile_filenames : list of str
+ List of tile filenames to pre-load
+ dataset_type : str, optional
+ Type of dataset ('MERIT', 'ETOPO', 'REMA'), by default 'MERIT'
+ verbose : bool, optional
+ Enable verbose logging, by default False
+
+ Attributes
+ ----------
+ tiles : dict
+ Dictionary mapping filenames to opened netCDF4.Dataset objects
+ tile_bounds : dict
+ Dictionary mapping filenames to (lat_min, lat_max, lon_min, lon_max) bounds
+ """
+
+ def __init__(
+ self,
+ data_dir: str,
+ tile_filenames: List[str],
+ dataset_type: str = "MERIT",
+ verbose: bool = False,
+ ):
+ self.data_dir = Path(data_dir)
+ self.dataset_type = dataset_type
+ self.verbose = verbose
+
+ # Cache dictionaries
+ self.tiles: Dict[str, nc.Dataset] = {}
+ self.tile_bounds: Dict[str, Tuple[float, float, float, float]] = {}
+ self.tile_lats: Dict[str, np.ndarray] = {}
+ self.tile_lons: Dict[str, np.ndarray] = {}
+
+ # ETOPO with empty tile list = lazy mode: tiles open on first access via
+ # get_etopo_data. MERIT keeps the existing eager pre-load behaviour.
+ if dataset_type == "ETOPO" and len(tile_filenames) == 0:
+ return
+
+ self._load_tiles(tile_filenames)
+
+ def _load_tiles(self, filenames: List[str]):
+ """Pre-load all tile files into memory."""
+ logger.info(f"Pre-loading {len(filenames)} topography tiles...")
+
+ for fn in filenames:
+ filepath = self.data_dir / fn
+
+ if not filepath.exists():
+ logger.warning(f"Tile file not found: {filepath}")
+ continue
+
+ try:
+ # Open NetCDF file under the shared HDF5 lock (HDF5 is not
+ # thread-safe on this system — see pycsa/core/io.py).
+ with _NETCDF_GLOBAL_LOCK:
+ ds = nc.Dataset(str(filepath), "r")
+ self.tiles[fn] = ds
+
+ # Cache coordinate arrays
+ lat = ds["lat"][:]
+ lon = ds["lon"][:]
+ self.tile_lats[fn] = lat
+ self.tile_lons[fn] = lon
+
+ # Cache bounds for quick lookup
+ self.tile_bounds[fn] = (
+ float(lat.min()),
+ float(lat.max()),
+ float(lon.min()),
+ float(lon.max()),
+ )
+
+ if self.verbose:
+ logger.debug(f"Loaded tile: {fn}")
+ logger.debug(
+ f" Bounds: lat[{lat.min():.2f}, {lat.max():.2f}], "
+ f"lon[{lon.min():.2f}, {lon.max():.2f}]"
+ )
+
+ except Exception as e:
+ logger.error(f"Failed to load tile {fn}: {e}")
+
+ def get_data_for_region(
+ self, lat_extent: np.ndarray, lon_extent: np.ndarray, merit_cg: int = 1
+ ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+ """
+ Extract topography data for a given lat/lon region.
+
+ This is designed to be a drop-in replacement for the current
+ read_merit_topo().get_topo() workflow.
+
+ Parameters
+ ----------
+ lat_extent : array-like
+ Latitude extent [lat_min, lat_max, ...]
+ lon_extent : array-like
+ Longitude extent [lon_min, lon_max, ...]
+ merit_cg : int, optional
+ Coarse-graining factor, by default 1
+
+ Returns
+ -------
+ lat : ndarray
+ Latitude coordinates
+ lon : ndarray
+ Longitude coordinates
+ topo : ndarray
+ Topography data (2D array)
+ """
+ lat_min = float(np.min(lat_extent))
+ lat_max = float(np.max(lat_extent))
+ lon_min = float(np.min(lon_extent))
+ lon_max = float(np.max(lon_extent))
+
+ # Handle dateline crossing — robust formula matching io.read_etopo_topo;
+ # the old `(lon_max - lon_min) > 180.0` test false-positived on western
+ # cells near the dateline (e.g. Aleutians).
+ crosses_dateline = compute_split_EW(lon_extent)
+ if crosses_dateline:
+ lon_min = (
+ max(np.where(lon_extent < 0.0, lon_extent + 360.0, lon_extent)) - 360.0
+ )
+ lon_max = min(np.where(lon_extent < 0.0, lon_extent + 360.0, lon_extent))
+
+ # Find tiles that overlap with this region
+ overlapping_tiles = self._find_overlapping_tiles(
+ lat_min, lat_max, lon_min, lon_max
+ )
+
+ if not overlapping_tiles:
+ logger.warning(
+ f"No tiles found for region: lat[{lat_min}, {lat_max}], lon[{lon_min}, {lon_max}]"
+ )
+ # Return empty arrays
+ return np.array([]), np.array([]), np.zeros((0, 0))
+
+ # Extract and merge data from overlapping tiles
+ lat_data, lon_data, topo_data = self._merge_tiles(
+ overlapping_tiles, lat_min, lat_max, lon_min, lon_max, crosses_dateline
+ )
+
+ # Apply coarse-graining if requested
+ if merit_cg > 1:
+ from pycsa.core import utils
+
+ # Adjust for high-latitude regions
+ iint = merit_cg
+ if lat_max < -85.0:
+ iint *= 5
+
+ # Coarse-grain using sliding window
+ lat_data = utils.sliding_window_view(
+ np.sort(lat_data), (iint,), (iint,)
+ ).mean(axis=-1)
+ lon_data = utils.sliding_window_view(
+ np.sort(lon_data), (iint,), (iint,)
+ ).mean(axis=-1)
+ topo_data = utils.sliding_window_view(
+ topo_data, (iint, iint), (iint, iint)
+ ).mean(axis=(-1, -2))[::-1, :]
+
+ return lat_data, lon_data, topo_data
+
+ def _find_overlapping_tiles(
+ self, lat_min: float, lat_max: float, lon_min: float, lon_max: float
+ ) -> List[str]:
+ """Find all tiles that overlap with the given region."""
+ overlapping = []
+
+ for fn, (
+ tile_lat_min,
+ tile_lat_max,
+ tile_lon_min,
+ tile_lon_max,
+ ) in self.tile_bounds.items():
+ # Check for overlap
+ lat_overlap = not (tile_lat_max < lat_min or tile_lat_min > lat_max)
+ lon_overlap = not (tile_lon_max < lon_min or tile_lon_min > lon_max)
+
+ if lat_overlap and lon_overlap:
+ overlapping.append(fn)
+
+ return overlapping
+
+ def _merge_tiles(
+ self,
+ tile_filenames: List[str],
+ lat_min: float,
+ lat_max: float,
+ lon_min: float,
+ lon_max: float,
+ crosses_dateline: bool,
+ ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+ """
+ Merge data from multiple tiles into a single contiguous array.
+
+ This handles the case where a cell region spans multiple MERIT/ETOPO tiles.
+ """
+ all_lats = []
+ all_lons = []
+ all_topos = []
+
+ for fn in tile_filenames:
+ ds = self.tiles[fn]
+ lat = self.tile_lats[fn]
+ lon = self.tile_lons[fn]
+
+ # Find indices within requested bounds
+ lat_mask = (lat >= lat_min) & (lat <= lat_max)
+ lon_mask = (lon >= lon_min) & (lon <= lon_max)
+
+ lat_idxs = np.where(lat_mask)[0]
+ lon_idxs = np.where(lon_mask)[0]
+
+ if len(lat_idxs) == 0 or len(lon_idxs) == 0:
+ continue
+
+ # Extract subset
+ lat_subset = lat[lat_idxs]
+ lon_subset = lon[lon_idxs]
+
+ # Handle elevation variable name (MERIT uses "Elevation", ETOPO may use different)
+ if "Elevation" in ds.variables:
+ elev_var = "Elevation"
+ elif "elevation" in ds.variables:
+ elev_var = "elevation"
+ elif "z" in ds.variables:
+ elev_var = "z"
+ else:
+ # Try to find any elevation-like variable
+ possible_names = ["topo", "topography", "height", "dem"]
+ elev_var = None
+ for name in possible_names:
+ if name in ds.variables:
+ elev_var = name
+ break
+ if elev_var is None:
+ logger.error(f"Could not find elevation variable in tile {fn}")
+ continue
+
+ with _NETCDF_GLOBAL_LOCK:
+ topo_subset = ds[elev_var][
+ lat_idxs[0] : lat_idxs[-1] + 1, lon_idxs[0] : lon_idxs[-1] + 1
+ ]
+
+ all_lats.append(lat_subset)
+ all_lons.append(lon_subset)
+ all_topos.append(topo_subset)
+
+ if not all_topos:
+ return np.array([]), np.array([]), np.zeros((0, 0))
+
+ # If only one tile, return directly
+ if len(all_topos) == 1:
+ return all_lats[0], all_lons[0], all_topos[0]
+
+ # Otherwise, need to merge multiple tiles
+ # For simplicity, concatenate and remove duplicates
+ merged_lat = np.unique(np.concatenate(all_lats))
+ merged_lon = np.unique(np.concatenate(all_lons))
+
+ # Create output array
+ merged_topo = np.zeros((len(merged_lat), len(merged_lon)))
+
+ # Fill from tiles (simple approach - could be optimized)
+ for i, lat_val in enumerate(merged_lat):
+ for j, lon_val in enumerate(merged_lon):
+ # Find which tile contains this point and extract value
+ for k, fn in enumerate(tile_filenames):
+ if (lat_val in all_lats[k]) and (lon_val in all_lons[k]):
+ lat_idx = np.where(all_lats[k] == lat_val)[0][0]
+ lon_idx = np.where(all_lons[k] == lon_val)[0][0]
+ merged_topo[i, j] = all_topos[k][lat_idx, lon_idx]
+ break
+
+ return merged_lat, merged_lon, merged_topo
+
+ # ------------------------------------------------------------------
+ # ETOPO path — byte-equivalent port of pycsa.core.io.read_etopo_topo
+ # ------------------------------------------------------------------
+ # The MERIT methods above (get_data_for_region, _find_overlapping_tiles,
+ # _merge_tiles) stay MERIT-specific. ETOPO has a fixed 15° tile grid and
+ # dateline handling that doesn't fit cleanly into bounds-based discovery,
+ # so the ETOPO path uses its own discovery + assembly mirroring io.py.
+
+ def _open_etopo_tile(self, fn: str) -> nc.Dataset:
+ """Open an ETOPO tile on first access; cache the handle thereafter.
+
+ Goes through _NETCDF_GLOBAL_LOCK because HDF5 is not thread-safe on
+ the target system. Once opened, the handle (and its lat/lon coordinate
+ arrays) stay cached for the lifetime of this TopographyTileCache.
+ """
+ if fn in self.tiles:
+ return self.tiles[fn]
+ filepath = str(self.data_dir / fn)
+ with _NETCDF_GLOBAL_LOCK:
+ ds = nc.Dataset(filepath, "r")
+ self.tiles[fn] = ds
+ # Coordinate arrays are small; cache so we don't re-read per cell.
+ self.tile_lats[fn] = ds["lat"][:]
+ self.tile_lons[fn] = ds["lon"][:]
+ return ds
+
+ @staticmethod
+ def _etopo_compute_idx(
+ vert: float, typ: str, direction: str, split_EW: bool
+ ) -> int:
+ """Look up which ETOPO tile-boundary index encloses ``vert``.
+
+ Mirrors pycsa.core.io.read_etopo_topo.__compute_idx (io.py:834-870).
+ """
+ fn_int = _ETOPO_FN_LON if direction == "lon" else _ETOPO_FN_LAT
+ where_idx = int(np.argmin(np.abs(fn_int - vert)))
+
+ if typ == "min":
+ if (vert - fn_int[where_idx]) < 0.0:
+ where_idx += -1 if direction == "lon" else 1
+ elif typ == "max":
+ if (vert - fn_int[where_idx]) > 0.0:
+ if direction == "lon":
+ if not split_EW:
+ where_idx += 1
+ else:
+ where_idx -= 1
+ if (where_idx == len(fn_int) - 1) and split_EW:
+ where_idx -= 1
+ return int(where_idx)
+
+ @staticmethod
+ def _etopo_get_fns(
+ lat_idx_rng: List[int], lon_idx_rng: List[int]
+ ) -> Tuple[List[str], int, int]:
+ """Build ETOPO filenames for a rectangular tile range.
+
+ Mirrors pycsa.core.io.read_etopo_topo.__get_fns (io.py:872-898).
+ Returns (filenames, lon_cnt, lat_cnt) where the counts are the
+ zero-based last enumerations (for __load_topo's row/col arithmetic).
+ """
+ fns: List[str] = []
+ lon_cnt = 0
+ lat_cnt = 0
+ for lat_cnt, lat_idx in enumerate(lat_idx_rng):
+ l_lat_bound = _ETOPO_FN_LAT[lat_idx]
+ for lon_cnt, lon_idx in enumerate(lon_idx_rng):
+ l_lon_bound = _ETOPO_FN_LON[lon_idx]
+ fns.append(_etopo_tile_filename(l_lat_bound, l_lon_bound))
+ return fns, lon_cnt, lat_cnt
+
+ @staticmethod
+ def _etopo_get_lon_idxs(
+ lon: np.ndarray,
+ lon_idx_rng: List[int],
+ n_col: int,
+ split_EW: bool,
+ lon_verts: np.ndarray,
+ ) -> Tuple[int, int]:
+ """Compute per-tile longitude slice indices.
+
+ Mirrors pycsa.core.io.read_etopo_topo.__get_lon_idxs (io.py:1052-1104).
+ """
+ l_lon_bound = _ETOPO_FN_LON[lon_idx_rng[n_col]]
+ r_idx = lon_idx_rng[n_col] + 1
+ if r_idx >= len(_ETOPO_FN_LON):
+ r_idx = 1 # 180° wraps to -165° (skip index 0 = -180° duplicate)
+ r_lon_bound = _ETOPO_FN_LON[r_idx]
+ lon_rng = r_lon_bound - l_lon_bound
+
+ lon_in_file = lon_verts[
+ ((lon_verts - l_lon_bound) >= 0) & ((lon_verts - l_lon_bound) <= lon_rng)
+ ]
+
+ if len(lon_in_file) == 0:
+ lon_high = int(np.argmin(np.abs(lon - r_lon_bound)))
+ lon_low = int(np.argmin(np.abs(lon - l_lon_bound)))
+ return lon_low, lon_high
+
+ if not split_EW:
+ if lon_in_file.max() == lon_verts.max():
+ lon_high = int(np.argmin(np.abs(lon - lon_in_file.max())))
+ else:
+ lon_high = int(np.argmin(np.abs(lon - r_lon_bound)))
+ if lon_in_file.min() == lon_verts.min():
+ lon_low = int(np.argmin(np.abs(lon - lon_in_file.min())))
+ else:
+ lon_low = int(np.argmin(np.abs(lon - l_lon_bound)))
+ return lon_low, lon_high
+
+ # split_EW = True (dateline crossing)
+ negative_lons = lon_verts[lon_verts < 0.0]
+ lon_high = int(np.argmin(np.abs(lon - r_lon_bound)))
+ lon_low = int(np.argmin(np.abs(lon - l_lon_bound)))
+ if len(negative_lons) > 0:
+ wrapped = np.where(lon_verts < 0.0, lon_verts + 360.0, lon_verts)
+ if lon_in_file.max() == wrapped.min():
+ lon_high = int(np.argmin(np.abs(lon - r_lon_bound)))
+ lon_low = int(np.argmin(np.abs(lon - lon_in_file.min())))
+ if lon_in_file.min() == (negative_lons.max() + 360.0 - 360.0):
+ lon_high = int(np.argmin(np.abs(lon - lon_in_file.max())))
+ lon_low = int(np.argmin(np.abs(lon - l_lon_bound)))
+ return lon_low, lon_high
+
+ def _etopo_load_topo(
+ self,
+ fns: List[str],
+ lon_cnt: int,
+ lat_cnt: int,
+ lat_idx_rng: List[int],
+ lon_idx_rng: List[int],
+ lat_verts: np.ndarray,
+ lon_verts: np.ndarray,
+ split_EW: bool,
+ ) -> Tuple[List[float], List[float], np.ndarray]:
+ """Assemble the regional topography array from per-tile slices.
+
+ Mirrors pycsa.core.io.read_etopo_topo.__load_topo (io.py:900-1050)
+ as a two-pass over ``fns`` — first pass computes the output shape,
+ second pass populates the array. Returns (lat_list, lon_list, topo).
+ """
+ # First pass: compute output shape (nc_lat, nc_lon).
+ n_col = 0
+ n_row = 0
+ nc_lon = 0
+ nc_lat = 0
+ for fn in fns:
+ ds = self._open_etopo_tile(fn)
+ lat = self.tile_lats[fn]
+ lon = self.tile_lons[fn]
+
+ lat_min_idx = np.argmin(
+ np.abs((lat - np.sign(lat) * 1e-4) - lat_verts.min())
+ )
+ lat_max_idx = np.argmin(
+ np.abs((lat + np.sign(lat) * 1e-4) - lat_verts.max())
+ )
+ lat_high = int(max(lat_min_idx, lat_max_idx))
+ lat_low = int(min(lat_min_idx, lat_max_idx))
+
+ lon_low, lon_high = self._etopo_get_lon_idxs(
+ lon, lon_idx_rng, n_col, split_EW, lon_verts
+ )
+
+ if n_row == 0:
+ nc_lon += lon_high - lon_low
+ if n_col == 0:
+ nc_lat += lat_high - lat_low
+
+ n_col += 1
+ if n_col == (lon_cnt + 1):
+ n_col = 0
+ n_row += 1
+
+ # Second pass: populate the array.
+ topo_arr = np.zeros((nc_lat, nc_lon))
+ cell_lat: List[float] = []
+ cell_lon: List[float] = []
+ n_col = 0
+ n_row = 0
+ lon_sz_old = 0
+ lat_sz_old = 0
+ for fn in fns:
+ ds = self.tiles[fn]
+ lat = self.tile_lats[fn]
+ lon = self.tile_lons[fn]
+
+ lat_min_idx = np.argmin(
+ np.abs((lat - np.sign(lat) * 1e-4) - lat_verts.min())
+ )
+ lat_max_idx = np.argmin(
+ np.abs((lat + np.sign(lat) * 1e-4) - lat_verts.max())
+ )
+ lat_high = int(max(lat_min_idx, lat_max_idx))
+ lat_low = int(min(lat_min_idx, lat_max_idx))
+
+ lon_low, lon_high = self._etopo_get_lon_idxs(
+ lon, lon_idx_rng, n_col, split_EW, lon_verts
+ )
+
+ with _NETCDF_GLOBAL_LOCK:
+ slab = ds["z"][lat_low:lat_high, lon_low:lon_high].data
+
+ curr_lon = lon[lon_low:lon_high].data.tolist()
+ if n_col == 0:
+ cell_lat += lat[lat_low:lat_high].data.tolist()
+ if n_row == 0:
+ cell_lon += curr_lon
+
+ lon_sz = lon_high - lon_low
+ lat_sz = lat_high - lat_low
+ topo_arr[
+ lat_sz_old : lat_sz_old + lat_sz,
+ lon_sz_old : lon_sz_old + lon_sz,
+ ] = slab
+
+ n_col += 1
+ lon_sz_old += lon_sz
+ if n_col == (lon_cnt + 1):
+ n_col = 0
+ lon_sz_old = 0
+ n_row += 1
+ lat_sz_old += lat_sz
+
+ return cell_lat, cell_lon, topo_arr
+
+ def get_etopo_data(
+ self,
+ lat_extent: np.ndarray,
+ lon_extent: np.ndarray,
+ etopo_cg: int = 1,
+ ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+ """Load ETOPO topography for a cell's lat/lon vertex extent.
+
+ Byte-equivalent to pycsa.core.io.read_etopo_topo.get_topo + __load_topo
+ (io.py:720-1050), but uses this cache's persistent file handles so the
+ same tile isn't re-opened across cells within a worker.
+
+ Parameters
+ ----------
+ lat_extent : array-like
+ Cell latitude vertices (1-D).
+ lon_extent : array-like
+ Cell longitude vertices (1-D), in [-180, 180).
+ etopo_cg : int, optional
+ Coarse-graining factor (stride). High southern latitudes
+ (lat_max < -85°) implicitly multiply this by 5 — see below.
+
+ Returns
+ -------
+ lat, lon, topo
+ 1-D coordinate arrays and the 2-D topography slab, sorted in
+ ascending lat/lon. ``lon`` is in [0, 360) when the cell crosses
+ the dateline; otherwise it stays in [-180, 180).
+ """
+ lat_verts = np.asarray(lat_extent)
+ lon_verts = np.asarray(lon_extent)
+
+ # Dateline detection (robust formula; see compute_split_EW).
+ lon_span = lon_verts.max() - lon_verts.min()
+ lon_verts_360 = np.where(lon_verts < 0.0, lon_verts + 360.0, lon_verts)
+ span_360 = lon_verts_360.max() - lon_verts_360.min()
+ split_EW = (span_360 < lon_span) and (lon_span > 180.0)
+
+ # Determine longitude tile range — three branches: global / dateline / normal.
+ if lon_span >= 360.0:
+ split_EW = False
+ lon_idx_rng = list(range(0, len(_ETOPO_FN_LON) - 1))
+ elif split_EW:
+ min_lon_360 = lon_verts_360.min()
+ max_lon_360 = lon_verts_360.max()
+ min_lon = min_lon_360 if min_lon_360 <= 180 else min_lon_360 - 360
+ max_lon = max_lon_360 if max_lon_360 <= 180 else max_lon_360 - 360
+ lon_min_idx = self._etopo_compute_idx(min_lon, "min", "lon", split_EW)
+ lon_max_idx = self._etopo_compute_idx(max_lon, "max", "lon", split_EW)
+ if lon_min_idx == lon_max_idx:
+ lon_idx_rng = [lon_min_idx]
+ if lon_min_idx >= len(_ETOPO_FN_LON) - 2:
+ lon_idx_rng.append(0)
+ else:
+ lon_idx_rng = list(range(lon_min_idx, len(_ETOPO_FN_LON) - 1)) + list(
+ range(0, lon_max_idx + 1)
+ )
+ else:
+ min_lon = lon_verts.min()
+ max_lon = lon_verts.max()
+ lon_min_idx = self._etopo_compute_idx(min_lon, "min", "lon", split_EW)
+ lon_max_idx = self._etopo_compute_idx(max_lon, "max", "lon", split_EW)
+ if lon_min_idx == lon_max_idx:
+ lon_max_idx += 1
+ lon_idx_rng = list(range(lon_min_idx, lon_max_idx))
+
+ # Latitude tile range — same logic across all longitude branches.
+ lat_min_tile_idx = self._etopo_compute_idx(
+ lat_verts.min(), "min", "lat", split_EW
+ )
+ lat_max_tile_idx = self._etopo_compute_idx(
+ lat_verts.max(), "max", "lat", split_EW
+ )
+ lat_idx_rng = list(range(lat_max_tile_idx, lat_min_tile_idx))
+
+ # Build filenames; load + assemble.
+ fns, lon_cnt, lat_cnt = self._etopo_get_fns(lat_idx_rng, lon_idx_rng)
+ cell_lat, cell_lon, topo_arr = self._etopo_load_topo(
+ fns,
+ lon_cnt,
+ lat_cnt,
+ lat_idx_rng,
+ lon_idx_rng,
+ lat_verts,
+ lon_verts,
+ split_EW,
+ )
+
+ # Wrap longitudes if dateline-crossing, then sort lat/lon and reorder topo.
+ lat_arr = np.array(cell_lat)
+ lon_arr = np.array(cell_lon)
+ if split_EW:
+ lon_arr = np.where(lon_arr < 0.0, lon_arr + 360.0, lon_arr)
+
+ lat_sort_idx = np.argsort(lat_arr)
+ lon_sort_idx = np.argsort(lon_arr)
+ lat_sorted = lat_arr[lat_sort_idx]
+ lon_sorted = lon_arr[lon_sort_idx]
+ topo_sorted = topo_arr[np.ix_(lat_sort_idx, lon_sort_idx)]
+
+ # Coarse-graining — io.py picks up a 5× multiplier for very-southern cells.
+ iint = etopo_cg
+ if iint > 1:
+ try:
+ out_lat = utils.sliding_window_view(lat_sorted, (iint,), (iint,)).mean(
+ axis=-1
+ )
+ out_lon = utils.sliding_window_view(lon_sorted, (iint,), (iint,)).mean(
+ axis=-1
+ )
+ out_topo = utils.sliding_window_view(
+ topo_sorted, (iint, iint), (iint, iint)
+ ).mean(axis=(-1, -2))
+ return out_lat, out_lon, out_topo
+ except (ValueError, MemoryError) as e:
+ logger.warning(
+ f"Coarse-graining failed ({e}); returning full resolution"
+ )
+ return lat_sorted, lon_sorted, topo_sorted
+
+ def close_all(self):
+ """Close all opened NetCDF files."""
+ for fn, ds in self.tiles.items():
+ try:
+ ds.close()
+ if self.verbose:
+ logger.debug(f"Closed tile: {fn}")
+ except Exception as e:
+ logger.error(f"Error closing tile {fn}: {e}")
+
+ self.tiles.clear()
+ self.tile_bounds.clear()
+ self.tile_lats.clear()
+ self.tile_lons.clear()
+
+ def __del__(self):
+ """Ensure files are closed when cache is destroyed."""
+ self.close_all()
+
+ def __enter__(self):
+ """Context manager entry."""
+ return self
+
+ def __exit__(self, exc_type, exc_val, exc_tb):
+ """Context manager exit - ensure files are closed."""
+ self.close_all()
+ return False
+
+
+def create_tile_cache_from_grid(
+ grid, params, padding: float = 0.5
+) -> TopographyTileCache:
+ """
+ Create a tile cache containing all tiles needed for a given grid.
+
+ This analyzes the grid to determine which tiles are needed, then
+ pre-loads them all at once.
+
+ Parameters
+ ----------
+ grid : pycsa.core.var.grid
+ ICON grid object with cell vertices
+ params : pycsa.core.var.params
+ Parameters object with path_merit, path_etopo, etc.
+ padding : float, optional
+ Extra padding in degrees to ensure tiles are loaded, by default 0.5
+
+ Returns
+ -------
+ TopographyTileCache
+ Initialized cache with all required tiles loaded
+ """
+ from pycsa.core import utils
+
+ # Determine global bounds of the grid
+ lat_min = np.min(grid.clat_vertices) - padding
+ lat_max = np.max(grid.clat_vertices) + padding
+ lon_min = np.min(grid.clon_vertices) - padding
+ lon_max = np.max(grid.clon_vertices) + padding
+
+ logger.info(
+ f"Grid spans: lat[{lat_min:.2f}, {lat_max:.2f}], lon[{lon_min:.2f}, {lon_max:.2f}]"
+ )
+
+ # Determine which tiles to load (using MERIT tile naming convention)
+ # TODO: Implement automatic tile discovery based on bounds
+ # For now, this is a placeholder - you'll need to implement the logic
+ # to determine required tile filenames based on the grid bounds
+
+ # Example: if using MERIT data with standard 30x30 degree tiles
+ tile_filenames = _get_merit_tiles_for_bounds(lat_min, lat_max, lon_min, lon_max)
+
+ logger.info(f"Loading {len(tile_filenames)} topography tiles for grid coverage")
+
+ # Create and return cache
+ return TopographyTileCache(
+ data_dir=params.path_merit,
+ tile_filenames=tile_filenames,
+ dataset_type="MERIT",
+ verbose=params.verbose if hasattr(params, "verbose") else False,
+ )
+
+
+def _get_merit_tiles_for_bounds(
+ lat_min: float, lat_max: float, lon_min: float, lon_max: float
+) -> List[str]:
+ """
+ Determine MERIT tile filenames needed to cover the given bounds.
+
+ MERIT tiles are 30x30 degrees and named like:
+ MERIT_N60-N90_W180-W150.nc4
+ """
+ # MERIT tile boundaries (standard grid)
+ merit_lat_bounds = np.array([90.0, 60.0, 30.0, 0.0, -30.0, -60.0, -90.0])
+ merit_lon_bounds = np.array(
+ [
+ -180.0,
+ -150.0,
+ -120.0,
+ -90.0,
+ -60.0,
+ -30.0,
+ 0.0,
+ 30.0,
+ 60.0,
+ 90.0,
+ 120.0,
+ 150.0,
+ 180.0,
+ ]
+ )
+
+ tile_filenames = []
+
+ # Find lat tile indices
+ lat_idx_min = np.searchsorted(merit_lat_bounds[::-1], lat_min, side="left")
+ lat_idx_max = np.searchsorted(merit_lat_bounds[::-1], lat_max, side="right")
+
+ # Find lon tile indices
+ lon_idx_min = np.searchsorted(merit_lon_bounds, lon_min, side="left")
+ lon_idx_max = np.searchsorted(merit_lon_bounds, lon_max, side="right")
+
+ def _get_nsew(val, coord_type):
+ """Get N/S/E/W tag for coordinate value."""
+ if coord_type == "lat":
+ return "N" if val >= 0 else "S"
+ else: # lon
+ return "E" if val >= 0 else "W"
+
+ # Generate filenames
+ for lat_idx in range(
+ max(0, lat_idx_min - 1), min(len(merit_lat_bounds) - 1, lat_idx_max + 1)
+ ):
+ l_lat = merit_lat_bounds[lat_idx]
+ r_lat = merit_lat_bounds[lat_idx + 1]
+ l_lat_tag = _get_nsew(l_lat, "lat")
+ r_lat_tag = _get_nsew(r_lat, "lat")
+
+ for lon_idx in range(
+ max(0, lon_idx_min - 1), min(len(merit_lon_bounds) - 1, lon_idx_max + 1)
+ ):
+ l_lon = merit_lon_bounds[lon_idx]
+ r_lon = merit_lon_bounds[lon_idx + 1]
+ l_lon_tag = _get_nsew(l_lon, "lon")
+ r_lon_tag = _get_nsew(r_lon, "lon")
+
+ # Check if this is REMA region (Antarctica)
+ if l_lat == -60.0 and r_lat == -90.0:
+ dataset_name = "REMA_BKG"
+ else:
+ dataset_name = "MERIT"
+
+ filename = f"{dataset_name}_{l_lat_tag}{abs(int(l_lat)):02d}-{r_lat_tag}{abs(int(r_lat)):02d}_{l_lon_tag}{abs(int(l_lon)):03d}-{r_lon_tag}{abs(int(r_lon)):03d}.nc4"
+ tile_filenames.append(filename)
+
+ return tile_filenames
+
+
+# ---------------------------------------------------------------------------
+# Per-worker cache lifecycle helpers
+# ---------------------------------------------------------------------------
+# The HPC main loop runs under Dask with processes=True, so each worker is a
+# separate process with its own module namespace. init_worker_cache is called
+# via client.run(...) once per memory batch to populate _WORKER_CACHE on each
+# worker; do_cell then reaches it via get_worker_cache(). This keeps NetCDF
+# file handles open across cells within a worker (the actual saving), without
+# trying to share state between processes (which would fail — nc.Dataset
+# handles aren't picklable).
+
+_WORKER_CACHE: Optional[TopographyTileCache] = None
+
+
+def init_worker_cache(data_dir: str, dataset_type: str = "ETOPO") -> bool:
+ """Initialise a lazy tile cache in the current worker process.
+
+ Intended to be called via `client.run(init_worker_cache, path_etopo)` at
+ the start of each memory batch. Idempotent: a second call with the same
+ arguments is a no-op so reinitialisation across batches is cheap.
+
+ Returns True so client.run reports {worker_addr: True, ...} on success.
+ """
+ global _WORKER_CACHE
+ if _WORKER_CACHE is not None and str(_WORKER_CACHE.data_dir) == str(Path(data_dir)):
+ return True
+ _WORKER_CACHE = TopographyTileCache(
+ data_dir=data_dir,
+ tile_filenames=[],
+ dataset_type=dataset_type,
+ verbose=False,
+ )
+ return True
+
+
+def get_worker_cache() -> TopographyTileCache:
+ """Return this worker's tile cache; raise if init_worker_cache wasn't called."""
+ if _WORKER_CACHE is None:
+ raise RuntimeError(
+ "TopographyTileCache not initialised on this worker. "
+ "Call init_worker_cache(data_dir) via client.run(...) first."
+ )
+ return _WORKER_CACHE
+
+
+def close_worker_cache() -> bool:
+ """Close NetCDF handles and drop the worker cache. Returns True."""
+ global _WORKER_CACHE
+ if _WORKER_CACHE is not None:
+ _WORKER_CACHE.close_all()
+ _WORKER_CACHE = None
+ return True
diff --git a/src/utils.py b/pycsa/core/utils.py
similarity index 79%
rename from src/utils.py
rename to pycsa/core/utils.py
index 4084165..42b2e01 100644
--- a/src/utils.py
+++ b/pycsa/core/utils.py
@@ -15,9 +15,23 @@ def pick_cell(
grid,
radius=1.0,
):
- """
- .. deprecated:: 0.90.0
+ """pick an ICON grid cell given (lon,lat) coorindates
+ Parameters
+ ----------
+ lat_ref : float
+ reference latitude coordinate in the cell to be picked
+ lon_ref : float
+ reference longitude coordinate in the cell to be picked
+ grid : class:`src.var.grid`
+ instance of an ICON grid
+ radius : float, optional
+ radius from `(lon_ref, lat_ref)` to search for `(clon,clat)`, by default 1.0
+
+ Returns
+ -------
+ _type_
+ _description_
"""
clat, clon = grid.clat, grid.clon
index = np.nonzero(
@@ -426,6 +440,7 @@ def get_lat_lon_segments(
topo_mask=None,
mask=None,
load_topo=False,
+ use_center=True,
):
"""
Populates an empty :class:`cell ` object given the vertices and underlying topography.
@@ -452,6 +467,9 @@ def get_lat_lon_segments(
2D Boolean mask to select for data points inside the non-quadrilateral grid cell, by default None
load_topo : bool, optional
explicitly replaces the topography attribute in the cell ``cell.topo`` with the data given in ``topo``, by default False
+ use_center : bool, optional
+ If True (default), use center of domain as projection origin (minimizes distortion)
+ If False, use corner of domain as projection origin (OLD behavior for testing)
"""
lat_max = get_closest_idx(lat_verts.max(), topo.lat) + padding
lat_min = get_closest_idx(lat_verts.min(), topo.lat) - padding
@@ -462,8 +480,15 @@ def get_lat_lon_segments(
cell.lat = np.copy(topo.lat[lat_min:lat_max])
cell.lon = np.copy(topo.lon[lon_min:lon_max])
- lon_origin = cell.lon[0]
- lat_origin = cell.lat[0]
+ # Choose projection origin based on use_center parameter
+ if use_center:
+ # NEW (default): Use midpoint of domain as projection center (minimizes distortion, especially at poles)
+ lon_origin = (cell.lon.min() + cell.lon.max()) / 2.0
+ lat_origin = (cell.lat.min() + cell.lat.max()) / 2.0
+ else:
+ # OLD: Use corner of domain as projection origin (for testing/comparison)
+ lon_origin = cell.lon[0]
+ lat_origin = cell.lat[0]
lat_in_m = latlon2m(cell.lat, lon_origin, latlon="lat")
lon_in_m = latlon2m(cell.lon, lat_origin, latlon="lon")
@@ -517,15 +542,35 @@ def get_lat_lon_segments(
if topo_mask is not None:
cell.topo *= topo_mask
+ # Convert vertices from degrees to planar coordinates (meters) for triangle masking
+ # This is critical at polar latitudes where degree-space and meter-space have different geometries
+ # We need to convert each vertex individually using the same projection origin as the grid
+
+ Rm = 6371000.0 # Earth radius in meters
+
+ # Convert latitude vertices (meridional distance from first grid point)
+ # Keep sign to preserve direction (north/south)
+ lat_ref = cell.lat[0] # Reference point (first grid latitude)
+ lat_verts_in_m = (np.radians(lat_verts) - np.radians(lat_ref)) * Rm
+
+ # Convert longitude vertices (zonal distance along parallel at lat_origin)
+ # Keep sign to preserve direction (east/west)
+ lon_ref = cell.lon[0] # Reference point (first grid longitude)
+ lon_verts_in_m = (
+ (np.radians(lon_verts) - np.radians(lon_ref))
+ * Rm
+ * np.cos(np.radians(lat_origin))
+ )
+
if padding > 0:
triangle = gen_triangle(
- lon_verts,
- lat_verts,
- x_rng=[cell.lon.min(), cell.lon.max()],
- y_rng=[cell.lat.min(), cell.lat.max()],
+ lon_verts_in_m,
+ lat_verts_in_m,
+ x_rng=[lon_in_m.min(), lon_in_m.max()],
+ y_rng=[lat_in_m.min(), lat_in_m.max()],
)
else:
- triangle = gen_triangle(lon_verts, lat_verts)
+ triangle = gen_triangle(lon_verts_in_m, lat_verts_in_m)
# crucial to update of the lat-lon data in the cell object AFTER the initialisation of the triangle object.
cell.lat = lat_in_m
@@ -547,43 +592,49 @@ def get_closest_idx(val, arr):
def latlon2m(arr, fix_pt, latlon):
- """Wrapper function to compute the distance of a list of values from a given fixed point (in meters).
+ """Compute along-axis distances (in meters) from the first element.
Parameters
----------
- arr : list
- list of values in degrees
+ arr : array-like
+ 1D list/array of coordinates in degrees (latitudes if ``latlon='lat'``,
+ longitudes if ``latlon='lon'``)
fix_pt : float
- given fixed point, e.g. the origin, in degrees
- latlon : str
- ``lat`` if the distance are to be computed in the latitudinal direction, ``lon`` otherwise.
+ Fixed coordinate in degrees:
+ - for ``latlon='lat'``: the fixed longitude at which meridional distances are evaluated
+ - for ``latlon='lon'``: the fixed latitude at which zonal (small-circle) distances are evaluated
+ latlon : {"lat", "lon"}
+ Which axis the distances are computed along.
Returns
-------
- float
- distance in meters
+ numpy.ndarray
+ Cumulative distances in meters starting at 0, monotonically non-decreasing.
"""
- arr = np.array(arr)
+ arr = np.asarray(arr, dtype=float)
assert arr.ndim == 1
- origin = arr[0]
- res = np.zeros_like(arr)
- res[0] = 0.0
-
- for cnt, idx in enumerate(range(1, len(arr))):
- cnt += 1
- if latlon == "lat":
- res[cnt] = __latlon2m_converter(fix_pt, fix_pt, origin, arr[idx])
- elif latlon == "lon":
- res[cnt] = __latlon2m_converter(origin, arr[idx], fix_pt, fix_pt)
- else:
- assert 0
+ Rm = 6371000.0 # mean Earth radius in meters
+
+ if latlon == "lat":
+ # Meridional arc length: great circle along a meridian
+ phi = np.radians(arr)
+ dphi = np.diff(phi, prepend=phi[0])
+ steps = np.abs(dphi) * Rm
+ elif latlon == "lon":
+ # Zonal distance along a parallel (small-circle) at latitude fix_pt
+ # Handle dateline by unwrapping longitudes first
+ lam = np.unwrap(np.radians(arr))
+ dlam = np.diff(lam, prepend=lam[0])
+ steps = np.abs(dlam) * Rm * np.cos(np.radians(fix_pt))
+ else:
+ raise ValueError("latlon must be 'lat' or 'lon'")
- return res * 1000
+ return np.cumsum(steps)
def __latlon2m_converter(lon1, lon2, lat1, lat2):
- """Helper function for lat-lon to meters conversion
+ """Helper function for great-circle distance between two lat-lon points.
Parameters
----------
@@ -599,7 +650,7 @@ def __latlon2m_converter(lon1, lon2, lat1, lat2):
Returns
-------
float
- distance between ``(lat1,lon1)`` and ``(lat2,lon2)`` in meters.
+ Great-circle distance between ``(lat1,lon1)`` and ``(lat2,lon2)`` in kilometers.
.. note:: Taken from https://stackoverflow.com/questions/19412462/getting-distance-between-two-points-based-on-latitude-longitude
@@ -794,3 +845,59 @@ def __stencil(gam):
stencil = (1.0 - gam) * stencil_iso + gam * stencil_aniso
return stencil
+
+
+def transfer_attributes(params, cls, prefix=""):
+ for key, value in vars(cls).items():
+ if len(prefix) > 0:
+ key = prefix + "_" + key
+
+ if not hasattr(params, key):
+ setattr(params, key, value)
+ elif getattr(params, key) == None:
+ setattr(params, key, value)
+
+
+def is_land(cell, simplex_lat, simplex_lon, topo, height_tol=0.5, percent_tol=0.95):
+
+ get_lat_lon_segments(
+ simplex_lat, simplex_lon, cell, topo, load_topo=True, filtered=False
+ )
+
+ if not (((cell.topo <= height_tol).sum() / cell.topo.size) > percent_tol):
+ return True
+ else:
+ return False
+
+
+def handle_latlon_expansion(
+ clat_vertices, clon_vertices, lat_expand=1.0, lon_expand=1.0
+):
+ clon_vertices = np.around(clon_vertices, 5)
+ clat_vertices = np.around(clat_vertices, 5)
+
+ # clon_vertices[np.where(np.abs(np.abs(clon_vertices) - 180.0) < 1e-5)] = 180.0
+ clon_vertices[np.where(clon_vertices == 180.0)] = (
+ np.sign(clon_vertices.min()) * 180.0
+ )
+ clon_vertices[np.where(clon_vertices == -180.0)] = (
+ np.sign(clon_vertices.max()) * 180.0
+ )
+
+ clat_vertices[np.argmax(clat_vertices)] += lat_expand
+ clon_vertices[np.argmax(clon_vertices)] += lon_expand
+
+ clat_vertices[np.argmin(clat_vertices)] -= lat_expand
+ clon_vertices[np.argmin(clon_vertices)] -= lon_expand
+
+ clon_vertices[np.where(clon_vertices < -180.0)] += 360.0
+ clon_vertices[np.where(clon_vertices > 180.0)] -= 360.0
+
+ clat_vertices = np.where(
+ clat_vertices < -90.0, clat_vertices + lat_expand, clat_vertices
+ )
+ clat_vertices = np.where(
+ clat_vertices > 90.0, clat_vertices - lat_expand, clat_vertices
+ )
+
+ return clat_vertices, clon_vertices
diff --git a/src/var.py b/pycsa/core/var.py
similarity index 93%
rename from src/var.py
rename to pycsa/core/var.py
index aadae16..0938d54 100644
--- a/src/var.py
+++ b/pycsa/core/var.py
@@ -3,7 +3,7 @@
"""
import numpy as np
-from src import utils, io
+from pycsa.core import utils, io
class grid(object):
@@ -22,6 +22,7 @@ def __init__(self):
self.clon = None
self.clon_vertices = None
self.links = None
+ self.cell_area = None
def apply_f(self, f):
"""
@@ -32,7 +33,7 @@ def apply_f(self, f):
f : ``function``
arbitrary function to be applied to class attributes, e.g. a radians-degrees converter.
"""
- self.non_convertibles = ["non_convertibles", "links"]
+ self.non_convertibles = ["non_convertibles", "links", "cell_area"]
for key, value in vars(self).items():
if key in self.non_convertibles:
pass
@@ -238,15 +239,19 @@ def get_attrs(self, fobj, freqs):
self.kks = fobj.m_i / (fobj.Ni)
self.lls = fobj.m_j / (fobj.Nj)
- self.kks, self.lls = np.meshgrid(self.kks, self.lls)
+ wla = self.wlat
+ wlo = self.wlon
- # self.kks = self.kks / self.kks.size
- # self.lls = self.lls / self.lls.size
+ kks = self.kks * 2.0 * np.pi
+ lls = self.lls * 2.0 * np.pi
- # self.clat = ma.getdata(df.variables['clat'][:])
- # clat_vertices = ma.getdata(df.variables['clat_vertices'][:])
- # clon = ma.getdata(df.variables['clon'][:])
- # clon_vertices = ma.getdata(df.variables['clon_vertices'][:])
+ kks = kks / wlo
+ lls = lls / wla
+
+ self.dk = np.diff(self.kks).mean()
+ self.dl = np.diff(self.lls).mean()
+
+ self.kks, self.lls = np.meshgrid(kks, lls)
def grid_kk_ll(self, fobj, dat):
"""
@@ -294,15 +299,15 @@ def __init__(self):
"""
# Define filenames
self.run_case = ""
- self.path = "../data/"
- self.fn_grid = self.path + "icon_compact.nc"
- self.fn_topo = self.path + "topo_compact.nc"
+ self.path_compact_grid = None
+ self.path_compact_topo = None
- self.output_fn = None
+ self.path_output = None
+ self.fn_output = None
self.enable_merit = True
self.merit_cg = 10
- self.merit_path = "/home/ray/Documents/orog_data/MERIT/"
+ self.path_merit = None
# Domain size
self.lat_extent = None
@@ -363,8 +368,8 @@ def self_test(self):
bool
True if test passed, False otherwise
"""
- if self.output_fn is None:
- self.output_fn = io.fn_gen(self)
+ if self.fn_output is None:
+ self.fn_output = io.fn_gen(self)
self.check_init()
diff --git a/pycsa/local_paths.py.template b/pycsa/local_paths.py.template
new file mode 100644
index 0000000..b5fd3bb
--- /dev/null
+++ b/pycsa/local_paths.py.template
@@ -0,0 +1,42 @@
+"""
+Template for local paths configuration.
+
+To use:
+1. Copy this file to local_paths.py: cp local_paths.py.template local_paths.py
+2. Edit local_paths.py with your actual paths
+3. Never commit local_paths.py (it's in .gitignore)
+
+Environment variables (optional):
+You can also set these as environment variables:
+- SPEC_APPX_DATA_DIR: Base directory for project data
+- SPEC_APPX_OUTPUT_DIR: Output directory
+- SPEC_APPX_MERIT_DIR: MERIT data directory
+- SPEC_APPX_REMA_DIR: REMA data directory
+- SPEC_APPX_ETOPO_DIR: ETOPO data directory
+"""
+
+import os
+from pathlib import Path
+from pycsa import var
+
+paths = var.obj()
+
+# Get base directories from environment or use defaults
+data_dir = os.getenv('SPEC_APPX_DATA_DIR', '/path/to/data')
+output_dir = os.getenv('SPEC_APPX_OUTPUT_DIR', '/path/to/outputs')
+merit_dir = os.getenv('SPEC_APPX_MERIT_DIR', '/path/to/MERIT')
+rema_dir = os.getenv('SPEC_APPX_REMA_DIR', '/path/to/REMA')
+etopo_dir = os.getenv('SPEC_APPX_ETOPO_DIR', '/path/to/etopo_15s')
+
+# Project data paths
+paths.compact_grid = os.path.join(data_dir, "icon_compact.nc")
+paths.compact_topo = os.path.join(data_dir, "topo_compact.nc")
+paths.icon_grid = os.path.join(data_dir, "icon_grid_0012_R02B04_G_linked.nc")
+
+# Output path
+paths.output = os.path.join(output_dir, "global_run/")
+
+# External data sources
+paths.merit = merit_dir
+paths.rema = rema_dir
+paths.etopo = etopo_dir
diff --git a/vis/__init__.py b/pycsa/plotting/__init__.py
similarity index 100%
rename from vis/__init__.py
rename to pycsa/plotting/__init__.py
diff --git a/vis/cart_plot.py b/pycsa/plotting/cart_plot.py
similarity index 97%
rename from vis/cart_plot.py
rename to pycsa/plotting/cart_plot.py
index 2587bce..7ebce86 100644
--- a/vis/cart_plot.py
+++ b/pycsa/plotting/cart_plot.py
@@ -17,7 +17,7 @@
)
-def lat_lon(topo, fs=(10, 6), int=1):
+def lat_lon(topo, fs=(10, 6), int=1, colorbar_margins=None):
"""
Does a simple Plate-Carre projection of a lat-lon topography data.
@@ -44,7 +44,9 @@ def lat_lon(topo, fs=(10, 6), int=1):
cmap="GnBu",
)
- cax = fig.add_axes([0.99, 0.22, 0.025, 0.55])
+ if colorbar_margins is None:
+ colorbar_margins = [0.99, 0.22, 0.025, 0.55]
+ cax = fig.add_axes(colorbar_margins)
fig.colorbar(im, cax=cax)
gl = ax.gridlines(
@@ -398,12 +400,12 @@ def lat_lon_icon(
fc="r",
alpha=0.2,
linewidth=1,
- transform=ccrs.Geodetic(),
+ transform=ccrs.PlateCarree(),
zorder=3,
)
ax.add_collection(coll)
- print("--> polygon collection done")
+ # print("--> polygon collection done")
if annotate_idxs:
ncells = kwargs["ncells"]
@@ -427,3 +429,4 @@ def lat_lon_icon(
# -- maximize and save the PNG file
if output_fig:
plt.savefig(fn, bbox_inches="tight", dpi=200)
+ plt.close()
diff --git a/vis/plotter.py b/pycsa/plotting/plotter.py
similarity index 94%
rename from vis/plotter.py
rename to pycsa/plotting/plotter.py
index db792ac..5535abe 100644
--- a/vis/plotter.py
+++ b/pycsa/plotting/plotter.py
@@ -36,7 +36,14 @@ def __init__(self, fig, nhi, nhj, cbar=True, set_label=True):
self.set_label = set_label
def phys_panel(
- self, axs, data, title="", extent=None, xlabel="", ylabel="", v_extent=None
+ self,
+ axs,
+ data,
+ title="",
+ extent=None,
+ xlabel="",
+ ylabel="",
+ v_extent=None,
):
"""
Plots a physical depiction of the input data.
@@ -157,10 +164,10 @@ def freq_panel(
if self.cbar:
self.fig.colorbar(im, ax=axs, fraction=0.2, pad=0.04, shrink=0.7)
- m_j = np.arange(-nhj / 2 + 1, nhj / 2 + 1)
+ m_j = np.arange(-nhj / 2 + 1, nhj / 2 + 1).astype(int)
ylocs = np.arange(0.5, nhj + 0.5, 1.0)
- m_i = np.arange(0, nhi)
+ m_i = np.arange(0, nhi).astype(int)
xlocs = np.arange(0.5, nhi + 0.5, 1.0)
axs.set_xticks(xlocs, m_i, rotation=-90)
@@ -168,9 +175,8 @@ def freq_panel(
axs.set_title(title)
if self.set_label:
- axs.set_ylabel(r"$m$", fontsize=12)
-
- axs.set_xlabel(r"$n$", fontsize=12)
+ axs.set_ylabel("m", fontsize=12, fontstyle="italic")
+ axs.set_xlabel("n", fontsize=12, fontstyle="italic")
# axs.set_aspect('equal')
# ref: https://stackoverflow.com/questions/20337664/cleanest-way-to-hide-every-nth-tick-label-in-matplotlib-colorbar
@@ -247,8 +253,8 @@ def fft_freq_panel(
axs.set_title(title)
if self.set_label:
- axs.set_xlabel(r"$k$ [m$^{-1}$]", fontsize=12)
- axs.set_ylabel(r"$l$ [m$^{-1}$]", fontsize=12)
+ axs.set_xlabel("k [1/m]", fontsize=12, fontstyle="italic")
+ axs.set_ylabel("l [1/m]", fontsize=12, fontstyle="italic")
if typ == "imag":
axs.set_aspect("equal")
@@ -268,6 +274,7 @@ def error_bar_plot(
fs=(10.0, 6.0),
ylabel="",
fontsize=8,
+ show_grid=True,
):
"""
Bar plot of errors.
@@ -298,6 +305,8 @@ def error_bar_plot(
y-axis label, by default ""
fontsize : int, optional
by default 8
+ show_grid : bool, optional
+ toggles grid in output, by default True
"""
data = pd.DataFrame(pmf_diff, index=idx_name, columns=["values"])
@@ -333,7 +342,8 @@ def error_bar_plot(
fontsize=fontsize,
)
- plt.grid()
+ if show_grid:
+ plt.grid()
plt.xlabel("first grid pair index", fontsize=fontsize + 3)
@@ -375,6 +385,7 @@ def error_bar_split_plot(
bs,
ts,
ts_ticks,
+ color,
fs=(3.5, 3.5),
title="",
output_fig=False,
@@ -396,10 +407,11 @@ def error_bar_split_plot(
ax2.set_ylim(0, bs)
ax1.set_ylim(ts[0], ts[1])
ax1.set_yticks(ts_ticks)
+ ax1.ticklabel_format(style="plain")
- bars1 = ax1.bar(XX.index, XX.values, color=("C0"))
- bars2 = ax2.bar(XX.index, XX.values, color=("C0", "C1", "C2", "r"))
- ax1.bar_label(bars1, padding=3)
+ bars1 = ax1.bar(XX.index, XX.values, color=color)
+ bars2 = ax2.bar(XX.index, XX.values, color=color)
+ ax1.bar_label(bars1, padding=3, fmt="%d")
ax2.bar_label(bars2, padding=3)
for tick in ax2.get_xticklabels():
@@ -542,7 +554,5 @@ def plot(self, Z, output_fig=True, output_fn="plot_3D", lbls=None, fs=(10, 10)):
plt.tight_layout()
if output_fig:
- plt.savefig(
- "../manuscript/%s.pdf" % output_fn, dpi=200, bbox_inches="tight"
- )
+ plt.savefig("./outputs/%s.pdf" % output_fn, dpi=200, bbox_inches="tight")
plt.show()
diff --git a/wrappers/__init__.py b/pycsa/wrappers/__init__.py
similarity index 100%
rename from wrappers/__init__.py
rename to pycsa/wrappers/__init__.py
diff --git a/wrappers/diagnostics.py b/pycsa/wrappers/diagnostics.py
similarity index 93%
rename from wrappers/diagnostics.py
rename to pycsa/wrappers/diagnostics.py
index 86fdd4d..2e51925 100644
--- a/wrappers/diagnostics.py
+++ b/pycsa/wrappers/diagnostics.py
@@ -1,17 +1,17 @@
"""
-Diagnostic wrapper module to ease setting up the CSAM building blocks
+Diagnostic wrapper module to ease setting up the CSA building blocks
"""
import numpy as np
-from src import physics
-from vis import plotter
+from pycsa.core import physics
+from pycsa.plotting import plotter
from copy import deepcopy
import matplotlib.pyplot as plt
class delaunay_metrics(object):
- """Helper class for evaluation of the CSAM on a Delaunay triangulated domain."""
+ """Helper class for evaluation of the CSA on a Delaunay triangulated domain."""
def __init__(self, params, tri, writer=None):
"""
@@ -67,7 +67,7 @@ def get_rel_err(self, triangle_pair):
Returns
-------
float
- the relative error of the CSAM on the Delaunay triangles against the FFT-computed reference
+ the relative error of the CSA on the Delaunay triangles against the FFT-computed reference
"""
self.update_pair(triangle_pair, store_error=False)
self.rel_err = self.__get_rel_diff(self.uw_sum, self.uw_ref)
@@ -168,10 +168,12 @@ def __write(self):
def __gen_percentage_errs(self):
"""Computes the relative and maximum errors in percentage"""
- max_idx = np.argmax(np.abs(self.pmf_refs))
- self.max_errs = self.__get_max_diff(
- self.pmf_sums, self.pmf_refs, np.array(self.pmf_refs[max_idx])
- )
+ if hasattr(self, "max_val"):
+ max_val = self.max_val
+ else:
+ max_idx = np.argmax(np.abs(self.pmf_refs))
+ max_val = self.pmf_refs[max_idx]
+ self.max_errs = self.__get_max_diff(self.pmf_sums, self.pmf_refs, max_val)
self.rel_errs = self.__get_rel_diff(self.pmf_sums, self.pmf_refs)
self.max_errs = np.array(self.max_errs) * 100
@@ -210,7 +212,7 @@ def __get_max_diff(arr, ref, max):
class diag_plotter(object):
- """Helper class to plot CSAM-computed data"""
+ """Helper class to plot CSA-computed data"""
def __init__(self, params, nhi, nhj):
"""
@@ -228,7 +230,7 @@ def __init__(self, params, nhi, nhj):
self.nhi = nhi
self.nhj = nhj
- self.output_dir = "../manuscript/"
+ self.output_dir = "./outputs/"
def show(
self,
@@ -252,7 +254,7 @@ def show(
sols : tuple
contains the data for plotting:
| (:class:`src.var.topo_cell` instance,
- | computed CSAM spectrum,
+ | computed CSA spectrum,
| computed idealised pseudo-momentum fluxes,
| the reconstructed physical data)
@@ -262,7 +264,7 @@ def show(
v_extent : list, optional
``[z_min, z_max]`` the vertical extent of the physical reconstruction, by default None
dfft_plot : bool, optional
- toggles whether a spectrum is the full FFT spectral space or the dense truncated CSAM spectrum, By default False, i.e. plot CSAM spectrum.
+ toggles whether a spectrum is the full FFT spectral space or the dense truncated CSA spectrum, By default False, i.e. plot CSA spectrum.
output_fig : bool, optional
toggles writing figure output, by default True
fs : tuple, optional
@@ -290,8 +292,8 @@ def show(
if ir_args is None:
if type(rect_idx) is int:
idxs_tag = "Cell %i" % rect_idx
- tag = "CSAM"
- fn = "plots_CSAM_%i" % rect_idx
+ tag = "CSA"
+ fn = "plots_CSA_%i" % rect_idx
elif len(rect_idx) == 2:
idxs_tag = "(%i,%i)" % (rect_idx[0], rect_idx[1])
tag = "FFT" if dfft_plot else "FA LSFF"
diff --git a/wrappers/interface.py b/pycsa/wrappers/interface.py
similarity index 89%
rename from wrappers/interface.py
rename to pycsa/wrappers/interface.py
index 39d2186..33b3a38 100644
--- a/wrappers/interface.py
+++ b/pycsa/wrappers/interface.py
@@ -1,10 +1,9 @@
"""
-Interface wrapper module to ease setting up the CSAM building blocks
+Interface wrapper module to ease setting up the CSA building blocks
"""
-
-from src import fourier, lin_reg, physics, reconstruction
-from src import utils, var
+from pycsa.core import fourier, lin_reg, physics, reconstruction
+from pycsa.core import utils, var
from copy import deepcopy
import numpy as np
@@ -31,7 +30,13 @@ def __init__(self, nhi, nhj, U, V, debug=False):
debug : bool, optional
debug flag, by default False
"""
- self.fobj = fourier.f_trans(nhi, nhj)
+ # Initialize buffer pool for memory-efficient array reuse
+ from pycsa.core.buffer_pool import BufferPool
+
+ self.buffer_pool = BufferPool()
+
+ # Initialize Fourier transformer with buffer pool
+ self.fobj = fourier.f_trans(nhi, nhj, buffer_pool=self.buffer_pool)
self.U = U
self.V = V
@@ -59,6 +64,8 @@ def sappx(self, cell, lmbda=0.1, scale=1.0, **kwargs):
lmbda,
kwargs.get("iter_solve", True),
kwargs.get("save_coeffs", False),
+ buffer_pool=self.buffer_pool,
+ use_sparse=kwargs.get("use_sparse", False),
)
if kwargs.get("save_am", False):
@@ -70,7 +77,11 @@ def sappx(self, cell, lmbda=0.1, scale=1.0, **kwargs):
if kwargs.get("refine", False):
cell.topo_m -= data_recons
am, data_recons = lin_reg.do(
- self.fobj, cell, lmbda, kwargs.get("iter_solve", True)
+ self.fobj,
+ cell,
+ lmbda,
+ kwargs.get("iter_solve", True),
+ buffer_pool=self.buffer_pool,
)
self.fobj.get_freq_grid(am)
@@ -362,7 +373,7 @@ def __init__(self, nhi, nhj, params, topo):
self.params = params
self.topo = topo
- def do(self, simplex_lat, simplex_lon, res_topo=None):
+ def do(self, simplex_lat, simplex_lon, res_topo=None, use_center=True):
"""Do the First Approximation step
Parameters
@@ -374,6 +385,8 @@ def do(self, simplex_lat, simplex_lon, res_topo=None):
_description_
res_topo : array-like, optional
residual orography, only required in iterative refinement, by default None
+ use_center : bool, optional
+ use centered planar projection (True) or corner-based (False), by default True
Returns
-------
@@ -381,7 +394,7 @@ def do(self, simplex_lat, simplex_lon, res_topo=None):
contains the data for plotting:
| (:class:`src.var.topo_cell` instance,
- | computed CSAM spectrum,
+ | computed CSA spectrum,
| computed idealised pseudo-momentum fluxes,
| the reconstructed physical data)
@@ -394,7 +407,12 @@ def do(self, simplex_lat, simplex_lon, res_topo=None):
taper_quad(self.params, simplex_lat, simplex_lon, cell_fa, self.topo)
else:
utils.get_lat_lon_segments(
- simplex_lat, simplex_lon, cell_fa, self.topo, rect=self.params.rect
+ simplex_lat,
+ simplex_lon,
+ cell_fa,
+ self.topo,
+ rect=self.params.rect,
+ use_center=use_center,
)
else:
cell_fa.topo = res_topo
@@ -406,6 +424,7 @@ def do(self, simplex_lat, simplex_lon, res_topo=None):
padding=self.params.padding,
rect=False,
mask=np.ones_like(res_topo).astype(bool),
+ use_center=use_center,
)
first_guess = get_pmf(self.nhi, self.nhj, self.params.U, self.params.V)
@@ -443,7 +462,7 @@ def __init__(self, nhi, nhj, params, topo, tri):
self.nhi, self.nhj = nhi, nhj
self.n_modes = params.n_modes
- def do(self, idx, ampls_fa, res_topo=None):
+ def do(self, idx, ampls_fa, res_topo=None, use_center=True):
"""Do the Second Approximation step
Parameters
@@ -454,6 +473,8 @@ def do(self, idx, ampls_fa, res_topo=None):
spectral modes identified in the first approximation step
res_topo : array-like, optional
residual orography, only required in iterative refinement, by default None
+ use_center : bool, optional
+ use centered planar projection (True) or corner-based (False), by default True
Returns
-------
@@ -461,7 +482,7 @@ def do(self, idx, ampls_fa, res_topo=None):
contains the data for plotting:
| (:class:`src.var.topo_cell` instance,
- | computed CSAM spectrum,
+ | computed CSA spectrum,
| computed idealised pseudo-momentum fluxes,
| the reconstructed physical data)
@@ -471,9 +492,9 @@ def do(self, idx, ampls_fa, res_topo=None):
"""
# make a copy of the spectrum obtained from the FA.
fq_cpy = np.copy(ampls_fa)
- fq_cpy[
- np.isnan(fq_cpy)
- ] = 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ fq_cpy[np.isnan(fq_cpy)] = (
+ 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ )
cell = var.topo_cell()
@@ -481,7 +502,9 @@ def do(self, idx, ampls_fa, res_topo=None):
simplex_lon = self.tri.tri_lon_verts[idx]
# use the non-quadrilateral self.topography
- utils.get_lat_lon_segments(simplex_lat, simplex_lon, cell, self.topo, rect=True)
+ utils.get_lat_lon_segments(
+ simplex_lat, simplex_lon, cell, self.topo, rect=True, use_center=use_center
+ )
save_am = True if self.params.recompute_rhs else False
@@ -489,7 +512,13 @@ def do(self, idx, ampls_fa, res_topo=None):
cell.topo = res_topo * cell.mask
utils.get_lat_lon_segments(
- simplex_lat, simplex_lon, cell, self.topo, rect=False, filtered=False
+ simplex_lat,
+ simplex_lon,
+ cell,
+ self.topo,
+ rect=False,
+ filtered=False,
+ use_center=use_center,
)
if self.params.taper_sa:
diff --git a/pyproject.toml b/pyproject.toml
index 966dee6..2728e3c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,3 +1,37 @@
+[project]
+name = "pyCSA"
+version = "0.95.1"
+
+dependencies = [
+ "Cartopy==0.25.0",
+ "dask[distributed]",
+ "h5py==3.15.1",
+ "matplotlib==3.10.7",
+ "netCDF4==1.7.3",
+ "noise==1.2.2",
+ "numba==0.62.1",
+ "numpy==2.2.6",
+ "pandas==2.3.3",
+ "scipy==1.15.3",
+ "tqdm>=4.66.0",
+]
+
+[project.optional-dependencies]
+test = [
+ "pytest>=7.0",
+ "pytest-cov>=4.0",
+]
+
+# Packaging
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["pycsa*"]
+
+
[tool.towncrier]
directory = "changelog.d"
filename = "CHANGELOG.rst"
@@ -31,4 +65,21 @@ showcontent = true
[[tool.towncrier.type]]
directory = "fixed"
name = "Fixed"
-showcontent = true
\ No newline at end of file
+showcontent = true
+
+# Pytest configuration
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_classes = ["Test*"]
+python_functions = ["test_*"]
+addopts = [
+ "-v",
+ "--tb=short",
+ "--strict-markers",
+]
+markers = [
+ "integration: integration tests (run full pipelines)",
+ "unit: unit tests (fast, isolated tests)",
+ "slow: slow tests (mark tests that take >10s)",
+]
\ No newline at end of file
diff --git a/runs/archive/delaunay_test.py b/runs/archive/delaunay_test.py
index eb05807..280f3eb 100644
--- a/runs/archive/delaunay_test.py
+++ b/runs/archive/delaunay_test.py
@@ -271,9 +271,9 @@
##############################################
fq_cpy = np.copy(freqs)
- fq_cpy[
- np.isnan(fq_cpy)
- ] = 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ fq_cpy[np.isnan(fq_cpy)] = (
+ 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ )
if params.debug:
total_power = fq_cpy.sum()
diff --git a/runs/archive/iterative_solver_test.py b/runs/archive/iterative_solver_test.py
index 8e78ec4..370f15f 100644
--- a/runs/archive/iterative_solver_test.py
+++ b/runs/archive/iterative_solver_test.py
@@ -13,7 +13,6 @@
from wrappers import interface
from vis import plotter, cart_plot
-
# %%
# from inputs.lam_run import params
# from inputs.selected_run import params
@@ -231,9 +230,9 @@
##############################################
fq_cpy = np.copy(freqs)
- fq_cpy[
- np.isnan(fq_cpy)
- ] = 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ fq_cpy[np.isnan(fq_cpy)] = (
+ 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ )
if params.debug:
total_power = fq_cpy.sum()
diff --git a/runs/chunk_consolidator.py b/runs/chunk_consolidator.py
new file mode 100644
index 0000000..354124c
--- /dev/null
+++ b/runs/chunk_consolidator.py
@@ -0,0 +1,70 @@
+# %%
+import numpy as np
+from tqdm import tqdm
+
+from pycsa.src import io, var
+from pycsa.inputs.icon_global_run import params
+
+chunk_start = 0
+n_cells = 20480
+chunk_sz = 100
+
+dat_path = params.path_output + "global_dataset/chunks/"
+out_path = params.path_output + "global_dataset/"
+out_fn = "icon_global_R2B4"
+
+global_dat = np.zeros((n_cells), dtype="object")
+
+cnt = 0
+for chunk in tqdm(range(chunk_start, n_cells, chunk_sz)):
+
+ sfx = "_" + str(chunk + chunk_sz)
+ fn = params.fn_output + sfx + ".nc"
+
+ writer = io.nc_writer(params, sfx)
+
+ if chunk + chunk_sz > n_cells:
+ chunk_end = n_cells
+ else:
+ chunk_end = chunk + chunk_sz
+
+ for ii in range(chunk, chunk_end):
+ struct = var.obj()
+ res = writer.read_dat(dat_path, fn, ii, struct)
+ global_dat[cnt] = struct
+ # print(cnt)
+ del struct
+
+ cnt += 1
+
+# print(cnt, chunk_end)
+print("\n==========")
+print("Collection done; writing output...")
+print("==========\n")
+assert (cnt) == chunk_end
+
+# %%
+from IPython import get_ipython
+
+ipython = get_ipython()
+
+if ipython is not None:
+ ipython.run_line_magic("load_ext", "autoreload")
+
+
+def autoreload():
+ if ipython is not None:
+ ipython.run_line_magic("autoreload", "2")
+
+
+# %%
+from pycsa.src import io
+
+autoreload()
+params.path_output = out_path
+global_writer = io.nc_writer(params, "")
+
+# for cnt, item in tqdm(enumerate(global_dat)):
+global_writer.duplicate_all(global_dat)
+
+# %%
diff --git a/runs/delaunay_runs.py b/runs/delaunay_runs.py
index 094557a..5928b9c 100644
--- a/runs/delaunay_runs.py
+++ b/runs/delaunay_runs.py
@@ -1,16 +1,11 @@
# %%
-import sys
-import os
-
-# set system path to find local modules
-sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
import numpy as np
-
-from src import io, var, utils, physics, delaunay
-from wrappers import interface, diagnostics
-from vis import plotter, cart_plot
import time
+from pycsa.core import io, var, utils, physics, delaunay
+from pycsa.wrappers import interface, diagnostics
+from pycsa.plotting import plotter, cart_plot
+
from IPython import get_ipython
ipython = get_ipython()
@@ -24,12 +19,11 @@ def autoreload():
ipython.run_line_magic("autoreload", "2")
-autoreload()
-
# %%
# from inputs.lam_run import params
from inputs.selected_run import params
+autoreload()
# from params.debug_run import params
from copy import deepcopy
@@ -44,11 +38,11 @@ def autoreload():
# read grid
reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
-reader.read_dat(params.fn_grid, grid)
+reader.read_dat(params.path_compact_grid, grid)
grid.apply_f(utils.rad2deg)
# writer object
-writer = io.writer(params.output_fn, params.rect_set, debug=params.debug_writer)
+writer = io.writer(params.fn_output, params.rect_set, debug=params.debug_writer)
# we only keep the topography that is inside this lat-lon extent.
lat_verts = np.array(params.lat_extent)
@@ -56,10 +50,11 @@ def autoreload():
# read topography
if not params.enable_merit:
- reader.read_dat(params.fn_topo, topo)
+ reader.read_dat(params.path_compact_topo, topo)
reader.read_topo(topo, topo, lon_verts, lat_verts)
else:
- reader.read_merit_topo(topo, params)
+ # reader.read_merit_topo(topo, params)
+ reader.read_etopo_topo(topo, params)
topo.topo[np.where(topo.topo < -500.0)] = -500.0
topo.gen_mgrids()
@@ -92,7 +87,7 @@ def autoreload():
fs=(12, 7),
highlight_indices=params.rect_set,
output_fig=True,
- fn="../manuscript/delaunay.pdf",
+ fn="./outputs/delaunay.pdf",
int=1,
raster=True,
)
@@ -420,7 +415,7 @@ def autoreload():
ylim=[-15, 15],
title="| FFT LRE | - | LSFF LRE |",
output_fig=True,
- fn="../manuscript/dfft_vs_lsff.pdf",
+ fn="./outputs/dfft_vs_lsff.pdf",
fontsize=12,
)
@@ -435,7 +430,7 @@ def autoreload():
ylim=[-100, 100],
output_fig=True,
title="percentage LRE",
- fn="../manuscript/lre_bar_ir.pdf",
+ fn="./outputs/lre_bar_ir.pdf",
fontsize=12,
comparison=np.array(rel_errs_orig) * 100,
)
@@ -462,7 +457,7 @@ def autoreload():
ylim=[-100, 100],
output_fig=True,
title="percentage LRE",
- fn="../manuscript/lre_bar_%s.pdf" % params.run_case,
+ fn="./outputs/lre_bar_%s.pdf" % params.run_case,
fontsize=12,
)
plotter.error_bar_plot(
@@ -475,7 +470,7 @@ def autoreload():
ylim=[-100, 100],
output_fig=True,
title="percentage MRE",
- fn="../manuscript/mre_bar_%s.pdf" % params.run_case,
+ fn="./outputs/mre_bar_%s.pdf" % params.run_case,
fontsize=12,
)
@@ -499,7 +494,7 @@ def autoreload():
fs=(12, 8),
highlight_indices=params.rect_set,
output_fig=True,
- fn="../manuscript/error_delaunay_%s.pdf" % params.run_case,
+ fn="./outputs/error_delaunay_%s.pdf" % params.run_case,
iint=1,
errors=errors,
alpha_max=0.6,
diff --git a/runs/icon_etopo_global.py b/runs/icon_etopo_global.py
new file mode 100644
index 0000000..04bdef4
--- /dev/null
+++ b/runs/icon_etopo_global.py
@@ -0,0 +1,1021 @@
+#!/usr/bin/env python3
+"""
+ICON ETOPO Global Processing Script
+
+IMPORTANT: Thread control environment variables must be set BEFORE numpy/numba import
+to prevent thread over-subscription with Dask workers.
+"""
+
+import os
+
+# ============================================================================
+# CRITICAL: Set thread limits BEFORE importing numpy/numba/scipy
+# This prevents thread over-subscription when using Dask with threads_per_worker > 1
+# ============================================================================
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+os.environ["NUMBA_NUM_THREADS"] = (
+ "1" # Critical: prevents Numba parallel=True conflicts
+)
+os.environ["VECLIB_MAXIMUM_THREADS"] = "1"
+
+import numpy as np
+import matplotlib
+
+matplotlib.use("Agg") # Use non-GUI backend for parallel processing
+import matplotlib.pyplot as plt
+from matplotlib.colors import TwoSlopeNorm
+import matplotlib.colors as mcolors
+from pathlib import Path
+import gc
+import logging
+from datetime import datetime
+
+from pycsa.core import io, var, utils, tile_cache
+from pycsa.wrappers import interface, diagnostics
+from pycsa.plotting import plotter
+
+# Initialize logger (will be configured in main)
+logger = logging.getLogger(__name__)
+
+
+def setup_logger(log_dir="logs"):
+ """
+ Set up logging configuration for ETOPO global run.
+
+ Parameters
+ ----------
+ log_dir : str
+ Directory for log files (default: "logs")
+
+ Returns
+ -------
+ Path
+ Path to the log file
+ """
+ # Create log directory
+ log_path = Path(log_dir)
+ log_path.mkdir(parents=True, exist_ok=True)
+
+ # Create timestamped log filename
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ log_file = log_path / f"icon_etopo_global_{timestamp}.log"
+
+ # Configure logger
+ logger = logging.getLogger(__name__)
+ logger.setLevel(logging.INFO)
+
+ # Remove any existing handlers
+ logger.handlers.clear()
+
+ # File handler - logs everything
+ file_handler = logging.FileHandler(log_file, mode="w")
+ file_handler.setLevel(logging.INFO)
+ file_formatter = logging.Formatter(
+ "%(asctime)s - %(levelname)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S"
+ )
+ file_handler.setFormatter(file_formatter)
+ logger.addHandler(file_handler)
+
+ # Also silence matplotlib and other libraries from console
+ logging.getLogger("matplotlib").setLevel(logging.WARNING)
+ logging.getLogger("distributed").setLevel(logging.WARNING)
+
+ return log_file
+
+
+def get_topo_colormap():
+ """
+ Create a topography colormap with blue for ocean (< 0m) and terrain colors for land (> 0m).
+ Transition occurs exactly at sea level (0m) with smooth blending.
+
+ For TwoSlopeNorm to work correctly, we need equal colors on each side:
+ 128 colors for ocean (< 0m) + 128 colors for land (> 0m) = 256 total
+ """
+ # Ocean colors (blue shades from deep to shallow)
+ ocean_colors = plt.cm.Blues_r(np.linspace(0.4, 0.95, 120))
+
+ # Smooth transition zone around sea level (8 colors on each side)
+ # Get the last ocean color and first land color
+ last_ocean = plt.cm.Blues_r(0.95)
+ first_land = plt.cm.terrain(0.25)
+
+ # Create smooth blend from ocean to land
+ transition_colors = np.zeros((16, 4))
+ for i in range(4): # RGBA channels
+ transition_colors[:, i] = np.linspace(last_ocean[i], first_land[i], 16)
+
+ # Land colors (terrain-like: green to brown to white)
+ land_colors = plt.cm.terrain(np.linspace(0.28, 1.0, 120))
+
+ # Combine: 120 ocean + 16 transition + 120 land = 256 total
+ # Transition centered at index 128 (sea level)
+ colors = np.vstack((ocean_colors, transition_colors, land_colors))
+ return mcolors.LinearSegmentedColormap.from_list("topo", colors)
+
+
+def plot_cell_diagnostics(c_idx, cell_sa, ampls_sa, dat_2D_sa, output_dir, params):
+ """
+ Create 3-panel diagnostic plot for a single cell.
+
+ Panel 1: Loaded topography (original ETOPO data within cell)
+ Panel 2: Reconstructed topography after second approximation
+ Panel 3: Computed spectrum
+
+ Parameters
+ ----------
+ c_idx : int
+ Cell index
+ cell_sa : topo_cell
+ Cell object after second approximation (contains original topo in cell.topo)
+ ampls_sa : ndarray
+ Amplitude spectrum from second approximation
+ dat_2D_sa : ndarray
+ Reconstructed topography from second approximation
+ output_dir : Path
+ Output directory for saving plots
+ params : params object
+ Parameters object
+ """
+ # Create figure with 3 panels
+ fig, axs = plt.subplots(1, 3, figsize=(18, 6))
+
+ # Get elevation extent for consistent color scaling
+ vmin = -200.0 # Always fix ocean floor at -500m (blue portion)
+ vmax = np.nanmax(cell_sa.topo)
+
+ # Ensure vmax is positive (land)
+ if vmax <= 0:
+ vmax = 100.0 # Force some land color even if all ocean
+
+ # Create custom colormap with blue for ocean, terrain colors for land
+ topo_cmap = get_topo_colormap()
+
+ # Create normalization centered at sea level (0m)
+ # This makes the colormap transition exactly at 0m
+ norm = TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax)
+
+ # Panel 1: Original topography within cell
+ topo_original = cell_sa.topo.copy()
+ topo_original[~cell_sa.mask] = np.nan
+
+ im1 = axs[0].imshow(
+ topo_original, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0].set_title(
+ f"Cell {c_idx}: Loaded Topography\nRange: [{vmin:.0f}, {vmax:.0f}] m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0].set_xlabel("Longitude index")
+ axs[0].set_ylabel("Latitude index")
+ cbar1 = plt.colorbar(im1, ax=axs[0], fraction=0.046, pad=0.04)
+ cbar1.set_label("Elevation [m]", rotation=270, labelpad=15)
+
+ # Panel 2: Reconstructed topography (masked)
+ dat_2D_masked = dat_2D_sa.copy()
+ dat_2D_masked[~cell_sa.mask] = np.nan
+
+ # Compute reconstruction error
+ diff = cell_sa.topo - dat_2D_sa
+ rmse = np.sqrt(np.mean(diff[cell_sa.mask] ** 2))
+ rel_rmse = rmse / (vmax - vmin) * 100
+
+ im2 = axs[1].imshow(
+ dat_2D_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[1].set_title(
+ f"Reconstructed (2nd Approx)\nRMSE: {rmse:.1f} m ({rel_rmse:.1f}%)",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[1].set_xlabel("Longitude index")
+ axs[1].set_ylabel("Latitude index")
+ cbar2 = plt.colorbar(im2, ax=axs[1], fraction=0.046, pad=0.04)
+ cbar2.set_label("Elevation [m]", rotation=270, labelpad=15)
+
+ # Panel 3: Amplitude spectrum in (k,l) wavenumber space
+ fig_obj = plotter.fig_obj(fig, params.nhi, params.nhj, cbar=True, set_label=True)
+ axs[2] = fig_obj.freq_panel(
+ axs[2], ampls_sa, title="Amplitude Spectrum", v_extent=None
+ )
+
+ plt.tight_layout()
+
+ # Save figure
+ output_path = output_dir / f"cell_{c_idx:05d}.png"
+ plt.savefig(output_path, dpi=150, bbox_inches="tight")
+ plt.close(fig)
+
+ # Explicit memory cleanup - delete ALL objects to prevent memory leaks
+ del fig, axs, fig_obj, im1, im2, topo_original, dat_2D_masked
+ del cbar1, cbar2, norm, topo_cmap, diff
+ gc.collect() # Force garbage collection after plotting
+
+ logger.info(f" Plot saved: {output_path}")
+
+
+def do_cell(
+ c_idx,
+ grid,
+ params,
+ reader,
+ writer,
+ chunk_output_dir,
+ clat_rad,
+ clon_rad,
+):
+ """
+ Process a single ICON grid cell with ETOPO topography.
+
+ Parameters
+ ----------
+ c_idx : int
+ Cell index in the grid
+ grid : grid object
+ ICON grid (in degrees)
+ params : params object
+ Parameters
+ reader : ncdata object
+ Data reader
+ writer : nc_writer object
+ NetCDF writer
+ chunk_output_dir : Path
+ Output directory for this chunk
+ clat_rad : ndarray
+ Cell center latitudes in radians
+ clon_rad : ndarray
+ Cell center longitudes in radians
+
+ Returns
+ -------
+ grp_struct
+ Result structure for NetCDF output
+ """
+
+ import sys
+ import traceback
+
+ try:
+ logger.info(f"[START] Processing cell {c_idx}")
+
+ topo = var.topo_cell()
+
+ lat_verts = grid.clat_vertices[c_idx]
+ lon_verts = grid.clon_vertices[c_idx]
+
+ # Determine lat/lon extents with appropriate expansion for data loading
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ # Load topography for this cell from the worker-local tile cache.
+ # The cache is initialised once per memory batch via init_worker_cache
+ # (see the per-batch loop below); handles stay open across cells in
+ # the same worker so we don't re-open the same ETOPO tile per cell.
+ cache = tile_cache.get_worker_cache()
+ topo.lat, topo.lon, topo.topo = cache.get_etopo_data(
+ lat_extent, lon_extent, etopo_cg=params.etopo_cg
+ )
+ split_EW = tile_cache.compute_split_EW(lon_extent)
+
+ # Clip deep bathymetry to -500m (same as test_etopo_pole_cells.py)
+ # This prevents issues with extreme ocean depths creating artifacts
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+ topo.gen_mgrids()
+
+ # Handle dateline crossing BEFORE processing vertices for CSA
+ # This must be done before handle_latlon_expansion() to ensure consistent coordinates
+ if split_EW:
+ lon_verts = lon_verts.copy() # Don't modify the grid object
+ lon_verts[lon_verts < 0.0] += 360.0
+
+ # Process vertices for CSA (after dateline correction!)
+ lat_verts, lon_verts = utils.handle_latlon_expansion(
+ lat_verts, lon_verts, lat_expand=0.0, lon_expand=0.0
+ )
+
+ # Set up cell center and vertices
+ clon = np.array([grid.clon[c_idx]])
+ clat = np.array([grid.clat[c_idx]])
+ clon_vertices = np.array([lon_verts])
+ clat_vertices = np.array([lat_verts])
+
+ ncells = 1
+ nv = clon_vertices[0].size
+
+ triangles = np.zeros((ncells, nv, 2))
+
+ for i in range(0, ncells, 1):
+ triangles[i, :, 0] = np.array(clon_vertices[i, :])
+ triangles[i, :, 1] = np.array(clat_vertices[i, :])
+
+ # Initialize cell objects for CSA algorithm
+ tri_idx = 0
+ cell = var.topo_cell()
+ tri = var.obj()
+
+ nhi = params.nhi
+ nhj = params.nhj
+
+ fa = interface.first_appx(nhi, nhj, params, topo)
+ sa = interface.second_appx(nhi, nhj, params, topo, tri)
+
+ tri.tri_lon_verts = triangles[:, :, 0]
+ tri.tri_lat_verts = triangles[:, :, 1]
+
+ simplex_lat = tri.tri_lat_verts[tri_idx]
+ simplex_lon = tri.tri_lon_verts[tri_idx]
+
+ if not utils.is_land(cell, simplex_lat, simplex_lon, topo):
+ logger.info(f"[OCEAN] Cell {c_idx} is ocean, skipping")
+ return writer.grp_struct(
+ c_idx, clat_rad[c_idx], clon_rad[c_idx], 0, None, grid.cell_area[c_idx]
+ )
+ else:
+ is_land = 1
+ logger.info(f"[LAND] Cell {c_idx} is land, processing...")
+
+ # First approximation
+ cell_fa, ampls_fa, uw_fa, dat_2D_fa = fa.do(
+ simplex_lat, simplex_lon, use_center=True
+ )
+
+ # Second approximation
+ if USE_MODE_SELECTION:
+ # COMPRESSED MODE: Use sa.do() to select top n_modes wavenumbers
+ # This is the original workflow with spectral compression
+ if params.recompute_rhs:
+ sols, _ = sa.do(tri_idx, ampls_fa, use_center=True)
+ else:
+ sols = sa.do(tri_idx, ampls_fa, use_center=True)
+ cell_sa, ampls_sa, uw_sa, dat_2D_sa = sols
+
+ # Exclude ocean from spectral analysis (same as FULL SPECTRUM mode)
+ ocean_mask = cell_sa.topo < -200.0
+ cell_sa.mask = cell_sa.mask & ~ocean_mask
+ cell_sa.get_masked(mask=cell_sa.mask)
+ else:
+ # FULL SPECTRUM MODE: Use ALL wavenumbers (no mode selection)
+ # This gives ~20% better RMSE but no compression
+ cell_sa = var.topo_cell()
+
+ # Step 1: Load topo with rectangular mask
+ utils.get_lat_lon_segments(
+ simplex_lat,
+ simplex_lon,
+ cell_sa,
+ topo,
+ rect=True,
+ filtered=True,
+ padding=0,
+ use_center=True,
+ )
+
+ # Step 2: Apply triangular mask
+ utils.get_lat_lon_segments(
+ simplex_lat,
+ simplex_lon,
+ cell_sa,
+ topo,
+ rect=False,
+ filtered=False,
+ padding=0,
+ use_center=True,
+ )
+
+ # Run SA with ALL wavenumbers
+ sa_pmf = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_sa, uw_sa, dat_2D_sa = sa_pmf.sappx(
+ cell_sa,
+ lmbda=params.lmbda_sa,
+ iter_solve=params.sa_iter_solve,
+ updt_analysis=True, # Populate cell_sa.analysis for NetCDF output
+ )
+
+ # Exclude ocean from spectral analysis for orographic gravity waves
+ # The atmosphere flows over ocean SURFACE (0m), not the seafloor
+ # Threshold: -200m distinguishes deep ocean from below-sea-level land
+ # - Most below-sea-level land features: -200m to 0m (Death Valley -86m, etc.)
+ # - Coastal ocean bathymetry: typically < -200m
+ ocean_mask = cell_sa.topo < -200.0
+ cell_sa.mask = cell_sa.mask & ~ocean_mask
+ cell_sa.get_masked(mask=cell_sa.mask)
+
+ # Store analysis results
+ result = writer.grp_struct(
+ c_idx,
+ clat_rad[c_idx],
+ clon_rad[c_idx],
+ is_land,
+ cell_sa.analysis,
+ grid.cell_area[c_idx],
+ )
+
+ # Generate 3-panel plot
+ if params.plot_output:
+ plot_cell_diagnostics(
+ c_idx, cell_sa, ampls_sa, dat_2D_sa, chunk_output_dir, params
+ )
+
+ logger.info(f"[DONE] Cell {c_idx} analysis complete")
+
+ # Explicit memory cleanup to help Dask workers
+ del (
+ topo,
+ cell_fa,
+ cell_sa,
+ ampls_fa,
+ ampls_sa,
+ uw_fa,
+ uw_sa,
+ dat_2D_fa,
+ dat_2D_sa,
+ )
+ del fa, sa, tri, cell, etopo_reader
+ gc.collect() # Force garbage collection
+
+ return result
+
+ except Exception as e:
+ # Catch ALL exceptions and log them before worker dies
+ error_msg = (
+ f"[FATAL ERROR] Cell {c_idx} crashed with {type(e).__name__}: {str(e)}"
+ )
+ logger.error(error_msg)
+ logger.error(traceback.format_exc())
+
+ # Print to stderr so it appears in worker logs
+ print(error_msg, file=sys.stderr)
+ print(traceback.format_exc(), file=sys.stderr)
+
+ # Re-raise to let Dask handle it
+ raise
+
+
+def estimate_cell_memory_gb(lat_deg):
+ """
+ Estimate memory requirements (in GB) for processing a cell based on its latitude.
+
+ At polar latitudes, cells cover a larger longitudinal range in degree-space,
+ requiring more topographic data points to be loaded with coarse-graining.
+
+ Parameters
+ ----------
+ lat_deg : float
+ Cell center latitude in degrees (-90 to 90)
+
+ Returns
+ -------
+ float
+ Estimated memory requirement in GB
+
+ Notes
+ -----
+ - Equatorial cells (~0°): ~10 GB sufficient
+ - Mid-latitude cells (~45°): ~10 GB
+ - High-latitude cells (~70°): ~25 GB
+ - Polar cells (~80-89°): ~60 GB required
+
+ Memory scales approximately with 1/cos(lat) due to meridian convergence,
+ but caps at ~60 GB for cells very close to the poles.
+ """
+ abs_lat = np.abs(lat_deg)
+
+ # Base memory requirement at equator
+ base_memory_gb = 10.0
+
+ # Scale factor based on latitude (empirical fit)
+ if abs_lat < 60.0:
+ # Below 60°, memory is fairly constant
+ scale_factor = 1.0
+ elif abs_lat < 85.0:
+ # Between 60° and 85°, use power law scaling
+ # At 70°: (1/0.342)^0.7 ≈ 2.5, giving 25 GB
+ # At 80°: (1/0.174)^0.7 ≈ 4.3, giving 43 GB
+ lat_rad = np.deg2rad(abs_lat)
+ cos_lat = np.cos(lat_rad)
+ scale_factor = (1.0 / cos_lat) ** 0.7
+ else:
+ # Above 85°, cap at 6x base (60 GB) to avoid unrealistic estimates
+ # Very close to poles, the ICON grid cells are smaller and don't
+ # actually require infinite memory despite cos(lat)→0
+ scale_factor = 6.0
+
+ return base_memory_gb * scale_factor
+
+
+def group_cells_by_memory(clat_rad, max_memory_per_batch_gb=240.0):
+ """
+ Group cells into batches with similar memory requirements.
+
+ Parameters
+ ----------
+ clat_rad : ndarray
+ Cell center latitudes in radians
+ max_memory_per_batch_gb : float
+ Maximum total memory available for a batch (default: 240 GB for 6 workers × 40 GB)
+
+ Returns
+ -------
+ list of dict
+ List of batch configurations, each containing:
+ - 'cell_indices': list of cell indices in this batch
+ - 'memory_per_cell_gb': average memory per cell in GB
+ - 'n_workers': recommended number of workers
+ - 'memory_per_worker_gb': recommended memory per worker
+ """
+ n_cells = len(clat_rad)
+ clat_deg = np.rad2deg(clat_rad)
+
+ # Estimate memory for each cell
+ cell_memory_gb = np.array([estimate_cell_memory_gb(lat) for lat in clat_deg])
+
+ # Sort cells by memory requirement (process high-memory cells first)
+ sorted_indices = np.argsort(cell_memory_gb)[::-1]
+
+ batches = []
+ current_batch_indices = []
+ current_batch_memory = []
+
+ for idx in sorted_indices:
+ mem = cell_memory_gb[idx]
+
+ # Check if adding this cell would exceed batch memory limit
+ if current_batch_indices:
+ avg_mem = np.mean(current_batch_memory + [mem])
+ # Ensure we can fit at least 1 worker with this memory
+ if avg_mem * len(current_batch_indices) > max_memory_per_batch_gb:
+ # Finalize current batch
+ avg_mem_current = np.mean(current_batch_memory)
+ # Use 30% safety margin for diskless NetCDF loading
+ safety_factor = 1.0
+ n_workers = max(
+ 1, int(max_memory_per_batch_gb / (avg_mem_current * safety_factor))
+ )
+ mem_per_worker = avg_mem_current * safety_factor
+
+ batches.append(
+ {
+ "cell_indices": sorted(
+ current_batch_indices
+ ), # Sort by original index order
+ "memory_per_cell_gb": avg_mem_current,
+ "n_workers": n_workers,
+ "memory_per_worker_gb": mem_per_worker,
+ }
+ )
+
+ # Start new batch
+ current_batch_indices = [idx]
+ current_batch_memory = [mem]
+ else:
+ current_batch_indices.append(idx)
+ current_batch_memory.append(mem)
+ else:
+ current_batch_indices.append(idx)
+ current_batch_memory.append(mem)
+
+ # Finalize last batch
+ if current_batch_indices:
+ avg_mem = np.mean(current_batch_memory)
+ # Use 30% safety margin for diskless NetCDF loading
+ safety_factor = 1.0
+ n_workers = max(1, int(max_memory_per_batch_gb / (avg_mem * safety_factor)))
+ mem_per_worker = avg_mem * safety_factor
+
+ batches.append(
+ {
+ "cell_indices": sorted(current_batch_indices),
+ "memory_per_cell_gb": avg_mem,
+ "n_workers": n_workers,
+ "memory_per_worker_gb": mem_per_worker,
+ }
+ )
+
+ return batches
+
+
+def parallel_wrapper(
+ grid, params, reader, writer, chunk_output_dir, clat_rad, clon_rad
+):
+ return lambda ii: do_cell(
+ ii, grid, params, reader, writer, chunk_output_dir, clat_rad, clon_rad
+ )
+
+
+from inputs.icon_global_run import params
+from dask.distributed import Client, progress
+import dask
+from tqdm import tqdm
+
+if __name__ == "__main__":
+ # ========================================================================
+ # CONFIGURATION SELECTOR
+ # ========================================================================
+ # Choose one: 'generic_laptop', 'dkrz_hpc', 'laptop_performance'
+ SYSTEM_CONFIG = "laptop_performance" # ← Edit this line to switch configs
+ # ========================================================================
+
+ # ========================================================================
+ # QUICK START GUIDE - Processing Specific Cell Ranges
+ # ========================================================================
+ # To process specific cell ranges (e.g., to regenerate corrupted chunks):
+ #
+ # 1. Scroll down to "CELL RANGE CONFIGURATION" section (around line 690)
+ # 2. Set cell_start and cell_end:
+ #
+ # Examples:
+ # cell_start = 0, cell_end = 100 → Process cells 0-99 only
+ # cell_start = 2900, cell_end = 3000 → Process cells 2900-2999 only
+ # cell_start = 0, cell_end = None → Process all cells from 0 to end
+ # cell_start = 3000, cell_end = None → Process from 3000 to end
+ #
+ # 3. Run the script - it will create appropriately named NetCDF files
+ #
+ # Note: Files are created in chunks of netcdf_chunk_size (default: 100)
+ # Example: cells 0-99 → icon_etopo_global_cells_00000-00099.nc
+ # ========================================================================
+
+ CONFIGS = {
+ "generic_laptop": {
+ "total_cores": 12, # Conservative: use 12 of 16 threads
+ "total_memory_gb": 12.0,
+ "netcdf_chunk_size": 100,
+ "threads_per_worker": 1, # Set to None for auto-compute
+ "memory_per_cpu_mb": None, # Will calculate dynamically
+ "description": "Generic laptop (16 threads, 16GB RAM)",
+ },
+ "dkrz_hpc": {
+ "total_cores": 250,
+ "total_memory_gb": 240.0,
+ "netcdf_chunk_size": 100,
+ "threads_per_worker": None, # Auto-compute based on worker memory
+ "memory_per_cpu_mb": None, # SLURM quota on interactive partition
+ "description": "DKRZ HPC interactive partition (standard memory node)",
+ },
+ "laptop_performance": {
+ "total_cores": 20, # Use 20 of 24 threads (leave 4 for background)
+ "total_memory_gb": 80.0,
+ "netcdf_chunk_size": 100,
+ "threads_per_worker": None, # Auto-compute based on worker memory
+ "memory_per_cpu_mb": None, # Will calculate dynamically
+ "description": "AMD Ryzen AI 9 HX 370 (24 threads, 94GB RAM)",
+ },
+ }
+
+ # Validate configuration selection
+ if SYSTEM_CONFIG not in CONFIGS:
+ raise ValueError(
+ f"Invalid SYSTEM_CONFIG '{SYSTEM_CONFIG}'. Choose from: {list(CONFIGS.keys())}"
+ )
+
+ config = CONFIGS[SYSTEM_CONFIG]
+
+ # Set up logging first
+ log_file = setup_logger(log_dir="logs")
+ print(f"Logging to: {log_file}")
+ print("=" * 80)
+ print(f"SYSTEM CONFIG: {SYSTEM_CONFIG}")
+ print(f" {config['description']}")
+ print(f" Cores: {config['total_cores']}, Memory: {config['total_memory_gb']} GB")
+ print("=" * 80)
+
+ # Override/add ETOPO-specific parameters
+ params.fn_output = "icon_etopo_global"
+ params.etopo_cg = (
+ 4 # Coarse-graining factor (1.8km at equator, ~0.9-1.8km at Drake Passage)
+ )
+
+ # Use traditional first approximation
+ params.dfft_first_guess = False
+ params.recompute_rhs = False
+
+ # Disable plotting by default (set to True if you want diagnostic plots for each cell)
+ params.plot_output = True
+
+ # ========================================================================
+ # SPECTRAL COMPRESSION TOGGLE
+ # ========================================================================
+ # Toggle between full spectrum vs compressed spectrum in second approximation:
+ #
+ # False (COMPRESSED - default): Use top n_modes=100 wavenumbers
+ # - Pros: 20x smaller NetCDF files, fast I/O, spectral compression feature
+ # - Cons: ~20% higher RMSE (e.g., 150.9m vs 121.0m for cell 3091)
+ #
+ # True (FULL SPECTRUM): Use ALL nhi*nhj=2048 wavenumbers
+ # - Pros: Best reconstruction quality (~20% lower RMSE)
+ # - Cons: 20x larger NetCDF files, no compression benefit
+ #
+ USE_FULL_SPECTRUM = False # Set to True to disable spectral compression
+
+ if USE_FULL_SPECTRUM:
+ logger.info(
+ "*** FULL SPECTRUM MODE: Using ALL wavenumbers (no compression) ***"
+ )
+ params.n_modes = params.nhi * params.nhj # 2048 modes
+ USE_MODE_SELECTION = False # Use all modes in SA
+ else:
+ logger.info("*** COMPRESSED SPECTRUM MODE: Using top 100 wavenumbers ***")
+ # params.n_modes already set to 100 in icon_global_run
+ USE_MODE_SELECTION = True # Select top n_modes in SA
+ # ========================================================================
+
+ if params.self_test():
+ params.print()
+
+ grid = var.grid()
+
+ # Read ICON grid
+ reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+ reader.read_dat(params.path_icon_grid, grid)
+
+ clat_rad = np.copy(grid.clat)
+ clon_rad = np.copy(grid.clon)
+
+ grid.apply_f(utils.rad2deg)
+
+ n_cells = grid.clat.size
+
+ # Create base output directory
+ base_output_dir = Path("outputs") / params.fn_output
+ base_output_dir.mkdir(parents=True, exist_ok=True)
+ logger.info(f"Base output directory: {base_output_dir}")
+
+ # ========================================================================
+ # DYNAMIC MEMORY ALLOCATION SETUP
+ # ========================================================================
+ # Instead of fixed worker configuration, we'll dynamically adjust based on
+ # the memory requirements of cells being processed (latitude-dependent)
+
+ import multiprocessing
+ import os
+
+ # Use configuration values
+ total_cores = config["total_cores"]
+ total_memory_gb = config["total_memory_gb"]
+ netcdf_chunk_size = config["netcdf_chunk_size"]
+
+ logger.info("=" * 80)
+ logger.info(f"RESOURCE CONFIGURATION: {SYSTEM_CONFIG}")
+ logger.info(f" Description: {config['description']}")
+ logger.info(f" Available cores: {total_cores}")
+ logger.info(f" Available memory: {total_memory_gb} GB")
+ logger.info(f" NetCDF chunk size: {netcdf_chunk_size} cells")
+
+ # Threading configuration display
+ if config["threads_per_worker"] is not None:
+ logger.info(
+ f" Threading mode: MANUAL (threads_per_worker = {config['threads_per_worker']})"
+ )
+ else:
+ logger.info(f" Threading mode: AUTO (will compute based on worker count)")
+
+ if config["memory_per_cpu_mb"] is not None:
+ logger.info(f" SLURM quota: {config['memory_per_cpu_mb']} MB per CPU")
+ logger.info("=" * 80)
+
+ # Group cells by memory requirements for dynamic worker allocation
+ logger.info(f"\nAnalyzing cells by latitude for dynamic memory allocation...")
+ memory_batches = group_cells_by_memory(
+ clat_rad, max_memory_per_batch_gb=total_memory_gb
+ )
+
+ logger.info(f"Created {len(memory_batches)} memory-based batches:")
+ for i, batch in enumerate(memory_batches):
+ logger.info(
+ f" Batch {i}: {len(batch['cell_indices'])} cells, "
+ f"{batch['memory_per_cell_gb']:.1f} GB/cell, "
+ f"{batch['n_workers']} workers × {batch['memory_per_worker_gb']:.1f} GB"
+ )
+
+ # We'll create Dask client dynamically for each memory batch
+ # Start with None (will be created per batch)
+ client = None
+ current_batch_idx = None
+
+ logger.info(f"Total cells in grid: {n_cells}")
+
+ # ========================================================================
+ # CELL RANGE CONFIGURATION
+ # ========================================================================
+ # Set cell_start and cell_end to process specific ranges
+ # Examples:
+ # cell_start = 0, cell_end = None → Process all cells (0 to n_cells-1)
+ # cell_start = 2900, cell_end = 3000 → Process cells 2900-2999 only
+ # cell_start = 0, cell_end = 100 → Process cells 0-99 only
+ cell_start = 0 # First cell to process (inclusive)
+ cell_end = None # Last cell to process (exclusive), None means process to end
+ # ========================================================================
+
+ # Validate and set cell_end
+ if cell_end is None:
+ cell_end = n_cells
+ else:
+ cell_end = min(cell_end, n_cells) # Don't exceed total cells
+
+ if cell_start >= cell_end:
+ raise ValueError(
+ f"Invalid cell range: cell_start ({cell_start}) >= cell_end ({cell_end})"
+ )
+
+ # Progress tracking
+ cells_to_process = cell_end - cell_start
+ total_netcdf_chunks = (
+ cells_to_process + netcdf_chunk_size - 1
+ ) // netcdf_chunk_size
+ logger.info(
+ f"\nProcessing cell range: {cell_start} to {cell_end-1} ({cells_to_process} cells)"
+ )
+ logger.info(
+ f" NetCDF chunks: {total_netcdf_chunks} files ({netcdf_chunk_size} cells each)\n"
+ )
+
+ # Statistics
+ total_land_cells = 0
+ total_ocean_cells = 0
+
+ # Configure task retries and logging (do this once)
+ import dask
+ import logging
+
+ dask.config.set({"distributed.scheduler.allowed-failures": 0})
+ logging.getLogger("distributed.worker.memory").setLevel(logging.ERROR)
+
+ # Create a mapping from cell_idx to memory batch index for quick lookup
+ cell_to_batch = {}
+ for batch_idx, batch in enumerate(memory_batches):
+ for cell_idx in batch["cell_indices"]:
+ cell_to_batch[cell_idx] = batch_idx
+
+ # ========================================================================
+ # SEQUENTIAL PROCESSING BY MEMORY BATCH
+ # ========================================================================
+ # Process memory batches sequentially (equatorial → mid-lat → polar)
+ # This allows easy restart: if script crashes, you know all previous
+ # memory batches are complete and can skip to the current batch.
+ # ========================================================================
+
+ logger.info("\n" + "=" * 80)
+ logger.info("PROCESSING STRATEGY: Sequential by Memory Batch")
+ logger.info("=" * 80)
+ for batch_idx, batch_config in enumerate(memory_batches):
+ logger.info(f"\n{'='*80}")
+ logger.info(
+ f"MEMORY BATCH {batch_idx}/{len(memory_batches)-1}: {len(batch_config['cell_indices'])} cells"
+ )
+ logger.info(f" Memory per cell: {batch_config['memory_per_cell_gb']:.1f} GB")
+ logger.info(f" Workers: {batch_config['n_workers']}")
+ logger.info(f"{'='*80}\n")
+
+ # Get all cells in this memory batch
+ batch_cell_indices = set(batch_config["cell_indices"])
+
+ # Create Dask client for this memory batch
+ n_workers = batch_config["n_workers"]
+ # Single-worker batches (high-memory polar cells) get the full machine
+ # memory; multi-worker batches share by config.
+ if n_workers == 1:
+ memory_per_worker = f"{int(total_memory_gb)}GB"
+ logger.info(
+ f" Single-worker mode: allowing full memory access ({total_memory_gb} GB)"
+ )
+ else:
+ memory_per_worker = f"{int(batch_config['memory_per_worker_gb'])}GB"
+ threads_per_worker = 1 # HDF5 not thread-safe
+
+ logger.info(f"Starting Dask client for memory batch {batch_idx}:")
+ logger.info(f" Workers: {n_workers} × {memory_per_worker}")
+ logger.info(f" Threads per worker: {threads_per_worker}")
+
+ client = Client(
+ threads_per_worker=threads_per_worker,
+ n_workers=n_workers,
+ processes=True,
+ memory_limit=memory_per_worker,
+ silence_logs="ERROR",
+ )
+ logger.info(f" Dashboard: {client.dashboard_link}\n")
+
+ # Initialise the per-worker tile cache. Each worker is a separate
+ # process, so this populates a module-level _WORKER_CACHE inside that
+ # process; do_cell then reaches it via tile_cache.get_worker_cache().
+ # The cache opens ETOPO tile files lazily on first access and keeps
+ # the handles for the rest of the worker's lifetime.
+ init_results = client.run(
+ tile_cache.init_worker_cache, params.path_etopo, "ETOPO"
+ )
+ logger.info(
+ f" Initialised tile cache on {sum(bool(v) for v in init_results.values())} "
+ f"of {len(init_results)} workers"
+ )
+
+ # Inner loop: NetCDF file creation (one file per netcdf_chunk_size cells)
+ # Only process NetCDF chunks that contain cells from this memory batch
+ for netcdf_chunk_idx, netcdf_chunk_start in enumerate(
+ tqdm(
+ range(cell_start, n_cells, netcdf_chunk_size),
+ desc=f"NetCDF chunks (batch {batch_idx})",
+ total=total_netcdf_chunks,
+ )
+ ):
+ netcdf_chunk_end = min(netcdf_chunk_start + netcdf_chunk_size, n_cells)
+
+ # Filter: only process cells in this NetCDF chunk that belong to current memory batch
+ cell_indices_in_chunk = []
+ for c_idx in range(netcdf_chunk_start, netcdf_chunk_end):
+ if c_idx in batch_cell_indices:
+ cell_indices_in_chunk.append(c_idx)
+
+ # Skip this NetCDF chunk if no cells belong to current memory batch
+ if not cell_indices_in_chunk:
+ continue
+
+ logger.info(
+ f"\n Processing NetCDF chunk {netcdf_chunk_idx}: cells {netcdf_chunk_start}-{netcdf_chunk_end-1}"
+ )
+ logger.info(f" Cells in this batch: {len(cell_indices_in_chunk)}")
+
+ # Create subdirectory for this NetCDF chunk's plots
+ chunk_output_dir = (
+ base_output_dir
+ / f"cells_{netcdf_chunk_start:05d}-{netcdf_chunk_end-1:05d}"
+ )
+ chunk_output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Writer object for this NetCDF chunk
+ sfx = f"_cells_{netcdf_chunk_start:05d}-{netcdf_chunk_end-1:05d}"
+ writer = io.nc_writer(params, sfx)
+
+ pw_run = parallel_wrapper(
+ grid, params, reader, writer, chunk_output_dir, clat_rad, clon_rad
+ )
+
+ # Process cells in smaller batches to avoid overwhelming scheduler
+ processing_batch_size = min(n_workers * 2, len(cell_indices_in_chunk))
+
+ for i in range(0, len(cell_indices_in_chunk), processing_batch_size):
+ batch_cells = cell_indices_in_chunk[i : i + processing_batch_size]
+
+ # Submit batch to Dask
+ lazy_results = []
+ for c_idx in batch_cells:
+ lazy_result = dask.delayed(pw_run)(c_idx)
+ lazy_results.append(lazy_result)
+
+ # Compute batch
+ results = dask.compute(*lazy_results)
+
+ # Write results to NetCDF file
+ for item in results:
+ writer.duplicate(item.c_idx, item)
+ if item.is_land:
+ total_land_cells += 1
+ else:
+ total_ocean_cells += 1
+
+ # Cleanup after each NetCDF chunk
+ if hasattr(reader, "close_cached_files"):
+ reader.close_cached_files()
+
+ gc.collect()
+
+ logger.info(
+ f" NetCDF chunk complete: {len(cell_indices_in_chunk)} cells processed"
+ )
+ logger.info(
+ f" Running totals - Land: {total_land_cells}, Ocean: {total_ocean_cells}"
+ )
+
+ # Close Dask client after finishing this memory batch
+ client.close()
+ logger.info(f"\n{'='*80}")
+ logger.info(f"MEMORY BATCH {batch_idx} COMPLETE")
+ logger.info(f" Processed {len(batch_cell_indices)} cells")
+ logger.info(f"{'='*80}\n")
+
+ # Cleanup: close all cached NetCDF files
+ logger.info("\n" + "=" * 80)
+ logger.info("PROCESSING COMPLETE")
+ logger.info("=" * 80)
+ logger.info(f"Total cells processed: {total_land_cells + total_ocean_cells}")
+ logger.info(f" Land cells: {total_land_cells}")
+ logger.info(f" Ocean cells: {total_ocean_cells}")
+ logger.info(f"\nNetCDF files created: {total_netcdf_chunks}")
+ logger.info(f" Location: {params.path_output}datasets/")
+ logger.info(f" Pattern: icon_etopo_global_cells_XXXXX-XXXXX.nc")
+ logger.info(f"\nTo merge into single file, run:")
+ logger.info(f" python3 -m runs.merge_netcdf_chunks")
+ logger.info("=" * 80)
+
+ if hasattr(reader, "close_cached_files"):
+ reader.close_cached_files()
+ logger.info("\n✓ Closed cached topography files")
+
+ # Final console message
+ print("=" * 80)
+ print(f"PROCESSING COMPLETE - Check log file: {log_file}")
+ print("=" * 80)
diff --git a/runs/icon_merit_global.py b/runs/icon_merit_global.py
new file mode 100644
index 0000000..9a583f6
--- /dev/null
+++ b/runs/icon_merit_global.py
@@ -0,0 +1,253 @@
+import numpy as np
+
+from pycsa.core import io, var, utils
+from pycsa.wrappers import interface, diagnostics
+from pycsa.plotting import cart_plot
+
+
+def do_cell(
+ c_idx,
+ grid,
+ params,
+ reader,
+ writer,
+):
+
+ print(c_idx)
+
+ topo = var.topo_cell()
+
+ lat_verts = grid.clat_vertices[c_idx]
+ lon_verts = grid.clon_vertices[c_idx]
+
+ # Determine lat/lon extents with appropriate expansion for data loading
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+ lat_verts, lon_verts = utils.handle_latlon_expansion(
+ lat_verts, lon_verts, lat_expand=0.0, lon_expand=0.0
+ )
+
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ # Load topography data for this cell
+ reader = reader.read_merit_topo(None, params, is_parallel=True)
+ reader.get_topo(topo)
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+ topo.gen_mgrids()
+
+ # Set up cell center and vertices
+ clon = np.array([grid.clon[c_idx]])
+ clat = np.array([grid.clat[c_idx]])
+ clon_vertices = np.array([lon_verts])
+ clat_vertices = np.array([lat_verts])
+
+ ncells = 1
+ nv = clon_vertices[0].size
+
+ # Handle dateline crossing
+ if reader.split_EW:
+ clon_vertices[clon_vertices < 0.0] += 360.0
+
+ triangles = np.zeros((ncells, nv, 2))
+
+ for i in range(0, ncells, 1):
+ triangles[i, :, 0] = np.array(clon_vertices[i, :])
+ triangles[i, :, 1] = np.array(clat_vertices[i, :])
+
+ if params.plot or params.plot_output:
+ output_fn = params.path_output + str(c_idx) + ".png"
+ cart_plot.lat_lon_icon(
+ topo,
+ triangles,
+ ncells=ncells,
+ clon=clon,
+ clat=clat,
+ title=c_idx,
+ fn=output_fn,
+ output_fig=True,
+ )
+
+ # Initialize cell objects for CSA algorithm
+ tri_idx = 0
+ cell = var.topo_cell()
+ tri = var.obj()
+
+ nhi = params.nhi
+ nhj = params.nhj
+
+ fa = interface.first_appx(nhi, nhj, params, topo)
+ sa = interface.second_appx(nhi, nhj, params, topo, tri)
+
+ dplot = diagnostics.diag_plotter(params, nhi, nhj)
+ dplot.output_dir = params.path_output
+
+ tri.tri_lon_verts = triangles[:, :, 0]
+ tri.tri_lat_verts = triangles[:, :, 1]
+
+ simplex_lat = tri.tri_lat_verts[tri_idx]
+ simplex_lon = tri.tri_lon_verts[tri_idx]
+
+ if not utils.is_land(cell, simplex_lat, simplex_lon, topo):
+ # writer.output(c_idx, clat_rad[c_idx], clon_rad[c_idx], 0)
+ print("--> skipping ocean cell")
+ return writer.grp_struct(c_idx, clat_rad[c_idx], clon_rad[c_idx], 0)
+ else:
+ is_land = 1
+
+ if params.dfft_first_guess:
+ # do tapering
+ if params.taper_fa:
+ interface.taper_quad(params, simplex_lat, simplex_lon, cell, topo)
+ else:
+ utils.get_lat_lon_segments(
+ simplex_lat, simplex_lon, cell, topo, rect=params.rect
+ )
+
+ dfft_run = interface.get_pmf(nhi, nhj, params.U, params.V)
+ ampls_fa, uw_fa, dat_2D_fa, kls_fa = dfft_run.dfft(cell)
+
+ cell_fa = cell
+
+ nhi = len(cell_fa.lon)
+ nhj = len(cell_fa.lat)
+
+ sa.nhi = nhi
+ sa.nhj = nhj
+ else:
+ cell_fa, ampls_fa, uw_fa, dat_2D_fa = fa.do(simplex_lat, simplex_lon)
+
+ sols = (cell_fa, ampls_fa, uw_fa, dat_2D_fa)
+
+ v_extent = [dat_2D_fa.min(), dat_2D_fa.max()]
+
+ if params.plot:
+ if params.dfft_first_guess:
+ dplot.show(
+ tri_idx,
+ sols,
+ kls=kls_fa,
+ v_extent=v_extent,
+ dfft_plot=True,
+ output_fig=False,
+ )
+ else:
+ dplot.show(c_idx, sols, v_extent=v_extent, output_fig=False)
+
+ if params.recompute_rhs:
+ sols, _ = sa.do(tri_idx, ampls_fa)
+ else:
+ sols = sa.do(tri_idx, ampls_fa)
+
+ cell, ampls_sa, uw_sa, dat_2D_sa = sols
+ v_extent = [dat_2D_sa.min(), dat_2D_sa.max()]
+
+ # writer.output(c_idx, clat_rad[c_idx], clon_rad[c_idx], is_land, cell.analysis)
+ result = writer.grp_struct(
+ c_idx, clat_rad[c_idx], clon_rad[c_idx], is_land, cell.analysis
+ )
+
+ if params.plot:
+ if params.dfft_first_guess:
+ dplot.show(
+ tri_idx,
+ sols,
+ kls=kls_fa,
+ v_extent=v_extent,
+ dfft_plot=True,
+ output_fig=False,
+ )
+ else:
+ dplot.show(c_idx, sols, v_extent=v_extent, output_fig=False)
+
+ print("--> analysis done")
+
+ return result
+
+
+def parallel_wrapper(grid, params, reader, writer):
+ return lambda ii: do_cell(ii, grid, params, reader, writer)
+
+
+from pycsa.inputs.icon_global_run import params
+from dask.distributed import Client, progress
+import dask
+from tqdm import tqdm
+
+if __name__ == "__main__":
+ if params.self_test():
+ params.print()
+
+ grid = var.grid()
+
+ # Read ICON grid
+ reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+ reader.read_dat(params.path_icon_grid, grid)
+
+ clat_rad = np.copy(grid.clat)
+ clon_rad = np.copy(grid.clon)
+
+ grid.apply_f(utils.rad2deg)
+
+ n_cells = grid.clat.size
+
+ # Configure Dask for parallel processing
+ # Use processes (not threads) to avoid NetCDF file locking issues
+ # Each worker gets 1 thread to avoid GIL contention
+ import multiprocessing
+
+ n_workers = min(multiprocessing.cpu_count() - 2, 20) # Leave 2 cores for system
+ print(f"Initializing Dask with {n_workers} workers...")
+
+ client = Client(
+ threads_per_worker=1,
+ n_workers=n_workers,
+ processes=True,
+ memory_limit="4GB", # Per worker
+ )
+ print(f"Dask dashboard available at: {client.dashboard_link}")
+
+ print(f"Total cells to process: {n_cells}")
+
+ chunk_sz = 10
+ chunk_start = 20400
+
+ # Progress tracking
+ total_chunks = (n_cells - chunk_start + chunk_sz - 1) // chunk_sz
+ print(
+ f"\nProcessing {n_cells - chunk_start} cells in {total_chunks} chunks of {chunk_sz}..."
+ )
+
+ for chunk_idx, chunk in enumerate(
+ tqdm(range(chunk_start, n_cells, chunk_sz), desc="Processing chunks")
+ ):
+ # Writer object for this chunk
+ sfx = "_" + str(chunk + chunk_sz)
+ writer = io.nc_writer(params, sfx)
+
+ pw_run = parallel_wrapper(grid, params, reader, writer)
+
+ lazy_results = []
+
+ if chunk + chunk_sz > n_cells:
+ chunk_end = n_cells
+ else:
+ chunk_end = chunk + chunk_sz
+
+ for c_idx in range(chunk, chunk_end):
+ lazy_result = dask.delayed(pw_run)(c_idx)
+ lazy_results.append(lazy_result)
+
+ results = dask.compute(*lazy_results)
+
+ for item in results:
+ writer.duplicate(item.c_idx, item)
+
+ # Cleanup: close all cached NetCDF files and shut down Dask client
+ print("\nCleaning up...")
+ if hasattr(reader, "close_cached_files"):
+ reader.close_cached_files()
+ print("✓ Closed cached topography files")
+
+ client.close()
+ print("✓ Shut down Dask client")
+ print("Processing complete!")
diff --git a/runs/icon_merit_regional.py b/runs/icon_merit_regional.py
new file mode 100644
index 0000000..43ed955
--- /dev/null
+++ b/runs/icon_merit_regional.py
@@ -0,0 +1,227 @@
+# %%
+# import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+from pycsa.src import io, var, utils, fourier, physics
+from pycsa.wrappers import interface
+from pycsa.vis import plotter, cart_plot
+
+from IPython import get_ipython
+
+ipython = get_ipython()
+
+if ipython is not None:
+ ipython.run_line_magic("load_ext", "autoreload")
+else:
+ print(ipython)
+
+
+def autoreload():
+ if ipython is not None:
+ ipython.run_line_magic("autoreload", "2")
+
+
+from sys import exit
+
+if __name__ != "__main__":
+ exit(0)
+# %%
+autoreload()
+from pycsa.inputs.icon_regional_run import params
+
+if params.self_test():
+ params.print()
+
+print(params.path_compact_topo)
+
+grid = var.grid()
+topo = var.topo_cell()
+
+# read grid
+reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+
+# writer object
+writer = io.nc_writer(params)
+
+reader.read_dat(params.path_compact_grid, grid)
+
+clat_rad = np.copy(grid.clat)
+clon_rad = np.copy(grid.clon)
+
+grid.apply_f(utils.rad2deg)
+
+# we only keep the topography that is inside this lat-lon extent.
+lat_verts = np.array(params.lat_extent)
+lon_verts = np.array(params.lon_extent)
+
+# read topography
+if not params.enable_merit:
+ reader.read_dat(params.fn_topo, topo)
+ reader.read_topo(topo, topo, lon_verts, lat_verts)
+else:
+ reader.read_merit_topo(topo, params)
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+
+topo.gen_mgrids()
+
+
+# %%
+
+# if params.run_full_land_model:
+# params.rect_set = delaunay.get_land_cells(tri, topo, height_tol=0.5)
+# print(params.rect_set)
+
+# params_orig = deepcopy(params)
+# writer.write_all_attrs(params)
+# writer.populate("decomposition", "rect_set", params.rect_set)
+
+clon = grid.clon
+clat = grid.clat
+clon_vertices = grid.clon_vertices
+clat_vertices = grid.clat_vertices
+
+ncells, nv = clon_vertices.shape[0], clon_vertices.shape[1]
+
+# -- print information to stdout
+print("Cells: %6d " % clon.size)
+
+# -- create the triangles
+clon_vertices = np.where(clon_vertices < -180.0, clon_vertices + 360.0, clon_vertices)
+clon_vertices = np.where(clon_vertices > 180.0, clon_vertices - 360.0, clon_vertices)
+
+triangles = np.zeros((ncells, nv, 2), np.float32)
+
+for i in range(0, ncells, 1):
+ triangles[i, :, 0] = np.array(clon_vertices[i, :])
+ triangles[i, :, 1] = np.array(clat_vertices[i, :])
+
+print("--> triangles done")
+
+cart_plot.lat_lon_icon(topo, triangles, ncells=ncells, clon=clon, clat=clat)
+
+
+# %%
+autoreload()
+idxs = []
+pmfs = []
+
+for tri_idx in params.tri_set:
+ # initialise cell object
+ cell = var.topo_cell()
+
+ simplex_lon = triangles[tri_idx, :, 0]
+ simplex_lat = triangles[tri_idx, :, 1]
+
+ utils.get_lat_lon_segments(simplex_lat, simplex_lon, cell, topo, rect=params.rect)
+
+ topo_orig = np.copy(cell.topo)
+
+ if params.dfft_first_guess:
+ nhi = len(cell.lon)
+ nhj = len(cell.lat)
+
+ first_guess = interface.get_pmf(nhi, nhj, params.U, params.V)
+ fobj_tri = fourier.f_trans(nhi, nhj)
+
+ #######################################################
+ # do fourier...
+
+ if not params.dfft_first_guess:
+ freqs, uw_pmf_freqs, dat_2D_fg0 = first_guess.sappx(cell, lmbda=0.0)
+
+ #######################################################
+ # do fourier using DFFT
+
+ if params.dfft_first_guess:
+ ampls, uw_pmf_freqs, dat_2D_fg0, kls = first_guess.dfft(cell)
+ freqs = np.copy(ampls)
+
+ print("uw_pmf_freqs_sum:", uw_pmf_freqs.sum())
+
+ fq_cpy = np.copy(freqs)
+
+ indices = []
+ max_ampls = []
+
+ for ii in range(params.n_modes):
+ max_idx = np.unravel_index(fq_cpy.argmax(), fq_cpy.shape)
+ indices.append(max_idx)
+ max_ampls.append(fq_cpy[max_idx])
+ max_val = fq_cpy[max_idx]
+ fq_cpy[max_idx] = 0.0
+
+ utils.get_lat_lon_segments(simplex_lat, simplex_lon, cell, topo, rect=False)
+
+ k_idxs = [pair[1] for pair in indices]
+ l_idxs = [pair[0] for pair in indices]
+
+ second_guess = interface.get_pmf(nhi, nhj, params.U, params.V)
+
+ if params.dfft_first_guess:
+ second_guess.fobj.set_kls(
+ k_idxs, l_idxs, recompute_nhij=True, components="real"
+ )
+ else:
+ second_guess.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+
+ freqs, uw, dat_2D_sg0 = second_guess.sappx(cell, lmbda=1e-1, updt_analysis=True)
+
+ cell.topo = topo_orig
+
+ writer.output(tri_idx, clat_rad[tri_idx], clon_rad[tri_idx], cell.analysis)
+
+ cell.uw = uw
+
+ if params.plot:
+ fs = (15, 9.0)
+ v_extent = [dat_2D_sg0.min(), dat_2D_sg0.max()]
+
+ fig, axs = plt.subplots(2, 2, figsize=fs)
+
+ fig_obj = plotter.fig_obj(
+ fig, second_guess.fobj.nhar_i, second_guess.fobj.nhar_j
+ )
+ axs[0, 0] = fig_obj.phys_panel(
+ axs[0, 0],
+ dat_2D_sg0,
+ title="T%i: Reconstruction" % tri_idx,
+ xlabel="longitude [km]",
+ ylabel="latitude [km]",
+ extent=[cell.lon.min(), cell.lon.max(), cell.lat.min(), cell.lat.max()],
+ v_extent=v_extent,
+ )
+
+ axs[0, 1] = fig_obj.phys_panel(
+ axs[0, 1],
+ cell.topo * cell.mask,
+ title="T%i: Reconstruction" % tri_idx,
+ xlabel="longitude [km]",
+ ylabel="latitude [km]",
+ extent=[cell.lon.min(), cell.lon.max(), cell.lat.min(), cell.lat.max()],
+ v_extent=v_extent,
+ )
+
+ if params.dfft_first_guess:
+ axs[1, 0] = fig_obj.fft_freq_panel(
+ axs[1, 0], freqs, kls[0], kls[1], typ="real"
+ )
+ axs[1, 1] = fig_obj.fft_freq_panel(
+ axs[1, 1], uw, kls[0], kls[1], title="PMF spectrum", typ="real"
+ )
+ else:
+ axs[1, 0] = fig_obj.freq_panel(axs[1, 0], freqs)
+ axs[1, 1] = fig_obj.freq_panel(axs[1, 1], uw, title="PMF spectrum")
+
+ plt.tight_layout()
+ plt.savefig("%sT%i.pdf" % (params.path_output, tri_idx))
+ plt.show()
+
+ ideal = physics.ideal_pmf(U=params.U, V=params.V)
+ uw_comp = ideal.compute_uw_pmf(cell.analysis)
+
+ idxs.append(tri_idx)
+ pmfs.append(uw_comp)
+
+
+# %%
diff --git a/runs/icon_usgs_test.py b/runs/icon_usgs_test.py
index a4dad71..1368457 100644
--- a/runs/icon_usgs_test.py
+++ b/runs/icon_usgs_test.py
@@ -1,23 +1,18 @@
# %%
-import sys
-
-# set system path to find local modules
-sys.path.append("..")
-
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
-from src import io, var, utils, fourier, physics
-from wrappers import interface
-from vis import plotter, cart_plot
+from pycsa.core import io, var, utils, fourier, physics
+from pycsa.wrappers import interface
+from pycsa.plotting import plotter, cart_plot
# %%
fn_grid = "../data/icon_compact.nc"
fn_topo = "../data/topo_compact.nc"
-lat_extent = [52.0, 64.0, 64.0]
-lon_extent = [-141.0, -158.0, -127.0]
+lat_extent = [48.0, 64.0, 64.0]
+lon_extent = [-148.0, -148.0, -112.0]
tri_set = [13, 104, 105, 106]
@@ -27,7 +22,7 @@
n_modes = 100
-U, V = 10.0, 0.1
+U, V = 10.0, 0.0
rect = True
@@ -101,10 +96,7 @@
simplex_lon = triangles[tri_idx, :, 0]
simplex_lat = triangles[tri_idx, :, 1]
- triangle = utils.triangle(simplex_lon, simplex_lat)
- utils.get_lat_lon_segments(
- simplex_lat, simplex_lon, cell, topo, triangle, rect=rect
- )
+ utils.get_lat_lon_segments(simplex_lat, simplex_lon, cell, topo, rect=rect)
topo_orig = np.copy(cell.topo)
@@ -142,9 +134,7 @@
max_val = fq_cpy[max_idx]
fq_cpy[max_idx] = 0.0
- utils.get_lat_lon_segments(
- simplex_lat, simplex_lon, cell, topo, triangle, rect=False
- )
+ utils.get_lat_lon_segments(simplex_lat, simplex_lon, cell, topo, rect=False)
k_idxs = [pair[1] for pair in indices]
l_idxs = [pair[0] for pair in indices]
diff --git a/runs/idealised_delaunay.py b/runs/idealised_delaunay.py
index 19cb56f..4945551 100644
--- a/runs/idealised_delaunay.py
+++ b/runs/idealised_delaunay.py
@@ -4,14 +4,8 @@
from matplotlib import pyplot as plt
from copy import deepcopy
-import sys
-import os
-
-# set system path to find local modules
-sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
-
-from src import utils, var
-from wrappers import interface, diagnostics
+from pycsa.core import utils, var
+from pycsa.wrappers import interface, diagnostics
from IPython import get_ipython
diff --git a/runs/idealised_isosceles.py b/runs/idealised_isosceles.py
index c12bff2..238e10a 100644
--- a/runs/idealised_isosceles.py
+++ b/runs/idealised_isosceles.py
@@ -1,16 +1,8 @@
# %%
-import sys
-import os
-
-# set system path to find local modules
-sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
-
import numpy as np
import matplotlib.pyplot as plt
-from src import var, utils
-from wrappers import interface
-from vis import plotter
+from pycsa import var, utils, interface, plotter
from copy import deepcopy
from IPython import get_ipython
@@ -140,8 +132,8 @@ def sinusoidal_basis(Ak, nk, Al, nl, sc, typ):
dat_arr = np.array([None] * num_experiments, dtype=object)
-#### helper function to run the CSAM algorithm
-def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
+#### helper function to run the CSA algorithm
+def csa_run(cell, n_modes, lmbda_fg, lmbda_sg):
first_guess = interface.get_pmf(nhi, nhj, U, V)
cell.get_masked(mask=np.ones_like(cell.topo).astype("bool"))
@@ -152,9 +144,9 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
freqs_fg, _, dat_2D_fg = first_guess.sappx(cell, lmbda=lmbda_fg, iter_solve=False)
fq_cpy = np.copy(freqs_fg)
- fq_cpy[
- np.isnan(fq_cpy)
- ] = 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ fq_cpy[np.isnan(fq_cpy)] = (
+ 0.0 # necessary. Otherwise, popping with fq_cpy.max() gives the np.nan entries first.
+ )
indices = []
max_ampls = []
@@ -202,11 +194,11 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
#### regularised lsff run
freqs_arr[2], _, dat_arr[2] = reg_lsff.sappx(cell, lmbda=lmbda_reg, iter_solve=False)
-#### optimal CSAM run
-freqs_arr[3], _, dat_arr[3] = csam_run(cell, sz, lmbda_fg, lmbda_sg)
+#### optimal CSA run
+freqs_arr[3], _, dat_arr[3] = csa_run(cell, sz, lmbda_fg, lmbda_sg)
-#### suboptimal CSAM run
-freqs_arr[4], _, dat_arr[4] = csam_run(cell, n_modes, lmbda_fg, lmbda_sg)
+#### suboptimal CSA run
+freqs_arr[4], _, dat_arr[4] = csa_run(cell, n_modes, lmbda_fg, lmbda_sg)
freqs_arr = np.array([np.nan_to_num(freq) for freq in freqs_arr])
@@ -236,7 +228,7 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
selected_sums = []
selected_sum_errs = []
-phys_lbls = ["reference", "pLSFF", "optCSAM", "subCSAM"]
+phys_lbls = ["reference", "pLSFF", "optCSA", "subCSA"]
spec_lbls = ["", "", "", ""]
for cnt, idx in enumerate(idxs):
@@ -268,7 +260,7 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
axs[1, 0].set_ylabel("$m$", fontsize=12)
# plt.tight_layout()
-plt.savefig("../manuscript/idealized_plots.pdf", bbox_inches="tight")
+plt.savefig("outputs/baseline_results/idealized_plots.pdf", bbox_inches="tight")
plt.show()
@@ -285,7 +277,7 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
fontsize=14,
fs=(3.5, 2.5),
output_fig=True,
- fn="../manuscript/l2_errs.pdf",
+ fn="outputs/baseline_results/l2_errs.pdf",
)
plotter.error_bar_abs_plot(
selected_sums,
@@ -296,7 +288,7 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
fontsize=14,
fs=(4.5, 2.5),
output_fig=True,
- fn="../manuscript/powers.pdf",
+ fn="outputs/baseline_results/powers.pdf",
)
@@ -353,7 +345,7 @@ def csam_run(cell, n_modes, lmbda_fg, lmbda_sg):
axs[2].set_ylabel("$m$", fontsize=12)
plt.tight_layout()
-plt.savefig("../manuscript/overfitting_issue.pdf", bbox_inches="tight")
+plt.savefig("outputs/baseline_results/overfitting_issue.pdf", bbox_inches="tight")
plt.show()
# %%
diff --git a/runs/merge_netcdf_chunks.py b/runs/merge_netcdf_chunks.py
new file mode 100644
index 0000000..8721b32
--- /dev/null
+++ b/runs/merge_netcdf_chunks.py
@@ -0,0 +1,261 @@
+"""
+Merge NetCDF chunk files into a single final NetCDF file.
+
+This script:
+1. Finds all icon_etopo_global_cells_*.nc files
+2. Validates that all expected chunks are present
+3. Merges them into icon_etopo_global_FINAL.nc
+4. Optionally removes intermediate chunk files
+
+Usage:
+ python3 -m runs.merge_netcdf_chunks [--cleanup] [--output OUTPUT_NAME]
+
+Options:
+ --cleanup Remove intermediate chunk files after successful merge
+ --output Output filename (default: icon_etopo_global_FINAL.nc)
+"""
+
+import netCDF4 as nc
+import numpy as np
+from pathlib import Path
+import re
+import argparse
+from tqdm import tqdm
+
+
+def find_chunk_files(datasets_dir):
+ """Find all NetCDF chunk files and extract their cell ranges."""
+ pattern = re.compile(r"icon_etopo_global_cells_(\d+)-(\d+)\.nc")
+
+ chunks = []
+ for filepath in sorted(datasets_dir.glob("icon_etopo_global_cells_*.nc")):
+ match = pattern.match(filepath.name)
+ if match:
+ start_cell = int(match.group(1))
+ end_cell = int(match.group(2))
+ chunks.append(
+ {
+ "filepath": filepath,
+ "start": start_cell,
+ "end": end_cell,
+ "size": end_cell - start_cell + 1,
+ }
+ )
+
+ return sorted(chunks, key=lambda x: x["start"])
+
+
+def validate_chunks(chunks, expected_total_cells=20480):
+ """Validate that chunks cover all cells without gaps or overlaps."""
+ if not chunks:
+ raise ValueError("No chunk files found!")
+
+ print(f"\nFound {len(chunks)} chunk files")
+ print(f" First chunk: cells {chunks[0]['start']}-{chunks[0]['end']}")
+ print(f" Last chunk: cells {chunks[-1]['start']}-{chunks[-1]['end']}")
+
+ # Check for gaps
+ for i in range(len(chunks) - 1):
+ current_end = chunks[i]["end"]
+ next_start = chunks[i + 1]["start"]
+ if current_end + 1 != next_start:
+ raise ValueError(
+ f"Gap detected: chunk ends at {current_end}, next starts at {next_start}"
+ )
+
+ # Check coverage
+ total_cells = chunks[-1]["end"] + 1 - chunks[0]["start"]
+ if chunks[0]["start"] != 0:
+ print(f"\n⚠ Warning: First chunk starts at cell {chunks[0]['start']}, not 0")
+
+ if total_cells < expected_total_cells:
+ print(f"\n⚠ Warning: Only {total_cells}/{expected_total_cells} cells covered")
+
+ print(f"\n✓ Validation passed: {total_cells} cells in {len(chunks)} chunks\n")
+ return True
+
+
+def merge_chunks(chunks, output_path, datasets_dir):
+ """Merge chunk files into a single NetCDF file."""
+
+ print(f"Merging {len(chunks)} chunks into: {output_path.name}")
+ print("=" * 80)
+
+ # Read first chunk to get global attributes and parameters
+ first_chunk = nc.Dataset(chunks[0]["filepath"], "r")
+
+ # Create output file
+ output_nc = nc.Dataset(output_path, "w", format="NETCDF4")
+
+ # Copy global attributes from first chunk
+ print("\nCopying global attributes...")
+ for attr_name in first_chunk.ncattrs():
+ setattr(output_nc, attr_name, getattr(first_chunk, attr_name))
+
+ # Create dimensions
+ nspec = (
+ first_chunk.dimensions["nspec"].size
+ if "nspec" in first_chunk.dimensions
+ else 100
+ )
+ output_nc.createDimension("nspec", nspec)
+
+ first_chunk.close()
+
+ # Merge all chunks
+ print(f"\nMerging chunks...")
+ total_land_cells = 0
+ total_ocean_cells = 0
+
+ for chunk in tqdm(chunks, desc="Processing chunks"):
+ src_nc = nc.Dataset(chunk["filepath"], "r")
+
+ # Iterate through all groups (cells) in this chunk
+ for group_name in src_nc.groups:
+ src_group = src_nc.groups[group_name]
+
+ # Create group in output
+ dst_group = output_nc.createGroup(group_name)
+
+ # Copy variables
+ for var_name in src_group.variables:
+ src_var = src_group.variables[var_name]
+
+ # Create variable in output
+ if src_var.dimensions:
+ dst_var = dst_group.createVariable(
+ var_name, src_var.datatype, src_var.dimensions
+ )
+ else:
+ dst_var = dst_group.createVariable(var_name, src_var.datatype)
+
+ # Copy data
+ dst_var[:] = src_var[:]
+
+ # Copy attributes
+ for attr_name in src_var.ncattrs():
+ setattr(dst_var, attr_name, getattr(src_var, attr_name))
+
+ # Track statistics
+ if "is_land" in src_group.variables:
+ if src_group.variables["is_land"][:]:
+ total_land_cells += 1
+ else:
+ total_ocean_cells += 1
+
+ src_nc.close()
+
+ output_nc.close()
+
+ print("\n" + "=" * 80)
+ print("MERGE COMPLETE")
+ print("=" * 80)
+ print(f"Output file: {output_path}")
+ print(f"File size: {output_path.stat().st_size / 1024 / 1024:.1f} MB")
+ print(f"\nCells merged:")
+ print(f" Land cells: {total_land_cells}")
+ print(f" Ocean cells: {total_ocean_cells}")
+ print(f" Total: {total_land_cells + total_ocean_cells}")
+ print("=" * 80)
+
+ return total_land_cells + total_ocean_cells
+
+
+def cleanup_chunks(chunks):
+ """Remove intermediate chunk files."""
+ print("\nCleaning up intermediate files...")
+ for chunk in tqdm(chunks, desc="Removing chunks"):
+ chunk["filepath"].unlink()
+ print(f"✓ Removed {len(chunks)} chunk files")
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Merge ICON ETOPO NetCDF chunk files")
+ parser.add_argument(
+ "--cleanup",
+ action="store_true",
+ help="Remove intermediate chunk files after merge",
+ )
+ parser.add_argument(
+ "--output",
+ type=str,
+ default="icon_etopo_global_FINAL.nc",
+ help="Output filename (default: icon_etopo_global_FINAL.nc)",
+ )
+ parser.add_argument(
+ "--datasets-dir",
+ type=str,
+ help="Directory containing chunk files (default: auto-detect)",
+ )
+
+ args = parser.parse_args()
+
+ # Find datasets directory
+ if args.datasets_dir:
+ datasets_dir = Path(args.datasets_dir)
+ else:
+ # Try to find it automatically
+ possible_paths = [
+ Path("outputs/global_run/datasets"),
+ Path("../outputs/global_run/datasets"),
+ Path("../../outputs/global_run/datasets"),
+ ]
+ datasets_dir = None
+ for path in possible_paths:
+ if path.exists():
+ datasets_dir = path
+ break
+
+ if datasets_dir is None:
+ print("Error: Could not find datasets directory")
+ print("Please specify with --datasets-dir")
+ return 1
+
+ print(f"Datasets directory: {datasets_dir}")
+
+ # Find chunk files
+ chunks = find_chunk_files(datasets_dir)
+ if not chunks:
+ print("Error: No chunk files found!")
+ print(f"Looking for: icon_etopo_global_cells_*.nc in {datasets_dir}")
+ return 1
+
+ # Validate
+ try:
+ validate_chunks(chunks)
+ except ValueError as e:
+ print(f"\n❌ Validation error: {e}")
+ print("\nChunk files found:")
+ for chunk in chunks:
+ print(f" {chunk['filepath'].name}: cells {chunk['start']}-{chunk['end']}")
+ return 1
+
+ # Merge
+ output_path = datasets_dir / args.output
+ if output_path.exists():
+ response = input(f"\n⚠ {output_path.name} already exists. Overwrite? [y/N] ")
+ if response.lower() != "y":
+ print("Merge cancelled")
+ return 0
+
+ try:
+ total_cells = merge_chunks(chunks, output_path, datasets_dir)
+ except Exception as e:
+ print(f"\n❌ Merge failed: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return 1
+
+ # Cleanup if requested
+ if args.cleanup:
+ response = input(f"\nRemove {len(chunks)} chunk files? [y/N] ")
+ if response.lower() == "y":
+ cleanup_chunks(chunks)
+
+ print(f"\n✓ Success! Merged file: {output_path}")
+ return 0
+
+
+if __name__ == "__main__":
+ exit(main())
diff --git a/runs/submit_etopo_global.sh b/runs/submit_etopo_global.sh
new file mode 100755
index 0000000..f0704b0
--- /dev/null
+++ b/runs/submit_etopo_global.sh
@@ -0,0 +1,55 @@
+#!/bin/bash
+#SBATCH --job-name=icon_etopo_global
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=128
+#SBATCH --mem=256G
+#SBATCH --time=48:00:00
+#SBATCH --output=logs/icon_etopo_%j.log
+#SBATCH --error=logs/icon_etopo_%j.err
+
+# SLURM submission script for ICON ETOPO global processing
+# Optimized for: 128 cores, 256 GB RAM single node
+
+echo "========================================="
+echo "ICON ETOPO Global Processing"
+echo "========================================="
+echo "Job ID: $SLURM_JOB_ID"
+echo "Node: $SLURM_NODELIST"
+echo "Cores: $SLURM_CPUS_PER_TASK"
+echo "Memory: 256 GB"
+echo "Start time: $(date)"
+echo "========================================="
+echo ""
+
+# Create logs directory if it doesn't exist
+mkdir -p logs
+
+# Load required modules (adjust for your HPC system)
+# module load anaconda3 # or your Python environment
+# module load netcdf4
+
+# Activate conda environment
+# source activate playground # or your environment name
+
+# Set OpenMP threads to 1 (we use Dask for parallelism)
+export OMP_NUM_THREADS=1
+export MKL_NUM_THREADS=1
+export NUMEXPR_NUM_THREADS=1
+
+# Increase file descriptor limits (NetCDF files)
+ulimit -n 4096
+
+# Run the HPC-optimized script
+echo "Starting ICON ETOPO processing..."
+python3 -m runs.icon_etopo_global
+
+exit_code=$?
+
+echo ""
+echo "========================================="
+echo "Job completed with exit code: $exit_code"
+echo "End time: $(date)"
+echo "========================================="
+
+exit $exit_code
diff --git a/runs/tapering_test.py b/runs/tapering_test.py
index da0f3f2..6bc20f4 100644
--- a/runs/tapering_test.py
+++ b/runs/tapering_test.py
@@ -1,14 +1,9 @@
# %%
-import sys
-
-# setting path
-sys.path.append("..")
-
import numpy as np
import matplotlib.pyplot as plt
-from src import io, var, utils, delaunay
-from vis import cart_plot, plotter
+from pycsa.core import io, var, utils, delaunay
+from pycsa.plotting import cart_plot, plotter
from copy import deepcopy
diff --git a/runs/validate_chunks.py b/runs/validate_chunks.py
new file mode 100644
index 0000000..5bb66e4
--- /dev/null
+++ b/runs/validate_chunks.py
@@ -0,0 +1,126 @@
+"""
+Quick validation script to check NetCDF chunk completeness.
+
+Usage:
+ python3 -m runs.validate_chunks [--datasets-dir PATH]
+"""
+
+from pathlib import Path
+import re
+import argparse
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Validate ICON ETOPO NetCDF chunks")
+ parser.add_argument(
+ "--datasets-dir",
+ type=str,
+ help="Directory containing chunk files (default: auto-detect)",
+ )
+ args = parser.parse_args()
+
+ # Find datasets directory
+ if args.datasets_dir:
+ datasets_dir = Path(args.datasets_dir)
+ else:
+ possible_paths = [
+ Path("outputs/global_run/datasets"),
+ Path("../outputs/global_run/datasets"),
+ Path("../../outputs/global_run/datasets"),
+ ]
+ datasets_dir = None
+ for path in possible_paths:
+ if path.exists():
+ datasets_dir = path
+ break
+
+ if datasets_dir is None:
+ print("❌ Could not find datasets directory")
+ return 1
+
+ print(f"Checking: {datasets_dir}\n")
+
+ # Find chunk files
+ pattern = re.compile(r"icon_etopo_global_cells_(\d+)-(\d+)\.nc")
+ chunks = []
+
+ for filepath in sorted(datasets_dir.glob("icon_etopo_global_cells_*.nc")):
+ match = pattern.match(filepath.name)
+ if match:
+ start_cell = int(match.group(1))
+ end_cell = int(match.group(2))
+ file_size = filepath.stat().st_size / 1024 # KB
+ chunks.append(
+ {
+ "filepath": filepath,
+ "start": start_cell,
+ "end": end_cell,
+ "size_kb": file_size,
+ }
+ )
+
+ chunks = sorted(chunks, key=lambda x: x["start"])
+
+ if not chunks:
+ print("❌ No chunk files found!")
+ print(f" Looking for: icon_etopo_global_cells_*.nc")
+ return 1
+
+ # Display summary
+ print(f"Found {len(chunks)} chunk files:")
+ print(f" First: cells {chunks[0]['start']}-{chunks[0]['end']}")
+ print(f" Last: cells {chunks[-1]['start']}-{chunks[-1]['end']}")
+
+ # Check for issues
+ issues = []
+
+ # Check for gaps
+ for i in range(len(chunks) - 1):
+ current_end = chunks[i]["end"]
+ next_start = chunks[i + 1]["start"]
+ if current_end + 1 != next_start:
+ issues.append(
+ f"Gap: chunk {i} ends at {current_end}, chunk {i+1} starts at {next_start}"
+ )
+
+ # Check start
+ if chunks[0]["start"] != 0:
+ issues.append(
+ f"First chunk doesn't start at 0 (starts at {chunks[0]['start']})"
+ )
+
+ # Check expected coverage
+ expected_cells = 20480
+ total_cells = chunks[-1]["end"] + 1 - chunks[0]["start"]
+
+ print(
+ f"\nCoverage: {total_cells}/{expected_cells} cells ({total_cells/expected_cells*100:.1f}%)"
+ )
+
+ if total_cells < expected_cells:
+ issues.append(f"Incomplete: only {total_cells}/{expected_cells} cells")
+
+ # Calculate total size
+ total_size_mb = sum(c["size_kb"] for c in chunks) / 1024
+ print(f"Total size: {total_size_mb:.1f} MB")
+
+ # Report
+ print("\n" + "=" * 60)
+ if issues:
+ print("⚠ ISSUES FOUND:")
+ for issue in issues:
+ print(f" - {issue}")
+ print("=" * 60)
+ return 1
+ else:
+ print("✓ ALL CHECKS PASSED")
+ print(" - No gaps in cell coverage")
+ print(" - All chunks present")
+ print("\nReady to merge with:")
+ print(" python3 -m runs.merge_netcdf_chunks")
+ print("=" * 60)
+ return 0
+
+
+if __name__ == "__main__":
+ exit(main())
diff --git a/scripts/check_etopo_sizes.sh b/scripts/check_etopo_sizes.sh
new file mode 100755
index 0000000..15e703b
--- /dev/null
+++ b/scripts/check_etopo_sizes.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+# Safer script - just checks file sizes
+# ETOPO 15s surface files should be 5-35 MB typically
+
+DATA_DIR="${1:-./data/etopo_15s}"
+
+echo "Checking ETOPO file sizes in: $DATA_DIR"
+echo "========================================="
+
+suspicious=()
+total=0
+
+for file in "$DATA_DIR"/*.nc; do
+ if [ -f "$file" ]; then
+ total=$((total + 1))
+ size=$(stat -f%z "$file" 2>/dev/null || stat -c%s "$file" 2>/dev/null)
+ size_mb=$((size / 1048576))
+ filename=$(basename "$file")
+
+ # ETOPO 15s tiles are typically 5-35 MB
+ if [ "$size" -lt 1000000 ]; then # Less than 1 MB is definitely wrong
+ echo "⚠️ SUSPICIOUS: $filename (${size_mb} MB - too small!)"
+ suspicious+=("$file")
+ elif [ "$size" -gt 50000000 ]; then # More than 50 MB is suspicious
+ echo "⚠️ SUSPICIOUS: $filename (${size_mb} MB - too large!)"
+ suspicious+=("$file")
+ else
+ echo "✓ OK: $filename (${size_mb} MB)"
+ fi
+ fi
+done
+
+echo ""
+echo "========================================="
+echo "Total files: $total"
+echo "Suspicious files: ${#suspicious[@]}"
+
+if [ ${#suspicious[@]} -gt 0 ]; then
+ echo ""
+ echo "Suspicious files to check/re-download:"
+ for file in "${suspicious[@]}"; do
+ size=$(stat -f%z "$file" 2>/dev/null || stat -c%s "$file" 2>/dev/null)
+ echo " - $(basename "$file") ($(($size / 1048576)) MB)"
+ done
+fi
diff --git a/scripts/check_slurm_resources.py b/scripts/check_slurm_resources.py
new file mode 100644
index 0000000..4298a4a
--- /dev/null
+++ b/scripts/check_slurm_resources.py
@@ -0,0 +1,87 @@
+#!/usr/bin/env python3
+"""
+Check SLURM resource allocation for the current job.
+"""
+
+import os
+import subprocess
+
+
+def get_slurm_allocation():
+ """Get SLURM resource allocation for current job."""
+
+ # Check if running under SLURM
+ job_id = os.environ.get("SLURM_JOB_ID")
+
+ if not job_id:
+ print("Not running in a SLURM job")
+ return None
+
+ print(f"SLURM Job ID: {job_id}")
+ print("=" * 60)
+
+ # Get info from environment variables
+ info = {
+ "Job ID": os.environ.get("SLURM_JOB_ID"),
+ "Job Name": os.environ.get("SLURM_JOB_NAME"),
+ "Partition": os.environ.get("SLURM_JOB_PARTITION"),
+ "Nodes": os.environ.get("SLURM_JOB_NUM_NODES"),
+ "CPUs per Task": os.environ.get("SLURM_CPUS_PER_TASK"),
+ "Total CPUs": os.environ.get("SLURM_NTASKS"),
+ "Memory per Node (MB)": os.environ.get("SLURM_MEM_PER_NODE"),
+ "Memory per CPU (MB)": os.environ.get("SLURM_MEM_PER_CPU"),
+ "CPUs on Node": os.environ.get("SLURM_CPUS_ON_NODE"),
+ "Tasks per Node": os.environ.get("SLURM_TASKS_PER_NODE"),
+ }
+
+ print("\nEnvironment Variables:")
+ for key, value in info.items():
+ if value:
+ print(f" {key:25s}: {value}")
+
+ # Calculate total memory
+ mem_per_node_mb = os.environ.get("SLURM_MEM_PER_NODE")
+ num_nodes = os.environ.get("SLURM_JOB_NUM_NODES", "1")
+
+ if mem_per_node_mb:
+ mem_mb = int(mem_per_node_mb)
+ mem_gb = mem_mb / 1024
+ total_mem_gb = mem_gb * int(num_nodes)
+ print(f"\n Total Memory Allocated : {total_mem_gb:.1f} GB ({mem_mb} MB)")
+
+ # Get more details using scontrol
+ try:
+ result = subprocess.run(
+ ["scontrol", "show", "job", job_id], capture_output=True, text=True
+ )
+
+ if result.returncode == 0:
+ output = result.stdout
+
+ # Parse key fields
+ for line in output.split("\n"):
+ if "MinMemoryNode=" in line:
+ # Extract memory
+ parts = line.split()
+ for part in parts:
+ if "MinMemoryNode=" in part:
+ mem_str = part.split("=")[1]
+ print(f"\n MinMemoryNode (scontrol) : {mem_str}")
+
+ if "NumCPUs=" in line:
+ parts = line.split()
+ for part in parts:
+ if part.startswith("NumCPUs="):
+ cpus = part.split("=")[1]
+ print(f" NumCPUs (scontrol) : {cpus}")
+
+ except Exception as e:
+ print(f"\nCouldn't get scontrol info: {e}")
+
+ print("=" * 60)
+
+ return info
+
+
+if __name__ == "__main__":
+ get_slurm_allocation()
diff --git a/scripts/diagnose_netcdf_issue.sh b/scripts/diagnose_netcdf_issue.sh
new file mode 100755
index 0000000..730ac72
--- /dev/null
+++ b/scripts/diagnose_netcdf_issue.sh
@@ -0,0 +1,194 @@
+#!/bin/bash
+# Diagnostic script for NetCDF/HDF errors on HPC
+# Usage: ./diagnose_netcdf_issue.sh /path/to/etopo_file.nc
+
+NETCDF_FILE="${1}"
+
+if [ -z "$NETCDF_FILE" ]; then
+ echo "Usage: $0 /path/to/netcdf_file.nc"
+ exit 1
+fi
+
+echo "========================================="
+echo "NetCDF/HDF Diagnostic Tool"
+echo "========================================="
+echo ""
+
+echo "File: $NETCDF_FILE"
+echo ""
+
+# 1. Check if file exists
+echo "1. File existence check:"
+if [ -f "$NETCDF_FILE" ]; then
+ echo " ✓ File exists"
+else
+ echo " ✗ File does not exist!"
+ exit 1
+fi
+echo ""
+
+# 2. Check file size
+echo "2. File size check:"
+FILE_SIZE=$(stat -c%s "$NETCDF_FILE" 2>/dev/null || stat -f%z "$NETCDF_FILE" 2>/dev/null)
+FILE_SIZE_MB=$((FILE_SIZE / 1048576))
+echo " Size: ${FILE_SIZE} bytes (${FILE_SIZE_MB} MB)"
+if [ "$FILE_SIZE" -lt 1000000 ]; then
+ echo " ⚠️ WARNING: File seems too small (< 1 MB), likely corrupted"
+elif [ "$FILE_SIZE" -gt 50000000 ]; then
+ echo " ⚠️ WARNING: File seems too large (> 50 MB), unusual for 15s tile"
+else
+ echo " ✓ File size seems reasonable"
+fi
+echo ""
+
+# 3. Check file permissions
+echo "3. File permissions check:"
+FILE_PERMS=$(ls -lh "$NETCDF_FILE" | awk '{print $1}')
+echo " Permissions: $FILE_PERMS"
+if [ -r "$NETCDF_FILE" ]; then
+ echo " ✓ File is readable"
+else
+ echo " ✗ File is NOT readable!"
+fi
+echo ""
+
+# 4. Check file type
+echo "4. File type check:"
+FILE_TYPE=$(file "$NETCDF_FILE" 2>/dev/null || echo "file command not available")
+echo " Type: $FILE_TYPE"
+if echo "$FILE_TYPE" | grep -qi "netcdf\|hdf"; then
+ echo " ✓ File appears to be NetCDF/HDF format"
+else
+ echo " ⚠️ WARNING: File may not be valid NetCDF/HDF"
+fi
+echo ""
+
+# 5. Check first few bytes (magic number)
+echo "5. File header check (magic number):"
+HEADER=$(xxd -l 16 -p "$NETCDF_FILE" 2>/dev/null | tr -d '\n')
+echo " First 16 bytes (hex): $HEADER"
+
+# NetCDF-3: starts with "CDF" (43 44 46)
+# NetCDF-4/HDF5: starts with HDF5 signature (89 48 44 46 0d 0a 1a 0a)
+if [[ "$HEADER" == 434446* ]]; then
+ echo " ✓ NetCDF-3 format detected"
+elif [[ "$HEADER" == 894844460d0a1a0a* ]]; then
+ echo " ✓ NetCDF-4/HDF5 format detected"
+else
+ echo " ✗ INVALID: Does not match NetCDF format signature!"
+ echo " This file is corrupted or not a NetCDF file"
+fi
+echo ""
+
+# 6. Check with ncdump (if available)
+echo "6. ncdump validation check:"
+if command -v ncdump &> /dev/null; then
+ if ncdump -h "$NETCDF_FILE" > /dev/null 2>&1; then
+ echo " ✓ File can be opened with ncdump"
+ echo ""
+ echo " Variables in file:"
+ ncdump -h "$NETCDF_FILE" | grep -E "^\s+(float|double|int|short|byte)" | head -10
+ else
+ echo " ✗ ncdump FAILED to open file"
+ echo ""
+ echo " Error output:"
+ ncdump -h "$NETCDF_FILE" 2>&1 | head -5
+ fi
+else
+ echo " ⚠️ ncdump not available (load netcdf module?)"
+fi
+echo ""
+
+# 7. Try Python netCDF4 library
+echo "7. Python netCDF4 library check:"
+if command -v python3 &> /dev/null; then
+ python3 << EOF
+import sys
+try:
+ import netCDF4 as nc
+ print(" ✓ netCDF4 module is available")
+ try:
+ ds = nc.Dataset("$NETCDF_FILE", "r")
+ print(" ✓ File opened successfully with Python netCDF4")
+ print(f" Variables: {list(ds.variables.keys())}")
+ ds.close()
+ except Exception as e:
+ print(f" ✗ Python netCDF4 FAILED to open file")
+ print(f" Error: {e}")
+ sys.exit(1)
+except ImportError:
+ print(" ⚠️ netCDF4 module not available in Python")
+ sys.exit(1)
+EOF
+else
+ echo " ⚠️ python3 not available"
+fi
+echo ""
+
+# 8. Check filesystem
+echo "8. Filesystem check:"
+FILESYSTEM=$(df -T "$NETCDF_FILE" 2>/dev/null | tail -1 | awk '{print $2}')
+MOUNT_POINT=$(df "$NETCDF_FILE" 2>/dev/null | tail -1 | awk '{print $NF}')
+echo " Filesystem type: $FILESYSTEM"
+echo " Mount point: $MOUNT_POINT"
+
+# Check if on /scratch (common on HPC)
+if [[ "$MOUNT_POINT" == *"scratch"* ]]; then
+ echo " ⚠️ File is on /scratch - check quota and purge policies"
+fi
+echo ""
+
+# 9. Check disk space
+echo "9. Disk space check:"
+df -h "$NETCDF_FILE" | tail -1
+echo ""
+
+# 10. Suggest fixes
+echo "========================================="
+echo "DIAGNOSTIC SUMMARY & SUGGESTIONS"
+echo "========================================="
+echo ""
+
+if [ "$FILE_SIZE" -lt 1000000 ]; then
+ echo "⚠️ LIKELY ISSUE: File is corrupted/incomplete (too small)"
+ echo ""
+ echo "SOLUTION:"
+ echo " 1. Delete the file:"
+ echo " rm '$NETCDF_FILE'"
+ echo ""
+ echo " 2. Re-download:"
+ filename=$(basename "$NETCDF_FILE")
+ echo " wget https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/15s/15s_surface_elev_netcdf/$filename"
+ echo ""
+elif ! echo "$HEADER" | grep -qE "^(434446|894844460d0a1a0a)"; then
+ echo "⚠️ LIKELY ISSUE: File is corrupted (invalid magic number)"
+ echo ""
+ echo "SOLUTION: Re-download the file (see above)"
+ echo ""
+else
+ echo "❓ File appears valid, but Python netCDF4 cannot open it."
+ echo ""
+ echo "Possible causes:"
+ echo " 1. HDF5 library version mismatch"
+ echo " 2. NetCDF4 compiled with different HDF5 than runtime"
+ echo " 3. File locking issues (multiple processes)"
+ echo " 4. Filesystem issues (NFS, /scratch)"
+ echo ""
+ echo "Try:"
+ echo " 1. Check loaded modules:"
+ echo " module list"
+ echo ""
+ echo " 2. Try reloading HDF5/NetCDF modules:"
+ echo " module purge"
+ echo " module load netcdf-c hdf5"
+ echo ""
+ echo " 3. Check if file is locked by another process:"
+ echo " lsof '$NETCDF_FILE'"
+ echo ""
+ echo " 4. Copy file to local /tmp and try opening:"
+ echo " cp '$NETCDF_FILE' /tmp/"
+ echo " # Then test with /tmp version"
+fi
+
+echo ""
+echo "========================================="
diff --git a/scripts/download_etopo_with_validation.sh b/scripts/download_etopo_with_validation.sh
new file mode 100755
index 0000000..03cca0b
--- /dev/null
+++ b/scripts/download_etopo_with_validation.sh
@@ -0,0 +1,258 @@
+#!/bin/bash
+# Enhanced ETOPO download script with validation
+# Checks remote file size and validates after download
+# Usage:
+# Download mode: ./download_etopo_with_validation.sh [output_dir]
+# Verify mode: ./download_etopo_with_validation.sh --verify [output_dir]
+
+set -e
+
+# Check for verify mode
+VERIFY_ONLY=false
+if [ "$1" = "--verify" ] || [ "$1" = "-v" ]; then
+ VERIFY_ONLY=true
+ OUTPUT_DIR="${2:-./data/etopo_15s}"
+else
+ OUTPUT_DIR="${1:-./data/etopo_15s}"
+fi
+
+DATA_TYPE="${ETOPO_DATA_TYPE:-surface}"
+if [ "$DATA_TYPE" = "bed" ]; then
+ BASE_URL="https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/15s/15s_bed_elev_netcdf"
+ FILE_SUFFIX="bed"
+else
+ BASE_URL="https://www.ngdc.noaa.gov/thredds/fileServer/global/ETOPO2022/15s/15s_surface_elev_netcdf"
+ FILE_SUFFIX="surface"
+fi
+
+mkdir -p "$OUTPUT_DIR"
+
+if [ "$VERIFY_ONLY" = true ]; then
+ echo "ETOPO 2022 15s Verification Mode"
+else
+ echo "ETOPO 2022 15s Download with Validation"
+fi
+echo "Data type: $DATA_TYPE"
+echo "Directory: $OUTPUT_DIR"
+echo "========================================"
+
+# Function to get remote file size
+get_remote_size() {
+ local url="$1"
+ # Use wget --spider to get headers only
+ local size=$(wget --spider --server-response "$url" 2>&1 | grep -i Content-Length | tail -1 | awk '{print $2}')
+ echo "$size"
+}
+
+# Function to get local file size
+get_local_size() {
+ local file="$1"
+ if [ -f "$file" ]; then
+ stat -f%z "$file" 2>/dev/null || stat -c%s "$file" 2>/dev/null
+ else
+ echo "0"
+ fi
+}
+
+# Function to verify a single tile (no download)
+verify_tile() {
+ local lat="$1"
+ local lon="$2"
+ local filename="ETOPO_2022_v1_15s_${lat}${lon}_${FILE_SUFFIX}.nc"
+ local filepath="${OUTPUT_DIR}/${filename}"
+ local url="${BASE_URL}/${filename}"
+
+ echo -n "Verifying ${lat}${lon}... "
+
+ # Check if file exists locally
+ local local_size=$(get_local_size "$filepath")
+
+ if [ "$local_size" = "0" ]; then
+ echo "✗ Missing"
+ return 1
+ fi
+
+ # Get remote size
+ local remote_size=$(get_remote_size "$url")
+
+ if [ -z "$remote_size" ] || [ "$remote_size" = "0" ]; then
+ echo "⚠️ Cannot verify (server unavailable)"
+ return 2
+ fi
+
+ # Compare sizes
+ if [ "$local_size" = "$remote_size" ]; then
+ echo "✓ Valid ($(($remote_size / 1048576)) MB)"
+ return 0
+ else
+ local local_mb=$(($local_size / 1048576))
+ local remote_mb=$(($remote_size / 1048576))
+ echo "✗ Size mismatch! Local: ${local_mb} MB, Expected: ${remote_mb} MB"
+ return 1
+ fi
+}
+
+# Function to download and validate a single tile
+download_tile() {
+ local lat="$1"
+ local lon="$2"
+ local filename="ETOPO_2022_v1_15s_${lat}${lon}_${FILE_SUFFIX}.nc"
+ local filepath="${OUTPUT_DIR}/${filename}"
+ local url="${BASE_URL}/${filename}"
+
+ # Check if file exists and get sizes
+ local local_size=$(get_local_size "$filepath")
+
+ echo -n "Checking ${lat}${lon}... "
+
+ # Get remote size
+ local remote_size=$(get_remote_size "$url")
+
+ if [ -z "$remote_size" ] || [ "$remote_size" = "0" ]; then
+ echo "⚠️ File not available on server"
+ return 1
+ fi
+
+ # Check if local file matches remote size
+ if [ "$local_size" = "$remote_size" ]; then
+ echo "✓ Already downloaded ($(($remote_size / 1048576)) MB)"
+ return 0
+ fi
+
+ # Download the file
+ echo "Downloading ($(($remote_size / 1048576)) MB)..."
+ if wget -c -O "$filepath" "$url" 2>&1 | grep -v "^--" | grep -v "^Saving" | grep -v "^Length"; then
+ # Verify download
+ local final_size=$(get_local_size "$filepath")
+ if [ "$final_size" = "$remote_size" ]; then
+ echo " ✓ Download verified"
+ return 0
+ else
+ echo " ✗ Size mismatch! Expected: $remote_size, Got: $final_size"
+ echo " Deleting incomplete file..."
+ rm -f "$filepath"
+ return 1
+ fi
+ else
+ echo " ✗ Download failed"
+ rm -f "$filepath"
+ return 1
+ fi
+}
+
+# All latitude/longitude combinations
+declare -a LATS=(N00 N15 N30 N45 N60 N75 N90 S15 S30 S45 S60 S75)
+declare -a LONS=(W180 W165 W150 W135 W120 W105 W090 W075 W060 W045 W030 W015 E000 E015 E030 E045 E060 E075 E090 E105 E120 E135 E150 E165)
+
+# Track statistics
+total_tiles=0
+valid=0
+invalid=0
+missing=0
+failed=0
+
+echo ""
+if [ "$VERIFY_ONLY" = true ]; then
+ echo "Verifying existing files..."
+else
+ echo "Starting download..."
+fi
+echo ""
+
+# Store corrupted files for optional deletion
+declare -a corrupted_files=()
+
+for lat in "${LATS[@]}"; do
+ for lon in "${LONS[@]}"; do
+ total_tiles=$((total_tiles + 1))
+
+ if [ "$VERIFY_ONLY" = true ]; then
+ # Verify mode
+ result=$(verify_tile "$lat" "$lon"; echo $?)
+ case $result in
+ 0)
+ valid=$((valid + 1))
+ ;;
+ 1)
+ invalid=$((invalid + 1))
+ filename="ETOPO_2022_v1_15s_${lat}${lon}_${FILE_SUFFIX}.nc"
+ filepath="${OUTPUT_DIR}/${filename}"
+ if [ -f "$filepath" ]; then
+ corrupted_files+=("$filepath")
+ else
+ missing=$((missing + 1))
+ fi
+ ;;
+ 2)
+ failed=$((failed + 1))
+ ;;
+ esac
+ else
+ # Download mode
+ if download_tile "$lat" "$lon"; then
+ valid=$((valid + 1))
+ else
+ failed=$((failed + 1))
+ fi
+ fi
+ done
+done
+
+echo ""
+echo "========================================"
+if [ "$VERIFY_ONLY" = true ]; then
+ echo "Verification Summary:"
+ echo " Total tiles checked: $total_tiles"
+ echo " Valid files: $valid"
+ echo " Invalid/corrupted: $invalid"
+ echo " Missing files: $missing"
+ echo " Could not verify: $failed"
+
+ if [ $invalid -gt 0 ]; then
+ echo ""
+ echo "⚠️ Found $invalid corrupted/invalid files"
+ echo ""
+ echo "Corrupted files:"
+ for file in "${corrupted_files[@]}"; do
+ echo " - $(basename "$file")"
+ done
+ echo ""
+ read -p "Delete corrupted files and re-download? (yes/no): " delete_confirm
+ if [ "$delete_confirm" = "yes" ]; then
+ for file in "${corrupted_files[@]}"; do
+ echo "Deleting: $(basename "$file")"
+ rm -f "$file"
+ done
+ echo ""
+ echo "Deleted $invalid corrupted files"
+ echo "Now re-run without --verify to download missing files:"
+ echo " $0 $OUTPUT_DIR"
+ fi
+ exit 1
+ elif [ $missing -gt 0 ]; then
+ echo ""
+ echo "⚠️ $missing files are missing"
+ echo "Run without --verify to download them:"
+ echo " $0 $OUTPUT_DIR"
+ exit 1
+ else
+ echo ""
+ echo "✓ All files verified successfully!"
+ exit 0
+ fi
+else
+ echo "Download Summary:"
+ echo " Total tiles attempted: $total_tiles"
+ echo " Successfully validated: $valid"
+ echo " Failed/Not available: $failed"
+ echo ""
+
+ if [ $failed -gt 0 ]; then
+ echo "⚠️ Some tiles failed to download."
+ echo "Re-run this script to retry failed downloads."
+ exit 1
+ else
+ echo "✓ All tiles downloaded and validated successfully!"
+ exit 0
+ fi
+fi
diff --git a/scripts/merge_icon_etopo_outputs.py b/scripts/merge_icon_etopo_outputs.py
new file mode 100644
index 0000000..d102f03
--- /dev/null
+++ b/scripts/merge_icon_etopo_outputs.py
@@ -0,0 +1,351 @@
+#!/usr/bin/env python3
+"""
+Merge ETOPO NetCDF Output Files
+
+This script merges all chunked NetCDF outputs from the ETOPO processing into a single file,
+ensuring that:
+1. All cell IDs (groups) are represented in the merged file
+2. Each cell has an 'is_land' attribute
+3. Missing cells are filled with ocean placeholders (is_land=0)
+"""
+
+import netCDF4
+import numpy as np
+from pathlib import Path
+from tqdm import tqdm
+import sys
+
+
+def get_expected_cell_range(files):
+ """
+ Determine the expected cell range from filenames.
+
+ Parameters
+ ----------
+ files : list of Path
+ List of NetCDF files
+
+ Returns
+ -------
+ tuple
+ (min_cell, max_cell) expected in the dataset
+ """
+ min_cell = float("inf")
+ max_cell = float("-inf")
+
+ for f in files:
+ parts = f.stem.split("_")
+ range_part = parts[-1] # e.g., '00000-00099'
+ start, end = map(int, range_part.split("-"))
+ min_cell = min(min_cell, start)
+ max_cell = max(max_cell, end)
+
+ return int(min_cell), int(max_cell)
+
+
+def collect_all_cells(files):
+ """
+ Collect all cell data from chunked NetCDF files.
+
+ Parameters
+ ----------
+ files : list of Path
+ List of NetCDF files to merge
+
+ Returns
+ -------
+ dict
+ Dictionary mapping cell_id (int) to cell data dict containing:
+ - is_land: int (0 or 1)
+ - clat: float (radians)
+ - clon: float (radians)
+ - cell_area: float or None (m^2)
+ - analysis: dict of arrays (only for land cells)
+ """
+ cell_data = {}
+
+ print("Reading cell data from NetCDF files...")
+ for nc_file in tqdm(files, desc="Processing files"):
+ try:
+ nc = netCDF4.Dataset(nc_file, "r")
+
+ # Iterate over all groups (cell IDs) in this file
+ for group_name in nc.groups.keys():
+ cell_id = int(group_name)
+ group = nc.groups[group_name]
+
+ # Extract cell data
+ is_land = int(group.variables["is_land"][:])
+ clat = float(group.variables["clat"][:])
+ clon = float(group.variables["clon"][:])
+
+ # Extract cell_area if available
+ cell_area = None
+ if "cell_area" in group.variables:
+ cell_area = float(group.variables["cell_area"][:])
+
+ cell_info = {
+ "is_land": is_land,
+ "clat": clat,
+ "clon": clon,
+ "cell_area": cell_area,
+ }
+
+ # For land cells, also extract analysis data
+ if is_land == 1:
+ cell_info["analysis"] = {}
+ for var_name in group.variables.keys():
+ if var_name not in ["is_land", "clat", "clon", "cell_area"]:
+ cell_info["analysis"][var_name] = group.variables[var_name][
+ :
+ ]
+
+ cell_data[cell_id] = cell_info
+
+ nc.close()
+
+ except Exception as e:
+ print(f"Error reading {nc_file.name}: {e}")
+ continue
+
+ return cell_data
+
+
+def create_merged_netcdf(cell_data, output_path, expected_min, expected_max):
+ """
+ Create merged NetCDF file with all cells.
+
+ Parameters
+ ----------
+ cell_data : dict
+ Dictionary of cell data from collect_all_cells()
+ output_path : Path
+ Output file path
+ expected_min : int
+ Expected minimum cell ID
+ expected_max : int
+ Expected maximum cell ID
+ """
+ print(f"\nCreating merged NetCDF file: {output_path}")
+
+ # Create new NetCDF file
+ nc_out = netCDF4.Dataset(output_path, "w", format="NETCDF4")
+
+ # Set global attributes
+ nc_out.title = "ICON ETOPO Global Topography - Merged Output"
+ nc_out.description = "Merged spectral analysis of ETOPO topography on ICON grid"
+ nc_out.source = "pycsa spectral approximation framework"
+
+ # Statistics counters
+ land_cells = 0
+ ocean_cells = 0
+ missing_cells = 0
+
+ print(f"Writing cells {expected_min} to {expected_max}...")
+
+ # Iterate through all expected cells
+ for cell_id in tqdm(range(expected_min, expected_max + 1), desc="Writing cells"):
+ # Create group for this cell
+ grp = nc_out.createGroup(str(cell_id))
+
+ if cell_id in cell_data:
+ # Cell exists in data
+ cell = cell_data[cell_id]
+ is_land = cell["is_land"]
+ clat = cell["clat"]
+ clon = cell["clon"]
+ cell_area = cell.get("cell_area", None)
+
+ if is_land:
+ land_cells += 1
+ else:
+ ocean_cells += 1
+
+ else:
+ # Missing cell - create ocean placeholder
+ print(f"Warning: Cell {cell_id} missing, creating ocean placeholder")
+ is_land = 0
+ clat = 0.0 # Placeholder
+ clon = 0.0 # Placeholder
+ cell_area = None
+ missing_cells += 1
+ ocean_cells += 1
+
+ # Write basic cell attributes (always present)
+ var_is_land = grp.createVariable("is_land", "i4")
+ var_is_land[:] = is_land
+
+ var_clat = grp.createVariable("clat", "f8")
+ var_clat[:] = clat
+ var_clat.units = "radians"
+ var_clat.long_name = "cell center latitude"
+
+ var_clon = grp.createVariable("clon", "f8")
+ var_clon[:] = clon
+ var_clon.units = "radians"
+ var_clon.long_name = "cell center longitude"
+
+ # Write cell_area if available
+ if cell_area is not None:
+ var_cell_area = grp.createVariable("cell_area", "f8")
+ var_cell_area[:] = cell_area
+ var_cell_area.units = "m^2"
+ var_cell_area.long_name = "Area of ICON grid cell"
+
+ # Write analysis data for land cells
+ if is_land and cell_id in cell_data:
+ analysis = cell_data[cell_id]["analysis"]
+ for var_name, var_data in analysis.items():
+ # Create variable with appropriate dimensions
+ if var_data.ndim == 0:
+ # Scalar variable (0-dimensional)
+ var = grp.createVariable(var_name, var_data.dtype)
+ var[:] = var_data
+ elif var_data.ndim == 1:
+ dim_name = f"dim_{var_name}"
+ grp.createDimension(dim_name, var_data.shape[0])
+ var = grp.createVariable(var_name, var_data.dtype, (dim_name,))
+ var[:] = var_data
+ elif var_data.ndim == 2:
+ dim0_name = f"dim0_{var_name}"
+ dim1_name = f"dim1_{var_name}"
+ grp.createDimension(dim0_name, var_data.shape[0])
+ grp.createDimension(dim1_name, var_data.shape[1])
+ var = grp.createVariable(
+ var_name, var_data.dtype, (dim0_name, dim1_name)
+ )
+ var[:] = var_data
+ else:
+ print(
+ f"Warning: Skipping variable {var_name} with unsupported dimensions: {var_data.ndim}"
+ )
+ continue
+
+ nc_out.close()
+
+ # Print statistics
+ print("\n" + "=" * 80)
+ print("MERGE COMPLETE")
+ print("=" * 80)
+ print(f"Output file: {output_path}")
+ print(f"Total cells: {expected_max - expected_min + 1}")
+ print(f" Land cells (is_land=1): {land_cells}")
+ print(f" Ocean cells (is_land=0): {ocean_cells}")
+ if missing_cells > 0:
+ print(f" Missing cells (filled with ocean): {missing_cells}")
+ print(
+ f"\nLand/Ocean ratio: {land_cells}/{ocean_cells} = {land_cells/ocean_cells:.3f}"
+ if ocean_cells > 0
+ else ""
+ )
+ print(f"Land percentage: {100*land_cells/(land_cells+ocean_cells):.2f}%")
+ print("=" * 80)
+
+
+def verify_merged_file(output_path, expected_min, expected_max):
+ """
+ Verify the merged NetCDF file has all cells with is_land attribute.
+
+ Parameters
+ ----------
+ output_path : Path
+ Path to merged NetCDF file
+ expected_min : int
+ Expected minimum cell ID
+ expected_max : int
+ Expected maximum cell ID
+
+ Returns
+ -------
+ bool
+ True if verification passes
+ """
+ print(f"\nVerifying merged file: {output_path}")
+
+ nc = netCDF4.Dataset(output_path, "r")
+
+ expected_cells = set(range(expected_min, expected_max + 1))
+ found_cells = set(int(g) for g in nc.groups.keys())
+
+ # Check all cells present
+ missing = expected_cells - found_cells
+ if missing:
+ print(f"ERROR: Missing cells: {sorted(missing)[:10]}... ({len(missing)} total)")
+ nc.close()
+ return False
+
+ # Check extra cells
+ extra = found_cells - expected_cells
+ if extra:
+ print(f"Warning: Extra cells: {sorted(extra)[:10]}... ({len(extra)} total)")
+
+ # Check is_land attribute and count land vs ocean
+ cells_without_is_land = []
+ land_count = 0
+ ocean_count = 0
+ for group_name in nc.groups.keys():
+ group = nc.groups[group_name]
+ if "is_land" not in group.variables:
+ cells_without_is_land.append(group_name)
+ else:
+ is_land_val = int(group.variables["is_land"][:])
+ if is_land_val == 1:
+ land_count += 1
+ else:
+ ocean_count += 1
+
+ if cells_without_is_land:
+ print(
+ f"ERROR: Cells without is_land attribute: {cells_without_is_land[:10]}... ({len(cells_without_is_land)} total)"
+ )
+ nc.close()
+ return False
+
+ nc.close()
+
+ print("✓ Verification PASSED")
+ print(f" All {len(expected_cells)} cells present")
+ print(f" All cells have 'is_land' attribute")
+ print(f" Land cells (is_land=1): {land_count}")
+ print(f" Ocean cells (is_land=0): {ocean_count}")
+ print(f" Land percentage: {100*land_count/(land_count+ocean_count):.2f}%")
+
+ return True
+
+
+if __name__ == "__main__":
+ # Configuration
+ input_dir = Path("datasets")
+ output_dir = Path("datasets")
+ output_filename = "icon_etopo_global_merged.nc"
+
+ # Find all input files
+ input_files = sorted(input_dir.glob("icon_etopo_global_cells_*.nc"))
+
+ if not input_files:
+ print(f"ERROR: No NetCDF files found in {input_dir}")
+ sys.exit(1)
+
+ print(f"Found {len(input_files)} NetCDF files to merge")
+
+ # Determine expected cell range
+ expected_min, expected_max = get_expected_cell_range(input_files)
+ print(
+ f"Expected cell range: {expected_min} to {expected_max} ({expected_max - expected_min + 1} cells)"
+ )
+
+ # Collect all cell data
+ cell_data = collect_all_cells(input_files)
+ print(f"Collected data for {len(cell_data)} cells")
+
+ # Create merged file
+ output_path = output_dir / output_filename
+ create_merged_netcdf(cell_data, output_path, expected_min, expected_max)
+
+ # Verify merged file
+ if verify_merged_file(output_path, expected_min, expected_max):
+ print(f"\n✓ Successfully created merged file: {output_path}")
+ print(f" Size: {output_path.stat().st_size / (1024**2):.1f} MB")
+ else:
+ print(f"\n✗ Verification failed for: {output_path}")
+ sys.exit(1)
diff --git a/scripts/plot_pacific_detail.py b/scripts/plot_pacific_detail.py
new file mode 100644
index 0000000..9171a87
--- /dev/null
+++ b/scripts/plot_pacific_detail.py
@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+"""
+Detailed Pacific region plot showing island cells more clearly.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap
+from pathlib import Path
+
+# Load data
+data = np.load("outputs/verification/verification_data.npz")
+clat_deg = data["clat_deg"]
+clon_deg = data["clon_deg"]
+land_fractions = data["land_fractions"]
+
+# Create colormap
+colors_gradient = [
+ "#0033aa",
+ "#0066cc",
+ "#3399ff",
+ "#66ccff",
+ "#99ff99",
+ "#66cc66",
+ "#339933",
+ "#006600",
+]
+cmap_land_ocean = LinearSegmentedColormap.from_list(
+ "land_ocean", colors_gradient, N=256
+)
+
+# Define Pacific regions
+regions = {
+ "Hawaii": (15, 25, -165, -150),
+ "Micronesia": (0, 15, 130, 170),
+ "Polynesia": (-30, 0, -180, -130),
+ "Indonesia": (-10, 10, 95, 140),
+}
+
+fig, axes = plt.subplots(2, 2, figsize=(16, 12))
+axes = axes.flatten()
+
+for idx, (name, (lat_min, lat_max, lon_min, lon_max)) in enumerate(regions.items()):
+ ax = axes[idx]
+
+ # Find cells in region
+ mask = (
+ (clat_deg >= lat_min)
+ & (clat_deg <= lat_max)
+ & (clon_deg >= lon_min)
+ & (clon_deg <= lon_max)
+ )
+
+ # Separate by land fraction
+ pure_ocean = mask & (land_fractions < 0.05)
+ has_land = mask & (land_fractions >= 0.05)
+
+ # Plot
+ if np.any(pure_ocean):
+ ax.scatter(
+ clon_deg[pure_ocean],
+ clat_deg[pure_ocean],
+ c="#E0F2F7",
+ s=80,
+ alpha=0.5,
+ edgecolors="gray",
+ linewidths=0.3,
+ label="Ocean (<5% land)",
+ )
+
+ if np.any(has_land):
+ sc = ax.scatter(
+ clon_deg[has_land],
+ clat_deg[has_land],
+ c=land_fractions[has_land],
+ cmap=cmap_land_ocean,
+ s=120,
+ alpha=0.95,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="black",
+ linewidths=0.8,
+ )
+
+ # Add cell numbers for high land fraction
+ high_land = has_land & (land_fractions > 0.3)
+ for cell_idx in np.where(high_land)[0]:
+ ax.text(
+ clon_deg[cell_idx],
+ clat_deg[cell_idx],
+ f"{100*land_fractions[cell_idx]:.0f}%",
+ fontsize=7,
+ ha="center",
+ va="center",
+ bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
+ )
+
+ # Format
+ ax.set_xlabel("Longitude [°]", fontsize=10)
+ ax.set_ylabel("Latitude [°]", fontsize=10)
+ ax.set_title(
+ f"{name} Region\n{np.sum(has_land)} cells with ≥5% land, "
+ f"{np.sum(pure_ocean)} pure ocean cells",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax.grid(True, alpha=0.3)
+ ax.set_xlim(lon_min, lon_max)
+ ax.set_ylim(lat_min, lat_max)
+
+ if idx == 0:
+ ax.legend(loc="best", fontsize=8)
+
+plt.tight_layout()
+
+# Add colorbar at the bottom
+cbar_ax = fig.add_axes([0.25, -0.02, 0.5, 0.02]) # [left, bottom, width, height]
+cbar = fig.colorbar(sc, cax=cbar_ax, orientation="horizontal")
+cbar.set_label("Land Fraction (0=Ocean, 1=Land)", fontsize=11)
+
+output_file = Path("outputs/verification/pacific_islands_detail.png")
+plt.savefig(output_file, dpi=200, bbox_inches="tight")
+print(f"Saved: {output_file}")
+
+# Print statistics
+print("\nPacific Island Statistics:")
+for name, (lat_min, lat_max, lon_min, lon_max) in regions.items():
+ mask = (
+ (clat_deg >= lat_min)
+ & (clat_deg <= lat_max)
+ & (clon_deg >= lon_min)
+ & (clon_deg <= lon_max)
+ )
+ has_land = mask & (land_fractions >= 0.05)
+
+ if np.any(has_land):
+ print(f"\n{name}:")
+ print(f" Cells with land: {np.sum(has_land)}")
+ print(f" Max land fraction: {np.max(land_fractions[has_land]):.1%}")
+ print(f" Mean land fraction: {np.mean(land_fractions[has_land]):.1%}")
diff --git a/scripts/plot_verification_improved.py b/scripts/plot_verification_improved.py
new file mode 100755
index 0000000..cf0658f
--- /dev/null
+++ b/scripts/plot_verification_improved.py
@@ -0,0 +1,360 @@
+#!/usr/bin/env python3
+"""
+Improved plotting script for ICON ETOPO verification data.
+Loads the saved verification data and creates enhanced visualizations.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap
+from pathlib import Path
+
+
+def load_verification_data():
+ """Load the verification data from npz file."""
+ data_file = Path("outputs/verification/verification_data.npz")
+
+ if not data_file.exists():
+ print(f"Error: {data_file} not found.")
+ print("Please run verify_icon_etopo_land_ocean.py first.")
+ return None
+
+ data = np.load(data_file)
+ print(f"Loaded verification data:")
+ print(f" Total cells: {data['n_cells']}")
+ print(f" Land cells: {data['land_count']}")
+ print(f" Ocean cells: {data['ocean_count']}")
+ print(f" ETOPO coarse-graining: {data['etopo_cg']}")
+ print()
+
+ return data
+
+
+def create_improved_plots(data, output_dir):
+ """Create improved visualization plots."""
+
+ clat_deg = data["clat_deg"]
+ clon_deg = data["clon_deg"]
+ land_cells = data["land_cells"]
+ ocean_cells = data["ocean_cells"]
+ land_fractions = data["land_fractions"]
+ land_count = data["land_count"]
+ ocean_count = data["ocean_count"]
+
+ output_dir = Path(output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Convert to Mollweide projection coordinates
+ lon_plot = np.deg2rad(clon_deg)
+ lon_plot[lon_plot > np.pi] -= 2 * np.pi
+ lat_plot = np.deg2rad(clat_deg)
+
+ # ========================================================================
+ # Figure 1: Multiple views with different thresholds
+ # ========================================================================
+ fig = plt.figure(figsize=(20, 12))
+
+ # Custom colormap from blue (ocean) to green (land)
+ colors_gradient = [
+ "#0033aa",
+ "#0066cc",
+ "#3399ff",
+ "#66ccff",
+ "#99ff99",
+ "#66cc66",
+ "#339933",
+ "#006600",
+ ]
+ cmap_land_ocean = LinearSegmentedColormap.from_list(
+ "land_ocean", colors_gradient, N=256
+ )
+
+ # Plot 1: Continuous land fraction (original)
+ ax1 = fig.add_subplot(231, projection="mollweide")
+ scatter1 = ax1.scatter(
+ lon_plot,
+ lat_plot,
+ c=land_fractions,
+ cmap=cmap_land_ocean,
+ s=5,
+ alpha=0.9,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="none",
+ )
+ cbar1 = plt.colorbar(
+ scatter1, ax=ax1, orientation="horizontal", pad=0.05, shrink=0.7
+ )
+ cbar1.set_label("Land Fraction", fontsize=10)
+ ax1.set_title(
+ f"Continuous Land Fraction\n(All gradations)", fontsize=11, fontweight="bold"
+ )
+ ax1.grid(True, alpha=0.3)
+
+ # Plot 2: Binary classification (>50% land = green, else blue)
+ ax2 = fig.add_subplot(232, projection="mollweide")
+ binary_colors = np.where(land_fractions > 0.5, "#228B22", "#1E90FF")
+ ax2.scatter(lon_plot, lat_plot, c=binary_colors, s=5, alpha=0.9, edgecolors="none")
+ ax2.set_title(
+ f"Binary: >50% Land = Green\nLand: {land_count}, Ocean: {ocean_count}",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax2.grid(True, alpha=0.3)
+
+ # Plot 3: Highlight mixed coastal cells (10-90% land)
+ ax3 = fig.add_subplot(233, projection="mollweide")
+ coastal_mask = (land_fractions > 0.1) & (land_fractions < 0.9)
+ pure_land_mask = land_fractions >= 0.9
+ pure_ocean_mask = land_fractions <= 0.1
+
+ # Plot pure ocean (light blue), pure land (green), coastal (red)
+ if np.any(pure_ocean_mask):
+ ax3.scatter(
+ lon_plot[pure_ocean_mask],
+ lat_plot[pure_ocean_mask],
+ c="#B0E0E6",
+ s=4,
+ alpha=0.5,
+ label="Pure Ocean (<10% land)",
+ )
+ if np.any(pure_land_mask):
+ ax3.scatter(
+ lon_plot[pure_land_mask],
+ lat_plot[pure_land_mask],
+ c="#90EE90",
+ s=4,
+ alpha=0.5,
+ label="Pure Land (>90% land)",
+ )
+ if np.any(coastal_mask):
+ ax3.scatter(
+ lon_plot[coastal_mask],
+ lat_plot[coastal_mask],
+ c="#FF6347",
+ s=8,
+ alpha=0.9,
+ label=f"Mixed Coastal (10-90% land)",
+ )
+
+ ax3.set_title(
+ f"Coastal/Mixed Cells Highlighted\n{np.sum(coastal_mask)} mixed cells",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax3.legend(loc="lower left", fontsize=8, markerscale=2)
+ ax3.grid(True, alpha=0.3)
+
+ # Plot 4: Grid structure (all cells same size/color)
+ ax4 = fig.add_subplot(234, projection="mollweide")
+ ax4.scatter(lon_plot, lat_plot, c="gray", s=2, alpha=0.6)
+ ax4.set_title(
+ f"ICON R2B4 Grid Structure\n{len(clat_deg)} cells total",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax4.grid(True, alpha=0.3)
+
+ # Plot 5: Only cells with ANY land (>5% threshold)
+ ax5 = fig.add_subplot(235, projection="mollweide")
+ any_land_mask = land_fractions > 0.05
+ if np.any(~any_land_mask):
+ ax5.scatter(
+ lon_plot[~any_land_mask],
+ lat_plot[~any_land_mask],
+ c="#1E90FF",
+ s=3,
+ alpha=0.3,
+ label="Pure Ocean",
+ )
+ if np.any(any_land_mask):
+ scatter5 = ax5.scatter(
+ lon_plot[any_land_mask],
+ lat_plot[any_land_mask],
+ c=land_fractions[any_land_mask],
+ cmap=cmap_land_ocean,
+ s=8,
+ alpha=0.9,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="none",
+ label="Has Land",
+ )
+ ax5.set_title(
+ f"Cells with >5% Land Highlighted\n{np.sum(any_land_mask)} cells with land",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax5.legend(loc="lower left", fontsize=8)
+ ax5.grid(True, alpha=0.3)
+
+ # Plot 6: Latitude distribution
+ ax6 = fig.add_subplot(236)
+ lat_bins = np.linspace(-90, 90, 37)
+
+ # Create histogram for different land fraction ranges
+ pure_ocean_hist, _ = np.histogram(clat_deg[land_fractions <= 0.1], bins=lat_bins)
+ coastal_hist, _ = np.histogram(clat_deg[coastal_mask], bins=lat_bins)
+ pure_land_hist, _ = np.histogram(clat_deg[land_fractions >= 0.9], bins=lat_bins)
+
+ bin_centers = (lat_bins[:-1] + lat_bins[1:]) / 2
+ width = 5
+
+ ax6.barh(
+ bin_centers,
+ pure_ocean_hist,
+ height=width,
+ color="#1E90FF",
+ alpha=0.6,
+ label="Pure Ocean (≤10% land)",
+ )
+ ax6.barh(
+ bin_centers,
+ coastal_hist,
+ height=width,
+ left=pure_ocean_hist,
+ color="#FF6347",
+ alpha=0.6,
+ label="Coastal (10-90% land)",
+ )
+ ax6.barh(
+ bin_centers,
+ pure_land_hist,
+ height=width,
+ left=pure_ocean_hist + coastal_hist,
+ color="#228B22",
+ alpha=0.6,
+ label="Pure Land (≥90% land)",
+ )
+
+ ax6.set_xlabel("Number of cells", fontsize=10)
+ ax6.set_ylabel("Latitude [degrees]", fontsize=10)
+ ax6.set_title("Cell Distribution by Latitude", fontsize=11, fontweight="bold")
+ ax6.legend(fontsize=8)
+ ax6.grid(True, alpha=0.3)
+
+ plt.tight_layout()
+
+ output_file = output_dir / "improved_verification_plots.png"
+ plt.savefig(output_file, dpi=150, bbox_inches="tight")
+ print(f"Saved: {output_file}")
+ plt.close()
+
+ # ========================================================================
+ # Figure 2: Pacific region zoom
+ # ========================================================================
+ fig2 = plt.figure(figsize=(16, 8))
+
+ # Define Pacific region
+ pacific_mask = (
+ (clat_deg >= -30)
+ & (clat_deg <= 30)
+ & (
+ ((clon_deg >= 120) & (clon_deg <= 180))
+ | ((clon_deg >= -180) & (clon_deg <= -100))
+ )
+ )
+
+ # Plot 1: Pacific overview with land fraction
+ ax1 = fig2.add_subplot(121)
+ scatter_pac = ax1.scatter(
+ clon_deg[pacific_mask],
+ clat_deg[pacific_mask],
+ c=land_fractions[pacific_mask],
+ cmap=cmap_land_ocean,
+ s=20,
+ alpha=0.9,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="gray",
+ linewidths=0.3,
+ )
+ cbar = plt.colorbar(scatter_pac, ax=ax1)
+ cbar.set_label("Land Fraction", fontsize=10)
+ ax1.set_xlabel("Longitude [degrees]", fontsize=10)
+ ax1.set_ylabel("Latitude [degrees]", fontsize=10)
+ ax1.set_title(
+ "Pacific Region: Land Fraction\n(Many islands are correctly detected)",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax1.grid(True, alpha=0.3)
+ ax1.set_xlim([120, -100])
+
+ # Plot 2: Pacific with only significant land (>20%)
+ ax2 = fig2.add_subplot(122)
+ pacific_ocean = pacific_mask & (land_fractions <= 0.2)
+ pacific_land = pacific_mask & (land_fractions > 0.2)
+
+ if np.any(pacific_ocean):
+ ax2.scatter(
+ clon_deg[pacific_ocean],
+ clat_deg[pacific_ocean],
+ c="#1E90FF",
+ s=10,
+ alpha=0.4,
+ label="Ocean (≤20% land)",
+ )
+ if np.any(pacific_land):
+ ax2.scatter(
+ clon_deg[pacific_land],
+ clat_deg[pacific_land],
+ c=land_fractions[pacific_land],
+ cmap=cmap_land_ocean,
+ s=30,
+ alpha=0.9,
+ vmin=0.2,
+ vmax=1.0,
+ edgecolors="black",
+ linewidths=0.5,
+ label="Land (>20% land)",
+ )
+
+ ax2.set_xlabel("Longitude [degrees]", fontsize=10)
+ ax2.set_ylabel("Latitude [degrees]", fontsize=10)
+ ax2.set_title(
+ f"Pacific: Cells with >20% Land\n{np.sum(pacific_land)} cells",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax2.legend(fontsize=9)
+ ax2.grid(True, alpha=0.3)
+ ax2.set_xlim([120, -100])
+
+ plt.tight_layout()
+
+ output_file2 = output_dir / "pacific_region_detail.png"
+ plt.savefig(output_file2, dpi=150, bbox_inches="tight")
+ print(f"Saved: {output_file2}")
+ plt.close()
+
+ # Print statistics
+ print("\n" + "=" * 80)
+ print("STATISTICS")
+ print("=" * 80)
+ print(f"Pure ocean cells (≤10% land): {np.sum(land_fractions <= 0.1)}")
+ print(f"Coastal/mixed cells (10-90% land): {np.sum(coastal_mask)}")
+ print(f"Pure land cells (≥90% land): {np.sum(land_fractions >= 0.9)}")
+ print()
+ print(f"Mean land fraction: {np.mean(land_fractions):.3f}")
+ print(f"Median land fraction: {np.median(land_fractions):.3f}")
+ print()
+ print(f"Pacific region cells: {np.sum(pacific_mask)}")
+ print(f"Pacific cells with >20% land: {np.sum(pacific_land)}")
+ print(f"Pacific land fraction: {np.mean(land_fractions[pacific_mask]):.3f}")
+ print("=" * 80)
+
+
+if __name__ == "__main__":
+ print("=" * 80)
+ print("IMPROVED VERIFICATION PLOTTING")
+ print("=" * 80)
+ print()
+
+ data = load_verification_data()
+
+ if data is not None:
+ output_dir = Path("outputs") / "verification"
+ create_improved_plots(data, output_dir)
+ print("\n✓ Improved plots created successfully!")
+ print(f" Location: {output_dir}")
diff --git a/scripts/verify_icon_etopo_land_ocean.py b/scripts/verify_icon_etopo_land_ocean.py
new file mode 100644
index 0000000..9ef8e50
--- /dev/null
+++ b/scripts/verify_icon_etopo_land_ocean.py
@@ -0,0 +1,642 @@
+#!/usr/bin/env python3
+"""
+Verify ETOPO Land/Ocean Cell Counts
+
+This script loads the ICON grid and ETOPO topography data, counts how many
+cells are land vs ocean, and creates comprehensive plots.
+
+Usage:
+ python verify_icon_etopo_land_ocean.py # Full verification + plotting
+ python verify_icon_etopo_land_ocean.py --plot-only # Load saved data and plot only
+"""
+
+import os
+
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+
+import sys
+import argparse
+import numpy as np
+import matplotlib.pyplot as plt
+from matplotlib.colors import TwoSlopeNorm, LinearSegmentedColormap
+import matplotlib.colors as mcolors
+from pathlib import Path
+
+
+def get_topo_colormap():
+ """
+ Create a topography colormap with blue for ocean (< 0m) and terrain colors for land (> 0m).
+ """
+ # Ocean colors (blue shades from deep to shallow)
+ ocean_colors = plt.cm.Blues_r(np.linspace(0.4, 0.95, 120))
+
+ # Smooth transition zone around sea level
+ last_ocean = plt.cm.Blues_r(0.95)
+ first_land = plt.cm.terrain(0.25)
+
+ # Create smooth blend from ocean to land
+ transition_colors = np.zeros((16, 4))
+ for i in range(4): # RGBA channels
+ transition_colors[:, i] = np.linspace(last_ocean[i], first_land[i], 16)
+
+ # Land colors (terrain-like: green to brown to white)
+ land_colors = plt.cm.terrain(np.linspace(0.28, 1.0, 120))
+
+ # Combine: 120 ocean + 16 transition + 120 land = 256 total
+ colors = np.vstack((ocean_colors, transition_colors, land_colors))
+ return mcolors.LinearSegmentedColormap.from_list("topo", colors)
+
+
+def count_land_ocean_cells(grid, params, reader):
+ """
+ Count how many cells in the ICON grid are land vs ocean based on ETOPO data.
+ Also computes land fraction for each cell for gradient visualization.
+
+ Parameters
+ ----------
+ grid : grid object
+ ICON grid (in degrees)
+ params : params object
+ Parameters with ETOPO settings
+ reader : ncdata object
+ Data reader
+
+ Returns
+ -------
+ tuple
+ (land_count, ocean_count, land_cells, ocean_cells, land_fractions)
+ land_cells and ocean_cells are lists of cell indices
+ land_fractions is array of land fraction [0-1] for each cell
+ """
+ n_cells = grid.clat.size
+ land_cells = []
+ ocean_cells = []
+ land_fractions = np.zeros(n_cells) # Store land fraction for each cell
+
+ print(f"Checking {n_cells} cells for land/ocean classification...")
+
+ for c_idx in range(n_cells):
+ if c_idx % 1000 == 0:
+ print(f" Processing cell {c_idx}/{n_cells}...")
+
+ topo = var.topo_cell()
+
+ lat_verts = grid.clat_vertices[c_idx]
+ lon_verts = grid.clon_vertices[c_idx]
+
+ # Determine lat/lon extents
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ # Load topography data
+ etopo_reader = reader.read_etopo_topo(None, params, is_parallel=True)
+ etopo_reader.get_topo(topo)
+
+ # Clip deep bathymetry to -500m
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+ topo.gen_mgrids()
+
+ # Handle dateline crossing
+ if etopo_reader.split_EW:
+ lon_verts = lon_verts.copy()
+ lon_verts[lon_verts < 0.0] += 360.0
+
+ # Process vertices for CSA
+ lat_verts, lon_verts = utils.handle_latlon_expansion(
+ lat_verts, lon_verts, lat_expand=0.0, lon_expand=0.0
+ )
+
+ # Initialize cell objects
+ tri_idx = 0
+ cell = var.topo_cell()
+ tri = var.obj()
+
+ # Set up triangles
+ clon_vertices = np.array([lon_verts])
+ clat_vertices = np.array([lat_verts])
+ ncells = 1
+ nv = clon_vertices[0].size
+
+ triangles = np.zeros((ncells, nv, 2))
+ triangles[0, :, 0] = clon_vertices[0, :]
+ triangles[0, :, 1] = clat_vertices[0, :]
+
+ tri.tri_lon_verts = triangles[:, :, 0]
+ tri.tri_lat_verts = triangles[:, :, 1]
+
+ simplex_lat = tri.tri_lat_verts[tri_idx]
+ simplex_lon = tri.tri_lon_verts[tri_idx]
+
+ # Check if land (binary classification)
+ is_land_cell = utils.is_land(cell, simplex_lat, simplex_lon, topo)
+
+ # Calculate land fraction (fraction of cell with elevation > 0m)
+ land_points = np.sum(cell.topo > 0.0)
+ total_points = cell.topo.size
+ land_fractions[c_idx] = land_points / total_points if total_points > 0 else 0.0
+
+ if is_land_cell:
+ land_cells.append(c_idx)
+ else:
+ ocean_cells.append(c_idx)
+
+ return len(land_cells), len(ocean_cells), land_cells, ocean_cells, land_fractions
+
+
+def create_comprehensive_plots(
+ clat_deg, clon_deg, land_cells, ocean_cells, land_fractions, output_dir
+):
+ """
+ Create comprehensive plots of land/ocean classification.
+
+ Parameters
+ ----------
+ clat_deg : array
+ Cell latitudes in degrees
+ clon_deg : array
+ Cell longitudes in degrees
+ land_cells : list
+ List of land cell indices
+ ocean_cells : list
+ List of ocean cell indices
+ land_fractions : array
+ Array of land fraction [0-1] for each cell
+ output_dir : Path
+ Output directory for plots
+ """
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ land_count = len(land_cells)
+ ocean_count = len(ocean_cells)
+
+ # Convert to Mollweide projection coordinates
+ lon_plot = np.deg2rad(clon_deg)
+ lon_plot[lon_plot > np.pi] -= 2 * np.pi
+ lat_plot = np.deg2rad(clat_deg)
+
+ # Custom colormap from blue (ocean) to green (land)
+ colors_gradient = [
+ "#0033aa",
+ "#0066cc",
+ "#3399ff",
+ "#66ccff",
+ "#99ff99",
+ "#66cc66",
+ "#339933",
+ "#006600",
+ ]
+ cmap_land_ocean = LinearSegmentedColormap.from_list(
+ "land_ocean", colors_gradient, N=256
+ )
+
+ # ========================================================================
+ # Figure 1: Multiple global views with different thresholds
+ # ========================================================================
+ print(" Creating global overview plots...")
+ fig = plt.figure(figsize=(20, 12))
+
+ # Plot 1: Continuous land fraction
+ ax1 = fig.add_subplot(231, projection="mollweide")
+ scatter1 = ax1.scatter(
+ lon_plot,
+ lat_plot,
+ c=land_fractions,
+ cmap=cmap_land_ocean,
+ s=5,
+ alpha=0.9,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="none",
+ )
+ cbar1 = plt.colorbar(
+ scatter1, ax=ax1, orientation="horizontal", pad=0.05, shrink=0.7
+ )
+ cbar1.set_label("Land Fraction", fontsize=10)
+ ax1.set_title(
+ f"Continuous Land Fraction\n(All gradations)", fontsize=11, fontweight="bold"
+ )
+ ax1.grid(True, alpha=0.3)
+
+ # Plot 2: Binary classification (>50% land = green, else blue)
+ ax2 = fig.add_subplot(232, projection="mollweide")
+ binary_colors = np.where(land_fractions > 0.5, "#228B22", "#1E90FF")
+ ax2.scatter(lon_plot, lat_plot, c=binary_colors, s=5, alpha=0.9, edgecolors="none")
+ ax2.set_title(
+ f"Binary: >50% Land = Green\nLand: {land_count}, Ocean: {ocean_count}",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax2.grid(True, alpha=0.3)
+
+ # Plot 3: Highlight mixed coastal cells (10-90% land)
+ ax3 = fig.add_subplot(233, projection="mollweide")
+ coastal_mask = (land_fractions > 0.1) & (land_fractions < 0.9)
+ pure_land_mask = land_fractions >= 0.9
+ pure_ocean_mask = land_fractions <= 0.1
+
+ if np.any(pure_ocean_mask):
+ ax3.scatter(
+ lon_plot[pure_ocean_mask],
+ lat_plot[pure_ocean_mask],
+ c="#B0E0E6",
+ s=4,
+ alpha=0.5,
+ label="Pure Ocean (<10% land)",
+ )
+ if np.any(pure_land_mask):
+ ax3.scatter(
+ lon_plot[pure_land_mask],
+ lat_plot[pure_land_mask],
+ c="#90EE90",
+ s=4,
+ alpha=0.5,
+ label="Pure Land (>90% land)",
+ )
+ if np.any(coastal_mask):
+ ax3.scatter(
+ lon_plot[coastal_mask],
+ lat_plot[coastal_mask],
+ c="#FF6347",
+ s=8,
+ alpha=0.9,
+ label=f"Mixed Coastal (10-90% land)",
+ )
+
+ ax3.set_title(
+ f"Coastal/Mixed Cells Highlighted\n{np.sum(coastal_mask)} mixed cells",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax3.legend(loc="lower left", fontsize=8, markerscale=2)
+ ax3.grid(True, alpha=0.3)
+
+ # Plot 4: Grid structure
+ ax4 = fig.add_subplot(234, projection="mollweide")
+ ax4.scatter(lon_plot, lat_plot, c="gray", s=2, alpha=0.6)
+ ax4.set_title(
+ f"ICON R2B4 Grid Structure\n{len(clat_deg)} cells total",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax4.grid(True, alpha=0.3)
+
+ # Plot 5: Only cells with ANY land (>5% threshold)
+ ax5 = fig.add_subplot(235, projection="mollweide")
+ any_land_mask = land_fractions > 0.05
+ if np.any(~any_land_mask):
+ ax5.scatter(
+ lon_plot[~any_land_mask],
+ lat_plot[~any_land_mask],
+ c="#1E90FF",
+ s=3,
+ alpha=0.3,
+ label="Pure Ocean",
+ )
+ if np.any(any_land_mask):
+ scatter5 = ax5.scatter(
+ lon_plot[any_land_mask],
+ lat_plot[any_land_mask],
+ c=land_fractions[any_land_mask],
+ cmap=cmap_land_ocean,
+ s=8,
+ alpha=0.9,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="none",
+ label="Has Land",
+ )
+ ax5.set_title(
+ f"Cells with >5% Land Highlighted\n{np.sum(any_land_mask)} cells with land",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax5.legend(loc="lower left", fontsize=8)
+ ax5.grid(True, alpha=0.3)
+
+ # Plot 6: Latitude distribution
+ ax6 = fig.add_subplot(236)
+ lat_bins = np.linspace(-90, 90, 37)
+
+ pure_ocean_hist, _ = np.histogram(clat_deg[land_fractions <= 0.1], bins=lat_bins)
+ coastal_hist, _ = np.histogram(clat_deg[coastal_mask], bins=lat_bins)
+ pure_land_hist, _ = np.histogram(clat_deg[land_fractions >= 0.9], bins=lat_bins)
+
+ bin_centers = (lat_bins[:-1] + lat_bins[1:]) / 2
+ width = 5
+
+ ax6.barh(
+ bin_centers,
+ pure_ocean_hist,
+ height=width,
+ color="#1E90FF",
+ alpha=0.6,
+ label="Pure Ocean (≤10% land)",
+ )
+ ax6.barh(
+ bin_centers,
+ coastal_hist,
+ height=width,
+ left=pure_ocean_hist,
+ color="#FF6347",
+ alpha=0.6,
+ label="Coastal (10-90% land)",
+ )
+ ax6.barh(
+ bin_centers,
+ pure_land_hist,
+ height=width,
+ left=pure_ocean_hist + coastal_hist,
+ color="#228B22",
+ alpha=0.6,
+ label="Pure Land (≥90% land)",
+ )
+
+ ax6.set_xlabel("Number of cells", fontsize=10)
+ ax6.set_ylabel("Latitude [degrees]", fontsize=10)
+ ax6.set_title("Cell Distribution by Latitude", fontsize=11, fontweight="bold")
+ ax6.legend(fontsize=8)
+ ax6.grid(True, alpha=0.3)
+
+ plt.tight_layout()
+
+ output_file = output_dir / "improved_verification_plots.png"
+ plt.savefig(output_file, dpi=150, bbox_inches="tight")
+ print(f" Saved: {output_file}")
+ plt.close()
+
+ # ========================================================================
+ # Figure 2: Pacific region details
+ # ========================================================================
+ print(" Creating Pacific region detail plots...")
+
+ regions = {
+ "Hawaii": (15, 25, -165, -150),
+ "Micronesia": (0, 15, 130, 170),
+ "Polynesia": (-30, 0, -180, -130),
+ "Indonesia": (-10, 10, 95, 140),
+ }
+
+ fig2, axes = plt.subplots(2, 2, figsize=(16, 12))
+ axes = axes.flatten()
+
+ for idx, (name, (lat_min, lat_max, lon_min, lon_max)) in enumerate(regions.items()):
+ ax = axes[idx]
+
+ # Find cells in region
+ mask = (
+ (clat_deg >= lat_min)
+ & (clat_deg <= lat_max)
+ & (clon_deg >= lon_min)
+ & (clon_deg <= lon_max)
+ )
+
+ # Separate by land fraction
+ pure_ocean = mask & (land_fractions < 0.05)
+ has_land = mask & (land_fractions >= 0.05)
+
+ # Plot
+ if np.any(pure_ocean):
+ ax.scatter(
+ clon_deg[pure_ocean],
+ clat_deg[pure_ocean],
+ c="#E0F2F7",
+ s=80,
+ alpha=0.5,
+ edgecolors="gray",
+ linewidths=0.3,
+ label="Ocean (<5% land)",
+ )
+
+ sc = None # Initialize scatter plot variable
+ if np.any(has_land):
+ sc = ax.scatter(
+ clon_deg[has_land],
+ clat_deg[has_land],
+ c=land_fractions[has_land],
+ cmap=cmap_land_ocean,
+ s=120,
+ alpha=0.95,
+ vmin=0.0,
+ vmax=1.0,
+ edgecolors="black",
+ linewidths=0.8,
+ )
+
+ # Add cell percentages for high land fraction
+ high_land = has_land & (land_fractions > 0.3)
+ for cell_idx in np.where(high_land)[0]:
+ ax.text(
+ clon_deg[cell_idx],
+ clat_deg[cell_idx],
+ f"{100*land_fractions[cell_idx]:.0f}%",
+ fontsize=7,
+ ha="center",
+ va="center",
+ bbox=dict(boxstyle="round,pad=0.3", facecolor="white", alpha=0.7),
+ )
+
+ # Format
+ ax.set_xlabel("Longitude [°]", fontsize=10)
+ ax.set_ylabel("Latitude [°]", fontsize=10)
+ ax.set_title(
+ f"{name} Region\n{np.sum(has_land)} cells with ≥5% land, "
+ f"{np.sum(pure_ocean)} pure ocean cells",
+ fontsize=11,
+ fontweight="bold",
+ )
+ ax.grid(True, alpha=0.3)
+ ax.set_xlim(lon_min, lon_max)
+ ax.set_ylim(lat_min, lat_max)
+
+ if idx == 0:
+ ax.legend(loc="best", fontsize=8)
+
+ plt.tight_layout()
+
+ # Add colorbar at the bottom (if we have scatter data)
+ if sc is not None:
+ cbar_ax = fig2.add_axes([0.25, -0.02, 0.5, 0.02])
+ cbar = fig2.colorbar(sc, cax=cbar_ax, orientation="horizontal")
+ cbar.set_label("Land Fraction (0=Ocean, 1=Land)", fontsize=11)
+
+ output_file2 = output_dir / "pacific_islands_detail.png"
+ plt.savefig(output_file2, dpi=200, bbox_inches="tight")
+ print(f" Saved: {output_file2}")
+ plt.close()
+
+ # Print statistics
+ print("\n" + "=" * 80)
+ print("STATISTICS")
+ print("=" * 80)
+ print(f"Pure ocean cells (≤10% land): {np.sum(land_fractions <= 0.1)}")
+ print(f"Coastal/mixed cells (10-90% land): {np.sum(coastal_mask)}")
+ print(f"Pure land cells (≥90% land): {np.sum(land_fractions >= 0.9)}")
+ print()
+ print(f"Mean land fraction: {np.mean(land_fractions):.3f}")
+ print(f"Median land fraction: {np.median(land_fractions):.3f}")
+ print()
+
+ # Pacific statistics
+ for name, (lat_min, lat_max, lon_min, lon_max) in regions.items():
+ mask = (
+ (clat_deg >= lat_min)
+ & (clat_deg <= lat_max)
+ & (clon_deg >= lon_min)
+ & (clon_deg <= lon_max)
+ )
+ has_land = mask & (land_fractions >= 0.05)
+
+ if np.any(has_land):
+ print(f"{name}:")
+ print(f" Cells with land: {np.sum(has_land)}")
+ print(f" Max land fraction: {np.max(land_fractions[has_land]):.1%}")
+ print(f" Mean land fraction: {np.mean(land_fractions[has_land]):.1%}")
+
+ print("=" * 80)
+
+
+def load_saved_data(data_file):
+ """Load previously saved verification data."""
+ if not data_file.exists():
+ print(f"Error: {data_file} not found.")
+ print("Please run verification first without --plot-only flag.")
+ sys.exit(1)
+
+ data = np.load(data_file)
+ print(f"Loaded verification data from: {data_file}")
+ print(f" Total cells: {data['n_cells']}")
+ print(f" Land cells: {data['land_count']}")
+ print(f" Ocean cells: {data['ocean_count']}")
+ print(f" ETOPO coarse-graining: {data['etopo_cg']}")
+ print()
+
+ return (
+ data["clat_deg"],
+ data["clon_deg"],
+ list(data["land_cells"]),
+ list(data["ocean_cells"]),
+ data["land_fractions"],
+ )
+
+
+if __name__ == "__main__":
+ # Parse command line arguments
+ parser = argparse.ArgumentParser(
+ description="Verify ETOPO land/ocean classification and create plots",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ python verify_icon_etopo_land_ocean.py # Full verification + plotting
+ python verify_icon_etopo_land_ocean.py --plot-only # Load saved data and plot only
+ """,
+ )
+ parser.add_argument(
+ "--plot-only",
+ action="store_true",
+ help="Only create plots from saved data (skip verification)",
+ )
+ args = parser.parse_args()
+
+ print("=" * 80)
+ print("ETOPO LAND/OCEAN VERIFICATION")
+ print("=" * 80)
+
+ output_dir = Path("outputs") / "verification"
+ data_file = output_dir / "verification_data.npz"
+
+ if args.plot_only:
+ # Plot-only mode: Load saved data
+ print("\nMode: PLOT ONLY (loading saved data)")
+ print("=" * 80)
+ clat_deg, clon_deg, land_cells, ocean_cells, land_fractions = load_saved_data(
+ data_file
+ )
+
+ else:
+ # Full verification mode
+ print("\nMode: FULL VERIFICATION (compute + save + plot)")
+ print("=" * 80)
+
+ # Import modules needed for verification
+ from pycsa.core import io, var, utils
+ from inputs.icon_global_run import params
+
+ # Load ICON grid
+ print("\nLoading ICON grid...")
+ grid = var.grid()
+ reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+ reader.read_dat(params.path_icon_grid, grid)
+
+ # Store radians for later use
+ clat_rad = np.copy(grid.clat)
+ clon_rad = np.copy(grid.clon)
+
+ # Convert to degrees for processing
+ grid.apply_f(utils.rad2deg)
+
+ n_cells = grid.clat.size
+ print(f" Total cells in grid: {n_cells}")
+
+ # Set ETOPO parameters
+ params.etopo_cg = 4 # Coarse-graining factor (matches processing used in icon_etopo_global.py)
+
+ # Count land/ocean cells
+ print("\nCounting land/ocean cells...")
+ land_count, ocean_count, land_cells, ocean_cells, land_fractions = (
+ count_land_ocean_cells(grid, params, reader)
+ )
+
+ # Print results
+ print("\n" + "=" * 80)
+ print("RESULTS")
+ print("=" * 80)
+ print(f"Total cells: {n_cells}")
+ print(f"Land cells (is_land=1): {land_count}")
+ print(f"Ocean cells (is_land=0): {ocean_count}")
+ print(
+ f"Land/Ocean ratio: {land_count}/{ocean_count} = {land_count/ocean_count:.3f}"
+ )
+ print(f"Land percentage: {100*land_count/(land_count+ocean_count):.2f}%")
+ print("=" * 80)
+
+ # Save plotting data for debugging
+ print("\nSaving verification data...")
+ output_dir.mkdir(parents=True, exist_ok=True)
+
+ # Convert grid coordinates to degrees for saving
+ clat_deg = np.rad2deg(clat_rad)
+ clon_deg = np.rad2deg(clon_rad)
+
+ # Save as compressed numpy file
+ np.savez_compressed(
+ data_file,
+ clat_deg=clat_deg,
+ clon_deg=clon_deg,
+ land_cells=np.array(land_cells),
+ ocean_cells=np.array(ocean_cells),
+ land_fractions=land_fractions,
+ n_cells=n_cells,
+ land_count=land_count,
+ ocean_count=ocean_count,
+ etopo_cg=params.etopo_cg,
+ )
+ print(f" Data saved: {data_file}")
+ print(
+ f" Contains: cell coordinates, land/ocean classifications, land fractions, and counts"
+ )
+
+ # Create comprehensive plots (both modes)
+ print("\nCreating comprehensive plots...")
+ create_comprehensive_plots(
+ clat_deg, clon_deg, land_cells, ocean_cells, land_fractions, output_dir
+ )
+
+ print("\n✓ Complete!")
+ print(f" Output directory: {output_dir}")
+ print(f" Plots created:")
+ print(f" - improved_verification_plots.png")
+ print(f" - pacific_islands_detail.png")
diff --git a/setup_paths.sh b/setup_paths.sh
new file mode 100755
index 0000000..413afb6
--- /dev/null
+++ b/setup_paths.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+# Setup script for local paths
+# Usage: source setup_paths.sh
+
+# Detect if we're on HPC or local machine
+if [[ -n "$SLURM_JOB_ID" ]] || [[ -n "$PBS_JOBID" ]] || [[ $(hostname) == *"hpc"* ]]; then
+ echo "Detected HPC environment"
+ export SPEC_APPX_ENV="HPC"
+
+ # HPC paths - UPDATE THESE FOR YOUR HPC
+ export SPEC_APPX_DATA_DIR="${HOME}/pyCSA/data"
+ export SPEC_APPX_OUTPUT_DIR="${HOME}/pyCSA/outputs"
+ export SPEC_APPX_MERIT_DIR="${HOME}/pyCSA/data/MERIT"
+ export SPEC_APPX_REMA_DIR="${HOME}/pyCSA/data/REMA"
+ export SPEC_APPX_ETOPO_DIR="${HOME}/pyCSA/data/etopo_15s/"
+else
+ echo "Detected local environment"
+ export SPEC_APPX_ENV="LOCAL"
+
+ # Local paths - UPDATE THESE FOR YOUR LOCAL MACHINE
+ export SPEC_APPX_DATA_DIR="${HOME}/pyCSA/data"
+ export SPEC_APPX_OUTPUT_DIR="${HOME}/pyCSA/outputs"
+ export SPEC_APPX_MERIT_DIR="${HOME}/pyCSA/data/MERIT"
+ export SPEC_APPX_REMA_DIR="${HOME}/pyCSA/data/REMA"
+ export SPEC_APPX_ETOPO_DIR="${HOME}/pyCSA/data/etopo_15s/"
+fi
+
+echo "Environment: $SPEC_APPX_ENV"
+echo "Data directory: $SPEC_APPX_DATA_DIR"
+echo "Output directory: $SPEC_APPX_OUTPUT_DIR"
+
+# Create local_paths.py if it doesn't exist
+if [ ! -f "pycsa/local_paths.py" ]; then
+ echo "Creating pycsa/local_paths.py from template..."
+ cp pycsa/local_paths.py.template pycsa/local_paths.py
+fi
diff --git a/src/io.py b/src/io.py
deleted file mode 100644
index 8f40cd8..0000000
--- a/src/io.py
+++ /dev/null
@@ -1,658 +0,0 @@
-"""
-Input/Output routines
-"""
-
-import netCDF4 as nc
-import numpy as np
-import h5py
-import os
-from datetime import datetime
-
-from src import utils
-
-
-class ncdata(object):
- """Helper class to read NetCDF4 topographic data"""
-
- def __init__(self, read_merit=False, padding=0, padding_tol=50):
- """
-
- Parameters
- ----------
- read_merit : bool, optional
- toggles between the `MERIT DEM `_ and `USGS GMTED 2010 `_ data files. By default False, i.e., read USGS GMTED 2010 data files.
- padding : int, optional
- number of data points to pad the loaded topography file, by default 0
- padding_tol : int, optional
- padding tolerance is added no matter the user-defined ``padding``, by default 50
- """
- self.read_merit = read_merit
- self.padding = padding_tol + padding
-
- def read_dat(self, fn, obj):
- """Reads data by attributes defined in the ``obj`` class.
-
- Parameters
- ----------
- fn : str
- filename
- obj : :class:`src.var.grid` or :class:`src.var.topo` or :class:`src.var.topo_cell`
- any data object in :mod:`src.var` accepting topography attributes
- """
- df = nc.Dataset(fn)
-
- for key, _ in vars(obj).items():
- if key in df.variables:
- setattr(obj, key, df.variables[key][:])
-
- df.close()
-
- def __get_truths(self, arr, vert_pts, d_pts):
- """Assembles Boolean array selecting for data points within a given lat-lon range, including padded boundary."""
- return (arr >= (vert_pts.min() - self.padding * d_pts)) & (
- arr <= vert_pts.max() + self.padding * d_pts
- )
-
- def read_topo(self, topo, cell, lon_vert, lat_vert):
- """Reads USGS GMTED 2010 dataset
-
- Parameters
- ----------
- topo : :class:`src.var.topo` or :class:`src.var.topo_cell`
- instance of a topography class containing the full regional or global topography loaded via :func:`src.io.read_dat`.
- cell : :class:`src.var.topo_cell`
- instance of a cell object
- lon_vert : list
- extent of the longitudinal coordinates encompassing the region to be loaded
- lat_vert : list
- extent of the latitudinal coordinates encompassing the region to be loaded
-
- .. note:: Loading the global topography in the ``topo`` argument may not be memory efficient. The notebook ``nc_compactifier.ipynb`` contains a script to extract a region of interest from the global GMTED 2010 dataset.
- """
- lon, lat, z = topo.lon, topo.lat, topo.topo
-
- nrecords = np.shape(z)[0]
-
- bool_arr = np.zeros_like(z).astype(bool)
- lat_arr = np.zeros_like(z)
- lon_arr = np.zeros_like(z)
-
- z = z[:, ::-1, :]
-
- for n in range(nrecords):
- lat_n = lat[n]
- lon_n = lon[n]
-
- dlat, dlon = np.diff(lat_n).mean(), np.diff(lon_n).mean()
-
- lon_nm, lat_nm = np.meshgrid(lon_n, lat_n)
-
- bool_arr[n] = self.__get_truths(lon_nm, lon_vert, dlon) & self.__get_truths(
- lat_nm, lat_vert, dlat
- )
-
- lat_arr[n] = lat_nm
- lon_arr[n] = lon_nm
-
- lon_res = lon_arr[bool_arr]
- lat_res = lat_arr[bool_arr]
- z_res = z[bool_arr].data
-
- # ---- processing of the lat,lon,topo to get the regular 2D grid for topography
- lon_uniq, lat_uniq = np.unique(lon_res), np.unique(
- lat_res
- ) # get unique values of lon,lat
- nla = len(lat_uniq)
- nlo = len(lon_uniq)
-
- lat_res_sort_idx = np.argsort(lat_res)
- lon_res_sort_idx = np.argsort(
- lon_res[lat_res_sort_idx].reshape(nla, nlo), axis=1
- )
- z_res = z_res[lat_res_sort_idx]
- z_res = np.take_along_axis(z_res.reshape(nla, nlo), lon_res_sort_idx, axis=1)
- topo_2D = z_res.reshape(nla, nlo)
-
- print("Data fetched...")
- cell.lon = lon_uniq
- cell.lat = lat_uniq
- cell.topo = topo_2D
-
- class read_merit_topo(object):
- """Subclass to read MERIT topographic data"""
-
- def __init__(self, cell, params, verbose=False):
- """Populates ``cell`` object instance with arguments from ``params``
-
- Parameters
- ----------
- cell : :class:`src.var.topo` or :class:`src.var.topo_cell`
- instance of an object with topograhy attribute
- params : :class:`src.var.params`
- user-defined run parameters
- verbose : bool, optional
- prints loading progression, by default False
- """
- self.dir = params.merit_path
- self.verbose = verbose
-
- self.fn_lon = np.array(
- [
- -180.0,
- -150.0,
- -120.0,
- -90.0,
- -60.0,
- -30.0,
- 0.0,
- 30.0,
- 60.0,
- 90.0,
- 120.0,
- 150.0,
- 180.0,
- ]
- )
- self.fn_lat = np.array([90.0, 60.0, 30.0, 0.0, -30.0, -60.0, -90.0])
-
- self.lat_verts = np.array(params.lat_extent)
- self.lon_verts = np.array(params.lon_extent)
-
- self.merit_cg = params.merit_cg
-
- lat_min_idx = self.__compute_idx(self.lat_verts.min(), "min", "lat")
- lat_max_idx = self.__compute_idx(self.lat_verts.max(), "max", "lat")
-
- lon_min_idx = self.__compute_idx(self.lon_verts.min(), "min", "lon")
- lon_max_idx = self.__compute_idx(self.lon_verts.max(), "max", "lon")
-
- fns, lon_cnt, lat_cnt = self.__get_fns(
- lat_min_idx, lat_max_idx, lon_min_idx, lon_max_idx
- )
-
- self.get_topo(cell, fns, lon_cnt, lat_cnt)
-
- def __compute_idx(self, vert, typ, direction):
- """Given a point ``vert``, look up which MERIT NetCDF file contains this point."""
- if direction == "lon":
- fn_int = self.fn_lon
- else:
- fn_int = self.fn_lat
-
- where_idx = np.argmin(np.abs(fn_int - vert))
-
- if self.verbose:
- print(fn_int, where_idx)
-
- if typ == "min":
- if (vert - fn_int[where_idx]) < 0.0:
- if direction == "lon":
- where_idx -= 1
- else:
- where_idx += 1
- elif typ == "max":
- if (vert - fn_int[where_idx]) > 0.0:
- if direction == "lon":
- where_idx += 1
- else:
- where_idx -= 1
-
- where_idx = int(where_idx)
-
- if self.verbose:
- print("where_idx, vert, fn_int[where_idx] for typ:")
- print(where_idx, vert, fn_int[where_idx], typ)
- print("")
-
- return where_idx
-
- def __get_fns(self, lat_min_idx, lat_max_idx, lon_min_idx, lon_max_idx):
- """Construct the full filenames required for the loading of the topographic data from the indices identified in :func:`src.io.ncdata.read_merit_topo.__compute_idx`"""
- fns = []
-
- for lat_cnt, lat_idx in enumerate(range(lat_max_idx, lat_min_idx)):
- l_lat_bound, r_lat_bound = (
- self.fn_lat[lat_idx],
- self.fn_lat[lat_idx + 1],
- )
- l_lat_tag, r_lat_tag = self.__get_NSEW(
- l_lat_bound, "lat"
- ), self.__get_NSEW(r_lat_bound, "lat")
-
- for lon_cnt, lon_idx in enumerate(range(lon_min_idx, lon_max_idx)):
- l_lon_bound, r_lon_bound = (
- self.fn_lon[lon_idx],
- self.fn_lon[lon_idx + 1],
- )
- l_lon_tag, r_lon_tag = self.__get_NSEW(
- l_lon_bound, "lon"
- ), self.__get_NSEW(r_lon_bound, "lon")
-
- name = "MERIT_%s%.2d-%s%.2d_%s%.3d-%s%.3d.nc4" % (
- l_lat_tag,
- np.abs(l_lat_bound),
- r_lat_tag,
- np.abs(r_lat_bound),
- l_lon_tag,
- np.abs(l_lon_bound),
- r_lon_tag,
- np.abs(r_lon_bound),
- )
-
- fns.append(name)
-
- return fns, lon_cnt, lat_cnt
-
- def get_topo(self, cell, fns, lon_cnt, lat_cnt, init=True, populate=True):
- """
- This method assembles a contiguous array in ``cell.topo`` containing the regional topography to be loaded.
-
- However, this full regional array is assembled from an array of block arrays. Each block array is loaded from a separated MERIT data file and varies in shape that is not known beforehand.
-
- Therefore, the ``get_topo`` method is run recursively:
- 1. The first run determines the shape of each constituting block array and subsequently the shape of the full regional array. An empty array in initialised.
- 2. The second run populates the empty array with the information of the block arrays obtained in the first run.
- """
- if (cell.topo is None) and (init):
- self.get_topo(cell, fns, lon_cnt, lat_cnt, init=False, populate=False)
-
- if not populate:
- nc_lon = 0
- nc_lat = 0
- else:
- n_col = 0
- n_row = 0
- lon_sz_old = 0
- lat_sz_old = 0
- cell.lat = []
- cell.lon = []
-
- for cnt, fn in enumerate(fns):
- test = nc.Dataset(self.dir + fn)
-
- lat = test["lat"]
- lat_min_idx = np.argmin(np.abs(lat - self.lat_verts.min()))
- lat_max_idx = np.argmin(np.abs(lat - self.lat_verts.max()))
-
- lat_high = np.max((lat_min_idx, lat_max_idx))
- lat_low = np.min((lat_min_idx, lat_max_idx))
-
- lon = test["lon"]
- lon_min_idx = np.argmin(np.abs(lon - (self.lon_verts.min())))
- lon_max_idx = np.argmin(np.abs(lon - (self.lon_verts.max())))
-
- lon_high = np.max((lon_min_idx, lon_max_idx))
- lon_low = np.min((lon_min_idx, lon_max_idx))
-
- if not populate:
- if cnt < (lon_cnt + 1):
- nc_lon += lon_high - lon_low
- if (cnt % (lat_cnt + 1)) == 0:
- nc_lat += lat_high - lat_low
- else:
- topo = test["Elevation"][lat_low:lat_high, lon_low:lon_high]
- if n_col == 0:
- cell.lat += lat[lat_low:lat_high].tolist()
- if n_row == 0:
- cell.lon += lon[lon_low:lon_high].tolist()
-
- lon_sz = lon_high - lon_low
- lat_sz = lat_high - lat_low
-
- cell.topo[
- n_row * lat_sz_old : n_row * lat_sz_old + lat_sz,
- n_col * lon_sz_old : n_col * lon_sz_old + lon_sz,
- ] = topo
-
- n_col += 1
- if n_col == (lon_cnt + 1):
- n_col = 0
- n_row += 1
- lat_sz_old = np.copy(lat_sz)
-
- lon_sz_old = np.copy(lon_sz)
-
- test.close()
-
- if not populate:
- cell.topo = np.zeros((nc_lat, nc_lon))
- else:
- iint = self.merit_cg
- # cell.lat = np.sort(cell.lat)[::iint]
- # cell.lon = np.sort(cell.lon)[::iint][:-1]
-
- cell.lat = utils.sliding_window_view(
- np.sort(cell.lat), (iint,), (iint,)
- ).mean(axis=-1)
- cell.lon = utils.sliding_window_view(
- np.sort(cell.lon), (iint,), (iint,)
- ).mean(axis=-1)
-
- cell.topo = utils.sliding_window_view(
- cell.topo, (iint, iint), (iint, iint)
- ).mean(axis=(-1, -2))[::-1, :]
-
- @staticmethod
- def __get_NSEW(vert, typ):
- """Method to determine `NSEW` in MERIT filename"""
- if typ == "lat":
- if vert >= 0.0:
- dir_tag = "N"
- else:
- dir_tag = "S"
- if typ == "lon":
- if vert >= 0.0:
- dir_tag = "E"
- else:
- dir_tag = "W"
-
- return dir_tag
-
-
-class writer(object):
- """
- HDF5 writer class
-
- Contains methods to create HDF5 file, create data sets and populate them with output variables.
-
- .. note:: This class was taken from an I/O routine originally written for the numerical flow solver used in `Chew et al. (2022) `_ and `Chew et al. (2023) `_.
- """
-
- def __init__(self, fn, idxs, sfx="", debug=False):
- """
- Creates an empty HDF5 file with filename ``fn`` and a group for each index in ``idxs``
-
- Parameters
- ----------
- fn : str
- filename
- idxs : list
- list of cell indices
- sfx : str, optional
- suffixes to the filename, by default ''
- debug : bool, optional
- debug flag, by default False
- """
-
- self.FORMAT = ".h5"
- self.OUTPUT_FOLDER = "../outputs/"
- self.OUTPUT_FILENAME = fn
- self.OUTPUT_FULLPATH = self.OUTPUT_FOLDER + self.OUTPUT_FILENAME
- self.SUFFIX = sfx
- self.DEBUG = debug
-
- self.IDXS = idxs
- self.PATHS = [
- # vars from the 'tri' object
- "tri_lat_verts",
- "tri_lon_verts",
- "tri_clats",
- "tri_clons",
- "points",
- "simplices",
- # vars from the 'cell' object
- "lon",
- "lat",
- "lon_grid",
- "lat_grid",
- # vars from the 'analysis' object
- "ampls",
- "kks",
- "lls",
- "recon",
- ]
-
- self.ATTRS = [
- # vars from the 'analysis' object
- "wlat",
- "wlon",
- ]
-
- if debug:
- self.PATHS = np.append(
- self.PATHS,
- [
- "mask",
- "topo_ref",
- "pmf_ref",
- "spectrum_ref",
- "spectrum_fg",
- "recon_fg",
- "pmf_fg",
- ],
- )
-
- self.io_create_file(self.IDXS)
-
- def io_create_file(self, paths):
- """
- Helper function to create file.
-
- Parameters
- ----------
- paths : list
- List of strings containing the name of the groups.
-
- Notes
- -----
- Currently, if the filename of the HDF5 file already exists, this function will append the existing filename with '_old' and create an empty HDF5 file with the same filename in its place.
-
- """
- # If directory does not exist, create it.
- if not os.path.exists(self.OUTPUT_FOLDER):
- os.mkdir(self.OUTPUT_FOLDER)
-
- # If file exists, rename it with old.
- if os.path.exists(self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT):
- os.rename(
- self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT,
- self.OUTPUT_FULLPATH + self.SUFFIX + "_old" + self.FORMAT,
- )
-
- file = h5py.File(self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT, "a")
- for path in paths:
- path = str(path)
- # check if groups have been created
- # if not created, create empty groups
- if not (path in file):
- file.create_group(path, track_order=True)
-
- file.close()
-
- def write_all(self, idx, *args):
- """Write all attributes and datasets of a given class instance to the group ``idx``.
-
- Parameters
- ----------
- idx : str or int
- group name to write the attributes or datasets
- """
- for arg in args:
- for attr in self.PATHS:
- if hasattr(arg, attr):
- self.populate(idx, attr, getattr(arg, attr))
-
- for attr in self.ATTRS:
- if hasattr(arg, attr):
- self.write_attr(idx, attr, getattr(arg, attr))
-
- def write_attr(self, idx, key, value):
- """Write HDF5 attributes for a group
-
- Parameters
- ----------
- idx : str or int
- group name to write the attributes
- key : str
- attribute name
- value : any
- attribute value that is accepted by HDF5
- """
- file = h5py.File(self.OUTPUT_FULLPATH + self.SUFFIX + self.FORMAT, "r+")
-
- try:
- file[str(idx)].attrs.create(str(key), value)
- except:
- file[str(idx)].attrs.create(
- str(key), repr(value), dtype="`.
-
- Parameters
- ----------
- fobj : :class:`src.fourier.f_trans` instance
- instance of the Fourier transformer class.
-
- Returns
- -------
- array-like
- 2D array corresponding to the ``M`` matrix.
- """
- Ncos = fobj.bf_cos
- Nsin = fobj.bf_sin
-
- coeff = np.hstack([Ncos, Nsin])
-
- del fobj.bf_cos
- del fobj.bf_sin
-
- if fobj.grad:
- coeff = np.vstack([coeff, coeff])
-
- return coeff
-
-
-def do(fobj, cell, lmbda=0.0, iter_solve=True, save_coeffs=False):
- """
- Does the linear regression
-
- Parameters
- ----------
- fobj : :class:`src.fourier.f_trans` instance
- instance of the Fourier transformer class.
- cell : :class:`src.var.topo_cell` instance
- cell object instance
- lmbda : float, optional
- regularisation parameter, by default 0.0
- iter_solve : bool, optional
- toggles between using direct or iterative solver, by default True
- save_coeffs : bool, optional
- skips the linear regression and just saves the generated ``M`` matrix for diagnostics and debugging, by default False
-
- Returns
- -------
- a_m : list
- list of Fourier amplitudes corresponding to the unknown vector in the linear problem
- data_recons : like
- vector-like topography reconstructed from ``a_m``
- """
- if fobj.grad:
- cell.get_grad()
- data = cell.grad_topo_m
- else:
- data = cell.topo_m
-
- coeff = get_coeffs(fobj)
-
- if save_coeffs:
- fobj.coeff = coeff
- return None, None
-
- # tot_coeff = coeff.shape[1]
-
- # E_tilda_lm = np.zeros((tot_coeff,tot_coeff))
-
- h_tilda_l = np.dot(coeff.T, data.reshape(-1, 1)).flatten()
-
- E_tilda_lm = np.dot(coeff.T, coeff)
-
- trace = np.trace(E_tilda_lm) / len(np.diag(E_tilda_lm)) * lmbda
- szc = E_tilda_lm.shape[0]
- for ttr in range(szc):
- E_tilda_lm[ttr, ttr] += trace
-
- if iter_solve:
- a_m, _ = gmres(E_tilda_lm, h_tilda_l)
- else:
- a_m = la.inv(E_tilda_lm).dot(h_tilda_l)
-
- # regular FFT considers normalization by total nu mber of datapoints N=100
- # so multiply the Fourier coefficients by N here
- # a_m = a_m#*len(data)
-
- data_recons = coeff.dot(a_m)
-
- return a_m, data_recons
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..d80072f
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,162 @@
+"""
+Shared pytest fixtures and utilities for pyCSA tests.
+"""
+
+# ---------------------------------------------------------------------------
+# Cartopy stub — let tests run in environments without cartopy installed.
+# pycsa.__init__ eagerly imports pycsa.plotting.cart_plot which imports
+# cartopy. The tests don't actually call any plotting functions, so a stub
+# is enough to satisfy the import chain. If real cartopy is installed, this
+# is a no-op.
+# ---------------------------------------------------------------------------
+try:
+ import cartopy # noqa: F401
+except ImportError:
+ import sys
+ import types
+
+ def _stub_pkg(name):
+ m = types.ModuleType(name)
+ m.__path__ = [] # marks as package so submodule imports work
+ sys.modules[name] = m
+ return m
+
+ def _stub_attrs(mod, *names):
+ for n in names:
+ setattr(mod, n, type(n, (), {}))
+
+ _stub_pkg("cartopy")
+ _crs = _stub_pkg("cartopy.crs")
+ _stub_attrs(_crs, "PlateCarree", "Mollweide", "Robinson", "Geodetic")
+ _stub_pkg("cartopy.mpl")
+ _ticker = _stub_pkg("cartopy.mpl.ticker")
+ _stub_attrs(
+ _ticker,
+ "LongitudeFormatter",
+ "LatitudeFormatter",
+ "LongitudeLocator",
+ "LatitudeLocator",
+ )
+ _stub_pkg("cartopy.feature")
+ _stub_pkg("cartopy.io")
+ _stub_pkg("cartopy.io.shapereader")
+
+import numpy as np
+import pytest
+from pathlib import Path
+
+
+@pytest.fixture
+def project_root():
+ """Return the project root directory."""
+ return Path(__file__).parent.parent
+
+
+@pytest.fixture
+def baseline_dir(project_root):
+ """Return the baseline results directory."""
+ return project_root / "outputs" / "baseline_results"
+
+
+@pytest.fixture
+def test_output_dir(project_root, tmp_path):
+ """Return a temporary directory for test outputs."""
+ return tmp_path
+
+
+def assert_arrays_close(actual, expected, rtol=1e-5, atol=1e-8, name="array"):
+ """
+ Assert that two numpy arrays are close within tolerance.
+
+ Parameters
+ ----------
+ actual : np.ndarray
+ The actual computed array
+ expected : np.ndarray
+ The expected baseline array
+ rtol : float
+ Relative tolerance
+ atol : float
+ Absolute tolerance
+ name : str
+ Name of the array for error messages
+ """
+ np.testing.assert_allclose(
+ actual,
+ expected,
+ rtol=rtol,
+ atol=atol,
+ err_msg=f"{name} does not match baseline within tolerance (rtol={rtol}, atol={atol})",
+ )
+
+
+def assert_values_close(actual, expected, rtol=1e-5, atol=1e-8, name="value"):
+ """
+ Assert that two scalar values are close within tolerance.
+
+ Parameters
+ ----------
+ actual : float
+ The actual computed value
+ expected : float
+ The expected baseline value
+ rtol : float
+ Relative tolerance
+ atol : float
+ Absolute tolerance
+ name : str
+ Name of the value for error messages
+ """
+ np.testing.assert_allclose(
+ actual,
+ expected,
+ rtol=rtol,
+ atol=atol,
+ err_msg=f"{name} = {actual} does not match baseline {expected} within tolerance",
+ )
+
+
+class BaselineComparison:
+ """Helper class for comparing test results against baseline."""
+
+ def __init__(self, rtol=1e-5, atol=1e-8):
+ """
+ Initialize baseline comparison.
+
+ Parameters
+ ----------
+ rtol : float
+ Relative tolerance for comparisons
+ atol : float
+ Absolute tolerance for comparisons
+ """
+ self.rtol = rtol
+ self.atol = atol
+ self.results = {}
+
+ def add_result(self, name, actual, expected):
+ """Add a result to compare."""
+ self.results[name] = {"actual": actual, "expected": expected, "passed": None}
+
+ def compare_all(self):
+ """Compare all added results and return summary."""
+ summary = {"passed": 0, "failed": 0, "failures": []}
+
+ for name, data in self.results.items():
+ try:
+ if isinstance(data["actual"], np.ndarray):
+ assert_arrays_close(
+ data["actual"], data["expected"], self.rtol, self.atol, name
+ )
+ else:
+ assert_values_close(
+ data["actual"], data["expected"], self.rtol, self.atol, name
+ )
+ self.results[name]["passed"] = True
+ summary["passed"] += 1
+ except AssertionError as e:
+ self.results[name]["passed"] = False
+ summary["failed"] += 1
+ summary["failures"].append({"name": name, "error": str(e)})
+
+ return summary
diff --git a/tests/debug/debug_etopo_single_cell.py b/tests/debug/debug_etopo_single_cell.py
new file mode 100644
index 0000000..ad809d0
--- /dev/null
+++ b/tests/debug/debug_etopo_single_cell.py
@@ -0,0 +1,714 @@
+"""
+Debug test for individual cells with verbose plotting and diagnostics.
+
+Usage:
+ # Edit CELL_INDICES list below, then run:
+ pytest tests/test_single_cell_debug.py -v -s
+
+This will create detailed plots and logs for debugging specific cell failures.
+"""
+
+import pytest
+import numpy as np
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from pathlib import Path
+import traceback
+import sys
+
+from pycsa.core import io, var, utils
+from pycsa.wrappers import interface
+
+# =============================================================================
+# CONFIGURE WHICH CELLS TO DEBUG HERE
+# =============================================================================
+CELL_INDICES = [
+ 1086, # FileNotFoundError: E180 tile (N90E180)
+ # 1027, # FileNotFoundError: E180 tile (N90E180)
+ # 1219, # FileNotFoundError: E180 tile (N75E180)
+]
+# =============================================================================
+
+
+@pytest.fixture(params=CELL_INDICES, ids=lambda x: f"cell_{x}")
+def cell_idx(request):
+ """Get cell index from parameter list."""
+ return request.param
+
+
+@pytest.fixture
+def output_dir(cell_idx):
+ """Create output directory for this specific cell."""
+ base_dir = Path(__file__).parent.parent / "outputs" / "cell_debug"
+ base_dir.mkdir(parents=True, exist_ok=True)
+
+ cell_dir = base_dir / f"cell_{cell_idx}"
+ cell_dir.mkdir(parents=True, exist_ok=True)
+
+ print(f"\n📁 Debug output directory: {cell_dir}")
+ return cell_dir
+
+
+@pytest.fixture
+def test_params():
+ """Create test parameters using ETOPO data."""
+ params = var.params()
+
+ # Import local paths
+ try:
+ from pycsa import local_paths
+
+ utils.transfer_attributes(params, local_paths.paths, prefix="path")
+ except ImportError as e:
+ pytest.skip(f"Could not import local_paths: {e}")
+
+ # Verify ETOPO path exists
+ if not hasattr(params, "path_etopo") or not Path(params.path_etopo).exists():
+ pytest.skip(f"ETOPO data path not found")
+
+ # Test region: Alaska (will be overridden per cell)
+ params.lat_extent = [48.0, 64.0, 64.0]
+ params.lon_extent = [-148.0, -148.0, -112.0]
+
+ # ETOPO coarse-graining factor
+ params.etopo_cg = 50
+
+ # CSA parameters
+ params.nhi = 24
+ params.nhj = 48
+ params.n_modes = 50
+ params.padding = 10
+
+ params.U, params.V = 10.0, 0.0
+ params.rect = True
+
+ # Enable verbose mode
+ params.plot = False
+ params.plot_output = False
+ params.debug = False
+ params.dfft_first_guess = False
+ params.refine = False
+ params.verbose = True
+
+ return params
+
+
+@pytest.fixture
+def test_grid(test_params):
+ """Load ICON grid."""
+ grid = var.grid()
+
+ try:
+ reader = io.ncdata()
+ reader.read_dat(test_params.path_icon_grid, grid)
+ except Exception as e:
+ pytest.skip(f"Could not load ICON grid: {e}")
+
+ # Convert to degrees
+ grid.apply_f(utils.rad2deg)
+
+ return grid
+
+
+def test_debug_cell(cell_idx, output_dir, test_params, test_grid):
+ """Debug a single cell with verbose output and plotting."""
+
+ print(f"\n{'='*70}")
+ print(f"DEBUGGING CELL {cell_idx}")
+ print(f"{'='*70}\n")
+
+ # Create log file
+ log_file = output_dir / "debug_log.txt"
+
+ def log_and_print(msg):
+ """Print and log message."""
+ print(msg)
+ with open(log_file, "a") as f:
+ f.write(msg + "\n")
+
+ log_and_print(f"Cell Index: {cell_idx}")
+ log_and_print(f"Output Directory: {output_dir}")
+ log_and_print("")
+
+ # Step 1: Get cell geometry
+ log_and_print("=" * 70)
+ log_and_print("STEP 1: Cell Geometry")
+ log_and_print("=" * 70)
+
+ try:
+ lat_verts = test_grid.clat_vertices[cell_idx]
+ lon_verts = test_grid.clon_vertices[cell_idx]
+ cell_lat = test_grid.clat[cell_idx]
+ cell_lon = test_grid.clon[cell_idx]
+
+ log_and_print(f"Cell center: lat={cell_lat:.4f}°, lon={cell_lon:.4f}°")
+ log_and_print(f"Vertices (lat): {lat_verts}")
+ log_and_print(f"Vertices (lon): {lon_verts}")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR getting cell geometry: {e}")
+ log_and_print(traceback.format_exc())
+ raise
+
+ # Step 2: Handle lat/lon expansion
+ log_and_print("=" * 70)
+ log_and_print("STEP 2: Lat/Lon Expansion")
+ log_and_print("=" * 70)
+
+ try:
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+ lat_verts_expanded, lon_verts_expanded = utils.handle_latlon_expansion(
+ lat_verts, lon_verts, lat_expand=0.0, lon_expand=0.0
+ )
+
+ log_and_print(f"Original vertices:")
+ log_and_print(f" lat: {lat_verts}")
+ log_and_print(f" lon: {lon_verts}")
+ log_and_print(f"")
+ log_and_print(f"Expanded extents:")
+ log_and_print(f" lat_extent: {lat_extent}")
+ log_and_print(f" lon_extent: {lon_extent}")
+ log_and_print(f"")
+ log_and_print(f"Expanded vertices:")
+ log_and_print(f" lat: {lat_verts_expanded}")
+ log_and_print(f" lon: {lon_verts_expanded}")
+ log_and_print("")
+
+ # Update params
+ test_params.lat_extent = lat_extent
+ test_params.lon_extent = lon_extent
+
+ except Exception as e:
+ log_and_print(f"ERROR in lat/lon expansion: {e}")
+ log_and_print(traceback.format_exc())
+ raise
+
+ # Step 3: Initialize ETOPO reader
+ log_and_print("=" * 70)
+ log_and_print("STEP 3: Initialize ETOPO Reader")
+ log_and_print("=" * 70)
+
+ try:
+ reader = io.ncdata(padding=test_params.padding)
+ topo = var.topo_cell()
+
+ log_and_print(f"Creating ETOPO reader with:")
+ log_and_print(f" padding: {test_params.padding}")
+ log_and_print(f" lat_extent: {test_params.lat_extent}")
+ log_and_print(f" lon_extent: {test_params.lon_extent}")
+ log_and_print(f" etopo_cg: {test_params.etopo_cg}")
+ log_and_print("")
+
+ etopo_reader = reader.read_etopo_topo(
+ None, test_params, is_parallel=True, verbose=True
+ )
+
+ log_and_print(f"ETOPO reader created successfully")
+ log_and_print(f" split_EW: {etopo_reader.split_EW}")
+ if hasattr(etopo_reader, "split_NS"):
+ log_and_print(f" split_NS: {etopo_reader.split_NS}")
+ if hasattr(etopo_reader, "file_cache"):
+ log_and_print(f" file_cache size: {len(etopo_reader.file_cache)}")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR initializing ETOPO reader: {e}")
+ log_and_print(traceback.format_exc())
+ raise
+
+ # Step 4: Load topography data
+ log_and_print("=" * 70)
+ log_and_print("STEP 4: Load Topography Data")
+ log_and_print("=" * 70)
+
+ try:
+ log_and_print("Calling etopo_reader.get_topo()...")
+ etopo_reader.get_topo(topo)
+
+ log_and_print(f"Topography loaded successfully!")
+ log_and_print(f" Shape: {topo.topo.shape}")
+ log_and_print(f" Min elevation: {np.min(topo.topo):.2f} m")
+ log_and_print(f" Max elevation: {np.max(topo.topo):.2f} m")
+ log_and_print(f" Mean elevation: {np.mean(topo.topo):.2f} m")
+ log_and_print(f" Lat shape: {topo.lat.shape}")
+ log_and_print(f" Lon shape: {topo.lon.shape}")
+ log_and_print(f" Lat range: [{np.min(topo.lat):.4f}, {np.max(topo.lat):.4f}]")
+ log_and_print(f" Lon range: [{np.min(topo.lon):.4f}, {np.max(topo.lon):.4f}]")
+ log_and_print("")
+
+ # Apply elevation floor
+ below_floor = np.sum(topo.topo < -500.0)
+ if below_floor > 0:
+ log_and_print(f"Applying elevation floor: {below_floor} points below -500m")
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+
+ topo.gen_mgrids()
+ log_and_print("Generated mesh grids")
+ log_and_print("")
+
+ # Save topography data for inspection
+ np.save(output_dir / "topo_elevation.npy", topo.topo)
+ np.save(output_dir / "topo_lat.npy", topo.lat)
+ np.save(output_dir / "topo_lon.npy", topo.lon)
+ log_and_print(f"Saved topography arrays to {output_dir}")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR loading topography: {e}")
+ log_and_print(traceback.format_exc())
+
+ # Try to get more debug info from the reader
+ if hasattr(etopo_reader, "__get_fns"):
+ try:
+ log_and_print("\nAttempting to get file info...")
+ # This might fail but could give us useful info
+ lat_idx_rng = getattr(etopo_reader, "lat_idx_rng", None)
+ lon_idx_rng = getattr(etopo_reader, "lon_idx_rng", None)
+ log_and_print(f" lat_idx_rng: {lat_idx_rng}")
+ log_and_print(f" lon_idx_rng: {lon_idx_rng}")
+ except:
+ pass
+
+ raise
+
+ # Step 5: Set up cell geometry for land check
+ log_and_print("=" * 70)
+ log_and_print("STEP 5: Cell Geometry Setup")
+ log_and_print("=" * 70)
+
+ try:
+ clon = np.array([test_grid.clon[cell_idx]])
+ clat = np.array([test_grid.clat[cell_idx]])
+ clon_vertices = np.array([lon_verts_expanded])
+ clat_vertices = np.array([lat_verts_expanded])
+
+ log_and_print(f"Cell geometry:")
+ log_and_print(f" clon: {clon}")
+ log_and_print(f" clat: {clat}")
+ log_and_print(f" clon_vertices: {clon_vertices}")
+ log_and_print(f" clat_vertices: {clat_vertices}")
+ log_and_print("")
+
+ ncells = 1
+ nv = clon_vertices[0].size
+
+ # Handle dateline crossing
+ if etopo_reader.split_EW:
+ log_and_print("Handling dateline crossing (split_EW=True)")
+ orig_clon_vertices = clon_vertices.copy()
+ clon_vertices[clon_vertices < 0.0] += 360.0
+ log_and_print(f" Before: {orig_clon_vertices}")
+ log_and_print(f" After: {clon_vertices}")
+ log_and_print("")
+
+ triangles = np.zeros((ncells, nv, 2))
+ for i in range(0, ncells, 1):
+ triangles[i, :, 0] = np.array(clon_vertices[i, :])
+ triangles[i, :, 1] = np.array(clat_vertices[i, :])
+
+ log_and_print(f"Triangle vertices:")
+ log_and_print(f" {triangles}")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR setting up cell geometry: {e}")
+ log_and_print(traceback.format_exc())
+ raise
+
+ # Step 6: Check if land
+ log_and_print("=" * 70)
+ log_and_print("STEP 6: Land/Ocean Check")
+ log_and_print("=" * 70)
+
+ try:
+ tri_idx = 0
+ cell = var.topo_cell()
+ tri = var.obj()
+
+ tri.tri_lon_verts = triangles[:, :, 0]
+ tri.tri_lat_verts = triangles[:, :, 1]
+ simplex_lat = tri.tri_lat_verts[tri_idx]
+ simplex_lon = tri.tri_lon_verts[tri_idx]
+
+ log_and_print(f"Simplex vertices for land check:")
+ log_and_print(f" simplex_lat: {simplex_lat}")
+ log_and_print(f" simplex_lon: {simplex_lon}")
+ log_and_print("")
+
+ # This is where the error happens in some cells
+ log_and_print("Calling utils.is_land()...")
+ is_land = utils.is_land(cell, simplex_lat, simplex_lon, topo)
+
+ log_and_print(f"is_land result: {is_land}")
+ log_and_print(
+ f"Cell lat shape: {cell.lat.shape if hasattr(cell, 'lat') and cell.lat is not None else 'None'}"
+ )
+ log_and_print(
+ f"Cell lon shape: {cell.lon.shape if hasattr(cell, 'lon') and cell.lon is not None else 'None'}"
+ )
+ log_and_print("")
+
+ if not is_land:
+ log_and_print("Cell is OCEAN - skipping CSA processing")
+ # Still plot the topography
+ plot_topography(
+ output_dir, topo, simplex_lat, simplex_lon, cell_idx, is_land=False
+ )
+ return
+
+ log_and_print("Cell is LAND - proceeding with CSA")
+
+ # Save cell data for inspection
+ if hasattr(cell, "lat") and cell.lat is not None:
+ np.save(output_dir / "cell_lat.npy", cell.lat)
+ np.save(output_dir / "cell_lon.npy", cell.lon)
+ if hasattr(cell, "topo") and cell.topo is not None:
+ np.save(output_dir / "cell_topo.npy", cell.topo)
+ log_and_print(f"Saved cell arrays to {output_dir}")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR in land check: {e}")
+ log_and_print(traceback.format_exc())
+
+ # Try to plot what we have so far
+ try:
+ plot_topography(
+ output_dir,
+ topo,
+ simplex_lat,
+ simplex_lon,
+ cell_idx,
+ is_land=None,
+ error=str(e),
+ )
+ except:
+ pass
+
+ raise
+
+ # Step 7: Get lat/lon segments
+ log_and_print("=" * 70)
+ log_and_print("STEP 7: Get Lat/Lon Segments")
+ log_and_print("=" * 70)
+
+ try:
+ log_and_print(f"Calling utils.get_lat_lon_segments()...")
+ log_and_print(f" simplex_lat: {simplex_lat}")
+ log_and_print(f" simplex_lon: {simplex_lon}")
+ log_and_print(f" rect: {test_params.rect}")
+ log_and_print("")
+
+ utils.get_lat_lon_segments(
+ simplex_lat, simplex_lon, cell, topo, rect=test_params.rect
+ )
+
+ log_and_print(f"Segments extracted successfully!")
+ log_and_print(f" cell.lat shape: {cell.lat.shape}")
+ log_and_print(f" cell.lon shape: {cell.lon.shape}")
+ log_and_print(f" cell.topo shape: {cell.topo.shape}")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR getting lat/lon segments: {e}")
+ log_and_print(traceback.format_exc())
+ raise
+
+ # Step 8: Run spectral approximation
+ log_and_print("=" * 70)
+ log_and_print("STEP 8: Spectral Approximation")
+ log_and_print("=" * 70)
+
+ try:
+ nhi = test_params.nhi
+ nhj = test_params.nhj
+
+ log_and_print(f"Running CSA with:")
+ log_and_print(f" nhi: {nhi}")
+ log_and_print(f" nhj: {nhj}")
+ log_and_print(f" U, V: {test_params.U}, {test_params.V}")
+ log_and_print(f" n_modes: {test_params.n_modes}")
+ log_and_print("")
+
+ pmf = interface.get_pmf(nhi, nhj, test_params.U, test_params.V)
+ ampls, uw_pmf, dat_2D = pmf.sappx(cell, lmbda=0.1)
+
+ # Filter out NaNs from spectrum
+ ampls_valid = ampls[~np.isnan(ampls)]
+
+ log_and_print(f"CSA complete!")
+ log_and_print(f" ampls shape: {ampls.shape}")
+ log_and_print(f" ampls total elements: {ampls.size}")
+ log_and_print(f" ampls valid (non-NaN): {len(ampls_valid)}")
+ if len(ampls_valid) > 0:
+ log_and_print(f" ampls max (valid): {np.max(ampls_valid):.6e}")
+ log_and_print(f" ampls sum (valid): {np.sum(ampls_valid):.6e}")
+ else:
+ log_and_print(f" ampls max: No valid values (all NaN)")
+ log_and_print("")
+
+ # Save spectrum
+ np.save(output_dir / "spectrum.npy", ampls)
+ log_and_print(f"Saved spectrum to {output_dir}/spectrum.npy")
+ log_and_print("")
+
+ except Exception as e:
+ log_and_print(f"ERROR in spectral approximation: {e}")
+ log_and_print(traceback.format_exc())
+ raise
+
+ # Step 9: Generate plots
+ log_and_print("=" * 70)
+ log_and_print("STEP 9: Generate Diagnostic Plots")
+ log_and_print("=" * 70)
+
+ try:
+ plot_topography(
+ output_dir,
+ topo,
+ simplex_lat,
+ simplex_lon,
+ cell_idx,
+ is_land=True,
+ cell=cell,
+ ampls=ampls,
+ )
+ log_and_print("✓ Generated diagnostic plots")
+ except Exception as e:
+ log_and_print(f"ERROR generating plots: {e}")
+ log_and_print(traceback.format_exc())
+
+ log_and_print("")
+ log_and_print("=" * 70)
+ log_and_print(f"DEBUG COMPLETE FOR CELL {cell_idx}")
+ log_and_print("=" * 70)
+ log_and_print(f"All outputs saved to: {output_dir}")
+ log_and_print("")
+
+ print(f"\n✓ Debug complete! Check {output_dir} for detailed outputs")
+
+
+def plot_topography(
+ output_dir,
+ topo,
+ simplex_lat,
+ simplex_lon,
+ cell_idx,
+ is_land=None,
+ cell=None,
+ ampls=None,
+ error=None,
+):
+ """Generate comprehensive topography plots."""
+
+ fig = plt.figure(figsize=(16, 12))
+
+ # Plot 1: Full topography with cell outline
+ ax1 = plt.subplot(2, 3, 1)
+ if topo.topo is not None and topo.topo.size > 0:
+ im1 = ax1.contourf(topo.lon, topo.lat, topo.topo, levels=50, cmap="terrain")
+ plt.colorbar(im1, ax=ax1, label="Elevation (m)")
+
+ # Overlay cell polygon
+ if simplex_lat is not None and simplex_lon is not None and len(simplex_lat) > 0:
+ # Close the polygon
+ poly_lat = np.append(simplex_lat, simplex_lat[0])
+ poly_lon = np.append(simplex_lon, simplex_lon[0])
+ ax1.plot(poly_lon, poly_lat, "r-", linewidth=2, label="Cell boundary")
+ ax1.legend()
+ else:
+ ax1.text(0.5, 0.5, "No topography data", ha="center", va="center")
+
+ ax1.set_xlabel("Longitude (°)")
+ ax1.set_ylabel("Latitude (°)")
+ ax1.set_title(f"Cell {cell_idx}: Full Topography")
+ ax1.grid(True, alpha=0.3)
+
+ # Plot 2: Topography 3D view
+ ax2 = plt.subplot(2, 3, 2, projection="3d")
+ if topo.topo is not None and topo.topo.size > 0:
+ # Downsample for 3D plotting if too large
+ stride = max(1, topo.topo.shape[0] // 50)
+ X, Y = np.meshgrid(topo.lon[::stride], topo.lat[::stride])
+ Z = topo.topo[::stride, ::stride]
+ ax2.plot_surface(X, Y, Z, cmap="terrain", alpha=0.8)
+ ax2.set_xlabel("Longitude (°)")
+ ax2.set_ylabel("Latitude (°)")
+ ax2.set_zlabel("Elevation (m)")
+ else:
+ ax2.text2D(
+ 0.5,
+ 0.5,
+ "No topography data",
+ transform=ax2.transAxes,
+ ha="center",
+ va="center",
+ )
+ ax2.set_title("3D View")
+
+ # Plot 3: Elevation histogram
+ ax3 = plt.subplot(2, 3, 3)
+ if topo.topo is not None and topo.topo.size > 0:
+ ax3.hist(topo.topo.flatten(), bins=50, edgecolor="black", alpha=0.7)
+ ax3.axvline(0, color="blue", linestyle="--", linewidth=2, label="Sea level")
+ ax3.axvline(
+ -500, color="red", linestyle="--", linewidth=2, label="Floor (-500m)"
+ )
+ ax3.set_xlabel("Elevation (m)")
+ ax3.set_ylabel("Count")
+ ax3.legend()
+ else:
+ ax3.text(0.5, 0.5, "No topography data", ha="center", va="center")
+ ax3.set_title("Elevation Distribution")
+ ax3.grid(True, alpha=0.3)
+
+ # Plot 4: Cell topography (if extracted)
+ ax4 = plt.subplot(2, 3, 4)
+ if (
+ cell is not None
+ and hasattr(cell, "topo")
+ and cell.topo is not None
+ and cell.topo.size > 0
+ ):
+ im4 = ax4.contourf(cell.lon, cell.lat, cell.topo, levels=50, cmap="terrain")
+ plt.colorbar(im4, ax=ax4, label="Elevation (m)")
+ ax4.set_xlabel("Longitude (°)")
+ ax4.set_ylabel("Latitude (°)")
+ ax4.set_title("Extracted Cell Topography")
+ ax4.grid(True, alpha=0.3)
+ else:
+ status = "OCEAN" if is_land == False else "ERROR" if error else "No cell data"
+ ax4.text(
+ 0.5, 0.5, status, ha="center", va="center", fontsize=14, fontweight="bold"
+ )
+ if error:
+ ax4.text(
+ 0.5,
+ 0.3,
+ f"Error: {error[:50]}...",
+ ha="center",
+ va="center",
+ fontsize=8,
+ color="red",
+ )
+ ax4.set_title("Cell Data")
+
+ # Plot 5: Spectrum (if available)
+ ax5 = plt.subplot(2, 3, 5)
+ if ampls is not None and ampls.size > 0:
+ # Plot non-NaN values
+ ampls_valid = ampls[~np.isnan(ampls)]
+ if len(ampls_valid) > 0:
+ # Find indices of valid values for proper x-axis
+ valid_indices = np.where(~np.isnan(ampls.flatten()))[0]
+ ax5.semilogy(valid_indices, ampls_valid, "o-", markersize=4)
+ ax5.set_xlabel("Mode index")
+ ax5.set_ylabel("Amplitude")
+ ax5.set_title(
+ f"Spectral Amplitudes ({len(ampls_valid)}/{ampls.size} valid)"
+ )
+ ax5.grid(True, alpha=0.3)
+ else:
+ ax5.text(
+ 0.5,
+ 0.5,
+ "No valid spectrum values\n(all NaN)",
+ ha="center",
+ va="center",
+ fontsize=10,
+ )
+ else:
+ ax5.text(0.5, 0.5, "No spectrum computed", ha="center", va="center")
+
+ # Plot 6: Summary info
+ ax6 = plt.subplot(2, 3, 6)
+ ax6.axis("off")
+
+ info_lines = [
+ f"Cell Index: {cell_idx}",
+ f"",
+ f"Topography Grid:",
+ f" Shape: {topo.topo.shape if topo.topo is not None else 'None'}",
+ (
+ f" Lat: [{np.min(topo.lat):.4f}, {np.max(topo.lat):.4f}]°"
+ if topo.lat is not None
+ else " Lat: None"
+ ),
+ (
+ f" Lon: [{np.min(topo.lon):.4f}, {np.max(topo.lon):.4f}]°"
+ if topo.lon is not None
+ else " Lon: None"
+ ),
+ f"",
+ f"Elevation:",
+ f" Min: {np.min(topo.topo):.1f} m" if topo.topo is not None else " Min: None",
+ f" Max: {np.max(topo.topo):.1f} m" if topo.topo is not None else " Max: None",
+ (
+ f" Mean: {np.mean(topo.topo):.1f} m"
+ if topo.topo is not None
+ else " Mean: None"
+ ),
+ f"",
+ f"Land Classification: {is_land if is_land is not None else 'Unknown'}",
+ ]
+
+ if cell is not None and hasattr(cell, "topo") and cell.topo is not None:
+ info_lines.extend(
+ [
+ f"",
+ f"Cell Data:",
+ f" Shape: {cell.topo.shape}",
+ f" Points: {cell.topo.size}",
+ ]
+ )
+
+ if ampls is not None:
+ ampls_valid = ampls[~np.isnan(ampls)]
+ info_lines.extend(
+ [
+ f"",
+ f"Spectrum:",
+ f" Total modes: {ampls.size}",
+ f" Valid modes: {len(ampls_valid)}",
+ ]
+ )
+ if len(ampls_valid) > 0:
+ info_lines.append(f" Max: {np.max(ampls_valid):.6e}")
+ else:
+ info_lines.append(f" Max: N/A (all NaN)")
+
+ if error:
+ info_lines.extend(
+ [
+ f"",
+ f"ERROR:",
+ f" {error[:60]}",
+ ]
+ )
+
+ info_text = "\n".join(info_lines)
+ ax6.text(
+ 0.1,
+ 0.9,
+ info_text,
+ transform=ax6.transAxes,
+ fontsize=9,
+ verticalalignment="top",
+ family="monospace",
+ )
+
+ plt.suptitle(f"Cell {cell_idx} Debug Plots", fontsize=16, fontweight="bold")
+ plt.tight_layout()
+ plt.savefig(output_dir / f"cell_{cell_idx}_debug.png", dpi=150, bbox_inches="tight")
+ plt.close()
+
+ print(f" ✓ Saved plot: {output_dir / f'cell_{cell_idx}_debug.png'}")
+
+
+if __name__ == "__main__":
+ # Run directly
+ print(f"Testing cells: {CELL_INDICES}")
+ pytest.main([__file__, "-v", "-s"])
diff --git a/tests/integration/test_delaunay_workflow.py b/tests/integration/test_delaunay_workflow.py
new file mode 100644
index 0000000..42342f3
--- /dev/null
+++ b/tests/integration/test_delaunay_workflow.py
@@ -0,0 +1,277 @@
+"""
+Integration test for Delaunay decomposition workflow (FIXED).
+
+Tests the full pipeline using the correct first_appx/second_appx API.
+"""
+
+import pytest
+import numpy as np
+from pathlib import Path
+from pycsa.core import io, var, utils, delaunay
+from pycsa.wrappers import interface, diagnostics
+
+
+@pytest.mark.integration
+class TestDelaunayWorkflow:
+ """Test Delaunay decomposition and triangle pair processing."""
+
+ @pytest.fixture
+ def data_dir(self):
+ """Return path to test data directory."""
+ return Path(__file__).parent.parent.parent / "data"
+
+ @pytest.fixture
+ def mock_params(self):
+ """Create mock params object for interface classes."""
+
+ class MockParams:
+ U = 10.0
+ V = 0.0
+ n_modes = 20
+ lmbda_fa = 1e-1
+ lmbda_sa = 1e-6
+ taper_ref = False
+ taper_fa = True
+ taper_sa = True
+ dfft_first_guess = False
+ rect = True
+ no_corrections = True
+ recompute_rhs = False
+ run_case = "TEST"
+ rect_set = [0, 2]
+ padding = 10
+ taper_art_it = 20
+ fa_iter_solve = False
+ sa_iter_solve = False
+ cg_spsp = False
+
+ return MockParams()
+
+ @pytest.fixture
+ def test_data(self, data_dir):
+ """Load test data (grid and topography)."""
+ grid_path = data_dir / "icon_compact_alaska.nc"
+ topo_path = data_dir / "topo_compact_alaska.nc"
+
+ if not grid_path.exists() or not topo_path.exists():
+ pytest.skip("Test data not available")
+
+ # Initialize data objects
+ grid = var.grid()
+ topo = var.topo_cell()
+
+ # Read data
+ reader = io.ncdata(padding=10, padding_tol=50)
+ reader.read_dat(str(grid_path), grid)
+ grid.apply_f(utils.rad2deg)
+
+ reader.read_dat(str(topo_path), topo)
+
+ # Define Alaska region
+ lat_verts = np.array([60.0, 64.0])
+ lon_verts = np.array([-148.0, -140.0])
+
+ # Extract topography for region
+ reader.read_topo(topo, topo, lon_verts, lat_verts)
+
+ # Clean up unrealistic values
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+
+ topo.gen_mgrids()
+
+ return grid, topo, reader
+
+ def test_delaunay_decomposition(self, test_data):
+ """Test Delaunay triangulation of domain."""
+ grid, topo, reader = test_data
+
+ # Perform Delaunay decomposition with small grid for testing
+ tri = delaunay.get_decomposition(topo, xnp=5, ynp=4, padding=reader.padding)
+
+ # Verify triangulation structure
+ assert hasattr(tri, "simplices"), "Triangulation missing simplices"
+ assert hasattr(tri, "points"), "Triangulation missing points"
+ assert tri.simplices is not None, "Simplices not computed"
+ assert tri.points is not None, "Points not computed"
+
+ # Check that we have triangles
+ assert len(tri.simplices) > 0, "No triangles created"
+
+ # Each triangle should have 3 vertices
+ assert tri.simplices.shape[1] == 3, "Triangles should have 3 vertices"
+
+ # Vertex indices should be valid
+ assert tri.simplices.min() >= 0, "Invalid vertex index"
+ assert tri.simplices.max() < len(tri.points), "Vertex index out of range"
+
+ # Check triangle vertex coordinates
+ assert hasattr(tri, "tri_lat_verts"), "Triangle lat vertices missing"
+ assert hasattr(tri, "tri_lon_verts"), "Triangle lon vertices missing"
+ assert len(tri.tri_lat_verts) == len(
+ tri.simplices
+ ), "Lat vertices count mismatch"
+ assert len(tri.tri_lon_verts) == len(
+ tri.simplices
+ ), "Lon vertices count mismatch"
+
+ # @pytest.mark.skip(reason="Requires complete params object - advanced test")
+ def test_first_appx_interface(self, test_data, mock_params):
+ """Test first approximation interface."""
+ grid, topo, reader = test_data
+
+ # Delaunay decomposition
+ tri = delaunay.get_decomposition(topo, xnp=5, ynp=4, padding=reader.padding)
+
+ rect_idx = 0
+ nhi = 12
+ nhj = 12
+
+ # Get reference cell
+ simplex_lat = tri.tri_lat_verts[rect_idx]
+ simplex_lon = tri.tri_lon_verts[rect_idx]
+
+ # Create first approximation object
+ fa = interface.first_appx(nhi, nhj, mock_params, topo)
+
+ # Run first approximation
+ cell_fa, ampls_fa, uw_fa, dat_2D_fa = fa.do(simplex_lat, simplex_lon)
+
+ # Verify results
+ assert cell_fa is not None, "Cell not returned"
+ assert ampls_fa is not None, "Amplitudes not computed"
+ assert uw_fa is not None, "PMF not computed"
+ assert dat_2D_fa is not None, "Reconstruction not computed"
+ assert ampls_fa.shape == (
+ nhj,
+ nhi,
+ ), f"Unexpected amplitude shape: {ampls_fa.shape}"
+
+ # @pytest.mark.skip(reason="Requires complete params object - advanced test")
+ def test_second_appx_interface(self, test_data, mock_params):
+ """Test second approximation interface."""
+ grid, topo, reader = test_data
+
+ # Delaunay decomposition
+ tri = delaunay.get_decomposition(topo, xnp=5, ynp=4, padding=reader.padding)
+
+ rect_idx = 0
+ nhi = 12
+ nhj = 12
+
+ # First approximation
+ simplex_lat = tri.tri_lat_verts[rect_idx]
+ simplex_lon = tri.tri_lon_verts[rect_idx]
+
+ fa = interface.first_appx(nhi, nhj, mock_params, topo)
+ cell_fa, ampls_fa, uw_fa, dat_2D_fa = fa.do(simplex_lat, simplex_lon)
+
+ # Second approximation
+ sa = interface.second_appx(nhi, nhj, mock_params, topo, tri)
+
+ # Process first triangle
+ idx = rect_idx
+ sols = sa.do(idx, ampls_fa)
+
+ cell, ampls_sa, uw_sa, dat_2D_sa = sols
+
+ # Verify results
+ assert cell is not None, "Cell not returned"
+ assert ampls_sa is not None, "Second approx amplitudes not computed"
+ assert uw_sa is not None, "PMF not computed"
+ assert dat_2D_sa is not None, "Reconstruction not computed"
+
+ # @pytest.mark.skip(reason="Requires complete params object - advanced test")
+ def test_triangle_pair_workflow(self, test_data, mock_params):
+ """Test complete triangle pair processing workflow."""
+ grid, topo, reader = test_data
+
+ # Delaunay decomposition
+ tri = delaunay.get_decomposition(topo, xnp=5, ynp=4, padding=reader.padding)
+
+ rect_idx = 0
+ nhi = 12
+ nhj = 12
+
+ # First approximation
+ simplex_lat = tri.tri_lat_verts[rect_idx]
+ simplex_lon = tri.tri_lon_verts[rect_idx]
+
+ fa = interface.first_appx(nhi, nhj, mock_params, topo)
+ cell_fa, ampls_fa, uw_fa, dat_2D_fa = fa.do(simplex_lat, simplex_lon)
+
+ # Second approximation on both triangles
+ sa = interface.second_appx(nhi, nhj, mock_params, topo, tri)
+
+ triangle_pair = []
+ for idx in [rect_idx, rect_idx + 1]:
+ cell, ampls_sa, uw_sa, dat_2D_sa = sa.do(idx, ampls_fa)
+ cell.uw = uw_sa
+ triangle_pair.append(cell)
+
+ # Verify triangle pair
+ assert len(triangle_pair) == 2, "Triangle pair should contain 2 triangles"
+ assert triangle_pair[0].topo is not None
+ assert triangle_pair[1].topo is not None
+ assert triangle_pair[0].analysis is not None
+ assert triangle_pair[1].analysis is not None
+
+
+@pytest.mark.integration
+class TestDelaunayDiagnostics:
+ """Test diagnostics for Delaunay workflow."""
+
+ @pytest.fixture
+ def mock_params(self):
+ """Create mock params."""
+
+ class MockParams:
+ run_case = "TEST"
+ rect_set = [0, 2]
+ padding = 10
+
+ return MockParams()
+
+ @pytest.fixture
+ def mock_triangle_pair(self):
+ """Create mock triangle pair for diagnostics testing."""
+ cell1 = var.topo_cell()
+ cell1.topo = np.random.randn(50, 50) * 100
+ cell1.lat = np.linspace(60, 61, 50)
+ cell1.lon = np.linspace(-150, -149, 50)
+ cell1.mask = np.ones((50, 50), dtype=bool)
+ cell1.uw = 1500.0
+
+ analysis1 = var.analysis()
+ analysis1.ampls = np.random.randn(12, 12) * 10
+ analysis1.recon = np.random.randn(50, 50) * 80
+ cell1.analysis = analysis1
+
+ cell2 = var.topo_cell()
+ cell2.topo = np.random.randn(50, 50) * 100
+ cell2.lat = np.linspace(60, 61, 50)
+ cell2.lon = np.linspace(-150, -149, 50)
+ cell2.mask = np.ones((50, 50), dtype=bool)
+ cell2.uw = 1200.0
+
+ analysis2 = var.analysis()
+ analysis2.ampls = np.random.randn(12, 12) * 10
+ analysis2.recon = np.random.randn(50, 50) * 80
+ cell2.analysis = analysis2
+
+ return [cell1, cell2]
+
+ @pytest.mark.skip(reason="Diagnostics API needs verification")
+ def test_diagnostics_basic(self, mock_params):
+ """Test basic diagnostics initialization."""
+
+ # Create mock triangulation
+ class MockTri:
+ simplices = np.array([[0, 1, 2], [1, 2, 3], [2, 3, 4]])
+
+ tri = MockTri()
+
+ diag = diagnostics.delaunay_metrics(mock_params, tri, writer=None)
+
+ # Just check it initializes without error
+ assert diag is not None
+ assert hasattr(diag, "rect_set")
diff --git a/tests/integration/test_idealised_delaunay.py b/tests/integration/test_idealised_delaunay.py
new file mode 100644
index 0000000..ccbad33
--- /dev/null
+++ b/tests/integration/test_idealised_delaunay.py
@@ -0,0 +1,348 @@
+"""
+Integration test for idealised Delaunay case with Perlin noise terrain.
+
+Tests CSA on synthetic terrain generated using Perlin noise,
+which provides more realistic multi-scale topography than pure sinusoids.
+"""
+
+import pytest
+import numpy as np
+from pycsa import var, utils, interface
+
+try:
+ import noise
+
+ NOISE_AVAILABLE = True
+except ImportError:
+ NOISE_AVAILABLE = False
+
+
+@pytest.mark.integration
+@pytest.mark.skipif(not NOISE_AVAILABLE, reason="noise package not available")
+class TestIdealisedDelaunay:
+ """Test CSA on Perlin noise synthetic terrain."""
+
+ @pytest.fixture
+ def perlin_terrain(self):
+ """Generate synthetic terrain using Perlin noise."""
+ res_x = res_y = 120 # Smaller for faster tests
+ scale_fac = 2000.0
+
+ shape = (res_x, res_y)
+ scale = 60.0
+ octaves = 6
+ persistence = 0.5
+ lacunarity = 2.0
+
+ world = np.zeros(shape)
+ for i in range(shape[0]):
+ for j in range(shape[1]):
+ world[i][j] = noise.pnoise2(
+ i / scale,
+ j / scale,
+ octaves=octaves,
+ persistence=persistence,
+ lacunarity=lacunarity,
+ repeatx=1024,
+ repeaty=1024,
+ base=42, # Fixed seed for reproducibility
+ )
+
+ world -= world.mean()
+ world /= world.max()
+ world *= scale_fac
+
+ return world, res_x, res_y, scale_fac
+
+ @pytest.fixture
+ def cosine_terrain(self):
+ """Generate simple cosine background terrain."""
+ res_x = res_y = 120
+ scale_fac = 2000.0
+
+ xx = np.linspace(0, 2.0 * np.pi * scale_fac, res_x)
+ X, Y = np.meshgrid(xx, xx)
+ kl = 1.0 / scale_fac
+
+ bg = -(scale_fac / 2.0) * (np.cos(kl * X + kl * Y))
+
+ return bg, res_x, res_y, scale_fac
+
+ def test_perlin_terrain_generation(self, perlin_terrain):
+ """Test that Perlin noise terrain is generated correctly."""
+ world, res_x, res_y, scale_fac = perlin_terrain
+
+ # Check shape
+ assert world.shape == (res_x, res_y), "Terrain shape incorrect"
+
+ # Check values are in expected range
+ assert np.abs(world).max() <= scale_fac, "Terrain values exceed scale factor"
+
+ # Check terrain has variation (not constant)
+ assert world.std() > 0, "Terrain has no variation"
+
+ # Check mean is close to zero (normalized)
+ assert np.abs(world.mean()) < 1.0, "Terrain mean not centered at zero"
+
+ def test_csa_on_perlin_terrain(self, perlin_terrain):
+ """Test CSA pipeline on Perlin noise terrain."""
+ world, res_x, res_y, scale_fac = perlin_terrain
+
+ # CSA parameters
+ U, V = 10.0, 0.0
+ nhi, nhj = 24, 48
+
+ # Initialize
+ grid = var.grid()
+ cell = var.topo_cell()
+ cell.topo = world
+
+ # Create isosceles triangle
+ vid = utils.isosceles(
+ grid,
+ cell,
+ ymax=2.0 * np.pi * scale_fac,
+ xmax=2.0 * np.pi * scale_fac,
+ res=res_x,
+ )
+
+ lat_v = grid.clat_vertices[vid, :]
+ lon_v = grid.clon_vertices[vid, :]
+
+ cell.gen_mgrids()
+
+ # Create triangle mask
+ triangle = utils.gen_triangle(lon_v, lat_v)
+ cell.get_masked(triangle=triangle)
+
+ cell.wlat = np.diff(cell.lat).mean()
+ cell.wlon = np.diff(cell.lon).mean()
+
+ # Run CSA
+ run = interface.get_pmf(nhi, nhj, U, V)
+ ampls, uw, recon = run.sappx(cell, lmbda=1e-3, iter_solve=False)
+
+ # Verify results
+ assert ampls is not None, "Amplitudes not computed"
+ assert ampls.shape == (nhj, nhi), f"Unexpected amplitude shape: {ampls.shape}"
+ assert not np.all(np.isnan(ampls)), "All amplitudes are NaN"
+
+ assert uw is not None, "PMF not computed"
+ # PMF can be scalar or array depending on configuration
+ if isinstance(uw, np.ndarray):
+ assert uw.size > 0, "PMF array is empty"
+ else:
+ assert isinstance(uw, (int, float, np.number)), "PMF should be numeric"
+
+ assert recon is not None, "Reconstruction not computed"
+ assert recon.shape == cell.topo.shape, "Reconstruction shape mismatch"
+
+ def test_csa_on_cosine_terrain(self, cosine_terrain):
+ """Test CSA on simple cosine terrain (should recover mode perfectly)."""
+ bg, res_x, res_y, scale_fac = cosine_terrain
+
+ # CSA parameters
+ U, V = 10.0, 0.0
+ nhi, nhj = 12, 24
+
+ # Initialize
+ grid = var.grid()
+ cell = var.topo_cell()
+ cell.topo = bg
+
+ # Create isosceles triangle
+ vid = utils.isosceles(
+ grid,
+ cell,
+ ymax=2.0 * np.pi * scale_fac,
+ xmax=2.0 * np.pi * scale_fac,
+ res=res_x,
+ )
+
+ lat_v = grid.clat_vertices[vid, :]
+ lon_v = grid.clon_vertices[vid, :]
+
+ cell.gen_mgrids()
+
+ # Create triangle mask
+ triangle = utils.gen_triangle(lon_v, lat_v)
+ cell.get_masked(triangle=triangle)
+
+ cell.wlat = np.diff(cell.lat).mean()
+ cell.wlon = np.diff(cell.lon).mean()
+
+ # Run CSA with regularization
+ run = interface.get_pmf(nhi, nhj, U, V)
+ ampls, uw, recon = run.sappx(cell, lmbda=1e-4, iter_solve=False)
+
+ # For a single cosine mode, we should have:
+ # - Most energy concentrated in one or a few modes
+ # - Good reconstruction quality
+
+ ampls_clean = np.nan_to_num(ampls)
+
+ # Check that we have non-zero amplitudes
+ assert np.any(ampls_clean != 0), "No modes recovered"
+
+ # Check that energy is concentrated (not uniform)
+ max_ampl = np.abs(ampls_clean).max()
+ mean_ampl = np.abs(ampls_clean).mean()
+ assert max_ampl > 3 * mean_ampl, "Energy should be concentrated in few modes"
+
+ def test_mode_selection_on_perlin_terrain(self, perlin_terrain):
+ """Test mode selection (top-N modes) on Perlin terrain."""
+ world, res_x, res_y, scale_fac = perlin_terrain
+
+ # CSA parameters
+ U, V = 10.0, 0.0
+ nhi, nhj = 24, 48
+ n_modes = 20
+
+ # Initialize
+ grid = var.grid()
+ cell = var.topo_cell()
+ cell.topo = world
+
+ # Create isosceles triangle
+ vid = utils.isosceles(
+ grid,
+ cell,
+ ymax=2.0 * np.pi * scale_fac,
+ xmax=2.0 * np.pi * scale_fac,
+ res=res_x,
+ )
+
+ lat_v = grid.clat_vertices[vid, :]
+ lon_v = grid.clon_vertices[vid, :]
+
+ cell.gen_mgrids()
+
+ triangle = utils.gen_triangle(lon_v, lat_v)
+ cell.get_masked(triangle=triangle)
+
+ cell.wlat = np.diff(cell.lat).mean()
+ cell.wlon = np.diff(cell.lon).mean()
+
+ # First approximation (get full spectrum)
+ first_appx = interface.get_pmf(nhi, nhj, U, V)
+ ampls_fa, uw_fa, recon_fa = first_appx.sappx(cell, lmbda=1e-2, iter_solve=False)
+
+ # Select top N modes
+ fq_cpy = np.copy(ampls_fa)
+ fq_cpy[np.isnan(fq_cpy)] = 0.0
+
+ indices = []
+ for ii in range(n_modes):
+ max_idx = np.unravel_index(fq_cpy.argmax(), fq_cpy.shape)
+ indices.append(max_idx)
+ max_val = fq_cpy[max_idx]
+ fq_cpy[max_idx] = 0.0
+
+ k_idxs = [pair[1] for pair in indices]
+ l_idxs = [pair[0] for pair in indices]
+
+ # Verify mode selection
+ assert len(k_idxs) == n_modes, "Incorrect number of k indices"
+ assert len(l_idxs) == n_modes, "Incorrect number of l indices"
+
+ # All indices should be within bounds
+ assert all(0 <= k < nhi for k in k_idxs), "k index out of bounds"
+ assert all(0 <= l < nhj for l in l_idxs), "l index out of bounds"
+
+ # Second approximation with selected modes
+ second_appx = interface.get_pmf(nhi, nhj, U, V)
+ second_appx.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+
+ ampls_sa, uw_sa, recon_sa = second_appx.sappx(
+ cell, lmbda=1e-5, updt_analysis=True, scale=1.0, iter_solve=False
+ )
+
+ # Verify second approximation
+ assert ampls_sa is not None, "Second approx failed"
+ assert not np.all(np.isnan(ampls_sa)), "Second approx all NaN"
+
+ # Second approximation should use fewer modes
+ ampls_sa_clean = np.nan_to_num(ampls_sa)
+ n_nonzero = np.sum(ampls_sa_clean != 0)
+ assert n_nonzero <= n_modes + 5, f"Too many modes in second approx: {n_nonzero}"
+
+ def test_deterministic_perlin_generation(self):
+ """Test that Perlin noise generation is deterministic with fixed seed."""
+
+ # Generate twice with same parameters
+ def generate_perlin():
+ res = 50
+ scale_fac = 1000.0
+ world = np.zeros((res, res))
+ for i in range(res):
+ for j in range(res):
+ world[i][j] = noise.pnoise2(
+ i / 30.0,
+ j / 30.0,
+ octaves=4,
+ persistence=0.5,
+ lacunarity=2.0,
+ repeatx=1024,
+ repeaty=1024,
+ base=42, # Fixed seed
+ )
+ return world
+
+ world1 = generate_perlin()
+ world2 = generate_perlin()
+
+ # Should be identical
+ np.testing.assert_array_equal(
+ world1, world2, err_msg="Perlin noise generation is not deterministic"
+ )
+
+ def test_reconstruction_quality(self, cosine_terrain):
+ """Test that reconstruction quality is reasonable for known terrain."""
+ bg, res_x, res_y, scale_fac = cosine_terrain
+
+ # CSA parameters
+ U, V = 10.0, 0.0
+ nhi, nhj = 24, 48
+
+ # Initialize
+ grid = var.grid()
+ cell = var.topo_cell()
+ cell.topo = bg
+
+ # Create isosceles triangle
+ vid = utils.isosceles(
+ grid,
+ cell,
+ ymax=2.0 * np.pi * scale_fac,
+ xmax=2.0 * np.pi * scale_fac,
+ res=res_x,
+ )
+
+ lat_v = grid.clat_vertices[vid, :]
+ lon_v = grid.clon_vertices[vid, :]
+
+ cell.gen_mgrids()
+
+ triangle = utils.gen_triangle(lon_v, lat_v)
+ cell.get_masked(triangle=triangle)
+
+ cell.wlat = np.diff(cell.lat).mean()
+ cell.wlon = np.diff(cell.lon).mean()
+
+ # Run CSA
+ run = interface.get_pmf(nhi, nhj, U, V)
+ ampls, uw, recon = run.sappx(cell, lmbda=1e-4, iter_solve=False)
+
+ # Compute reconstruction error
+ # Only compare where mask is True
+ original_masked = cell.topo * cell.mask
+ recon_masked = recon * cell.mask
+
+ # Relative L2 error
+ l2_error = np.linalg.norm(original_masked - recon_masked) / np.linalg.norm(
+ original_masked
+ )
+
+ # For a simple cosine, reconstruction should be good
+ # (not perfect due to triangular domain and regularization)
+ assert l2_error < 0.5, f"Reconstruction error too high: {l2_error:.3f}"
diff --git a/tests/integration/test_idealised_isosceles.py b/tests/integration/test_idealised_isosceles.py
new file mode 100644
index 0000000..e3037be
--- /dev/null
+++ b/tests/integration/test_idealised_isosceles.py
@@ -0,0 +1,271 @@
+"""
+Integration test for idealised isosceles triangle case.
+
+This test runs the full CSA pipeline on synthetic terrain with an isosceles
+triangular domain and compares results against baseline values from the
+published JAMES paper.
+"""
+
+import numpy as np
+import pytest
+from pycsa import var, utils, interface
+from copy import deepcopy
+
+
+class TestIdealisedIsosceles:
+ """Test suite for the idealised isosceles triangle case."""
+
+ @pytest.fixture
+ def baseline_results(self):
+ """Baseline numerical results from the JAMES paper."""
+ return {
+ "num_modes": 22,
+ "amplitudes": np.array(
+ [
+ 1243.29667409,
+ 1110972.57606147,
+ 1861.67185697,
+ 1243.32433928,
+ 1146.82593374,
+ 1110972.57606147,
+ ]
+ ),
+ "l2_errors": np.array(
+ [
+ 0.0,
+ 164291.56804783,
+ 115.71273229,
+ 85.67668202,
+ 111.37226442,
+ 164291.56804783,
+ ]
+ ),
+ "percentage_errors": np.array(
+ [0.0, 89256.997, 49.737, 0.002, 7.759, 89256.997]
+ ),
+ }
+
+ @pytest.fixture
+ def synthetic_terrain(self):
+ """Generate the synthetic terrain with known spectral content."""
+ np.random.seed(777)
+
+ # Generate random spectral modes
+ sz = 25
+ nk = np.random.randint(0, 12, size=sz)
+ nl = np.random.randint(-5, 7, size=sz)
+
+ for ii in range(sz):
+ if nk[ii] == 0 and nl[ii] < 0:
+ nk[ii] += np.random.randint(1, 11)
+ pts = [item for item in zip(nk, nl)]
+ pts = np.array(list(set(pts)))
+
+ nk = pts[:, 0]
+ nl = pts[:, 1]
+ sz = len(pts)
+
+ Ak = np.random.random(size=sz) * 100.0
+ Al = np.random.random(size=sz) * 100.0
+ sck = np.random.randint(0, 2, size=sz)
+ scl = np.random.randint(0, 2, size=sz)
+
+ return {
+ "nk": nk,
+ "nl": nl,
+ "Ak": Ak,
+ "Al": Al,
+ "sck": sck,
+ "scl": scl,
+ "sz": sz,
+ "pts": pts,
+ }
+
+ @pytest.fixture
+ def isosceles_cell(self, synthetic_terrain):
+ """Create an isosceles triangle cell with synthetic topography."""
+ nhi = 12
+ nhj = 12
+
+ # Initialize triangle
+ grid = var.grid()
+ cell = var.topo_cell()
+ vid = utils.isosceles(grid, cell)
+
+ lat_v = grid.clat_vertices[vid, :]
+ lon_v = grid.clon_vertices[vid, :]
+
+ cell.gen_mgrids()
+
+ # Fill with synthetic topography
+ cell.topo = np.zeros_like(cell.lat_grid)
+
+ def sinusoidal_basis(Ak, nk, Al, nl, sc):
+ nk_scaled = 2.0 * np.pi * nk / cell.lon.max()
+ nl_scaled = 2.0 * np.pi * nl / cell.lat.max()
+
+ if sc == 0:
+ bf = Ak * np.cos(nk_scaled * cell.lon_grid + nl_scaled * cell.lat_grid)
+ else:
+ bf = Al * np.sin(nk_scaled * cell.lon_grid + nl_scaled * cell.lat_grid)
+
+ return bf
+
+ terrain = synthetic_terrain
+ for ii in range(terrain["sz"]):
+ cell.topo += sinusoidal_basis(
+ terrain["Ak"][ii],
+ terrain["nk"][ii],
+ terrain["Al"][ii],
+ terrain["nl"][ii],
+ terrain["sck"][ii],
+ )
+
+ # Define triangle mask
+ triangle = utils.gen_triangle(lon_v, lat_v)
+ cell.get_masked(triangle=triangle)
+
+ cell.wlat = np.diff(cell.lat).mean()
+ cell.wlon = np.diff(cell.lon).mean()
+
+ return cell, triangle, terrain["sz"]
+
+ def test_spectral_approximation(
+ self, isosceles_cell, synthetic_terrain, baseline_results
+ ):
+ """Test that CSA pipeline runs and produces consistent results."""
+ cell, triangle, sz = isosceles_cell
+ terrain = synthetic_terrain
+
+ nhi = 12
+ nhj = 12
+ n_modes = 14
+ lmbda_reg = 8.0 * 1e-5
+ lmbda_fg = 1e-1
+ lmbda_sg = 1e-6
+
+ # Artificial winds (not used in idealised test)
+ U, V = 1.0, 1.0
+
+ # Build reference spectrum from known terrain components
+ freqs_ref = np.zeros((nhi, nhj))
+ cnt = 0
+ for pt in terrain["pts"]:
+ kk, ll = pt
+ ll += 5 # Offset as in original script
+ freqs_ref[ll, kk] = terrain["Ak"][cnt]
+ cnt += 1
+
+ # Run pure LSFF
+ pure_lsff = interface.get_pmf(nhi, nhj, U, V)
+ freqs_plsff, _, _ = pure_lsff.sappx(
+ cell, lmbda=0.0, iter_solve=False, save_am=True
+ )
+
+ # Run regularized LSFF
+ reg_lsff = interface.get_pmf(nhi, nhj, U, V)
+ freqs_rlsff, _, _ = reg_lsff.sappx(cell, lmbda=lmbda_reg, iter_solve=False)
+
+ # Run CSA (first approximation + mode selection + second approximation)
+ first_guess = interface.get_pmf(nhi, nhj, U, V)
+
+ # First approximation on quadrilateral domain
+ cell_fa = deepcopy(cell)
+ cell_fa.get_masked(mask=np.ones_like(cell.topo).astype("bool"))
+ cell_fa.wlat = np.diff(cell_fa.lat).mean()
+ cell_fa.wlon = np.diff(cell_fa.lon).mean()
+
+ freqs_fg, _, _ = first_guess.sappx(cell_fa, lmbda=lmbda_fg, iter_solve=False)
+
+ # Select top N modes
+ fq_cpy = np.copy(freqs_fg)
+ fq_cpy[np.isnan(fq_cpy)] = 0.0
+
+ indices = []
+ for ii in range(n_modes):
+ max_idx = np.unravel_index(fq_cpy.argmax(), fq_cpy.shape)
+ indices.append(max_idx)
+ fq_cpy[max_idx] = 0.0
+
+ k_idxs = [pair[1] for pair in indices]
+ l_idxs = [pair[0] for pair in indices]
+
+ # Second approximation on triangular domain
+ second_guess = interface.get_pmf(nhi, nhj, U, V)
+ second_guess.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+
+ cell_sa = deepcopy(cell)
+ cell_sa.get_masked(triangle=triangle)
+ cell_sa.wlat = np.diff(cell_sa.lat).mean()
+ cell_sa.wlon = np.diff(cell_sa.lon).mean()
+
+ freqs_csa, _, _ = second_guess.sappx(
+ cell_sa, lmbda=lmbda_sg, updt_analysis=True, scale=1.0, iter_solve=False
+ )
+
+ # Clean up NaN values
+ freqs_plsff = np.nan_to_num(freqs_plsff)
+ freqs_rlsff = np.nan_to_num(freqs_rlsff)
+ freqs_csa = np.nan_to_num(freqs_csa)
+ freqs_ref = np.nan_to_num(freqs_ref)
+
+ # Compute L2 errors against reference
+ err_plsff = np.linalg.norm(freqs_plsff - freqs_ref)
+ err_rlsff = np.linalg.norm(freqs_rlsff - freqs_ref)
+ err_csa = np.linalg.norm(freqs_csa - freqs_ref)
+
+ # Compare against baseline with reasonable tolerance
+ # The baseline L2 errors are: [0, 164291.57, 115.71, 85.68, 111.37, 164291.57]
+ # Where indices are: [ref, pLSFF, rLSFF, optCSA, subCSA, quad]
+ # We're running subCSA (n_modes=14), so compare against baseline[4] = 111.37
+
+ # For now, just check that computations run and produce reasonable values
+ assert err_plsff > 1000, "Pure LSFF should have large error (overfits)"
+ assert err_rlsff > 0, "Regularized LSFF should have some error"
+ assert err_csa > 0, "CSA should have some error"
+ assert err_csa < err_plsff, "CSA should perform better than pure LSFF"
+
+ # Check that we're in the right ballpark (within factor of 2)
+ assert (
+ 50 < err_csa < 250
+ ), f"CSA L2 error {err_csa:.2f} should be ~111 (baseline)"
+
+ # Amplitude sums should be positive
+ sum_plsff = freqs_plsff.sum()
+ sum_rlsff = freqs_rlsff.sum()
+ sum_csa = freqs_csa.sum()
+
+ assert sum_plsff > 0, "Pure LSFF amplitude sum should be positive"
+ assert sum_rlsff > 0, "Regularized LSFF amplitude sum should be positive"
+ assert sum_csa > 0, "CSA amplitude sum should be positive"
+
+ def test_mode_count(self, synthetic_terrain, baseline_results):
+ """Test that the correct number of unique modes are generated."""
+ sz = synthetic_terrain["sz"]
+
+ # Should match baseline number of unique modes
+ assert (
+ sz == baseline_results["num_modes"]
+ ), f"Expected {baseline_results['num_modes']} unique modes, got {sz}"
+
+ def test_deterministic_terrain_generation(self):
+ """Test that terrain generation is deterministic with fixed seed."""
+ np.random.seed(777)
+
+ # Generate terrain twice with same seed
+ sz1 = 25
+ nk1 = np.random.randint(0, 12, size=sz1)
+ nl1 = np.random.randint(-5, 7, size=sz1)
+
+ np.random.seed(777)
+
+ sz2 = 25
+ nk2 = np.random.randint(0, 12, size=sz2)
+ nl2 = np.random.randint(-5, 7, size=sz2)
+
+ np.testing.assert_array_equal(
+ nk1, nk2, err_msg="Terrain generation is not deterministic"
+ )
+ np.testing.assert_array_equal(
+ nl1, nl2, err_msg="Terrain generation is not deterministic"
+ )
diff --git a/tests/test_dynamic_memory.py b/tests/test_dynamic_memory.py
new file mode 100644
index 0000000..499605f
--- /dev/null
+++ b/tests/test_dynamic_memory.py
@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""
+Test script for dynamic memory allocation based on cell latitude.
+
+This verifies that:
+1. Memory estimation function works correctly
+2. Cells are properly grouped by memory requirements
+3. Configuration makes sense for different hardware setups
+"""
+
+import numpy as np
+from pycsa.core import io, var, utils
+
+# Import the new functions
+import sys
+
+sys.path.insert(0, "/home/ray/git-projects/spec_appx/runs")
+from icon_etopo_global import estimate_cell_memory_gb, group_cells_by_memory
+
+
+def test_memory_estimation():
+ """Test that memory estimation scales appropriately with latitude."""
+ print("=" * 80)
+ print("TEST 1: Memory Estimation Function")
+ print("=" * 80)
+
+ test_latitudes = [0, 30, 45, 60, 70, 75, 80, 85, 89]
+
+ print("\nMemory requirements by latitude:")
+ print(f"{'Latitude':<12} {'Memory (GB)':<15} {'Scale Factor':<15}")
+ print("-" * 42)
+
+ base_mem = estimate_cell_memory_gb(0)
+ for lat in test_latitudes:
+ mem_gb = estimate_cell_memory_gb(lat)
+ scale = mem_gb / base_mem
+ print(f"{lat:>3}° {mem_gb:>6.1f} GB {scale:>5.2f}x")
+
+ # Verify expectations
+ assert estimate_cell_memory_gb(0) == 10.0, "Equatorial cells should need 10 GB"
+ assert (
+ estimate_cell_memory_gb(85) >= 50.0
+ ), "Polar cells (~85°) should need >= 50 GB"
+ print("\n✓ Memory estimation function passes basic tests")
+
+
+def test_cell_grouping():
+ """Test that cells are properly grouped by memory requirements."""
+ print("\n" + "=" * 80)
+ print("TEST 2: Cell Grouping by Memory")
+ print("=" * 80)
+
+ # Load actual ICON grid to get realistic cell latitudes
+ print("\nLoading ICON grid...")
+ from inputs.icon_global_run import params
+
+ grid = var.grid()
+ reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+ reader.read_dat(params.path_icon_grid, grid)
+
+ clat_rad = grid.clat
+ n_cells = len(clat_rad)
+
+ print(f"Loaded {n_cells} cells")
+ print(
+ f"Latitude range: {np.rad2deg(clat_rad.min()):.1f}° to {np.rad2deg(clat_rad.max()):.1f}°"
+ )
+
+ # Test for laptop configuration (60 GB total)
+ print("\n--- LAPTOP CONFIGURATION (60 GB total) ---")
+ batches_laptop = group_cells_by_memory(clat_rad, max_memory_per_batch_gb=60.0)
+
+ print(f"\nCreated {len(batches_laptop)} memory batches:")
+ total_cells_batched = 0
+ for i, batch in enumerate(batches_laptop):
+ n = len(batch["cell_indices"])
+ total_cells_batched += n
+ print(
+ f" Batch {i}: {n:>6} cells, "
+ f"{batch['memory_per_cell_gb']:>5.1f} GB/cell, "
+ f"{batch['n_workers']:>2} workers × {batch['memory_per_worker_gb']:>5.1f} GB = "
+ f"{batch['n_workers'] * batch['memory_per_worker_gb']:>6.1f} GB total"
+ )
+
+ assert (
+ total_cells_batched == n_cells
+ ), f"All cells should be batched (got {total_cells_batched}, expected {n_cells})"
+ print(f"\n✓ All {n_cells} cells properly batched")
+
+ # Test for HPC configuration (240 GB total)
+ print("\n--- HPC CONFIGURATION (240 GB total) ---")
+ batches_hpc = group_cells_by_memory(clat_rad, max_memory_per_batch_gb=240.0)
+
+ print(f"\nCreated {len(batches_hpc)} memory batches:")
+ total_cells_batched = 0
+ for i, batch in enumerate(batches_hpc):
+ n = len(batch["cell_indices"])
+ total_cells_batched += n
+ print(
+ f" Batch {i}: {n:>6} cells, "
+ f"{batch['memory_per_cell_gb']:>5.1f} GB/cell, "
+ f"{batch['n_workers']:>2} workers × {batch['memory_per_worker_gb']:>5.1f} GB = "
+ f"{batch['n_workers'] * batch['memory_per_worker_gb']:>6.1f} GB total"
+ )
+
+ assert (
+ total_cells_batched == n_cells
+ ), f"All cells should be batched (got {total_cells_batched}, expected {n_cells})"
+ print(f"\n✓ All {n_cells} cells properly batched")
+
+ # Verify that HPC has better parallelism (more workers on average)
+ avg_workers_laptop = np.mean([b["n_workers"] for b in batches_laptop])
+ avg_workers_hpc = np.mean([b["n_workers"] for b in batches_hpc])
+
+ print(f"\nAverage workers per batch:")
+ print(f" Laptop: {avg_workers_laptop:.1f}")
+ print(f" HPC: {avg_workers_hpc:.1f}")
+
+ assert (
+ avg_workers_hpc > avg_workers_laptop
+ ), "HPC should have more workers on average"
+ print("✓ HPC configuration properly utilizes more workers")
+
+
+def test_specific_cells():
+ """Test memory estimation for specific problematic cells."""
+ print("\n" + "=" * 80)
+ print("TEST 3: Specific Cell Memory Requirements")
+ print("=" * 80)
+
+ from inputs.icon_global_run import params
+
+ grid = var.grid()
+ reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+ reader.read_dat(params.path_icon_grid, grid)
+
+ clat_rad = grid.clat
+ clat_deg = np.rad2deg(clat_rad)
+
+ # Test cell 16384 (known to need 60 GB)
+ test_cell_idx = 16384
+ if test_cell_idx < len(clat_deg):
+ cell_lat = clat_deg[test_cell_idx]
+ estimated_mem = estimate_cell_memory_gb(cell_lat)
+
+ print(f"\nCell {test_cell_idx}:")
+ print(f" Latitude: {cell_lat:.2f}°")
+ print(f" Estimated memory: {estimated_mem:.1f} GB")
+ print(f" Actual requirement (from tests): 60 GB")
+
+ if estimated_mem >= 50.0:
+ print(" ✓ Estimation is in the right ballpark")
+ else:
+ print(
+ f" ⚠ Estimation may be too low (got {estimated_mem:.1f} GB, expected >= 50 GB)"
+ )
+
+ # Show top 10 most memory-intensive cells
+ cell_memory_gb = np.array([estimate_cell_memory_gb(lat) for lat in clat_deg])
+ top_indices = np.argsort(cell_memory_gb)[-10:][::-1]
+
+ print(f"\nTop 10 most memory-intensive cells:")
+ print(f"{'Cell Index':<12} {'Latitude':<12} {'Est. Memory':<15}")
+ print("-" * 39)
+ for idx in top_indices:
+ print(f"{idx:<12} {clat_deg[idx]:>7.2f}° {cell_memory_gb[idx]:>6.1f} GB")
+
+
+if __name__ == "__main__":
+ try:
+ test_memory_estimation()
+ test_cell_grouping()
+ test_specific_cells()
+
+ print("\n" + "=" * 80)
+ print("ALL TESTS PASSED ✓")
+ print("=" * 80)
+
+ except Exception as e:
+ print(f"\n❌ TEST FAILED: {e}")
+ import traceback
+
+ traceback.print_exc()
+ sys.exit(1)
diff --git a/tests/test_etopo_edge_cases.py b/tests/test_etopo_edge_cases.py
new file mode 100644
index 0000000..23950c0
--- /dev/null
+++ b/tests/test_etopo_edge_cases.py
@@ -0,0 +1,230 @@
+#!/usr/bin/env python3
+"""
+ETOPO Edge Case Tests - Similar to test_merit_edge_cases.py
+
+Tests critical latitude/longitude boundaries where tile loading might fail.
+Includes visualization of edge cases like dateline and prime meridian.
+"""
+
+import sys
+import numpy as np
+
+# Force reload
+for mod in list(sys.modules.keys()):
+ if "pycsa" in mod:
+ del sys.modules[mod]
+
+from pycsa.core import io, var
+from pycsa.plotting import cart_plot
+import matplotlib.pyplot as plt
+
+
+def test_and_plot_region(lat_extent, lon_extent, description, plot=True):
+ """Test and optionally plot a specific region."""
+ print(f"\nTest: {description}")
+ print(f" Latitude: {lat_extent}")
+ print(f" Longitude: {lon_extent}")
+
+ class Params:
+ def __init__(self):
+ self.path_etopo = "/home/ray/git-projects/spec_appx/data/etopo_15s/"
+ self.lat_extent = lat_extent
+ self.lon_extent = lon_extent
+ self.etopo_cg = 8
+
+ params = Params()
+ cell = var.topo_cell()
+
+ try:
+ loader = io.ncdata.read_etopo_topo(cell, params, verbose=False)
+
+ print(f" ✓ Loaded successfully")
+ print(f" Shape: {cell.topo.shape}")
+ print(f" Lat range: [{cell.lat.min():.2f}, {cell.lat.max():.2f}]")
+ print(f" Lon range: [{cell.lon.min():.2f}, {cell.lon.max():.2f}]")
+ print(f" Elev range: [{cell.topo.min():.0f}, {cell.topo.max():.0f}] m")
+
+ # Plot if requested
+ if plot:
+ cell.gen_mgrids()
+ plt.figure(figsize=(12, 6))
+ ax = plt.subplot(111)
+
+ im = ax.contourf(
+ cell.lon_grid, cell.lat_grid, cell.topo, levels=20, cmap="terrain"
+ )
+ plt.colorbar(im, ax=ax, label="Elevation (m)")
+
+ ax.set_xlabel("Longitude (°)")
+ ax.set_ylabel("Latitude (°)")
+ ax.set_title(description)
+ ax.grid(True, alpha=0.3)
+
+ # Add dateline/meridian markers
+ if (
+ lon_extent[0] <= -180 <= lon_extent[1]
+ or lon_extent[0] <= 180 <= lon_extent[1]
+ ):
+ ax.axvline(
+ 180, color="red", linestyle="--", alpha=0.5, label="Dateline"
+ )
+ ax.axvline(-180, color="red", linestyle="--", alpha=0.5)
+ if lon_extent[0] <= 0 <= lon_extent[1]:
+ ax.axvline(
+ 0, color="blue", linestyle="--", alpha=0.5, label="Prime Meridian"
+ )
+
+ ax.legend()
+
+ # Save plot
+ filename = f"outputs/etopo_edge_case_{description.replace(' ', '_').replace('(', '').replace(')', '').replace('°', 'deg')}.png"
+ plt.savefig(filename, dpi=150, bbox_inches="tight")
+ print(f" Plot saved: {filename}")
+ plt.close()
+
+ return True, cell
+
+ except Exception as e:
+ print(f" ✗ FAILED: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return False, None
+
+
+def run_edge_case_tests():
+ """Run comprehensive edge case tests."""
+ print("=" * 80)
+ print("ETOPO EDGE CASE COMPREHENSIVE TEST SUITE")
+ print("=" * 80)
+ print()
+
+ results = []
+
+ # Test 1: Prime Meridian crossing (0° longitude)
+ print("\n" + "=" * 80)
+ print("TEST 1: PRIME MERIDIAN CROSSING")
+ print("=" * 80)
+ success, cell = test_and_plot_region(
+ lat_extent=[-30.0, 60.0],
+ lon_extent=[-30.0, 30.0],
+ description="Prime Meridian (-30 to 30°E)",
+ plot=True,
+ )
+ results.append(("Prime Meridian", success))
+
+ # Test 2: Dateline crossing (180° longitude)
+ print("\n" + "=" * 80)
+ print("TEST 2: DATELINE CROSSING")
+ print("=" * 80)
+ success, cell = test_and_plot_region(
+ lat_extent=[-30.0, 60.0],
+ lon_extent=[150.0, -150.0], # Crosses dateline
+ description="Dateline Crossing (150°E to 150°W)",
+ plot=True,
+ )
+ results.append(("Dateline", success))
+
+ # Test 3: Full global
+ print("\n" + "=" * 80)
+ print("TEST 3: FULL GLOBAL")
+ print("=" * 80)
+ success, cell = test_and_plot_region(
+ lat_extent=[-90.0, 90.0],
+ lon_extent=[-180.0, 180.0],
+ description="Full Global",
+ plot=True,
+ )
+ results.append(("Full Global", success))
+
+ # Test 4: Himalayas region (multi-tile)
+ print("\n" + "=" * 80)
+ print("TEST 4: HIMALAYAS REGION (Multi-tile)")
+ print("=" * 80)
+ success, cell = test_and_plot_region(
+ lat_extent=[15.0, 45.0],
+ lon_extent=[75.0, 105.0],
+ description="Himalayas (15-45°N, 75-105°E)",
+ plot=True,
+ )
+ if success and cell.topo.max() > 5000:
+ print(f" ✓ High peaks found: {cell.topo.max():.0f}m")
+ max_idx = np.unravel_index(np.argmax(cell.topo), cell.topo.shape)
+ print(
+ f" Location: ({cell.lat[max_idx[0]]:.2f}°N, {cell.lon[max_idx[1]]:.2f}°E)"
+ )
+ results.append(("Himalayas", success))
+
+ # Test 5: Andes region
+ print("\n" + "=" * 80)
+ print("TEST 5: ANDES REGION (Multi-tile)")
+ print("=" * 80)
+ success, cell = test_and_plot_region(
+ lat_extent=[-45.0, -15.0],
+ lon_extent=[-75.0, -60.0],
+ description="Andes (45-15°S, 75-60°W)",
+ plot=True,
+ )
+ if success and cell.topo.max() > 4000:
+ print(f" ✓ High peaks found: {cell.topo.max():.0f}m")
+ results.append(("Andes", success))
+
+ # Test 6: Pacific dateline region (multiple tiles across dateline)
+ print("\n" + "=" * 80)
+ print("TEST 6: PACIFIC DATELINE (Multiple tiles)")
+ print("=" * 80)
+ success, cell = test_and_plot_region(
+ lat_extent=[0.0, 45.0],
+ lon_extent=[165.0, -165.0],
+ description="Pacific Dateline (165°E to 165°W)",
+ plot=True,
+ )
+ results.append(("Pacific Dateline", success))
+
+ # Summary
+ print("\n" + "=" * 80)
+ print("EDGE CASE TEST SUMMARY")
+ print("=" * 80)
+
+ passed = sum(1 for _, r in results if r)
+ total = len(results)
+
+ for desc, result in results:
+ status = "✓ PASS" if result else "✗ FAIL"
+ print(f" {status}: {desc}")
+
+ print()
+ print(f"Total: {passed}/{total} tests passed")
+
+ if passed == total:
+ print("\n✓✓✓ ALL EDGE CASE TESTS PASSED ✓✓✓")
+ print("\nPlots saved in outputs/ directory")
+ return True
+ else:
+ print(f"\n✗✗✗ {total - passed} TEST(S) FAILED ✗✗✗")
+ return False
+
+
+if __name__ == "__main__":
+ # Create outputs directory if it doesn't exist
+ import os
+
+ os.makedirs("outputs", exist_ok=True)
+
+ success = run_edge_case_tests()
+
+ print("\n" + "=" * 80)
+ print("ETOPO LOADER STATUS")
+ print("=" * 80)
+ print("✓ Dateline bug FIXED - can load lon_extent = [-180, 180]")
+ print("✓ Tile assembly bug FIXED - all latitude bands now load correctly")
+ print("✓ Edge cases working - prime meridian, dateline, full global")
+ print()
+ print("Note: Coarse-graining (CG) affects peak elevations:")
+ print(" - CG=1-2: Best accuracy (~8500m for Everest)")
+ print(" - CG=4: Good accuracy (~7000m)")
+ print(" - CG=8: Moderate (~6000m) - used in these tests")
+ print(" - CG=16: Heavy smoothing (~4500m)")
+ print("=" * 80)
+
+ sys.exit(0 if success else 1)
diff --git a/tests/test_etopo_global_plot.py b/tests/test_etopo_global_plot.py
new file mode 100755
index 0000000..702cf65
--- /dev/null
+++ b/tests/test_etopo_global_plot.py
@@ -0,0 +1,361 @@
+#!/usr/bin/env python3
+"""
+Test script to load ALL ETOPO data and plot it on a globe.
+
+This script validates that:
+1. The ETOPO loader can handle large extent regions, including full global coverage
+2. Coarse-graining works correctly to speed up loading and plotting
+3. The cart_plotter can visualize large datasets on a globe
+4. Data values are reasonable (elevation ranges)
+
+Author: Test Suite
+Date: 2025-10-22
+Updated: Fixed to support full global extent
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import time
+from pathlib import Path
+
+# Import CSA modules
+from pycsa.core import io, var
+from pycsa.plotting import cart_plot
+
+
+def create_global_params(etopo_cg=8):
+ """
+ Create parameters for global ETOPO data loading.
+
+ Parameters
+ ----------
+ etopo_cg : int, optional
+ Coarse-graining factor (default: 8)
+ - 1: Full resolution (~463m at equator) - VERY SLOW, huge memory
+ - 2: ~926m - Still very slow
+ - 4: ~1.85km - Moderate speed
+ - 8: ~3.70km - Good balance for global plots
+ - 16: ~7.4km - Fast, good for testing
+
+ Returns
+ -------
+ params : object
+ Parameter object with required attributes
+ """
+
+ class Params:
+ def __init__(self):
+ # Path to ETOPO data directory
+ self.path_etopo = "/home/ray/git-projects/spec_appx/data/etopo_15s/"
+
+ # Full global extent: entire world
+ self.lat_extent = [-90.0, 90.0]
+ self.lon_extent = [-180.0, 180.0]
+
+ # Coarse-graining factor to speed up loading
+ self.etopo_cg = etopo_cg
+
+ return Params()
+
+
+def test_global_etopo_load_and_plot():
+ """
+ Main test function: Load global ETOPO data and plot on globe.
+ """
+ print("=" * 80)
+ print("GLOBAL ETOPO DATA LOADING AND PLOTTING TEST")
+ print("=" * 80)
+ print()
+
+ # Configuration
+ coarse_grain_factor = 8 # 8x8 averaging for reasonable speed
+ plot_stride = 1 # Use all loaded data points for plotting
+
+ print(f"Configuration:")
+ print(f" - Region: Full Global (-90 to 90°N, -180 to 180°E)")
+ print(f" - Coverage: 100% of Earth's surface")
+ print(f" - Coarse-graining: {coarse_grain_factor}x{coarse_grain_factor}")
+ print(f" - Effective resolution: ~{0.463 * coarse_grain_factor:.2f} km at equator")
+ print(f" - Plot stride: every {plot_stride} point(s)")
+ print()
+
+ # Step 1: Create parameters
+ print("Step 1: Creating parameters...")
+ params = create_global_params(etopo_cg=coarse_grain_factor)
+
+ # Verify data directory exists
+ data_path = Path(params.path_etopo)
+ if not data_path.exists():
+ print(f"ERROR: ETOPO data directory not found: {data_path}")
+ print("Please ensure ETOPO data is downloaded and path is correct.")
+ return False
+ print(f" - Data directory: {data_path}")
+ print(f" - Directory exists: {data_path.exists()}")
+ print()
+
+ # Step 2: Initialize topo_cell object
+ print("Step 2: Initializing topo_cell object...")
+ cell = var.topo_cell()
+ print(" - topo_cell object created")
+ print()
+
+ # Step 3: Load ETOPO data
+ print("Step 3: Loading ETOPO data...")
+ print(
+ " (This will load all tiles for full global coverage - may take a few minutes even with coarse-graining)"
+ )
+ start_time = time.time()
+
+ try:
+ loader = io.ncdata.read_etopo_topo(
+ cell, params, verbose=True, is_parallel=False # Show progress
+ )
+ load_time = time.time() - start_time
+ print()
+ print(f" - Loading completed in {load_time:.2f} seconds")
+ print()
+
+ except Exception as e:
+ print(f"ERROR during loading: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return False
+
+ # Step 4: Validate loaded data
+ print("Step 4: Validating loaded data...")
+ print(f" - Latitude array shape: {cell.lat.shape}")
+ print(f" - Longitude array shape: {cell.lon.shape}")
+ print(f" - Topography array shape: {cell.topo.shape}")
+ print()
+ print(f" - Latitude range: [{cell.lat.min():.4f}, {cell.lat.max():.4f}] degrees")
+ print(f" - Longitude range: [{cell.lon.min():.4f}, {cell.lon.max():.4f}] degrees")
+ print()
+ print(f" - Elevation range: [{cell.topo.min():.1f}, {cell.topo.max():.1f}] meters")
+ print(f" - Mean elevation: {cell.topo.mean():.1f} meters")
+ print(f" - Median elevation: {np.median(cell.topo):.1f} meters")
+ print()
+
+ # Sanity checks
+ checks_passed = True
+
+ # Check data shapes
+ expected_lat_points = len(cell.lat)
+ expected_lon_points = len(cell.lon)
+ if cell.topo.shape != (expected_lat_points, expected_lon_points):
+ print(f" WARNING: Unexpected topo shape!")
+ checks_passed = False
+ else:
+ print(f" ✓ Topography shape matches lat/lon dimensions")
+
+ # Check elevation ranges (should be realistic)
+ if cell.topo.min() < -11500 or cell.topo.max() > 9000:
+ print(f" WARNING: Elevation values outside expected range!")
+ print(f" (Expected: ~-11000m to ~8850m)")
+ checks_passed = False
+ else:
+ print(f" ✓ Elevation values within expected range")
+
+ # Check for NaN or infinite values
+ if np.isnan(cell.topo).any():
+ print(f" WARNING: Found NaN values in topography data!")
+ checks_passed = False
+ else:
+ print(f" ✓ No NaN values found")
+
+ if np.isinf(cell.topo).any():
+ print(f" WARNING: Found infinite values in topography data!")
+ checks_passed = False
+ else:
+ print(f" ✓ No infinite values found")
+
+ print()
+
+ if not checks_passed:
+ print(" Some validation checks failed!")
+ return False
+
+ # Step 5: Optionally clip ocean cells before plotting
+ print("Step 5: Optionally clip ocean cells before plotting...")
+ import os
+
+ clip_ocean = True # Default: clip ocean cells to -500m
+ # Allow override via environment variable or function argument in future
+
+ if cell.topo is None:
+ print("ERROR: cell.topo is None. ETOPO data did not load correctly.")
+ print("Skipping plotting and summary.")
+ return False
+
+ land_mask = cell.topo > 0
+ ocean_mask = cell.topo <= 0
+ total_points = cell.topo.size
+ land_points = np.sum(land_mask)
+ ocean_points = np.sum(ocean_mask)
+
+ if clip_ocean:
+ # Clip all ocean cells to -500m for land-only orography test
+ cell.topo[ocean_mask] = -500.0
+ print(" - Ocean cells clipped to -500m for land orography test.")
+ else:
+ print(" - Ocean cells retain original bathymetry (full range).")
+
+ # Step 6: Generate meshgrid for plotting
+ print("Step 6: Generating meshgrid for plotting...")
+ cell.gen_mgrids()
+ print(f" - lon_grid shape: {cell.lon_grid.shape}")
+ print(f" - lat_grid shape: {cell.lat_grid.shape}")
+ print()
+
+ # Step 7: Create plot
+ print("Step 7: Creating global plot...")
+ print(" - Using cartopy PlateCarree projection")
+ print(" - This may take a moment to render...")
+ print()
+
+ try:
+ # Call the plotting function
+ cart_plot.lat_lon(
+ cell,
+ fs=(14, 8), # Larger figure for global view
+ int=plot_stride,
+ colorbar_margins=[0.92, 0.22, 0.035, 0.55], # More visible colorbar
+ )
+ print(" - Plot displayed successfully!")
+ print()
+
+ except Exception as e:
+ print(f"ERROR during plotting: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return False
+
+ # Step 8: Summary statistics
+ print("Step 8: Summary statistics...")
+ # Use the already-clipped topo for stats
+ print(f" - Total data points: {total_points:,}")
+ print(f" - Land points: {land_points:,} ({100*land_points/total_points:.1f}%)")
+ print(f" - Ocean points: {ocean_points:,} ({100*ocean_points/total_points:.1f}%)")
+ print()
+ print(f" - Mean land elevation: {cell.topo[land_mask].mean():.1f} m")
+ if not clip_ocean:
+ print(f" - Mean ocean depth: {cell.topo[ocean_mask].mean():.1f} m")
+ print()
+ print(f" - Highest point: {cell.topo.max():.1f} m (should be near Mt. Everest)")
+ print(
+ f" - Lowest point: {cell.topo.min():.1f} m (should be near Mariana Trench or -500m if clipped)"
+ )
+ print()
+
+ # Step 8: Report success
+ print("=" * 80)
+ print("TEST COMPLETED SUCCESSFULLY!")
+ print("=" * 80)
+ print()
+ print("Summary:")
+ print(f" - Loaded {total_points:,} elevation data points")
+ print(f" - Load time: {load_time:.2f} seconds")
+ print(f" - Data quality: PASSED all validation checks")
+ print(f" - Visualization: SUCCESS")
+ print()
+
+ return True
+
+
+def test_different_coarse_graining_factors():
+ """
+ Test loading with different coarse-graining factors.
+ This helps understand the speed/quality tradeoff.
+ """
+ print("=" * 80)
+ print("TESTING DIFFERENT COARSE-GRAINING FACTORS")
+ print("=" * 80)
+ print()
+
+ # Test with progressively coarser graining
+ test_factors = [16, 12, 8]
+
+ for cg_factor in test_factors:
+ print(f"\n{'='*60}")
+ print(f"Testing with coarse-graining factor: {cg_factor}")
+ print(f"Effective resolution: ~{0.463 * cg_factor:.2f} km at equator")
+ print(f"{'='*60}\n")
+
+ params = create_global_params(etopo_cg=cg_factor)
+ cell = var.topo_cell()
+
+ start_time = time.time()
+ try:
+ loader = io.ncdata.read_etopo_topo(cell, params, verbose=False)
+ load_time = time.time() - start_time
+
+ print(f" Load time: {load_time:.2f} seconds")
+ print(f" Grid size: {cell.topo.shape}")
+ print(f" Memory usage: ~{cell.topo.nbytes / 1e6:.1f} MB")
+ print(
+ f" Elevation range: [{cell.topo.min():.1f}, {cell.topo.max():.1f}] m"
+ )
+
+ except Exception as e:
+ print(f" ERROR: {e}")
+
+ print()
+
+
+if __name__ == "__main__":
+ import sys
+
+ # Run the main global test
+ success = test_global_etopo_load_and_plot()
+
+ if success:
+ print(
+ "\nAll tests passed! The ETOPO loader successfully loaded global coverage."
+ )
+ print()
+ print("=" * 80)
+ print("RECOMMENDED APPROACH FOR FULL GLOBAL COVERAGE")
+ print("=" * 80)
+ print()
+ print(
+ "The dateline handling has been improved, but for best elevation accuracy"
+ )
+ print("with full global coverage, use the two-hemisphere approach:")
+ print()
+ print(" # Load Western Hemisphere")
+ print(" params_west = create_global_params()")
+ print(" params_west.lon_extent = [-180.0, 0.0]")
+ print(" cell_west = var.topo_cell()")
+ print(" loader_west = io.ncdata.read_etopo_topo(cell_west, params_west)")
+ print()
+ print(" # Load Eastern Hemisphere")
+ print(" params_east = create_global_params()")
+ print(" params_east.lon_extent = [0.0, 180.0]")
+ print(" cell_east = var.topo_cell()")
+ print(" loader_east = io.ncdata.read_etopo_topo(cell_east, params_east)")
+ print()
+ print(" # Combine")
+ print(" cell_global = var.topo_cell()")
+ print(" cell_global.lon = np.concatenate([cell_west.lon, cell_east.lon])")
+ print(" cell_global.lat = cell_west.lat # Same for both")
+ print(
+ " cell_global.topo = np.concatenate([cell_west.topo, cell_east.topo], axis=1)"
+ )
+ print()
+ print("This approach preserves elevation accuracy better than loading")
+ print("all 288 tiles in a single operation.")
+ print("=" * 80)
+
+ # Optionally run coarse-graining comparison (only if running interactively)
+ if sys.stdin.isatty():
+ user_input = input("\nRun coarse-graining comparison test? (y/n): ")
+ if user_input.lower() == "y":
+ test_different_coarse_graining_factors()
+ else:
+ print(
+ "\nNote: Run interactively to test different coarse-graining factors."
+ )
+ else:
+ print("\nTest failed! Please check the errors above.")
+ sys.exit(1)
diff --git a/tests/test_etopo_parallel_benchmark.py b/tests/test_etopo_parallel_benchmark.py
new file mode 100644
index 0000000..ce0602b
--- /dev/null
+++ b/tests/test_etopo_parallel_benchmark.py
@@ -0,0 +1,686 @@
+"""
+Comprehensive benchmark test for ETOPO data processing with Dask parallelization.
+
+This test:
+1. Uses ETOPO input data instead of MERIT
+2. Processes 320 cells using 16+ cores
+3. Verifies Dask is working correctly
+4. Saves diagnostic outputs (topography plots, spectra)
+"""
+
+import pytest
+import numpy as np
+import time
+import os
+from pathlib import Path
+import matplotlib
+
+matplotlib.use("Agg") # Non-interactive backend
+import matplotlib.pyplot as plt
+from datetime import datetime
+
+from pycsa.core import io, var, utils
+from pycsa.wrappers import interface, diagnostics
+from pycsa.plotting import cart_plot
+
+# Dask imports
+from dask.distributed import Client, as_completed
+import dask
+
+
+class TestETOPOParallelBenchmark:
+ """Benchmark test for parallel ETOPO processing."""
+
+ @pytest.fixture(scope="class")
+ def output_dir(self, tmp_path_factory):
+ """Create output directory for test results."""
+ # Use a permanent directory instead of tmp for inspection
+ base_dir = Path(__file__).parent.parent / "outputs" / "benchmark_etopo"
+ base_dir.mkdir(parents=True, exist_ok=True)
+
+ # Create timestamped subdirectory
+ timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+ test_dir = base_dir / f"run_{timestamp}"
+ test_dir.mkdir(parents=True, exist_ok=True)
+
+ print(f"\n📁 Output directory: {test_dir}")
+ return test_dir
+
+ @pytest.fixture(scope="class")
+ def test_params(self):
+ """Create test parameters using ETOPO data."""
+ params = var.params()
+
+ # Import local paths
+ try:
+ from pycsa import local_paths
+
+ utils.transfer_attributes(params, local_paths.paths, prefix="path")
+ except ImportError as e:
+ print(f"ERROR: Could not import local_paths: {e}")
+ raise
+
+ # Verify ETOPO path exists
+ if not hasattr(params, "path_etopo") or not Path(params.path_etopo).exists():
+ pytest.skip(
+ f"ETOPO data path not found: {params.path_etopo if hasattr(params, 'path_etopo') else 'not set'}"
+ )
+
+ # Test region: Alaska (good for testing, has varied topography)
+ params.lat_extent = [48.0, 64.0, 64.0]
+ params.lon_extent = [-148.0, -148.0, -112.0]
+
+ # ETOPO coarse-graining factor
+ params.etopo_cg = 50
+
+ # CSA parameters
+ params.nhi = 24
+ params.nhj = 48
+ params.n_modes = 50
+ params.padding = 10
+
+ params.U, params.V = 10.0, 0.0
+ params.rect = True
+
+ # Disable plotting during cell processing (we'll plot diagnostics separately)
+ params.plot = False
+ params.plot_output = False
+
+ params.debug = False
+ params.dfft_first_guess = False
+ params.refine = False
+ params.verbose = False
+
+ return params
+
+ @pytest.fixture(scope="class")
+ def test_grid(self, test_params):
+ """Load a subset of ICON grid for testing."""
+ grid = var.grid()
+
+ # Read ICON grid
+ try:
+ reader = io.ncdata()
+ reader.read_dat(test_params.path_icon_grid, grid)
+ except Exception as e:
+ pytest.skip(f"Could not load ICON grid: {e}")
+
+ # Convert to degrees
+ grid.apply_f(utils.rad2deg)
+
+ return grid
+
+ def test_dask_initialization(self, output_dir):
+ """Test 1: Verify Dask initializes correctly with 16+ cores."""
+ import multiprocessing
+
+ n_workers = min(multiprocessing.cpu_count() - 2, 20)
+ assert n_workers >= 16, f"Not enough cores available: {n_workers} (need 16+)"
+
+ print(f"\n🚀 Initializing Dask with {n_workers} workers...")
+
+ client = Client(
+ threads_per_worker=1,
+ n_workers=n_workers,
+ processes=True,
+ memory_limit="4GB",
+ )
+
+ # Verify client is running
+ assert client.status == "running", "Dask client not running!"
+
+ # Verify workers
+ workers = client.scheduler_info()["workers"]
+ assert len(workers) >= 16, f"Only {len(workers)} workers started (expected 16+)"
+
+ print(f"✓ Dask running with {len(workers)} workers")
+ print(f"✓ Dashboard: {client.dashboard_link}")
+
+ # Save Dask info to output
+ with open(output_dir / "dask_info.txt", "w") as f:
+ f.write(f"Dask Benchmark Test\n")
+ f.write(f"===================\n\n")
+ f.write(f"Workers: {len(workers)}\n")
+ f.write(f"Threads per worker: 1\n")
+ f.write(f"Memory limit per worker: 4GB\n")
+ f.write(f"Dashboard: {client.dashboard_link}\n")
+ f.write(f"\nWorker details:\n")
+ for worker_id, worker_info in workers.items():
+ f.write(f" {worker_id}: {worker_info['memory_limit'] / 1e9:.1f}GB\n")
+
+ client.close()
+ print("✓ Dask client closed cleanly")
+
+ def test_etopo_file_caching(self, test_params, output_dir):
+ """Test 2: Verify ETOPO file caching works correctly."""
+ print("\n📦 Testing ETOPO file caching...")
+
+ # Create a test cell
+ test_cell = var.topo_cell()
+
+ # Initialize ETOPO reader with caching
+ reader = io.ncdata(padding=test_params.padding)
+ etopo_reader = reader.read_etopo_topo(
+ test_cell, test_params, verbose=True, is_parallel=True
+ )
+
+ # Verify cache exists
+ assert hasattr(
+ etopo_reader, "file_cache"
+ ), "ETOPO reader missing file_cache attribute!"
+ assert hasattr(
+ etopo_reader, "_get_cached_file"
+ ), "ETOPO reader missing _get_cached_file method!"
+ assert hasattr(
+ etopo_reader, "close_cached_files"
+ ), "ETOPO reader missing close_cached_files method!"
+
+ # Load data (this should populate the cache)
+ etopo_reader.get_topo(test_cell)
+
+ # Verify data was loaded
+ assert test_cell.topo is not None, "Topography not loaded!"
+ assert test_cell.lon is not None, "Longitude not loaded!"
+ assert test_cell.lat is not None, "Latitude not loaded!"
+
+ # Verify cache was used
+ cache_size = len(etopo_reader.file_cache)
+ print(f"✓ File cache contains {cache_size} open files")
+ assert cache_size > 0, "File cache is empty (caching not working!)"
+
+ # Load same region again - should reuse cache
+ test_cell2 = var.topo_cell()
+ etopo_reader.get_topo(test_cell2)
+
+ # Cache size should not have increased
+ cache_size_after = len(etopo_reader.file_cache)
+ assert (
+ cache_size_after == cache_size
+ ), f"Cache size increased ({cache_size} -> {cache_size_after}), files not being reused!"
+
+ print(f"✓ File cache correctly reused (size unchanged: {cache_size})")
+
+ # Clean up
+ etopo_reader.close_cached_files()
+ assert (
+ len(etopo_reader.file_cache) == 0
+ ), "Cache not cleared after close_cached_files()!"
+ print("✓ Cache cleared successfully")
+
+ # Save cache info
+ with open(output_dir / "cache_info.txt", "w") as f:
+ f.write("ETOPO File Caching Test\n")
+ f.write("=======================\n\n")
+ f.write(f"Cache size (unique files): {cache_size}\n")
+ f.write(f"Cache reuse verified: Yes\n")
+ f.write(f"Cache cleanup verified: Yes\n")
+
+ def test_parallel_320_cells(self, test_params, test_grid, output_dir):
+ """Test 3: Process 320 cells in parallel with full diagnostics."""
+ print(f"\n🔬 Processing 320 cells in parallel...")
+
+ n_test_cells = 320
+ total_cells = test_grid.clat.size
+
+ # Make sure we have enough cells
+ if total_cells < n_test_cells:
+ pytest.skip(f"Grid only has {total_cells} cells, need {n_test_cells}")
+
+ # Select cells to process (spread across the grid)
+ cell_indices = np.linspace(0, total_cells - 1, n_test_cells, dtype=int)
+
+ # Initialize Dask
+ import multiprocessing
+
+ n_workers = min(multiprocessing.cpu_count() - 2, 20)
+ print(f" Starting Dask with {n_workers} workers...")
+
+ client = Client(
+ threads_per_worker=1,
+ n_workers=n_workers,
+ processes=True,
+ memory_limit="4GB",
+ )
+ print(f" Dashboard: {client.dashboard_link}")
+
+ # Initialize reader with ETOPO
+ reader = io.ncdata(
+ padding=test_params.padding, padding_tol=(60 - test_params.padding)
+ )
+
+ # Store pre-computation info
+ clat_rad = np.copy(test_grid.clat)
+ clon_rad = np.copy(test_grid.clon)
+
+ # Scatter large objects to workers (avoid serialization overhead)
+ print(f"\n Scattering grid data to workers...")
+ grid_future = client.scatter(test_grid, broadcast=True)
+ params_future = client.scatter(test_params, broadcast=True)
+ clat_rad_future = client.scatter(clat_rad, broadcast=True)
+ clon_rad_future = client.scatter(clon_rad, broadcast=True)
+
+ # Diagnostic storage
+ processing_times = []
+ cell_results = []
+ error_cells = []
+
+ # Progress tracking
+ from tqdm import tqdm
+
+ print(f"\n Processing {n_test_cells} cells...")
+ start_time = time.time()
+
+ # Process cells
+ futures = []
+ for c_idx in cell_indices:
+ future = client.submit(
+ self._process_single_cell,
+ c_idx,
+ grid_future,
+ params_future,
+ reader,
+ clat_rad_future,
+ clon_rad_future,
+ )
+ futures.append((c_idx, future))
+
+ # Collect results with progress bar
+ for c_idx, future in tqdm(futures, desc="Processing cells"):
+ try:
+ result = future.result(timeout=120) # 2 min timeout per cell
+ if result is not None:
+ cell_results.append(result)
+ if "error" not in result:
+ processing_times.append(result["processing_time"])
+ else:
+ error_cells.append(result)
+ if len(error_cells) <= 3: # Only print first 3 errors
+ print(f"\n Cell {c_idx} error: {result['error']}")
+ except Exception as e:
+ print(f"\n Warning: Cell {c_idx} timed out: {e}")
+ error_cells.append({"c_idx": c_idx, "error": f"Timeout: {e}"})
+
+ total_time = time.time() - start_time
+
+ # Close cached files
+ if hasattr(reader, "close_cached_files"):
+ reader.close_cached_files()
+
+ # Shut down Dask
+ client.close()
+
+ # Analysis
+ n_total = len(cell_results)
+ n_errors = len(error_cells)
+ valid_results = [r for r in cell_results if "error" not in r]
+ n_successful = len(valid_results)
+ n_land = sum(1 for r in valid_results if r.get("is_land", False))
+ n_ocean = sum(1 for r in valid_results if r.get("is_land") == False)
+ success_rate = 100 * n_successful / n_test_cells
+
+ # Separate land and ocean processing times
+ land_times = [
+ r["processing_time"] for r in valid_results if r.get("is_land") == True
+ ]
+ ocean_times = [
+ r["processing_time"] for r in valid_results if r.get("is_land") == False
+ ]
+
+ print(f"\n📊 Results:")
+ print(f" Total time: {total_time:.1f}s")
+ print(f" Cells processed: {n_successful}/{n_test_cells} ({success_rate:.1f}%)")
+ if n_successful > 0:
+ print(f" - Land cells: {n_land} ({100*n_land/n_successful:.0f}%)")
+ print(
+ f" - Ocean cells: {n_ocean} ({100*n_ocean/n_successful:.0f}%) [skipped CSA]"
+ )
+ print(f" Errors/failures: {n_errors}")
+
+ if land_times:
+ print(f"\n Land cell timing (CSA processed):")
+ print(f" Avg: {np.mean(land_times):.2f}s")
+ print(f" Min: {np.min(land_times):.2f}s")
+ print(f" Max: {np.max(land_times):.2f}s")
+
+ if ocean_times:
+ print(f"\n Ocean cell timing (skipped):")
+ print(f" Avg: {np.mean(ocean_times):.3f}s")
+
+ if processing_times:
+ print(f"\n Overall throughput: {n_successful / total_time:.1f} cells/sec")
+ if land_times:
+ print(
+ f" Land-only throughput: {n_land / sum(land_times):.1f} cells/sec"
+ )
+
+ # Assertions (relaxed for initial benchmarking)
+ # Note: Success rate depends on grid coverage of test region
+ assert (
+ success_rate >= 60
+ ), f"Success rate too low: {success_rate:.1f}% (expected ≥60%)"
+ if processing_times:
+ assert (
+ np.mean(processing_times) < 10
+ ), f"Average processing time too high: {np.mean(processing_times):.1f}s"
+
+ # Print error summary if needed
+ if n_errors > 0:
+ print(
+ f"\n⚠️ Warning: {n_errors} cells had errors. Check outputs/benchmark_etopo/*/errors.txt for details"
+ )
+
+ # Save results
+ self._save_benchmark_results(
+ output_dir,
+ valid_results,
+ processing_times,
+ total_time,
+ n_test_cells,
+ error_cells,
+ )
+
+ # Generate diagnostic plots
+ self._generate_diagnostic_plots(output_dir, cell_results, test_params)
+
+ print(f"\n✓ Benchmark complete! Results saved to {output_dir}")
+
+ @staticmethod
+ def _process_single_cell(c_idx, grid, params, reader, clat_rad, clon_rad):
+ """Process a single cell (executed by Dask worker)."""
+ try:
+ start_time = time.time()
+
+ # Create cell object
+ topo = var.topo_cell()
+
+ # Get cell vertices
+ lat_verts = grid.clat_vertices[c_idx]
+ lon_verts = grid.clon_vertices[c_idx]
+
+ # Handle lat/lon expansion
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+ lat_verts, lon_verts = utils.handle_latlon_expansion(
+ lat_verts, lon_verts, lat_expand=0.0, lon_expand=0.0
+ )
+
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ # Load ETOPO topography data
+ etopo_reader = reader.read_etopo_topo(None, params, is_parallel=True)
+ etopo_reader.get_topo(topo)
+
+ # Apply elevation floor
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+ topo.gen_mgrids()
+
+ # Set up cell geometry
+ clon = np.array([grid.clon[c_idx]])
+ clat = np.array([grid.clat[c_idx]])
+ clon_vertices = np.array([lon_verts])
+ clat_vertices = np.array([lat_verts])
+
+ ncells = 1
+ nv = clon_vertices[0].size
+
+ # Handle dateline crossing
+ if etopo_reader.split_EW:
+ clon_vertices[clon_vertices < 0.0] += 360.0
+
+ triangles = np.zeros((ncells, nv, 2))
+ for i in range(0, ncells, 1):
+ triangles[i, :, 0] = np.array(clon_vertices[i, :])
+ triangles[i, :, 1] = np.array(clat_vertices[i, :])
+
+ # Check if land
+ tri_idx = 0
+ cell = var.topo_cell()
+ tri = var.obj()
+
+ tri.tri_lon_verts = triangles[:, :, 0]
+ tri.tri_lat_verts = triangles[:, :, 1]
+ simplex_lat = tri.tri_lat_verts[tri_idx]
+ simplex_lon = tri.tri_lon_verts[tri_idx]
+
+ is_land = utils.is_land(cell, simplex_lat, simplex_lon, topo)
+
+ if not is_land:
+ return {
+ "c_idx": c_idx,
+ "is_land": False,
+ "processing_time": time.time() - start_time,
+ }
+
+ # Run CSA (simplified - just first approximation for benchmark)
+ nhi = params.nhi
+ nhj = params.nhj
+
+ utils.get_lat_lon_segments(
+ simplex_lat, simplex_lon, cell, topo, rect=params.rect
+ )
+
+ # Run spectral approximation
+ pmf = interface.get_pmf(nhi, nhj, params.U, params.V)
+ ampls, uw_pmf, dat_2D = pmf.sappx(cell, lmbda=0.1)
+
+ processing_time = time.time() - start_time
+
+ # Filter out NaNs from spectrum for meaningful statistics
+ ampls_valid = ampls[~np.isnan(ampls)]
+ spectrum_max = (
+ float(np.max(ampls_valid)) if len(ampls_valid) > 0 else np.nan
+ )
+ n_valid_modes = len(ampls_valid)
+
+ return {
+ "c_idx": c_idx,
+ "is_land": True,
+ "processing_time": processing_time,
+ "topo_shape": topo.topo.shape,
+ "topo_min": float(np.min(topo.topo)),
+ "topo_max": float(np.max(topo.topo)),
+ "spectrum_max": spectrum_max,
+ "n_modes": ampls.size,
+ "n_valid_modes": n_valid_modes,
+ "lat_extent": params.lat_extent,
+ "lon_extent": params.lon_extent,
+ }
+
+ except Exception as e:
+ import traceback
+
+ return {
+ "c_idx": c_idx,
+ "is_land": None,
+ "processing_time": time.time() - start_time,
+ "error": str(e),
+ "traceback": traceback.format_exc(),
+ }
+
+ def _save_benchmark_results(
+ self,
+ output_dir,
+ cell_results,
+ processing_times,
+ total_time,
+ n_test_cells,
+ error_cells,
+ ):
+ """Save benchmark results to file."""
+ with open(output_dir / "benchmark_results.txt", "w") as f:
+ f.write("ETOPO Parallel Processing Benchmark\n")
+ f.write("=" * 50 + "\n\n")
+
+ f.write(f"Test Configuration:\n")
+ f.write(f" Total cells attempted: {n_test_cells}\n")
+ f.write(f" Successful cells: {len(cell_results)}\n")
+ f.write(f" Error/failed cells: {len(error_cells)}\n")
+ f.write(f"\n")
+
+ f.write(f"Timing Results:\n")
+ f.write(f" Total time: {total_time:.2f}s\n")
+ f.write(f" Average per cell: {np.mean(processing_times):.2f}s\n")
+ f.write(f" Median per cell: {np.median(processing_times):.2f}s\n")
+ f.write(f" Min per cell: {np.min(processing_times):.2f}s\n")
+ f.write(f" Max per cell: {np.max(processing_times):.2f}s\n")
+ f.write(f" Throughput: {len(cell_results) / total_time:.2f} cells/sec\n")
+ f.write(f"\n")
+
+ # Land/ocean statistics
+ land_cells = sum(1 for r in cell_results if r.get("is_land"))
+ ocean_cells = sum(1 for r in cell_results if r.get("is_land") == False)
+ f.write(f"Cell Statistics:\n")
+ f.write(f" Land cells: {land_cells}\n")
+ f.write(f" Ocean cells: {ocean_cells}\n")
+
+ # Error summary
+ if error_cells:
+ f.write(f"\nErrors:\n")
+ error_types = {}
+ for err in error_cells:
+ err_msg = err.get("error", "Unknown error")
+ # Group by error type (first line of error)
+ err_type = err_msg.split("\n")[0][:100]
+ error_types[err_type] = error_types.get(err_type, 0) + 1
+
+ for err_type, count in sorted(
+ error_types.items(), key=lambda x: x[1], reverse=True
+ ):
+ f.write(f" {count}x: {err_type}\n")
+
+ # Save detailed error log
+ if error_cells:
+ with open(output_dir / "errors.txt", "w") as f:
+ f.write(f"Detailed Error Log ({len(error_cells)} errors)\n")
+ f.write("=" * 70 + "\n\n")
+ for i, err in enumerate(error_cells[:10]): # First 10 errors
+ f.write(f"Error {i+1}: Cell {err.get('c_idx', 'unknown')}\n")
+ f.write(f"{'-' * 70}\n")
+ f.write(f"{err.get('error', 'No error message')}\n")
+ if "traceback" in err:
+ f.write(f"\nTraceback:\n{err['traceback']}\n")
+ f.write(f"\n{'=' * 70}\n\n")
+ if len(error_cells) > 10:
+ f.write(
+ f"\n... and {len(error_cells) - 10} more errors (see benchmark_results.txt for summary)\n"
+ )
+
+ print(f" ✓ Saved benchmark results")
+
+ def _generate_diagnostic_plots(self, output_dir, cell_results, params):
+ """Generate diagnostic plots from results."""
+ print("\n Generating diagnostic plots...")
+
+ # Filter land cells only
+ land_results = [r for r in cell_results if r["is_land"]]
+
+ if len(land_results) < 5:
+ print(" Skipping plots (not enough land cells)")
+ return
+
+ # Plot 1: Processing time distribution
+ fig, axes = plt.subplots(2, 2, figsize=(12, 10))
+
+ times = [r["processing_time"] for r in cell_results]
+ axes[0, 0].hist(times, bins=30, edgecolor="black", alpha=0.7)
+ axes[0, 0].set_xlabel("Processing Time (s)")
+ axes[0, 0].set_ylabel("Count")
+ axes[0, 0].set_title("Processing Time Distribution")
+ axes[0, 0].axvline(
+ np.mean(times),
+ color="red",
+ linestyle="--",
+ label=f"Mean: {np.mean(times):.2f}s",
+ )
+ axes[0, 0].legend()
+
+ # Plot 2: Topography elevation ranges
+ topo_mins = [r["topo_min"] for r in land_results]
+ topo_maxs = [r["topo_max"] for r in land_results]
+ axes[0, 1].scatter(topo_mins, topo_maxs, alpha=0.5)
+ axes[0, 1].set_xlabel("Min Elevation (m)")
+ axes[0, 1].set_ylabel("Max Elevation (m)")
+ axes[0, 1].set_title("Topography Elevation Ranges")
+ axes[0, 1].grid(True, alpha=0.3)
+
+ # Plot 3: Spectrum amplitudes
+ spectrum_maxs = [
+ r["spectrum_max"] for r in land_results if not np.isnan(r["spectrum_max"])
+ ]
+ if len(spectrum_maxs) > 0:
+ axes[1, 0].hist(spectrum_maxs, bins=30, edgecolor="black", alpha=0.7)
+ else:
+ axes[1, 0].text(
+ 0.5, 0.5, "No valid spectrum data", ha="center", va="center"
+ )
+ axes[1, 0].set_xlabel("Max Spectrum Amplitude")
+ axes[1, 0].set_ylabel("Count")
+ axes[1, 0].set_title("Spectral Amplitude Distribution")
+
+ # Plot 4: Topography grid sizes
+ topo_sizes = [r["topo_shape"][0] * r["topo_shape"][1] for r in land_results]
+ axes[1, 1].hist(topo_sizes, bins=30, edgecolor="black", alpha=0.7)
+ axes[1, 1].set_xlabel("Grid Points")
+ axes[1, 1].set_ylabel("Count")
+ axes[1, 1].set_title("Loaded Topography Grid Sizes")
+
+ plt.tight_layout()
+ plt.savefig(
+ output_dir / "diagnostics_summary.png", dpi=150, bbox_inches="tight"
+ )
+ plt.close()
+
+ print(f" ✓ Saved diagnostics_summary.png")
+
+ # Save a few example topography samples
+ n_samples = min(6, len(land_results))
+ sample_cells = np.random.choice(len(land_results), n_samples, replace=False)
+
+ fig, axes = plt.subplots(2, 3, figsize=(15, 10))
+ axes = axes.flatten()
+
+ for idx, sample_idx in enumerate(sample_cells):
+ result = land_results[sample_idx]
+ ax = axes[idx]
+
+ # Just show basic info since we don't have the actual topo data
+ spectrum_str = (
+ f"{result['spectrum_max']:.2e}"
+ if not np.isnan(result["spectrum_max"])
+ else "N/A"
+ )
+ n_valid = result.get("n_valid_modes", "?")
+ n_total = result.get("n_modes", "?")
+
+ info_text = (
+ f"Cell {result['c_idx']}\n"
+ f"Grid: {result['topo_shape']}\n"
+ f"Elev: [{result['topo_min']:.0f}, {result['topo_max']:.0f}]m\n"
+ f"Spectrum max: {spectrum_str}\n"
+ f"Valid modes: {n_valid}/{n_total}\n"
+ f"Time: {result['processing_time']:.2f}s"
+ )
+ ax.text(
+ 0.5,
+ 0.5,
+ info_text,
+ ha="center",
+ va="center",
+ fontsize=10,
+ family="monospace",
+ )
+ ax.set_xlim(0, 1)
+ ax.set_ylim(0, 1)
+ ax.axis("off")
+
+ plt.suptitle("Sample Cell Results", fontsize=14, fontweight="bold")
+ plt.tight_layout()
+ plt.savefig(output_dir / "sample_cells.png", dpi=150, bbox_inches="tight")
+ plt.close()
+
+ print(f" ✓ Saved sample_cells.png")
+
+
+if __name__ == "__main__":
+ # Run the test directly
+ pytest.main([__file__, "-v", "-s"])
diff --git a/tests/test_etopo_pole_cells.py b/tests/test_etopo_pole_cells.py
new file mode 100644
index 0000000..bd0cd32
--- /dev/null
+++ b/tests/test_etopo_pole_cells.py
@@ -0,0 +1,1345 @@
+"""
+Test script to compare old (corner-based) vs. new (centered) planar projection.
+
+Tests 10 pre-selected polar cells (5 Arctic, 5 Antarctic) to evaluate improvement
+in pyCSA RMSE when using centered projection instead of corner-based projection.
+"""
+
+import numpy as np
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+from matplotlib.colors import TwoSlopeNorm
+import matplotlib.colors as mcolors
+from pathlib import Path
+
+from pycsa.core import io, var, utils
+from pycsa.wrappers import interface
+from scipy import interpolate
+
+# Pre-selected cell indices from ICON grid
+# Users can comment/uncomment cells to test different scenarios
+# Focus on EXTREME POLAR cells where projection distortion is maximum
+
+POLAR_CELLS = [
+ # ========================================================================
+ # ARCTIC CELLS (Greenland, 80-82°N)
+ # ========================================================================
+ # Moderate latitude - smaller projection differences expected
+ # 3091, # Arctic: 80.35°N, -92.11°E - Greenland
+ # 3105, # Arctic: 79.77°N, -65.63°E - Greenland
+ # 3107, # Arctic: 79.77°N, -78.37°E - Greenland
+ # 3108, # Arctic: 81.28°N, -57.03°E - Greenland
+ # 3109, # Arctic: 82.56°N, -45.32°E - Greenland
+ # ========================================================================
+ # EXTREME ANTARCTIC CELLS (87-89°S)
+ # ========================================================================
+ # These cells are within 1-3 degrees of the South Pole where corner
+ # projection creates MAXIMUM distortion. This is where centered projection
+ # should show the biggest improvement!
+ # MOST EXTREME: -88.90°S (within 1.1° of South Pole!)
+ 17408, # Antarctic: -88.90°S, -108.00°E - Interior plateau, 100% land, elev=2699m
+ 16384, # Antarctic: -88.90°S, 180.00°E - Interior plateau, 100% land, elev=2761m
+ 18432, # Antarctic: -88.90°S, -36.00°E - Interior plateau, 100% land, elev=2649m
+ 15360, # Antarctic: -88.90°S, 108.00°E - Interior plateau, 100% land, elev=2941m
+ 19456, # Antarctic: -88.90°S, 36.00°E - Interior plateau, 100% land, elev=2835m
+ # VERY EXTREME: -88.07°S
+ 15362, # Antarctic: -88.07°S, 108.00°E - Interior plateau, 100% land, elev=3055m
+ 16386, # Antarctic: -88.07°S, 180.00°E - Interior plateau, 100% land, elev=2754m
+ 16387,
+ 17410, # Antarctic: -88.07°S, -108.00°E - Interior plateau, 100% land, elev=2554m
+ 19458, # Antarctic: -88.07°S, 36.00°E - Interior plateau, 100% land, elev=2882m
+ 18434, # Antarctic: -88.07°S, -36.00°E - Interior plateau, 100% land, elev=2445m
+ # EXTREME: -87.21°S
+ 15361, # Antarctic: -87.21°S, 129.75°E - Interior plateau, 100% land, elev=3023m
+ 15363, # Antarctic: -87.21°S, 86.25°E - Interior plateau, 100% land, elev=3105m
+ 16387, # Antarctic: -87.21°S, 158.25°E - Interior plateau, 100% land, elev=2698m
+ 17409, # Antarctic: -87.21°S, -86.25°E - Interior plateau, 100% land, elev=2384m
+ 19457, # Antarctic: -87.21°S, 57.75°E - Interior plateau, 100% land, elev=3059m
+ # ========================================================================
+ # LESS EXTREME ANTARCTIC CELLS (85-86°S)
+ # ========================================================================
+ # Still very high latitude but slightly less extreme than above
+ # 15364, # Antarctic: -85.39°S, 135.26°E - Interior plateau, 100% land, elev=2896m
+ # 15369, # Antarctic: -86.34°S, 90.55°E - Interior plateau, 100% land, elev=3214m
+ # 15370, # Antarctic: -85.75°S, 108.00°E - Interior plateau, 100% land, elev=3109m
+ # 15371, # Antarctic: -86.34°S, 125.45°E - Interior plateau, 100% land, elev=2987m
+ # 15372, # Antarctic: -85.39°S, 80.74°E - Interior plateau, 100% land, elev=3328m
+]
+
+# Equatorial/mid-latitude cells - to test if centered projection helps more here
+# Will be populated dynamically to find land cells near equator
+EQUATORIAL_CELLS_CANDIDATES = list(range(0, 25000)) # Will filter for equatorial land
+# EQUATORIAL_CELLS = [340, 992, 1015] # To be filled in
+
+
+def get_topo_colormap():
+ """Create topography colormap with blue for ocean, terrain for land."""
+ ocean_colors = plt.cm.Blues_r(np.linspace(0.4, 0.95, 120))
+ last_ocean = plt.cm.Blues_r(0.95)
+ first_land = plt.cm.terrain(0.25)
+
+ transition_colors = np.zeros((16, 4))
+ for i in range(4):
+ transition_colors[:, i] = np.linspace(last_ocean[i], first_land[i], 16)
+
+ land_colors = plt.cm.terrain(np.linspace(0.28, 1.0, 120))
+ colors = np.vstack((ocean_colors, transition_colors, land_colors))
+ return mcolors.LinearSegmentedColormap.from_list("topo", colors)
+
+
+def interpolate_to_reference_grid(data_2D, source_cell, target_cell):
+ """
+ Interpolate 2D data from source planar grid to target planar grid.
+
+ This is needed when comparing CSA outputs from different projection methods
+ (corner vs centered) against a common reference topography.
+
+ Parameters
+ ----------
+ data_2D : ndarray
+ 2D data on source grid (e.g., CSA reconstruction)
+ source_cell : topo_cell
+ Cell with source planar coordinates (lat, lon in meters)
+ target_cell : topo_cell
+ Cell with target planar coordinates (lat, lon in meters)
+
+ Returns
+ -------
+ ndarray
+ Data interpolated onto target grid, same shape as target_cell.topo
+ """
+ # Create source grid coordinates (meshgrid of lat/lon in meters)
+ source_lon_grid, source_lat_grid = np.meshgrid(source_cell.lon, source_cell.lat)
+
+ # Create target grid coordinates
+ target_lon_grid, target_lat_grid = np.meshgrid(target_cell.lon, target_cell.lat)
+
+ # Flatten source coordinates and data
+ source_points = np.column_stack([source_lon_grid.ravel(), source_lat_grid.ravel()])
+ source_values = data_2D.ravel()
+
+ # Flatten target coordinates
+ target_points = np.column_stack([target_lon_grid.ravel(), target_lat_grid.ravel()])
+
+ # Interpolate using griddata (linear interpolation)
+ interpolated_values = interpolate.griddata(
+ source_points,
+ source_values,
+ target_points,
+ method="linear",
+ fill_value=0.0, # Fill any out-of-bounds points with 0
+ )
+
+ # Reshape back to 2D grid
+ interpolated_2D = interpolated_values.reshape(target_cell.topo.shape)
+
+ return interpolated_2D
+
+
+def create_cell_with_projection(lat_verts, lon_verts, topo, use_center=True, rect=True):
+ """
+ Create cell using production code path (utils.get_lat_lon_segments).
+
+ Parameters
+ ----------
+ lat_verts, lon_verts : array
+ Vertex coordinates in degrees (processed by handle_latlon_expansion)
+ topo : topo_cell
+ Topography object
+ use_center : bool
+ If True, use center of domain as projection origin (NEW method)
+ If False, use corner of domain as projection origin (OLD method)
+ rect : bool
+ If True, use rectangular mask (for FA)
+ If False, use triangular mask (for SA)
+
+ Returns
+ -------
+ cell : topo_cell
+ Configured cell object
+ """
+ cell = var.topo_cell()
+
+ # Use production code path - this includes all preprocessing!
+ if rect:
+ # FA: Create rectangular cell with filtered topography
+ utils.get_lat_lon_segments(
+ lat_verts,
+ lon_verts,
+ cell,
+ topo,
+ rect=True,
+ filtered=True, # Remove features < 5km
+ padding=0,
+ use_center=use_center,
+ )
+ else:
+ # SA: Create triangular cell
+ # Production calls this twice on the same cell: first rect=True to load topo,
+ # then rect=False to apply triangular mask
+ # We'll do the same
+ utils.get_lat_lon_segments(
+ lat_verts,
+ lon_verts,
+ cell,
+ topo,
+ rect=True,
+ filtered=True,
+ padding=0,
+ use_center=use_center,
+ )
+ # Now apply triangular mask
+ utils.get_lat_lon_segments(
+ lat_verts,
+ lon_verts,
+ cell,
+ topo,
+ rect=False,
+ filtered=False,
+ padding=0,
+ use_center=use_center,
+ )
+
+ print(f" use_center={use_center}, rect={rect}")
+ print(
+ f" Mask: {cell.mask.sum()} / {cell.mask.size} points ({100*cell.mask.sum()/cell.mask.size:.1f}%)"
+ )
+ print(f" cell.lat range: [{cell.lat.min():.1f}, {cell.lat.max():.1f}] m")
+ print(f" cell.lon range: [{cell.lon.min():.1f}, {cell.lon.max():.1f}] m")
+
+ return cell
+
+
+def run_full_csa(cell, params, use_mode_selection=False):
+ """
+ Run full CSA algorithm (first + second approximation) on a cell.
+
+ Parameters
+ ----------
+ cell : topo_cell
+ Cell object with topography
+ params : params object
+ Parameters
+ use_mode_selection : bool, optional
+ If True, select top n_modes wavenumbers in SA (spectral compression)
+ If False, use ALL wavenumbers in SA (full spectrum, better RMSE)
+ Default: False (full spectrum)
+
+ Returns
+ -------
+ tuple : (ampls_fa, ampls_sa, dat_2D_sa, rmse_fa, rmse_sa)
+ """
+ # First approximation
+ fa = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_fa, uw_fa, dat_2D_fa = fa.sappx(
+ cell, lmbda=params.lmbda_fa, iter_solve=params.fa_iter_solve
+ )
+
+ # Compute first approximation RMSE
+ diff_fa = cell.topo - dat_2D_fa
+ mask = cell.mask if hasattr(cell, "mask") else np.ones_like(cell.topo, dtype=bool)
+ rmse_fa = np.sqrt(np.mean(diff_fa[mask] ** 2))
+
+ # Second approximation
+ if use_mode_selection:
+ # COMPRESSED MODE: Select top n_modes wavenumbers
+ # Extract top modes from FA spectrum
+ fq_cpy = np.copy(ampls_fa)
+ fq_cpy[np.isnan(fq_cpy)] = 0.0
+
+ indices = []
+ modes_cnt = 0
+ while modes_cnt < params.n_modes:
+ max_idx = np.unravel_index(fq_cpy.argmax(), fq_cpy.shape)
+ indices.append(max_idx)
+ fq_cpy[max_idx] = 0.0
+ modes_cnt += 1
+
+ k_idxs = [pair[1] for pair in indices]
+ l_idxs = [pair[0] for pair in indices]
+
+ # Create new PMF with selected modes only
+ sa = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ sa.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+ ampls_sa, uw_sa, dat_2D_sa = sa.sappx(
+ cell, lmbda=params.lmbda_sa, iter_solve=params.sa_iter_solve
+ )
+ else:
+ # FULL SPECTRUM MODE: Use ALL wavenumbers
+ sa = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_sa, uw_sa, dat_2D_sa = sa.sappx(
+ cell, lmbda=params.lmbda_sa, iter_solve=params.sa_iter_solve
+ )
+
+ # Compute second approximation RMSE
+ diff_sa = cell.topo - dat_2D_sa
+ rmse_sa = np.sqrt(np.mean(diff_sa[mask] ** 2))
+
+ return ampls_fa, ampls_sa, dat_2D_sa, rmse_fa, rmse_sa
+
+
+def plot_single_method(
+ c_idx,
+ lat,
+ topo_orig,
+ recon_fa,
+ recon_sa,
+ rmse_fa,
+ rmse_sa,
+ mask,
+ output_dir,
+ method_name,
+):
+ """
+ Create 5-panel plot for a single projection method.
+
+ Panels:
+ 1. Reference topography
+ 2. First Approximation reconstruction
+ 3. Second Approximation reconstruction
+ 4. First Approximation error map (absolute error)
+ 5. Second Approximation error map (absolute error)
+
+ Parameters
+ ----------
+ c_idx : int
+ Cell index
+ lat : float
+ Cell latitude in degrees
+ topo_orig : ndarray
+ Reference topography
+ recon_fa : ndarray
+ First approximation reconstruction
+ recon_sa : ndarray
+ Second approximation reconstruction
+ rmse_fa : float
+ First approximation RMSE
+ rmse_sa : float
+ Second approximation RMSE
+ mask : ndarray
+ Boolean mask for triangular cell
+ output_dir : Path
+ Output directory
+ method_name : str
+ 'OLD' or 'NEW' for labeling
+ """
+ fig, axs = plt.subplots(2, 3, figsize=(20, 12))
+
+ # Mask the reconstructions for visualization (show only triangular cell)
+ recon_fa_masked = np.ma.masked_where(~mask, recon_fa)
+ recon_sa_masked = np.ma.masked_where(~mask, recon_sa)
+ topo_orig_masked = np.ma.masked_where(~mask, topo_orig)
+
+ vmin = topo_orig[mask].min()
+ vmax = topo_orig[mask].max()
+
+ topo_cmap = get_topo_colormap()
+ norm = TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax)
+
+ method_label = "Corner-based" if method_name == "OLD" else "Centered"
+
+ # Panel 1: Reference topography
+ im1 = axs[0, 0].imshow(
+ topo_orig_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0, 0].set_title(
+ f"Cell {c_idx} at {lat:.1f}°: Reference Topo\nRange: [{vmin:.0f}, {vmax:.0f}] m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0, 0].set_xlabel("Longitude index")
+ axs[0, 0].set_ylabel("Latitude index")
+ plt.colorbar(im1, ax=axs[0, 0], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 2: First Approximation
+ im2 = axs[0, 1].imshow(
+ recon_fa_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0, 1].set_title(
+ f"{method_name} ({method_label}): 1st Approx\nRMSE: {rmse_fa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0, 1].set_xlabel("Longitude index")
+ axs[0, 1].set_ylabel("Latitude index")
+ plt.colorbar(im2, ax=axs[0, 1], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 3: Second Approximation
+ im3 = axs[0, 2].imshow(
+ recon_sa_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0, 2].set_title(
+ f"{method_name} ({method_label}): 2nd Approx\nRMSE: {rmse_sa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0, 2].set_xlabel("Longitude index")
+ axs[0, 2].set_ylabel("Latitude index")
+ plt.colorbar(im3, ax=axs[0, 2], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 4: First Approximation Error Map
+ error_fa = np.abs(topo_orig - recon_fa)
+ error_fa_masked = np.ma.masked_where(~mask, error_fa)
+ error_max_fa = error_fa[mask].max()
+
+ im4 = axs[1, 0].imshow(
+ error_fa_masked,
+ origin="lower",
+ cmap="Reds",
+ vmin=0,
+ vmax=error_max_fa,
+ aspect="auto",
+ )
+ axs[1, 0].set_title(
+ f"1st Approx: Absolute Error\nMax: {error_max_fa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[1, 0].set_xlabel("Longitude index")
+ axs[1, 0].set_ylabel("Latitude index")
+ plt.colorbar(im4, ax=axs[1, 0], fraction=0.046, pad=0.04).set_label(
+ "Absolute Error [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 5: Second Approximation Error Map
+ error_sa = np.abs(topo_orig - recon_sa)
+ error_sa_masked = np.ma.masked_where(~mask, error_sa)
+ error_max_sa = error_sa[mask].max()
+
+ im5 = axs[1, 1].imshow(
+ error_sa_masked,
+ origin="lower",
+ cmap="Reds",
+ vmin=0,
+ vmax=error_max_sa,
+ aspect="auto",
+ )
+ axs[1, 1].set_title(
+ f"2nd Approx: Absolute Error\nMax: {error_max_sa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[1, 1].set_xlabel("Longitude index")
+ axs[1, 1].set_ylabel("Latitude index")
+ plt.colorbar(im5, ax=axs[1, 1], fraction=0.046, pad=0.04).set_label(
+ "Absolute Error [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 6: Statistics summary (text panel)
+ axs[1, 2].axis("off")
+ stats_text = f"""
+ Method: {method_name} ({method_label})
+ Cell: {c_idx}
+ Latitude: {lat:.2f}°
+
+ Topography Range:
+ Min: {vmin:.1f} m
+ Max: {vmax:.1f} m
+
+ 1st Approximation:
+ RMSE: {rmse_fa:.1f} m
+ Max Error: {error_max_fa:.1f} m
+ Mean Error: {error_fa[mask].mean():.1f} m
+
+ 2nd Approximation:
+ RMSE: {rmse_sa:.1f} m
+ Max Error: {error_max_sa:.1f} m
+ Mean Error: {error_sa[mask].mean():.1f} m
+
+ Improvement (FA → SA):
+ RMSE: {rmse_fa - rmse_sa:.1f} m
+ Reduction: {((rmse_fa - rmse_sa)/rmse_fa*100):.1f}%
+ """
+ axs[1, 2].text(
+ 0.1,
+ 0.5,
+ stats_text,
+ fontsize=10,
+ family="monospace",
+ verticalalignment="center",
+ transform=axs[1, 2].transAxes,
+ )
+
+ plt.tight_layout()
+ output_path = (
+ output_dir / f"{method_name.lower()}_cell_{c_idx}_lat_{lat:.1f}deg.png"
+ )
+ plt.savefig(output_path, dpi=150, bbox_inches="tight")
+ plt.close(fig)
+
+ print(f" Plot saved: {output_path}")
+
+
+def plot_comparison(
+ c_idx,
+ lat,
+ topo_orig,
+ recon_old_fa,
+ recon_old_sa,
+ recon_new_fa,
+ recon_new_sa,
+ rmse_old_fa,
+ rmse_old_sa,
+ rmse_new_fa,
+ rmse_new_sa,
+ mask,
+ output_dir,
+):
+ """
+ Create 6-panel comparison plot (FA and SA for both methods).
+
+ All data is on the same grid (centered projection reference).
+ OLD method reconstructions have been interpolated to this reference grid.
+ """
+ fig, axs = plt.subplots(2, 3, figsize=(20, 12))
+
+ # Mask the reconstructions for visualization (show only triangular cell)
+ recon_old_fa_masked = np.ma.masked_where(~mask, recon_old_fa)
+ recon_old_sa_masked = np.ma.masked_where(~mask, recon_old_sa)
+ recon_new_fa_masked = np.ma.masked_where(~mask, recon_new_fa)
+ recon_new_sa_masked = np.ma.masked_where(~mask, recon_new_sa)
+ topo_orig_masked = np.ma.masked_where(~mask, topo_orig)
+
+ vmin = topo_orig[mask].min()
+ vmax = topo_orig[mask].max()
+
+ topo_cmap = get_topo_colormap()
+ norm = TwoSlopeNorm(vmin=vmin, vcenter=0.0, vmax=vmax)
+
+ # Panel 1: Reference topography (centered projection)
+ im1 = axs[0, 0].imshow(
+ topo_orig_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0, 0].set_title(
+ f"Cell {c_idx} at {lat:.1f}°: Reference (Centered)\nRange: [{vmin:.0f}, {vmax:.0f}] m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0, 0].set_xlabel("Longitude index")
+ axs[0, 0].set_ylabel("Latitude index")
+ plt.colorbar(im1, ax=axs[0, 0], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 2: OLD - First Approximation
+ im2 = axs[0, 1].imshow(
+ recon_old_fa_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0, 1].set_title(
+ f"OLD (Corner): 1st Approx\nRMSE: {rmse_old_fa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0, 1].set_xlabel("Longitude index")
+ axs[0, 1].set_ylabel("Latitude index")
+ plt.colorbar(im2, ax=axs[0, 1], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 3: OLD - Second Approximation
+ im3 = axs[0, 2].imshow(
+ recon_old_sa_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[0, 2].set_title(
+ f"OLD (Corner): 2nd Approx\nRMSE: {rmse_old_sa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ )
+ axs[0, 2].set_xlabel("Longitude index")
+ axs[0, 2].set_ylabel("Latitude index")
+ plt.colorbar(im3, ax=axs[0, 2], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 4: Error map (FA)
+ error_old_fa = np.abs(topo_orig - recon_old_fa)
+ error_new_fa = np.abs(topo_orig - recon_new_fa)
+ error_diff_fa = error_old_fa - error_new_fa
+ error_diff_fa_masked = np.ma.masked_where(~mask, error_diff_fa)
+ error_max_fa = max(
+ np.abs(error_diff_fa[mask].min()), np.abs(error_diff_fa[mask].max())
+ )
+
+ im4 = axs[1, 0].imshow(
+ error_diff_fa_masked,
+ origin="lower",
+ cmap="RdYlGn",
+ vmin=-error_max_fa,
+ vmax=error_max_fa,
+ aspect="auto",
+ )
+ imp_fa = (
+ ((rmse_old_fa - rmse_new_fa) / rmse_old_fa * 100) if rmse_old_fa > 0 else 0.0
+ )
+ axs[1, 0].set_title(
+ f"1st Approx Improvement\nGreen=Better | Imp: {imp_fa:.1f}%",
+ fontsize=11,
+ fontweight="bold",
+ color="green" if imp_fa > 0 else "red",
+ )
+ axs[1, 0].set_xlabel("Longitude index")
+ axs[1, 0].set_ylabel("Latitude index")
+ plt.colorbar(im4, ax=axs[1, 0], fraction=0.046, pad=0.04).set_label(
+ "Error Reduction [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 5: NEW - First Approximation
+ im5 = axs[1, 1].imshow(
+ recon_new_fa_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ axs[1, 1].set_title(
+ f"NEW (Centered): 1st Approx\nRMSE: {rmse_new_fa:.1f} m",
+ fontsize=11,
+ fontweight="bold",
+ color="green",
+ )
+ axs[1, 1].set_xlabel("Longitude index")
+ axs[1, 1].set_ylabel("Latitude index")
+ plt.colorbar(im5, ax=axs[1, 1], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ # Panel 6: NEW - Second Approximation
+ im6 = axs[1, 2].imshow(
+ recon_new_sa_masked, origin="lower", cmap=topo_cmap, norm=norm, aspect="auto"
+ )
+ imp_sa = (
+ ((rmse_old_sa - rmse_new_sa) / rmse_old_sa * 100) if rmse_old_sa > 0 else 0.0
+ )
+ axs[1, 2].set_title(
+ f"NEW (Centered): 2nd Approx\nRMSE: {rmse_new_sa:.1f} m | Imp: {imp_sa:.1f}%",
+ fontsize=11,
+ fontweight="bold",
+ color="green",
+ )
+ axs[1, 2].set_xlabel("Longitude index")
+ axs[1, 2].set_ylabel("Latitude index")
+ plt.colorbar(im6, ax=axs[1, 2], fraction=0.046, pad=0.04).set_label(
+ "Elevation [m]", rotation=270, labelpad=15
+ )
+
+ plt.tight_layout()
+ output_path = output_dir / f"comparison_cell_{c_idx}_lat_{lat:.1f}deg.png"
+ plt.savefig(output_path, dpi=150, bbox_inches="tight")
+ plt.close(fig)
+
+ print(f" Plot saved: {output_path}")
+ return imp_fa, imp_sa
+
+
+def main():
+ """
+ Main test function.
+
+ Tests OLD (corner-based) vs NEW (centered) planar projection methods.
+
+ KEY METHODOLOGY:
+ - Creates a SHARED REFERENCE topography using centered projection (geometrically accurate)
+ - OLD method: Runs CSA on corner-projection grid, then interpolates to reference grid
+ - NEW method: Runs CSA on centered-projection grid (same as reference, no interpolation)
+ - Both methods compared against the SAME reference for fair comparison
+ """
+ # ========================================================================
+ # USER CONFIGURATION - MODIFY THESE VALUES
+ # ========================================================================
+
+ # PROJECTION METHOD TOGGLE
+ # Options: 'BOTH', 'OLD', 'NEW'
+ # - 'BOTH': Compare OLD (corner-based) vs NEW (centered) methods side-by-side
+ # - 'OLD': Run only OLD (corner-based) projection method
+ # - 'NEW': Run only NEW (centered) projection method
+ RUN_METHOD = "NEW" # Change to 'OLD' or 'NEW' to run single method
+
+ # TOPOGRAPHY COARSENING FACTOR
+ # Higher values = coarser topography (faster, less memory)
+ # Typical values: 1 (full resolution), 2, 4, 8
+ ETOPO_CG = 12
+
+ # SPECTRAL COMPRESSION TOGGLE
+ # Toggle between full spectrum vs compressed spectrum in second approximation:
+ #
+ # False (FULL SPECTRUM - default for this test): Use ALL wavenumbers
+ # - Pros: Best reconstruction quality
+ # - Cons: No compression benefit, larger output
+ #
+ # True (COMPRESSED): Use top n_modes=100 wavenumbers
+ # - Pros: Spectral compression (20x smaller)
+ # - Cons: ~20% higher RMSE
+ USE_MODE_SELECTION = True # Set to True to test compressed mode
+
+ # ========================================================================
+ # END USER CONFIGURATION
+ # ========================================================================
+
+ print("=" * 80)
+ print("CENTERED PROJECTION TEST: Old vs. New Planar Projection")
+ print("Testing polar cells (Arctic + Antarctic) at extreme latitudes")
+ if RUN_METHOD == "BOTH":
+ print("Both methods compared against SHARED REFERENCE (centered projection)")
+ elif RUN_METHOD == "OLD":
+ print("Running ONLY OLD (corner-based) projection method")
+ elif RUN_METHOD == "NEW":
+ print("Running ONLY NEW (centered) projection method")
+ else:
+ raise ValueError(
+ f"Invalid RUN_METHOD='{RUN_METHOD}'. Must be 'BOTH', 'OLD', or 'NEW'"
+ )
+ print("=" * 80)
+
+ # Setup parameters
+ from inputs.icon_global_run import params
+
+ params.fn_output = "centered_projection_test"
+ params.etopo_cg = ETOPO_CG
+ params.dfft_first_guess = False
+ params.recompute_rhs = False
+ params.plot_output = False
+
+ # CSA parameters
+ params.lmbda_fa = 1e-2
+ params.lmbda_sa = 1e-1
+ params.fa_iter_solve = True
+ params.sa_iter_solve = True
+
+ if USE_MODE_SELECTION:
+ print(f"*** COMPRESSED MODE: Using top {params.n_modes} wavenumbers ***")
+ else:
+ print(
+ f"*** FULL SPECTRUM MODE: Using ALL {params.nhi * params.nhj} wavenumbers ***"
+ )
+
+ if not params.self_test():
+ print("ERROR: Parameters failed self-test")
+ return
+
+ # Create output directory
+ output_dir = Path("outputs/planar_test")
+ output_dir.mkdir(parents=True, exist_ok=True)
+ print(f"\nOutput directory: {output_dir}")
+
+ # Load ICON grid
+ print("\nLoading ICON grid...")
+ grid = var.grid()
+ reader = io.ncdata(padding=params.padding, padding_tol=(60 - params.padding))
+ reader.read_dat(params.path_icon_grid, grid)
+
+ clat_rad = np.copy(grid.clat)
+ clon_rad = np.copy(grid.clon)
+ grid.apply_f(utils.rad2deg)
+
+ # Use pre-selected extreme polar cells
+ # These cells are at -88.90°S to -87.21°S (within 1-3° of South Pole)
+ # where corner projection creates maximum distortion
+ ALL_TEST_CELLS = POLAR_CELLS
+
+ if len(ALL_TEST_CELLS) == 0:
+ print("\nERROR: No test cells found. Exiting.")
+ return
+
+ print(f"\nTesting {len(ALL_TEST_CELLS)} polar cells (Arctic + Antarctic)")
+
+ # Results storage
+ results = []
+
+ # Test each cell
+ for c_idx in ALL_TEST_CELLS:
+ actual_lat = grid.clat[c_idx]
+ actual_lon = grid.clon[c_idx]
+
+ print(f"\n{'='*80}")
+ print(
+ f"Testing cell {c_idx} at latitude {actual_lat:.2f}°, longitude {actual_lon:.2f}°"
+ )
+ print(f"{'='*80}")
+
+ # Get cell vertices
+ lat_verts = grid.clat_vertices[c_idx]
+ lon_verts = grid.clon_vertices[c_idx]
+ lat_extent, lon_extent = utils.handle_latlon_expansion(
+ lat_verts, lon_verts, lat_expand=0.0, lon_expand=0.0
+ )
+
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ # Load topography
+ print(f" Loading topography...")
+ topo = var.topo_cell()
+ etopo_reader = reader.read_etopo_topo(None, params, is_parallel=True)
+ etopo_reader.get_topo(topo)
+ topo.topo[np.where(topo.topo < -500.0)] = -500.0
+ topo.gen_mgrids()
+
+ # Handle dateline crossing BEFORE processing vertices (like production code)
+ if etopo_reader.split_EW:
+ lon_verts = lon_verts.copy() # Don't modify the grid object
+ lon_verts[lon_verts < 0.0] += 360.0
+
+ # Process vertices exactly like production code (using dateline-corrected lon_verts!)
+ lat_verts_processed, lon_verts_processed = utils.handle_latlon_expansion(
+ lat_verts,
+ lon_verts, # Use corrected vertices, not grid originals
+ lat_expand=0.0,
+ lon_expand=0.0,
+ )
+
+ print(
+ f" Vertices (degrees): lat={lat_verts_processed}, lon={lon_verts_processed}"
+ )
+
+ # ================================================================
+ # CREATE SHARED REFERENCE CELL (Centered Projection - Ground Truth)
+ # ================================================================
+ # This is the canonical reference topography that BOTH methods will be compared against.
+ # Using centered projection (use_center=True) because it's more geometrically accurate,
+ # especially at polar latitudes where corner projection introduces maximum distortion.
+ print(f" Creating shared reference cell (centered projection)...")
+ cell_reference = create_cell_with_projection(
+ lat_verts_processed,
+ lon_verts_processed,
+ topo,
+ use_center=True,
+ rect=False, # Triangular mask for final comparison
+ )
+ print(
+ f" REFERENCE: {cell_reference.mask.sum()} masked points, "
+ f"topo range: [{cell_reference.topo[cell_reference.mask].min():.1f}, "
+ f"{cell_reference.topo[cell_reference.mask].max():.1f}] m"
+ )
+
+ # Initialize variables for optional methods
+ rmse_old_fa, rmse_old_sa = None, None
+ rmse_new_fa, rmse_new_sa = None, None
+ dat_2D_old_fa_interp, dat_2D_old_sa_interp = None, None
+ dat_2D_new_fa, dat_2D_new_sa = None, None
+
+ # TEST 1: OLD projection (corner-based)
+ if RUN_METHOD in ["BOTH", "OLD"]:
+ print(f" Running CSA with OLD projection (corner-based)...")
+
+ # FA: Rectangular domain
+ print(
+ f" [FA] Creating cell with OLD (corner) projection + rectangular mask..."
+ )
+ cell_old_fa = create_cell_with_projection(
+ lat_verts_processed,
+ lon_verts_processed,
+ topo,
+ use_center=False,
+ rect=True,
+ )
+
+ # Run FA
+ fa_old = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_old_fa, uw_old_fa, dat_2D_old_fa = fa_old.sappx(
+ cell_old_fa, lmbda=params.lmbda_fa, iter_solve=params.fa_iter_solve
+ )
+
+ # SA: Triangular domain
+ print(
+ f" [SA] Creating cell with OLD (corner) projection + triangular mask..."
+ )
+ cell_old_sa = create_cell_with_projection(
+ lat_verts_processed,
+ lon_verts_processed,
+ topo,
+ use_center=False,
+ rect=False,
+ )
+
+ # Run SA
+ if USE_MODE_SELECTION:
+ # COMPRESSED MODE: Select top n_modes wavenumbers from FA
+ ampls_old_fa_copy = np.copy(ampls_old_fa)
+ ampls_old_fa_copy[np.isnan(ampls_old_fa_copy)] = 0.0
+ indices = []
+ for _ in range(params.n_modes):
+ max_idx = np.unravel_index(
+ ampls_old_fa_copy.argmax(), ampls_old_fa_copy.shape
+ )
+ indices.append(max_idx)
+ ampls_old_fa_copy[max_idx] = 0.0
+ k_idxs = [pair[1] for pair in indices]
+ l_idxs = [pair[0] for pair in indices]
+ sa_old = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ sa_old.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+ ampls_old_sa, uw_old_sa, dat_2D_old_sa = sa_old.sappx(
+ cell_old_sa, lmbda=params.lmbda_sa, iter_solve=params.sa_iter_solve
+ )
+ else:
+ # FULL SPECTRUM MODE: Use all wavenumbers
+ sa_old = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_old_sa, uw_old_sa, dat_2D_old_sa = sa_old.sappx(
+ cell_old_sa, lmbda=params.lmbda_sa, iter_solve=params.sa_iter_solve
+ )
+
+ # Interpolate OLD method outputs from corner-projection grid to reference grid
+ print(f" Interpolating OLD method outputs to reference grid...")
+ dat_2D_old_fa_interp = interpolate_to_reference_grid(
+ dat_2D_old_fa, cell_old_sa, cell_reference
+ )
+ dat_2D_old_sa_interp = interpolate_to_reference_grid(
+ dat_2D_old_sa, cell_old_sa, cell_reference
+ )
+
+ # Compute RMSE against shared reference (centered projection)
+ diff_fa = cell_reference.topo - dat_2D_old_fa_interp
+ diff_sa = cell_reference.topo - dat_2D_old_sa_interp
+ rmse_old_fa = np.sqrt(np.mean(diff_fa[cell_reference.mask] ** 2))
+ rmse_old_sa = np.sqrt(np.mean(diff_sa[cell_reference.mask] ** 2))
+
+ print(
+ f" OLD - 1st Approx RMSE (vs shared reference): {rmse_old_fa:.1f} m"
+ )
+ print(
+ f" OLD - 2nd Approx RMSE (vs shared reference): {rmse_old_sa:.1f} m"
+ )
+
+ # TEST 2: NEW projection (centered)
+ if RUN_METHOD in ["BOTH", "NEW"]:
+ print(f" Running CSA with NEW projection (centered)...")
+
+ # FA: Rectangular domain
+ print(
+ f" [FA] Creating cell with NEW (centered) projection + rectangular mask..."
+ )
+ cell_new_fa = create_cell_with_projection(
+ lat_verts_processed,
+ lon_verts_processed,
+ topo,
+ use_center=True,
+ rect=True,
+ )
+
+ # Run FA
+ fa_new = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_new_fa, uw_new_fa, dat_2D_new_fa = fa_new.sappx(
+ cell_new_fa, lmbda=params.lmbda_fa, iter_solve=params.fa_iter_solve
+ )
+
+ # SA: Triangular domain
+ print(
+ f" [SA] Creating cell with NEW (centered) projection + triangular mask..."
+ )
+ cell_new_sa = create_cell_with_projection(
+ lat_verts_processed,
+ lon_verts_processed,
+ topo,
+ use_center=True,
+ rect=False,
+ )
+
+ # Run SA
+ if USE_MODE_SELECTION:
+ # COMPRESSED MODE: Select top n_modes wavenumbers from FA
+ ampls_new_fa_copy = np.copy(ampls_new_fa)
+ ampls_new_fa_copy[np.isnan(ampls_new_fa_copy)] = 0.0
+ indices = []
+ for _ in range(params.n_modes):
+ max_idx = np.unravel_index(
+ ampls_new_fa_copy.argmax(), ampls_new_fa_copy.shape
+ )
+ indices.append(max_idx)
+ ampls_new_fa_copy[max_idx] = 0.0
+ k_idxs = [pair[1] for pair in indices]
+ l_idxs = [pair[0] for pair in indices]
+ sa_new = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ sa_new.fobj.set_kls(k_idxs, l_idxs, recompute_nhij=False)
+ ampls_new_sa, uw_new_sa, dat_2D_new_sa = sa_new.sappx(
+ cell_new_sa, lmbda=params.lmbda_sa, iter_solve=params.sa_iter_solve
+ )
+ else:
+ # FULL SPECTRUM MODE: Use all wavenumbers
+ sa_new = interface.get_pmf(params.nhi, params.nhj, params.U, params.V)
+ ampls_new_sa, uw_new_sa, dat_2D_new_sa = sa_new.sappx(
+ cell_new_sa, lmbda=params.lmbda_sa, iter_solve=params.sa_iter_solve
+ )
+
+ # Compute RMSE against shared reference (no interpolation needed - same grid!)
+ # Note: cell_new_sa and cell_reference both use centered projection,
+ # so they're on the same planar grid and can be compared directly
+ diff_fa = cell_reference.topo - dat_2D_new_fa
+ diff_sa = cell_reference.topo - dat_2D_new_sa
+ rmse_new_fa = np.sqrt(np.mean(diff_fa[cell_reference.mask] ** 2))
+ rmse_new_sa = np.sqrt(np.mean(diff_sa[cell_reference.mask] ** 2))
+
+ print(
+ f" NEW - 1st Approx RMSE (vs shared reference): {rmse_new_fa:.1f} m"
+ )
+ print(
+ f" NEW - 2nd Approx RMSE (vs shared reference): {rmse_new_sa:.1f} m"
+ )
+
+ # Compute improvements (only if BOTH methods were run)
+ if RUN_METHOD == "BOTH":
+ imp_fa = (
+ ((rmse_old_fa - rmse_new_fa) / rmse_old_fa * 100)
+ if rmse_old_fa > 0
+ else 0.0
+ )
+ imp_sa = (
+ ((rmse_old_sa - rmse_new_sa) / rmse_old_sa * 100)
+ if rmse_old_sa > 0
+ else 0.0
+ )
+ print(f" IMPROVEMENT - 1st Approx: {imp_fa:.1f}%")
+ print(f" IMPROVEMENT - 2nd Approx: {imp_sa:.1f}%")
+
+ # Generate comparison plot using shared reference topography
+ # Note: All reconstructions are now on the reference grid (centered projection)
+ print(f" Generating comparison plot...")
+ plot_comparison(
+ c_idx,
+ actual_lat,
+ cell_reference.topo, # Shared reference (centered projection)
+ dat_2D_old_fa_interp,
+ dat_2D_old_sa_interp, # OLD method (interpolated to reference grid)
+ dat_2D_new_fa,
+ dat_2D_new_sa, # NEW method (already on reference grid)
+ rmse_old_fa,
+ rmse_old_sa,
+ rmse_new_fa,
+ rmse_new_sa,
+ cell_reference.mask,
+ output_dir, # Use reference mask
+ )
+ elif RUN_METHOD == "OLD":
+ imp_fa = 0.0
+ imp_sa = 0.0
+ print(f" Generating visualization plot for OLD method...")
+ plot_single_method(
+ c_idx,
+ actual_lat,
+ cell_reference.topo, # Reference topography
+ dat_2D_old_fa_interp,
+ dat_2D_old_sa_interp, # OLD method reconstructions
+ rmse_old_fa,
+ rmse_old_sa, # RMSE values
+ cell_reference.mask,
+ output_dir, # Mask and output
+ method_name="OLD",
+ )
+ elif RUN_METHOD == "NEW":
+ imp_fa = 0.0
+ imp_sa = 0.0
+ print(f" Generating visualization plot for NEW method...")
+ plot_single_method(
+ c_idx,
+ actual_lat,
+ cell_reference.topo, # Reference topography
+ dat_2D_new_fa,
+ dat_2D_new_sa, # NEW method reconstructions
+ rmse_new_fa,
+ rmse_new_sa, # RMSE values
+ cell_reference.mask,
+ output_dir, # Mask and output
+ method_name="NEW",
+ )
+
+ # Store results with region tag
+ if actual_lat > 75.0:
+ region = "ARCTIC"
+ elif actual_lat < -75.0:
+ region = "ANTARCTIC"
+ else:
+ region = "MID-LATITUDE"
+
+ # Only store results if we have data to store
+ if RUN_METHOD == "BOTH":
+ results.append(
+ {
+ "cell_idx": c_idx,
+ "lat": actual_lat,
+ "lon": actual_lon,
+ "region": region,
+ "rmse_old_fa": rmse_old_fa,
+ "rmse_old_sa": rmse_old_sa,
+ "rmse_new_fa": rmse_new_fa,
+ "rmse_new_sa": rmse_new_sa,
+ "imp_fa": imp_fa,
+ "imp_sa": imp_sa,
+ }
+ )
+ elif RUN_METHOD == "OLD":
+ results.append(
+ {
+ "cell_idx": c_idx,
+ "lat": actual_lat,
+ "lon": actual_lon,
+ "region": region,
+ "rmse_old_fa": rmse_old_fa,
+ "rmse_old_sa": rmse_old_sa,
+ "rmse_new_fa": None,
+ "rmse_new_sa": None,
+ "imp_fa": None,
+ "imp_sa": None,
+ }
+ )
+ elif RUN_METHOD == "NEW":
+ results.append(
+ {
+ "cell_idx": c_idx,
+ "lat": actual_lat,
+ "lon": actual_lon,
+ "region": region,
+ "rmse_old_fa": None,
+ "rmse_old_sa": None,
+ "rmse_new_fa": rmse_new_fa,
+ "rmse_new_sa": rmse_new_sa,
+ "imp_fa": None,
+ "imp_sa": None,
+ }
+ )
+
+ # Separate results by region
+ arctic_results = [r for r in results if r["region"] == "ARCTIC"]
+ antarctic_results = [r for r in results if r["region"] == "ANTARCTIC"]
+ mid_lat_results = [r for r in results if r["region"] == "MID-LATITUDE"]
+
+ # Print summary
+ print(f"\n{'='*80}")
+ print("SUMMARY OF RESULTS")
+ print(f"{'='*80}")
+
+ # Helper function to format RMSE values (handle None)
+ def fmt_rmse(val):
+ return f"{val:>10.1f}" if val is not None else f"{'N/A':>10}"
+
+ def fmt_imp(val):
+ return f"{val:>7.1f}%" if val is not None else f"{'N/A':>8}"
+
+ if arctic_results:
+ print("\nARCTIC CELLS (lat > 75°N):")
+ print(
+ f"{'Cell':>6} {'Lat':>8} {'Lon':>8} {'OLD FA':>10} {'NEW FA':>10} {'Imp FA':>8} {'OLD SA':>10} {'NEW SA':>10} {'Imp SA':>8}"
+ )
+ print(f"{'-'*80}")
+ for r in arctic_results:
+ print(
+ f"{r['cell_idx']:>6d} {r['lat']:>8.2f} {r['lon']:>8.2f} "
+ f"{fmt_rmse(r['rmse_old_fa'])} {fmt_rmse(r['rmse_new_fa'])} {fmt_imp(r['imp_fa'])} "
+ f"{fmt_rmse(r['rmse_old_sa'])} {fmt_rmse(r['rmse_new_sa'])} {fmt_imp(r['imp_sa'])}"
+ )
+ if RUN_METHOD == "BOTH":
+ avg_arctic_fa = np.mean(
+ [r["imp_fa"] for r in arctic_results if r["imp_fa"] is not None]
+ )
+ avg_arctic_sa = np.mean(
+ [r["imp_sa"] for r in arctic_results if r["imp_sa"] is not None]
+ )
+ print(f" {'Arctic Average - 1st Approx:':>58} {avg_arctic_fa:>7.1f}%")
+ print(f" {'Arctic Average - 2nd Approx:':>58} {avg_arctic_sa:>7.1f}%")
+
+ if antarctic_results:
+ print("\nANTARCTIC CELLS (lat < -75°S):")
+ print(
+ f"{'Cell':>6} {'Lat':>8} {'Lon':>8} {'OLD FA':>10} {'NEW FA':>10} {'Imp FA':>8} {'OLD SA':>10} {'NEW SA':>10} {'Imp SA':>8}"
+ )
+ print(f"{'-'*80}")
+ for r in antarctic_results:
+ print(
+ f"{r['cell_idx']:>6d} {r['lat']:>8.2f} {r['lon']:>8.2f} "
+ f"{fmt_rmse(r['rmse_old_fa'])} {fmt_rmse(r['rmse_new_fa'])} {fmt_imp(r['imp_fa'])} "
+ f"{fmt_rmse(r['rmse_old_sa'])} {fmt_rmse(r['rmse_new_sa'])} {fmt_imp(r['imp_sa'])}"
+ )
+ if RUN_METHOD == "BOTH":
+ avg_antarctic_fa = np.mean(
+ [r["imp_fa"] for r in antarctic_results if r["imp_fa"] is not None]
+ )
+ avg_antarctic_sa = np.mean(
+ [r["imp_sa"] for r in antarctic_results if r["imp_sa"] is not None]
+ )
+ print(
+ f" {'Antarctic Average - 1st Approx:':>58} {avg_antarctic_fa:>7.1f}%"
+ )
+ print(
+ f" {'Antarctic Average - 2nd Approx:':>58} {avg_antarctic_sa:>7.1f}%"
+ )
+
+ if mid_lat_results:
+ print("\nMID-LATITUDE CELLS (|lat| < 75°):")
+ print(
+ f"{'Cell':>6} {'Lat':>8} {'Lon':>8} {'OLD FA':>10} {'NEW FA':>10} {'Imp FA':>8} {'OLD SA':>10} {'NEW SA':>10} {'Imp SA':>8}"
+ )
+ print(f"{'-'*80}")
+ for r in mid_lat_results:
+ print(
+ f"{r['cell_idx']:>6d} {r['lat']:>8.2f} {r['lon']:>8.2f} "
+ f"{fmt_rmse(r['rmse_old_fa'])} {fmt_rmse(r['rmse_new_fa'])} {fmt_imp(r['imp_fa'])} "
+ f"{fmt_rmse(r['rmse_old_sa'])} {fmt_rmse(r['rmse_new_sa'])} {fmt_imp(r['imp_sa'])}"
+ )
+ if RUN_METHOD == "BOTH":
+ avg_mid_lat_fa = np.mean(
+ [r["imp_fa"] for r in mid_lat_results if r["imp_fa"] is not None]
+ )
+ avg_mid_lat_sa = np.mean(
+ [r["imp_sa"] for r in mid_lat_results if r["imp_sa"] is not None]
+ )
+ print(
+ f" {'Mid-Latitude Average - 1st Approx:':>58} {avg_mid_lat_fa:>7.1f}%"
+ )
+ print(
+ f" {'Mid-Latitude Average - 2nd Approx:':>58} {avg_mid_lat_sa:>7.1f}%"
+ )
+
+ # Calculate overall averages (only for BOTH mode)
+ if RUN_METHOD == "BOTH":
+ avg_imp_fa = np.mean([r["imp_fa"] for r in results if r["imp_fa"] is not None])
+ avg_imp_sa = np.mean([r["imp_sa"] for r in results if r["imp_sa"] is not None])
+ print(f"\n{'OVERALL Average - 1st Approx:':>60} {avg_imp_fa:>7.1f}%")
+ print(f"{'OVERALL Average - 2nd Approx:':>60} {avg_imp_sa:>7.1f}%")
+
+ print(f"\n{'='*80}")
+ print(f"All plots saved to: {output_dir}")
+ print(f"{'='*80}")
+
+ # Save results to file
+ results_file = output_dir / "results_summary.txt"
+ with open(results_file, "w") as f:
+ f.write("CENTERED PROJECTION TEST RESULTS\n")
+ f.write("=" * 80 + "\n\n")
+ f.write(f"Testing {len(results)} cells:\n")
+ f.write(f" Arctic cells (lat > 75°N): {len(arctic_results)}\n")
+ f.write(f" Antarctic cells (lat < -75°S): {len(antarctic_results)}\n")
+ f.write(f" Mid-latitude cells (|lat| < 75°): {len(mid_lat_results)}\n\n")
+
+ if RUN_METHOD == "BOTH":
+ f.write(
+ f"Comparing OLD (corner-based) vs NEW (centered) planar projection\n"
+ )
+ f.write(
+ f"Running FULL pyCSA: First Approximation + Second Approximation\n\n"
+ )
+ f.write(
+ f"IMPORTANT: Both methods are compared against the SAME reference topography\n"
+ )
+ f.write(f" (centered projection, geometrically accurate).\n")
+ f.write(
+ f" OLD method reconstructions interpolated to reference grid.\n\n"
+ )
+ elif RUN_METHOD == "OLD":
+ f.write(f"Testing OLD (corner-based) planar projection ONLY\n")
+ f.write(
+ f"Running FULL pyCSA: First Approximation + Second Approximation\n\n"
+ )
+ elif RUN_METHOD == "NEW":
+ f.write(f"Testing NEW (centered) planar projection ONLY\n")
+ f.write(
+ f"Running FULL pyCSA: First Approximation + Second Approximation\n\n"
+ )
+
+ # Helper function for file writing
+ def fmt_rmse_file(val):
+ return f"{val:>10.1f}" if val is not None else f"{'N/A':>10}"
+
+ def fmt_imp_file(val):
+ return f"{val:>7.1f}%" if val is not None else f"{'N/A':>8}"
+
+ if arctic_results:
+ f.write("ARCTIC CELLS (lat > 75°N):\n")
+ f.write(
+ f"{'Cell':>6} {'Lat':>8} {'Lon':>8} {'OLD FA':>10} {'NEW FA':>10} {'Imp FA':>8} {'OLD SA':>10} {'NEW SA':>10} {'Imp SA':>8}\n"
+ )
+ f.write("-" * 80 + "\n")
+ for r in arctic_results:
+ f.write(
+ f"{r['cell_idx']:>6d} {r['lat']:>8.2f} {r['lon']:>8.2f} "
+ f"{fmt_rmse_file(r['rmse_old_fa'])} {fmt_rmse_file(r['rmse_new_fa'])} {fmt_imp_file(r['imp_fa'])} "
+ f"{fmt_rmse_file(r['rmse_old_sa'])} {fmt_rmse_file(r['rmse_new_sa'])} {fmt_imp_file(r['imp_sa'])}\n"
+ )
+ if RUN_METHOD == "BOTH":
+ avg_arctic_fa = np.mean(
+ [r["imp_fa"] for r in arctic_results if r["imp_fa"] is not None]
+ )
+ avg_arctic_sa = np.mean(
+ [r["imp_sa"] for r in arctic_results if r["imp_sa"] is not None]
+ )
+ f.write(
+ f" {'Arctic Average - 1st Approx:':>58} {avg_arctic_fa:>7.1f}%\n"
+ )
+ f.write(
+ f" {'Arctic Average - 2nd Approx:':>58} {avg_arctic_sa:>7.1f}%\n\n"
+ )
+ else:
+ f.write("\n")
+
+ if antarctic_results:
+ f.write("ANTARCTIC CELLS (lat < -75°S):\n")
+ f.write(
+ f"{'Cell':>6} {'Lat':>8} {'Lon':>8} {'OLD FA':>10} {'NEW FA':>10} {'Imp FA':>8} {'OLD SA':>10} {'NEW SA':>10} {'Imp SA':>8}\n"
+ )
+ f.write("-" * 80 + "\n")
+ for r in antarctic_results:
+ f.write(
+ f"{r['cell_idx']:>6d} {r['lat']:>8.2f} {r['lon']:>8.2f} "
+ f"{fmt_rmse_file(r['rmse_old_fa'])} {fmt_rmse_file(r['rmse_new_fa'])} {fmt_imp_file(r['imp_fa'])} "
+ f"{fmt_rmse_file(r['rmse_old_sa'])} {fmt_rmse_file(r['rmse_new_sa'])} {fmt_imp_file(r['imp_sa'])}\n"
+ )
+ if RUN_METHOD == "BOTH":
+ avg_antarctic_fa = np.mean(
+ [r["imp_fa"] for r in antarctic_results if r["imp_fa"] is not None]
+ )
+ avg_antarctic_sa = np.mean(
+ [r["imp_sa"] for r in antarctic_results if r["imp_sa"] is not None]
+ )
+ f.write(
+ f" {'Antarctic Average - 1st Approx:':>58} {avg_antarctic_fa:>7.1f}%\n"
+ )
+ f.write(
+ f" {'Antarctic Average - 2nd Approx:':>58} {avg_antarctic_sa:>7.1f}%\n\n"
+ )
+ else:
+ f.write("\n")
+
+ if mid_lat_results:
+ f.write("MID-LATITUDE CELLS (|lat| < 75°):\n")
+ f.write(
+ f"{'Cell':>6} {'Lat':>8} {'Lon':>8} {'OLD FA':>10} {'NEW FA':>10} {'Imp FA':>8} {'OLD SA':>10} {'NEW SA':>10} {'Imp SA':>8}\n"
+ )
+ f.write("-" * 80 + "\n")
+ for r in mid_lat_results:
+ f.write(
+ f"{r['cell_idx']:>6d} {r['lat']:>8.2f} {r['lon']:>8.2f} "
+ f"{fmt_rmse_file(r['rmse_old_fa'])} {fmt_rmse_file(r['rmse_new_fa'])} {fmt_imp_file(r['imp_fa'])} "
+ f"{fmt_rmse_file(r['rmse_old_sa'])} {fmt_rmse_file(r['rmse_new_sa'])} {fmt_imp_file(r['imp_sa'])}\n"
+ )
+ if RUN_METHOD == "BOTH":
+ avg_mid_lat_fa = np.mean(
+ [r["imp_fa"] for r in mid_lat_results if r["imp_fa"] is not None]
+ )
+ avg_mid_lat_sa = np.mean(
+ [r["imp_sa"] for r in mid_lat_results if r["imp_sa"] is not None]
+ )
+ f.write(
+ f" {'Mid-Latitude Average - 1st Approx:':>58} {avg_mid_lat_fa:>7.1f}%\n"
+ )
+ f.write(
+ f" {'Mid-Latitude Average - 2nd Approx:':>58} {avg_mid_lat_sa:>7.1f}%\n\n"
+ )
+ else:
+ f.write("\n")
+
+ f.write("-" * 80 + "\n")
+ if RUN_METHOD == "BOTH":
+ avg_imp_fa = np.mean(
+ [r["imp_fa"] for r in results if r["imp_fa"] is not None]
+ )
+ avg_imp_sa = np.mean(
+ [r["imp_sa"] for r in results if r["imp_sa"] is not None]
+ )
+ f.write(f"{'OVERALL Average - 1st Approx:':>60} {avg_imp_fa:>7.1f}%\n")
+ f.write(f"{'OVERALL Average - 2nd Approx:':>60} {avg_imp_sa:>7.1f}%\n")
+
+ print(f"\nResults summary saved to: {results_file}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/tests/test_icon_etopo_validation.py b/tests/test_icon_etopo_validation.py
new file mode 100644
index 0000000..c0d35b1
--- /dev/null
+++ b/tests/test_icon_etopo_validation.py
@@ -0,0 +1,760 @@
+"""
+Test ICON grid cells against real-world ETOPO topography.
+
+This module validates that ICON grid cells and their associated ETOPO topography
+data correctly correspond to real-world geographical features. This ensures that
+coordinate transformations, data loading, and spatial mapping are functioning correctly.
+
+Test categories:
+1. Mountains: Verify high elevation features (Himalayas, Andes, Alps, etc.)
+2. Lakes: Verify inland water bodies (Great Lakes, Lake Baikal, etc.)
+3. Oceans/Gulfs: Verify marine features (Pacific, Gulf of Mexico, etc.)
+4. Coasts: Verify land-ocean transitions
+5. Edge cases: Dateline, poles, tile boundaries
+"""
+
+import pytest
+import numpy as np
+from pathlib import Path
+import matplotlib.pyplot as plt
+from typing import Tuple, Dict, List, Optional
+
+from pycsa.core import io, var, utils
+from pycsa import local_paths
+
+
+class GeographicFeature:
+ """Represents a known geographic feature for validation."""
+
+ def __init__(
+ self,
+ name: str,
+ lat_range: Tuple[float, float],
+ lon_range: Tuple[float, float],
+ feature_type: str,
+ validation_func,
+ description: str = "",
+ ):
+ """
+ Initialize a geographic feature.
+
+ Args:
+ name: Feature name (e.g., "Himalayas", "Lake Superior")
+ lat_range: (min_lat, max_lat) in degrees
+ lon_range: (min_lon, max_lon) in degrees
+ feature_type: One of "mountain", "lake", "ocean", "gulf", "coast"
+ validation_func: Function that validates topography matches feature
+ description: Human-readable description
+ """
+ self.name = name
+ self.lat_range = lat_range
+ self.lon_range = lon_range
+ self.feature_type = feature_type
+ self.validation_func = validation_func
+ self.description = description
+
+ def get_center(self) -> Tuple[float, float]:
+ """Return (center_lat, center_lon) of feature."""
+ lat_center = np.mean(self.lat_range)
+ lon_center = np.mean(self.lon_range)
+ return lat_center, lon_center
+
+ def validate(self, topo_cell: var.topo_cell) -> Dict:
+ """
+ Validate that topography matches this geographic feature.
+
+ Returns:
+ Dict with keys: 'passed', 'message', 'stats'
+ """
+ return self.validation_func(topo_cell, self)
+
+
+# Validation functions for different feature types
+def validate_mountain(topo_cell: var.topo_cell, feature: GeographicFeature) -> Dict:
+ """Validate mountain features have high elevations."""
+ max_elev = topo_cell.topo.max()
+ min_expected = 3000 # meters
+
+ # Different mountain ranges have different heights
+ if "Himalayas" in feature.name or "Karakoram" in feature.name:
+ min_expected = 5000 # Should have peaks > 5km
+ elif "Andes" in feature.name or "Alps" in feature.name:
+ min_expected = 3500
+ elif "Rockies" in feature.name or "Appalachian" in feature.name:
+ min_expected = 2000
+
+ passed = max_elev >= min_expected
+ message = (
+ f"{feature.name}: max elevation {max_elev:.0f}m (expected >{min_expected}m)"
+ )
+
+ stats = {
+ "max_elevation": max_elev,
+ "mean_elevation": topo_cell.topo.mean(),
+ "min_elevation": topo_cell.topo.min(),
+ "std_elevation": topo_cell.topo.std(),
+ "high_terrain_fraction": (topo_cell.topo > 1000).sum() / topo_cell.topo.size,
+ }
+
+ return {"passed": passed, "message": message, "stats": stats}
+
+
+def validate_lake(topo_cell: var.topo_cell, feature: GeographicFeature) -> Dict:
+ """Validate lake features have appropriate water elevation."""
+ # Lakes regions include surrounding terrain, so we check:
+ # 1. Minimum elevation should be near expected lake level
+ # 2. Should have some low-elevation areas (the actual lake)
+
+ min_elev = topo_cell.topo.min()
+ mean_elev = topo_cell.topo.mean()
+
+ # Count how much of the area is near the expected lake elevation
+ # Special cases for different lakes
+ if "Titicaca" in feature.name:
+ expected_lake_elev = 3812 # meters
+ tolerance = 300 # Allow surrounding mountains
+ elif "Baikal" in feature.name:
+ expected_lake_elev = 456 # meters
+ tolerance = 500 # Mountainous region
+ elif "Great Lakes" in feature.name or "Superior" in feature.name:
+ expected_lake_elev = 183 # meters
+ tolerance = 200 # Relatively flat region
+ else:
+ expected_lake_elev = 100 # Generic lake
+ tolerance = 300
+
+ # Check that minimum elevation is close to lake level (below it due to lake depth)
+ lake_depth_margin = 500 # Lakes can be deep
+ min_expected = expected_lake_elev - lake_depth_margin
+ max_expected = expected_lake_elev + tolerance
+
+ # Count fraction of area near lake elevation (within tolerance)
+ near_lake_level = np.abs(topo_cell.topo - expected_lake_elev) < tolerance
+ lake_fraction = near_lake_level.sum() / topo_cell.topo.size
+
+ # Validate: minimum should be below/near lake level, and some area should be at lake level
+ has_low_areas = min_elev < expected_lake_elev + 100
+ has_lake_level_areas = lake_fraction > 0.05 # At least 5% at lake level
+
+ passed = has_low_areas and has_lake_level_areas
+ message = (
+ f"{feature.name}: min elev {min_elev:.0f}m, mean {mean_elev:.0f}m, "
+ f"{lake_fraction:.1%} near lake level ~{expected_lake_elev}m"
+ )
+
+ stats = {
+ "mean_elevation": mean_elev,
+ "min_elevation": min_elev,
+ "max_elevation": topo_cell.topo.max(),
+ "std_elevation": topo_cell.topo.std(),
+ "expected_lake_elevation": expected_lake_elev,
+ "fraction_near_lake_level": lake_fraction,
+ "has_low_areas": has_low_areas,
+ "has_lake_level_areas": has_lake_level_areas,
+ }
+
+ return {"passed": passed, "message": message, "stats": stats}
+
+
+def validate_ocean(topo_cell: var.topo_cell, feature: GeographicFeature) -> Dict:
+ """Validate ocean features have negative (below sea level) elevations."""
+ # Oceans should be mostly below sea level
+ water_fraction = (topo_cell.topo < 0).sum() / topo_cell.topo.size
+ mean_depth = (
+ -topo_cell.topo[topo_cell.topo < 0].mean() if (topo_cell.topo < 0).any() else 0
+ )
+
+ min_water_fraction = 0.80 # At least 80% should be water
+
+ # Deep ocean should have significant depth
+ if "Pacific" in feature.name or "Atlantic" in feature.name:
+ min_expected_depth = 3000 # Deep ocean
+ else:
+ min_expected_depth = 100 # Shallow seas/gulfs
+
+ passed = water_fraction >= min_water_fraction and mean_depth >= min_expected_depth
+ message = (
+ f"{feature.name}: water fraction {water_fraction:.1%}, "
+ f"mean depth {mean_depth:.0f}m (expected >{min_expected_depth}m)"
+ )
+
+ stats = {
+ "water_fraction": water_fraction,
+ "mean_depth": mean_depth,
+ "max_depth": -topo_cell.topo.min(),
+ "mean_elevation": topo_cell.topo.mean(),
+ "land_fraction": (topo_cell.topo >= 0).sum() / topo_cell.topo.size,
+ }
+
+ return {"passed": passed, "message": message, "stats": stats}
+
+
+def validate_gulf(topo_cell: var.topo_cell, feature: GeographicFeature) -> Dict:
+ """Validate gulf/bay features have mostly water with some coastline."""
+ # Gulfs should be mostly water but may have significant land depending on region bounds
+ water_fraction = (topo_cell.topo < 0).sum() / topo_cell.topo.size
+ mean_water_depth = (
+ -topo_cell.topo[topo_cell.topo < 0].mean() if (topo_cell.topo < 0).any() else 0
+ )
+
+ # Adjust thresholds based on specific gulf
+ if "Persian Gulf" in feature.name:
+ min_water_fraction = 0.70 # Fairly shallow, wide gulf
+ min_expected_depth = 30 # Persian Gulf is shallow
+ else:
+ min_water_fraction = 0.50 # At least 50% should be water
+ min_expected_depth = 50 # Should have some depth
+
+ passed = (
+ water_fraction >= min_water_fraction and mean_water_depth >= min_expected_depth
+ )
+
+ message = (
+ f"{feature.name}: water fraction {water_fraction:.1%}, "
+ f"mean depth {mean_water_depth:.0f}m (expected >{min_expected_depth}m)"
+ )
+
+ stats = {
+ "water_fraction": water_fraction,
+ "land_fraction": (topo_cell.topo >= 0).sum() / topo_cell.topo.size,
+ "mean_water_depth": mean_water_depth,
+ "mean_elevation": topo_cell.topo.mean(),
+ "elevation_range": topo_cell.topo.max() - topo_cell.topo.min(),
+ "min_expected_depth": min_expected_depth,
+ }
+
+ return {"passed": passed, "message": message, "stats": stats}
+
+
+def validate_coast(topo_cell: var.topo_cell, feature: GeographicFeature) -> Dict:
+ """Validate coastal features have both land and water."""
+ # Coasts should have significant mix of land and water
+ water_fraction = (topo_cell.topo < 0).sum() / topo_cell.topo.size
+ land_fraction = (topo_cell.topo >= 0).sum() / topo_cell.topo.size
+
+ # Coast should have reasonable mix (20-80% water)
+ min_water = 0.20
+ max_water = 0.80
+
+ passed = min_water <= water_fraction <= max_water
+ message = (
+ f"{feature.name}: water {water_fraction:.1%}, land {land_fraction:.1%} "
+ f"(expected {min_water:.0%}-{max_water:.0%} water)"
+ )
+
+ stats = {
+ "water_fraction": water_fraction,
+ "land_fraction": land_fraction,
+ "mean_elevation": topo_cell.topo.mean(),
+ "elevation_range": topo_cell.topo.max() - topo_cell.topo.min(),
+ "std_elevation": topo_cell.topo.std(),
+ }
+
+ return {"passed": passed, "message": message, "stats": stats}
+
+
+# Define known geographic features for testing
+GEOGRAPHIC_FEATURES = [
+ # Mountains
+ GeographicFeature(
+ "Himalayas",
+ (27.0, 30.0),
+ (85.0, 90.0),
+ "mountain",
+ validate_mountain,
+ "World's highest mountain range (Everest, K2)",
+ ),
+ GeographicFeature(
+ "Andes (Peru)",
+ (-15.0, -10.0),
+ (-77.0, -72.0),
+ "mountain",
+ validate_mountain,
+ "Andes mountain range in Peru",
+ ),
+ GeographicFeature(
+ "Alps",
+ (45.5, 47.5),
+ (6.0, 11.0),
+ "mountain",
+ validate_mountain,
+ "European Alps (Mont Blanc)",
+ ),
+ GeographicFeature(
+ "Rockies (Colorado)",
+ (38.0, 41.0),
+ (-108.0, -105.0),
+ "mountain",
+ validate_mountain,
+ "Rocky Mountains in Colorado",
+ ),
+ # Lakes
+ GeographicFeature(
+ "Lake Superior",
+ (46.5, 48.5),
+ (-89.0, -85.0),
+ "lake",
+ validate_lake,
+ "Largest Great Lake by area",
+ ),
+ GeographicFeature(
+ "Lake Baikal",
+ (51.5, 55.5),
+ (103.5, 109.5),
+ "lake",
+ validate_lake,
+ "World's deepest lake in Siberia",
+ ),
+ GeographicFeature(
+ "Lake Titicaca",
+ (-16.5, -15.0),
+ (-69.5, -68.5),
+ "lake",
+ validate_lake,
+ "High-altitude lake in Andes (Peru/Bolivia border)",
+ ),
+ # Oceans
+ GeographicFeature(
+ "Pacific Ocean (mid)",
+ (10.0, 15.0),
+ (-160.0, -150.0),
+ "ocean",
+ validate_ocean,
+ "Central Pacific Ocean",
+ ),
+ GeographicFeature(
+ "Atlantic Ocean (mid)",
+ (25.0, 30.0),
+ (-50.0, -40.0),
+ "ocean",
+ validate_ocean,
+ "Central Atlantic Ocean",
+ ),
+ # Gulfs and Bays
+ GeographicFeature(
+ "Gulf of Mexico",
+ (27.0, 29.5),
+ (-94.0, -89.0),
+ "gulf",
+ validate_gulf,
+ "Gulf of Mexico central region with coastal areas",
+ ),
+ GeographicFeature(
+ "Persian Gulf",
+ (26.0, 28.0),
+ (50.0, 52.0),
+ "gulf",
+ validate_gulf,
+ "Persian Gulf between Iran and Arabia",
+ ),
+ # Coasts
+ GeographicFeature(
+ "California Coast",
+ (35.0, 37.0),
+ (-122.0, -120.0),
+ "coast",
+ validate_coast,
+ "California coastline near Monterey",
+ ),
+ GeographicFeature(
+ "Mediterranean Coast (Spain)",
+ (40.0, 42.0),
+ (1.0, 3.0),
+ "coast",
+ validate_coast,
+ "Spanish Mediterranean coast",
+ ),
+]
+
+
+class TestICONETOPOValidation:
+ """Validate ICON grid cells against ETOPO topography."""
+
+ @pytest.fixture(scope="class")
+ def setup(self):
+ """Setup test parameters and data structures."""
+ params = var.params()
+ utils.transfer_attributes(params, local_paths.paths, prefix="path")
+ params.etopo_cg = 4 # Use coarse-graining for faster tests
+ params.padding = 0
+
+ # Load ICON grid
+ grid = var.grid()
+ reader = io.ncdata(padding=params.padding, padding_tol=60)
+ reader.read_dat(params.path_icon_grid, grid)
+ grid.apply_f(utils.rad2deg)
+
+ return {"params": params, "grid": grid, "reader": reader}
+
+ def load_region_topography(
+ self,
+ setup: Dict,
+ lat_range: Tuple[float, float],
+ lon_range: Tuple[float, float],
+ ) -> var.topo_cell:
+ """
+ Load topography for a specific lat/lon region.
+
+ Args:
+ setup: Test setup dictionary with params and reader
+ lat_range: (min_lat, max_lat) in degrees
+ lon_range: (min_lon, max_lon) in degrees
+
+ Returns:
+ topo_cell with loaded topography data
+ """
+ params = setup["params"]
+ reader = setup["reader"]
+
+ # Set region extents
+ params.lat_extent = list(lat_range)
+ params.lon_extent = list(lon_range)
+
+ # Load topography
+ topo = var.topo_cell()
+ etopo_reader = reader.read_etopo_topo(
+ None, params, is_parallel=True, verbose=False
+ )
+ etopo_reader.get_topo(topo)
+ etopo_reader.close_cached_files()
+
+ # Generate mesh grids
+ topo.gen_mgrids()
+
+ return topo
+
+ def load_cell_topography(
+ self, setup: Dict, cell_idx: int
+ ) -> Tuple[var.topo_cell, np.ndarray, np.ndarray]:
+ """
+ Load topography for a specific ICON grid cell.
+
+ Args:
+ setup: Test setup dictionary
+ cell_idx: ICON grid cell index
+
+ Returns:
+ (topo_cell, lat_vertices, lon_vertices)
+ """
+ params = setup["params"]
+ grid = setup["grid"]
+ reader = setup["reader"]
+
+ # Get cell vertices
+ lat_verts = grid.clat_vertices[cell_idx]
+ lon_verts = grid.clon_vertices[cell_idx]
+
+ # Handle edge cases (dateline, poles)
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ # Load topography
+ topo = var.topo_cell()
+ etopo_reader = reader.read_etopo_topo(
+ None, params, is_parallel=True, verbose=False
+ )
+ etopo_reader.get_topo(topo)
+ etopo_reader.close_cached_files()
+
+ topo.gen_mgrids()
+
+ return topo, lat_verts, lon_verts
+
+ def test_topography_data_quality_basic(self, setup):
+ """Test that loaded topography has valid data structure."""
+ # Load a simple region (central Pacific)
+ topo = self.load_region_topography(setup, (10.0, 20.0), (-160.0, -150.0))
+
+ # Basic structure checks
+ assert topo.topo is not None, "No topography loaded"
+ assert (
+ topo.lat is not None and topo.lon is not None
+ ), "Missing coordinate arrays"
+ assert topo.topo.shape[0] == len(topo.lat), "Latitude dimension mismatch"
+ assert topo.topo.shape[1] == len(topo.lon), "Longitude dimension mismatch"
+
+ # Check for NaN values
+ nan_count = np.sum(np.isnan(topo.topo))
+ assert nan_count == 0, f"Found {nan_count} NaN values in topography"
+
+ # Sanity check elevation range (Earth surface)
+ assert (
+ topo.topo.min() >= -12000
+ ), f"Elevation too low: {topo.topo.min()}m (deepest ocean ~-11km)"
+ assert (
+ topo.topo.max() <= 9000
+ ), f"Elevation too high: {topo.topo.max()}m (Everest ~8.8km)"
+
+ print(
+ f"✓ Data quality check passed: shape={topo.topo.shape}, "
+ f"elev=[{topo.topo.min():.0f}, {topo.topo.max():.0f}]m"
+ )
+
+ @pytest.mark.parametrize("feature", GEOGRAPHIC_FEATURES, ids=lambda f: f.name)
+ def test_geographic_feature(self, setup, feature: GeographicFeature):
+ """Test that a specific geographic feature validates correctly."""
+ print(f"\nTesting: {feature.name} ({feature.feature_type})")
+ print(f" Location: lat={feature.lat_range}, lon={feature.lon_range}")
+ print(f" Description: {feature.description}")
+
+ # Load topography for this region
+ topo = self.load_region_topography(setup, feature.lat_range, feature.lon_range)
+
+ # Validate against feature
+ result = feature.validate(topo)
+
+ # Print statistics
+ print(f" {result['message']}")
+ for key, value in result["stats"].items():
+ if isinstance(value, float):
+ print(f" {key}: {value:.2f}")
+ else:
+ print(f" {key}: {value}")
+
+ # Assert validation passed
+ assert result[
+ "passed"
+ ], f"{feature.name} validation failed: {result['message']}"
+ print(f" ✓ Validation PASSED")
+
+ def test_cell_near_himalayas(self, setup):
+ """Test loading a cell near the Himalayas and verify high elevations."""
+ grid = setup["grid"]
+
+ # Find cell near Himalayas (28°N, 87°E - near Everest)
+ cell_idx = utils.pick_cell(lat_ref=28.0, lon_ref=87.0, grid=grid, radius=1.0)
+ assert cell_idx is not None, "Could not find cell near Himalayas"
+
+ print(f"\nTesting ICON cell {cell_idx} near Himalayas")
+
+ # Load cell topography
+ topo, lat_verts, lon_verts = self.load_cell_topography(setup, cell_idx)
+
+ print(
+ f" Cell vertices: lat={np.rad2deg(lat_verts)}, lon={np.rad2deg(lon_verts)}"
+ )
+ print(f" Topography shape: {topo.topo.shape}")
+ print(
+ f" Elevation: [{topo.topo.min():.0f}, {topo.topo.max():.0f}]m, mean={topo.topo.mean():.0f}m"
+ )
+
+ # Verify high elevations
+ assert (
+ topo.topo.max() > 4000
+ ), f"Expected high peaks in Himalayas, got {topo.topo.max():.0f}m"
+ assert (
+ topo.topo.mean() > 2000
+ ), f"Expected high mean elevation, got {topo.topo.mean():.0f}m"
+
+ print(f" ✓ Himalayan cell validation PASSED")
+
+ def test_cell_in_pacific_ocean(self, setup):
+ """Test loading a cell in the Pacific Ocean and verify it's water."""
+ grid = setup["grid"]
+
+ # Find cell in Pacific (15°N, 155°W)
+ cell_idx = utils.pick_cell(lat_ref=15.0, lon_ref=-155.0, grid=grid, radius=1.0)
+ assert cell_idx is not None, "Could not find cell in Pacific"
+
+ print(f"\nTesting ICON cell {cell_idx} in Pacific Ocean")
+
+ # Load cell topography
+ topo, lat_verts, lon_verts = self.load_cell_topography(setup, cell_idx)
+
+ print(
+ f" Cell vertices: lat={np.rad2deg(lat_verts)}, lon={np.rad2deg(lon_verts)}"
+ )
+ print(f" Topography shape: {topo.topo.shape}")
+ print(
+ f" Elevation: [{topo.topo.min():.0f}, {topo.topo.max():.0f}]m, mean={topo.topo.mean():.0f}m"
+ )
+
+ # Verify it's ocean
+ water_fraction = (topo.topo < 0).sum() / topo.topo.size
+ print(f" Water fraction: {water_fraction:.1%}")
+
+ assert (
+ water_fraction > 0.95
+ ), f"Expected mostly water in Pacific, got {water_fraction:.1%}"
+ assert (
+ topo.topo.mean() < -1000
+ ), f"Expected deep ocean, got mean depth {-topo.topo.mean():.0f}m"
+
+ print(f" ✓ Pacific Ocean cell validation PASSED")
+
+ def test_cell_on_california_coast(self, setup):
+ """Test loading a coastal cell and verify land-water mix."""
+ grid = setup["grid"]
+
+ # Find cell on California coast (36°N, 122°W)
+ cell_idx = utils.pick_cell(lat_ref=36.0, lon_ref=-122.0, grid=grid, radius=1.0)
+ assert cell_idx is not None, "Could not find cell on California coast"
+
+ print(f"\nTesting ICON cell {cell_idx} on California coast")
+
+ # Load cell topography
+ topo, lat_verts, lon_verts = self.load_cell_topography(setup, cell_idx)
+
+ print(
+ f" Cell vertices: lat={np.rad2deg(lat_verts)}, lon={np.rad2deg(lon_verts)}"
+ )
+ print(f" Topography shape: {topo.topo.shape}")
+ print(f" Elevation: [{topo.topo.min():.0f}, {topo.topo.max():.0f}]m")
+
+ # Verify it's coastal (mix of land and water)
+ water_fraction = (topo.topo < 0).sum() / topo.topo.size
+ land_fraction = (topo.topo >= 0).sum() / topo.topo.size
+
+ print(f" Water fraction: {water_fraction:.1%}")
+ print(f" Land fraction: {land_fraction:.1%}")
+
+ # Coast should have both land and water
+ assert (
+ 0.10 < water_fraction < 0.90
+ ), f"Expected coastal mix, got {water_fraction:.1%} water"
+
+ print(f" ✓ Coastal cell validation PASSED")
+
+ def test_multiple_cells_consistency(self, setup):
+ """Test that multiple cells across different regions load consistently."""
+ grid = setup["grid"]
+
+ # Test cells at various locations
+ test_locations = [
+ (0.0, 0.0, "Equator/Prime Meridian"),
+ (45.0, 0.0, "Mid-latitude Europe"),
+ (0.0, 180.0, "Equator/Dateline"),
+ (-30.0, 150.0, "Australia region"),
+ (60.0, -100.0, "Northern Canada"),
+ ]
+
+ results = []
+ for lat, lon, description in test_locations:
+ cell_idx = utils.pick_cell(lat_ref=lat, lon_ref=lon, grid=grid, radius=1.0)
+ if cell_idx is None:
+ print(f" ⚠ Could not find cell at {description} ({lat}, {lon})")
+ continue
+
+ try:
+ topo, lat_verts, lon_verts = self.load_cell_topography(setup, cell_idx)
+
+ result = {
+ "location": description,
+ "cell_idx": cell_idx,
+ "lat": lat,
+ "lon": lon,
+ "shape": topo.topo.shape,
+ "elev_min": topo.topo.min(),
+ "elev_max": topo.topo.max(),
+ "elev_mean": topo.topo.mean(),
+ "has_nan": np.isnan(topo.topo).any(),
+ "success": True,
+ }
+ results.append(result)
+
+ print(
+ f" ✓ Cell {cell_idx} ({description}): "
+ f"shape={topo.topo.shape}, elev=[{topo.topo.min():.0f}, {topo.topo.max():.0f}]m"
+ )
+
+ except Exception as e:
+ print(f" ✗ Cell {cell_idx} ({description}) FAILED: {str(e)}")
+ results.append(
+ {
+ "location": description,
+ "cell_idx": cell_idx,
+ "success": False,
+ "error": str(e),
+ }
+ )
+
+ # Verify all succeeded
+ success_count = sum(1 for r in results if r["success"])
+ print(f"\n Summary: {success_count}/{len(results)} cells loaded successfully")
+
+ assert success_count == len(
+ results
+ ), f"Some cells failed to load: {len(results) - success_count} failures"
+
+ # Verify no NaN values in any cell
+ nan_count = sum(1 for r in results if r.get("has_nan", False))
+ assert nan_count == 0, f"Found NaN values in {nan_count} cells"
+
+
+class TestICONETOPOVisualization:
+ """Optional visualization tests for debugging (requires matplotlib)."""
+
+ @pytest.fixture(scope="class")
+ def setup(self):
+ """Setup test parameters and data structures."""
+ params = var.params()
+ utils.transfer_attributes(params, local_paths.paths, prefix="path")
+ params.etopo_cg = 4
+ params.padding = 0
+
+ grid = var.grid()
+ reader = io.ncdata(padding=params.padding, padding_tol=60)
+ reader.read_dat(params.path_icon_grid, grid)
+ grid.apply_f(utils.rad2deg)
+
+ return {"params": params, "grid": grid, "reader": reader}
+
+ def test_visualize_feature(self, setup):
+ """Visualize a geographic feature for debugging.
+
+ Run with: pytest -v -s -k visualization
+ """
+ # Pick a feature to visualize (Himalayas)
+ feature = GEOGRAPHIC_FEATURES[5] # Himalayas
+
+ # Load topography
+ params = setup["params"]
+ reader = setup["reader"]
+
+ params.lat_extent = list(feature.lat_range)
+ params.lon_extent = list(feature.lon_range)
+
+ topo = var.topo_cell()
+ etopo_reader = reader.read_etopo_topo(
+ None, params, is_parallel=True, verbose=True
+ )
+ etopo_reader.get_topo(topo)
+ etopo_reader.close_cached_files()
+ topo.gen_mgrids()
+
+ # Create visualization
+ fig, axes = plt.subplots(1, 2, figsize=(14, 6))
+
+ # Plot 1: Raw topography
+ im1 = axes[0].imshow(topo.topo, origin="lower", cmap="terrain", aspect="auto")
+ axes[0].set_title(f"{feature.name} - Raw Topography")
+ axes[0].set_xlabel(f"Longitude index")
+ axes[0].set_ylabel(f"Latitude index")
+ plt.colorbar(im1, ax=axes[0], label="Elevation (m)")
+
+ # Plot 2: Contour plot with coordinates
+ levels = 20
+ cs = axes[1].contourf(
+ topo.lon_grid, topo.lat_grid, topo.topo, levels=levels, cmap="terrain"
+ )
+ axes[1].set_title(f"{feature.name} - Contour Plot")
+ axes[1].set_xlabel("Longitude (°)")
+ axes[1].set_ylabel("Latitude (°)")
+ plt.colorbar(cs, ax=axes[1], label="Elevation (m)")
+
+ plt.tight_layout()
+
+ # Save figure
+ output_dir = Path(__file__).parent.parent / "outputs" / "test_visualizations"
+ output_dir.mkdir(parents=True, exist_ok=True)
+ output_path = output_dir / f"validation_{feature.name.replace(' ', '_')}.png"
+ plt.savefig(output_path, dpi=150, bbox_inches="tight")
+ print(f"\nSaved visualization to: {output_path}")
+
+ plt.show()
+
+
+if __name__ == "__main__":
+ # Allow running tests directly
+ pytest.main([__file__, "-v", "-s"])
diff --git a/tests/test_merit_edge_cases.py b/tests/test_merit_edge_cases.py
new file mode 100755
index 0000000..9ba513c
--- /dev/null
+++ b/tests/test_merit_edge_cases.py
@@ -0,0 +1,543 @@
+#!/usr/bin/env python3
+"""
+Edge case test script for MERIT topography data loading.
+
+This script tests the MERIT loader on challenging regions to validate:
+1. MERIT-REMA interface at -60° latitude (Antarctic boundary)
+2. Dateline crossing at ±180° longitude
+3. North Pole high-latitude region
+4. Prime Meridian crossing at 0° longitude
+5. Equator crossing at 0° latitude
+6. Multiple boundary crossings simultaneously
+
+These are the trickiest cases for global data loaders!
+
+Author: Test Suite
+Date: 2025-10-22
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import time
+from pathlib import Path
+import sys
+
+from pycsa.core import io, var
+from pycsa.plotting import cart_plot
+
+
+def test_region(name, lat_extent, lon_extent, merit_cg=50, description=""):
+ """
+ Test loading a specific region.
+
+ Parameters
+ ----------
+ name : str
+ Region name for display
+ lat_extent : list
+ [lat_min, lat_max]
+ lon_extent : list
+ [lon_min, lon_max]
+ merit_cg : int
+ Coarse-graining factor
+ description : str
+ Description of what makes this region tricky
+
+ Returns
+ -------
+ dict
+ Results dictionary with success status and statistics
+ """
+ print("=" * 80)
+ print(f"TEST: {name}")
+ print("=" * 80)
+ print()
+ print(f"Region Configuration:")
+ print(
+ f" Latitude: {lat_extent[0]:7.2f}° to {lat_extent[1]:7.2f}° (span: {lat_extent[1]-lat_extent[0]:.2f}°)"
+ )
+ print(
+ f" Longitude: {lon_extent[0]:7.2f}° to {lon_extent[1]:7.2f}° (span: {abs(lon_extent[1]-lon_extent[0]):.2f}°)"
+ )
+ print(f" Coarse-graining: {merit_cg}x{merit_cg}")
+ print()
+ if description:
+ print(f"Why this is tricky:")
+ print(f" {description}")
+ print()
+
+ # Create parameters
+ class Params:
+ def __init__(self):
+ self.path_merit = "/home/ray/Documents/orog_data/MERIT/"
+ self.path_rema = "/home/ray/Documents/orog_data/REMA/"
+ self.lat_extent = lat_extent
+ self.lon_extent = lon_extent
+ self.merit_cg = merit_cg
+
+ params = Params()
+
+ # Check data paths
+ if not Path(params.path_merit).exists():
+ print(f"ERROR: MERIT data not found at {params.path_merit}")
+ return {"success": False, "error": "Data path not found"}
+
+ # Load data
+ print("Loading MERIT data...")
+ cell = var.topo_cell()
+ start_time = time.time()
+
+ try:
+ loader = io.ncdata.read_merit_topo(cell, params, verbose=False)
+ load_time = time.time() - start_time
+ print(f"✓ Loaded in {load_time:.2f} seconds")
+ print()
+
+ except Exception as e:
+ print(f"✗ ERROR during loading: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return {"success": False, "error": str(e)}
+
+ # Apply data cleaning
+ n_clipped = np.sum(cell.topo < -500.0)
+ cell.topo[cell.topo < -500.0] = -500.0
+
+ # Validate data
+ print("Data Validation:")
+ print(f" Shape: {cell.topo.shape}")
+ print(f" Lat range: [{cell.lat.min():.4f}, {cell.lat.max():.4f}]°")
+ print(f" Lon range: [{cell.lon.min():.4f}, {cell.lon.max():.4f}]°")
+ print(f" Elevation: [{cell.topo.min():.1f}, {cell.topo.max():.1f}] m")
+ print(f" Mean elevation: {cell.topo.mean():.1f} m")
+ if n_clipped > 0:
+ print(f" Clipped {n_clipped:,} points below -500m")
+
+ # Check for issues
+ has_nan = np.isnan(cell.topo).any()
+ has_inf = np.isinf(cell.topo).any()
+
+ if has_nan:
+ print(f" ✗ WARNING: Contains NaN values!")
+ else:
+ print(f" ✓ No NaN values")
+
+ if has_inf:
+ print(f" ✗ WARNING: Contains infinite values!")
+ else:
+ print(f" ✓ No infinite values")
+
+ # Statistics
+ land_mask = cell.topo > 0
+ ocean_mask = cell.topo <= 0
+ land_pct = 100 * np.sum(land_mask) / cell.topo.size
+ ocean_pct = 100 * np.sum(ocean_mask) / cell.topo.size
+
+ print(f" Land/Ocean: {land_pct:.1f}% / {ocean_pct:.1f}%")
+ print()
+
+ # Plot
+ print("Creating plot...")
+ try:
+ cell.gen_mgrids()
+
+ # Adjust figure size based on region aspect ratio
+ lat_span = lat_extent[1] - lat_extent[0]
+ lon_span = abs(lon_extent[1] - lon_extent[0])
+ aspect = lon_span / max(lat_span, 1.0)
+
+ if aspect > 2:
+ figsize = (16, 8)
+ elif aspect < 0.5:
+ figsize = (8, 12)
+ else:
+ figsize = (12, 8)
+
+ cart_plot.lat_lon(cell, fs=figsize, int=1)
+ print(f"✓ Plot displayed")
+ print()
+
+ except Exception as e:
+ print(f"✗ ERROR during plotting: {e}")
+ import traceback
+
+ traceback.print_exc()
+ return {"success": False, "error": f"Plotting failed: {e}"}
+
+ # Success!
+ success = not (has_nan or has_inf)
+
+ results = {
+ "success": success,
+ "name": name,
+ "load_time": load_time,
+ "shape": cell.topo.shape,
+ "elevation_range": (cell.topo.min(), cell.topo.max()),
+ "mean_elevation": cell.topo.mean(),
+ "land_pct": land_pct,
+ "has_nan": has_nan,
+ "has_inf": has_inf,
+ }
+
+ if success:
+ print(f"✓ {name}: PASSED")
+ else:
+ print(f"⚠ {name}: COMPLETED WITH WARNINGS")
+ print()
+
+ return results
+
+
+def run_all_edge_case_tests():
+ """
+ Run all edge case tests.
+
+ Returns
+ -------
+ list
+ List of test results
+ """
+ print("=" * 80)
+ print("MERIT EDGE CASE COMPREHENSIVE TEST SUITE")
+ print("=" * 80)
+ print()
+ print("Testing the trickiest regions for global data loaders:")
+ print(" 1. MERIT-REMA interface at -60° latitude")
+ print(" 2. International dateline crossing at ±180° longitude")
+ print(" 3. North Pole high-latitude region")
+ print(" 4. Prime Meridian crossing at 0° longitude")
+ print(" 5. Equator crossing")
+ print(" 6. Multiple boundary crossings")
+ print()
+ input("Press Enter to start tests...")
+ print()
+
+ results = []
+
+ # Test 1: MERIT-REMA Interface at EXACTLY -60° (South Orkney Islands!)
+ # This is THE island you remember - sits right on the boundary!
+ results.append(
+ test_region(
+ name="MERIT-REMA Boundary (South Orkney Islands)",
+ lat_extent=[-61.5, -59.5], # Tight 2° centered on South Orkney at -60.5°
+ lon_extent=[
+ -47.0,
+ -44.0,
+ ], # Narrow 3° window over South Orkney Islands at -45.5°
+ merit_cg=10, # Finer resolution to catch the small islands
+ description="Tests EXACTLY the -60° latitude boundary with South Orkney Islands!\n"
+ " These islands sit RIGHT ON the MERIT-REMA transition at 60.5°S.\n"
+ " Perfect test case for seamless dataset integration.",
+ )
+ )
+
+ # Test 1b: MERIT-REMA Interface (Antarctic Peninsula - broader view)
+ results.append(
+ test_region(
+ name="MERIT-REMA Interface (Antarctic Peninsula)",
+ lat_extent=[-70.0, -55.0], # Crosses -60° boundary, broader range
+ lon_extent=[-65.0, -55.0], # Narrow 10° window over Antarctic Peninsula
+ merit_cg=30,
+ description="Crosses the -60° latitude boundary over Antarctic Peninsula.\n"
+ " Broader view of the MERIT-REMA transition zone.\n"
+ " Tests seamless data integration between datasets.",
+ )
+ )
+
+ # Test 2: Dateline Crossing - Kamchatka Peninsula (Russia, has land)
+ results.append(
+ test_region(
+ name="Dateline Crossing (Kamchatka Peninsula)",
+ lat_extent=[50.0, 62.0], # Kamchatka Peninsula latitude
+ lon_extent=[175.0, -175.0], # Narrow 10° window crossing dateline
+ merit_cg=30,
+ description="Crosses the international dateline at ±180° longitude.\n"
+ " Focuses on Kamchatka Peninsula (volcanoes, mountains).\n"
+ " Tests handling of longitude wraparound over land.",
+ )
+ )
+
+ # Test 3: North Pole Region - Greenland focus (has major topography)
+ results.append(
+ test_region(
+ name="North Pole Region (Greenland)",
+ lat_extent=[75.0, 85.0], # High Arctic, northern Greenland
+ lon_extent=[-50.0, -20.0], # Narrow window over Greenland ice sheet
+ merit_cg=40,
+ description="High latitude region near North Pole.\n"
+ " Focuses on northern Greenland (ice sheet with elevation).\n"
+ " Tests polar convergence and high-latitude handling.",
+ )
+ )
+
+ # Test 4: Prime Meridian Crossing - UK/France coast (small, fast, over land)
+ results.append(
+ test_region(
+ name="Prime Meridian Crossing (UK-France)",
+ lat_extent=[49.0, 52.0], # English Channel area, tight lat range
+ lon_extent=[-3.0, 3.0], # Narrow 6° window crossing 0° longitude
+ merit_cg=20,
+ description="Crosses the Prime Meridian at 0° longitude.\n"
+ " Focuses on UK-France region (Dover, Calais area).\n"
+ " Tests transition from negative to positive longitude over land.",
+ )
+ )
+
+ # Test 5: Equator Crossing - Mount Kenya area (has elevation features)
+ results.append(
+ test_region(
+ name="Equator Crossing (Mount Kenya)",
+ lat_extent=[-2.0, 2.0], # Narrow 4° crossing equator
+ lon_extent=[36.0, 38.0], # Tight 2° window on Mt. Kenya
+ merit_cg=20,
+ description="Crosses the Equator at 0° latitude.\n"
+ " Focuses on Mount Kenya (5199m, sits on equator!).\n"
+ " Tests hemisphere transition over dramatic topography.",
+ )
+ )
+
+ # Test 6: Tierra del Fuego - near MERIT-REMA boundary
+ results.append(
+ test_region(
+ name="Tierra del Fuego (Near Antarctic Boundary)",
+ lat_extent=[-56.0, -53.0], # Southernmost South America
+ lon_extent=[-70.0, -65.0], # Cape Horn area
+ merit_cg=25,
+ description="Southernmost tip of South America, near -60° boundary.\n"
+ " Tests high southern latitude (stays in MERIT, doesn't cross to REMA).\n"
+ " Drake Passage area with complex coastline.",
+ )
+ )
+
+ # Test 7: Bering Strait - dateline + high latitude (Alaska-Russia)
+ results.append(
+ test_region(
+ name="Bering Strait (Dateline + High Latitude)",
+ lat_extent=[64.0, 68.0], # Bering Strait, tight range
+ lon_extent=[177.0, -177.0], # Narrow 6° crossing dateline
+ merit_cg=25,
+ description="Bering Strait region between Alaska and Russia.\n"
+ " Tests BOTH dateline crossing AND high latitude.\n"
+ " Includes Bering Strait islands and coastlines.",
+ )
+ )
+
+ # Test 8: South Pole Region (Pure REMA) - smaller window
+ results.append(
+ test_region(
+ name="South Pole Region (Marie Byrd Land)",
+ lat_extent=[-85.0, -75.0], # Deep Antarctica
+ lon_extent=[-150.0, -100.0], # Narrower 50° window over Marie Byrd Land
+ merit_cg=60, # Higher CG for speed
+ description="Interior Antarctica (pure REMA data).\n"
+ " Focuses on Marie Byrd Land (West Antarctica, mountains).\n"
+ " Tests REMA dataset at extreme southern latitude.",
+ )
+ )
+
+ return results
+
+
+def print_summary(results):
+ """Print summary of all test results."""
+ print()
+ print("=" * 80)
+ print("EDGE CASE TEST SUMMARY")
+ print("=" * 80)
+ print()
+
+ passed = sum(1 for r in results if r.get("success", False))
+ total = len(results)
+
+ print(f"Tests Passed: {passed}/{total}")
+ print()
+
+ print(f"{'Test Name':<45} {'Status':<10} {'Time (s)':<10} {'Shape':<15}")
+ print("-" * 80)
+
+ for r in results:
+ if r.get("success"):
+ status = "✓ PASS"
+ elif "error" in r:
+ status = "✗ FAIL"
+ else:
+ status = "⚠ WARN"
+
+ name = r.get("name", "Unknown")[:44]
+ time_str = f"{r.get('load_time', 0):.2f}" if "load_time" in r else "N/A"
+ shape = str(r.get("shape", "N/A"))
+
+ print(f"{name:<45} {status:<10} {time_str:<10} {shape:<15}")
+
+ print()
+
+ if passed == total:
+ print("🎉 ALL EDGE CASE TESTS PASSED!")
+ print()
+ print("The MERIT loader correctly handles:")
+ print(" ✓ MERIT-REMA interface at -60° latitude")
+ print(" ✓ International dateline crossing (±180° longitude)")
+ print(" ✓ North and South Pole regions")
+ print(" ✓ Prime Meridian crossing (0° longitude)")
+ print(" ✓ Equator crossing (0° latitude)")
+ print(" ✓ Multiple simultaneous boundary crossings")
+ print()
+ print("The implementation is robust and production-ready! 🚀")
+ else:
+ print(f"⚠ {total - passed} test(s) had issues. Review details above.")
+
+ print()
+ return passed == total
+
+
+if __name__ == "__main__":
+ import argparse
+
+ parser = argparse.ArgumentParser(
+ description="Test MERIT data loader on edge cases and tricky regions"
+ )
+ parser.add_argument(
+ "--quick",
+ action="store_true",
+ help="Run quick test (only 3 most critical regions)",
+ )
+ parser.add_argument(
+ "--test",
+ type=str,
+ choices=[
+ "merit-rema",
+ "south-orkney",
+ "dateline",
+ "north-pole",
+ "prime-meridian",
+ "equator",
+ "tierra-del-fuego",
+ "bering",
+ "south-pole",
+ ],
+ help="Run only a specific test",
+ )
+
+ args = parser.parse_args()
+
+ if args.test:
+ # Run single test
+ test_configs = {
+ "merit-rema": {
+ "name": "MERIT-REMA Boundary (South Orkney Islands)",
+ "lat_extent": [-61.5, -59.5],
+ "lon_extent": [-47.0, -44.0],
+ "merit_cg": 10,
+ "description": "Tests EXACTLY -60° boundary with South Orkney Islands",
+ },
+ "south-orkney": {
+ "name": "MERIT-REMA Boundary (South Orkney Islands)",
+ "lat_extent": [-61.5, -59.5],
+ "lon_extent": [-47.0, -44.0],
+ "merit_cg": 10,
+ "description": "Tests EXACTLY -60° boundary with South Orkney Islands",
+ },
+ "dateline": {
+ "name": "Dateline Crossing (Kamchatka)",
+ "lat_extent": [50.0, 62.0],
+ "lon_extent": [175.0, -175.0],
+ "merit_cg": 30,
+ "description": "Tests ±180° longitude over Kamchatka Peninsula",
+ },
+ "north-pole": {
+ "name": "North Pole (Greenland)",
+ "lat_extent": [75.0, 85.0],
+ "lon_extent": [-50.0, -20.0],
+ "merit_cg": 40,
+ "description": "Tests high Arctic over northern Greenland",
+ },
+ "prime-meridian": {
+ "name": "Prime Meridian (UK-France)",
+ "lat_extent": [49.0, 52.0],
+ "lon_extent": [-3.0, 3.0],
+ "merit_cg": 20,
+ "description": "Tests 0° longitude crossing over UK-France",
+ },
+ "equator": {
+ "name": "Equator (Mount Kenya)",
+ "lat_extent": [-2.0, 2.0],
+ "lon_extent": [36.0, 38.0],
+ "merit_cg": 20,
+ "description": "Tests 0° latitude over Mount Kenya",
+ },
+ "tierra-del-fuego": {
+ "name": "Tierra del Fuego",
+ "lat_extent": [-56.0, -53.0],
+ "lon_extent": [-70.0, -65.0],
+ "merit_cg": 25,
+ "description": "Tests southern tip of South America",
+ },
+ "bering": {
+ "name": "Bering Strait",
+ "lat_extent": [64.0, 68.0],
+ "lon_extent": [177.0, -177.0],
+ "merit_cg": 25,
+ "description": "Tests dateline + high latitude over strait",
+ },
+ "south-pole": {
+ "name": "South Pole (Marie Byrd Land)",
+ "lat_extent": [-85.0, -75.0],
+ "lon_extent": [-150.0, -100.0],
+ "merit_cg": 60,
+ "description": "Tests pure REMA over West Antarctica",
+ },
+ }
+
+ config = test_configs[args.test]
+ result = test_region(**config)
+ success = result.get("success", False)
+ sys.exit(0 if success else 1)
+
+ elif args.quick:
+ # Run only 3 most critical tests
+ print("Running QUICK edge case tests (3 most critical regions)...\n")
+
+ results = []
+
+ # 1. MERIT-REMA interface at EXACT boundary (most critical!)
+ results.append(
+ test_region(
+ name="MERIT-REMA Boundary (South Orkney Islands)",
+ lat_extent=[-61.5, -59.5],
+ lon_extent=[-47.0, -44.0],
+ merit_cg=10,
+ description="EXACTLY -60° boundary with South Orkney Islands at 60.5°S",
+ )
+ )
+
+ # 2. Dateline crossing
+ results.append(
+ test_region(
+ name="Dateline Crossing (Kamchatka)",
+ lat_extent=[50.0, 62.0],
+ lon_extent=[175.0, -175.0],
+ merit_cg=30,
+ description="±180° longitude over Kamchatka Peninsula",
+ )
+ )
+
+ # 3. North Pole
+ results.append(
+ test_region(
+ name="North Pole (Greenland)",
+ lat_extent=[75.0, 85.0],
+ lon_extent=[-50.0, -20.0],
+ merit_cg=40,
+ description="High Arctic over northern Greenland",
+ )
+ )
+
+ success = print_summary(results)
+ sys.exit(0 if success else 1)
+
+ else:
+ # Run all tests
+ results = run_all_edge_case_tests()
+ success = print_summary(results)
+ sys.exit(0 if success else 1)
diff --git a/tests/test_tile_cache_etopo_equivalence.py b/tests/test_tile_cache_etopo_equivalence.py
new file mode 100644
index 0000000..3f0eb0e
--- /dev/null
+++ b/tests/test_tile_cache_etopo_equivalence.py
@@ -0,0 +1,161 @@
+"""Byte-equivalence test for TopographyTileCache.get_etopo_data vs read_etopo_topo.
+
+The cache's ETOPO path is a port of pycsa.core.io.read_etopo_topo.get_topo. This
+test loads representative ICON cells via both paths and asserts the returned
+(lat, lon, topo) arrays are identical. Run with:
+
+ pytest tests/test_tile_cache_etopo_equivalence.py -v
+
+Skips automatically if data/etopo_15s/ is missing.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+from pycsa.core import io as pcio, utils, var
+from pycsa import local_paths
+from pycsa.core.tile_cache import TopographyTileCache, compute_split_EW
+
+ETOPO_DIR = Path(local_paths.paths.etopo)
+ICON_GRID = local_paths.paths.icon_grid
+
+
+pytestmark = pytest.mark.skipif(
+ not ETOPO_DIR.exists() or not Path(ICON_GRID).exists(),
+ reason="ETOPO tiles or ICON grid not available locally",
+)
+
+
+# Representative cells covering each branch of the ETOPO loader.
+# Each tuple is (c_idx, description).
+TEST_CELLS = [
+ (1086, "typical non-dateline mid-latitude (lat ~76°N)"),
+ (2311, "Aleutians — false-positive dateline (all-negative lons near -176°)"),
+ (1074, "genuine dateline crossing (split_EW=True, lat ~80°N)"),
+ (17408, "extreme south polar (lat -88.90°S, exercises lat_idx_rng generation)"),
+]
+
+
+@pytest.fixture(scope="module")
+def grid():
+ """Load the ICON grid once and reuse across cells."""
+ g = var.grid()
+ pcio.ncdata().read_dat(ICON_GRID, g)
+ return g
+
+
+@pytest.fixture(scope="module")
+def params():
+ """Minimal params object with what read_etopo_topo needs."""
+ p = var.obj()
+ p.path_etopo = str(ETOPO_DIR) + "/"
+ p.etopo_cg = 4 # matches the default coarse-graining used by the global run
+ p.lat_extent = np.array([0.0, 0.0]) # placeholder; set per-cell
+ p.lon_extent = np.array([0.0, 0.0])
+ return p
+
+
+def _load_via_reader(grid, params, c_idx):
+ """Reference path: pycsa.core.io.read_etopo_topo."""
+ lat_verts = np.degrees(grid.clat_vertices[c_idx])
+ lon_verts = np.degrees(grid.clon_vertices[c_idx])
+ lat_extent, lon_extent = utils.handle_latlon_expansion(lat_verts, lon_verts)
+ params.lat_extent = lat_extent
+ params.lon_extent = lon_extent
+
+ topo = var.topo_cell()
+ reader = pcio.ncdata().read_etopo_topo(None, params, is_parallel=True)
+ reader.get_topo(topo)
+ return topo, reader.split_EW, lat_extent, lon_extent
+
+
+def _load_via_cache(cache, params, lat_extent, lon_extent):
+ """Candidate path: TopographyTileCache.get_etopo_data."""
+ lat, lon, topo = cache.get_etopo_data(
+ lat_extent, lon_extent, etopo_cg=params.etopo_cg
+ )
+ return lat, lon, topo
+
+
+@pytest.fixture(scope="module")
+def cache():
+ """Build a single lazy ETOPO cache used across all cells."""
+ return TopographyTileCache(
+ data_dir=str(ETOPO_DIR),
+ tile_filenames=[],
+ dataset_type="ETOPO",
+ verbose=False,
+ )
+
+
+def test_worker_cache_lifecycle(grid, params):
+ """init_worker_cache / get_worker_cache / close_worker_cache happy path.
+
+ This mirrors what do_cell does inside a Dask worker process: the main
+ loop calls client.run(init_worker_cache, ...), then each cell's do_cell
+ call retrieves the cache via get_worker_cache().
+ """
+ from pycsa.core import tile_cache as tc
+
+ # No cache should be initialised yet (or from a prior test).
+ tc.close_worker_cache()
+ with pytest.raises(RuntimeError):
+ tc.get_worker_cache()
+
+ assert tc.init_worker_cache(str(ETOPO_DIR), "ETOPO") is True
+ cache = tc.get_worker_cache()
+ assert cache.dataset_type == "ETOPO"
+
+ # Idempotency: second init with same dir should be a no-op (same object).
+ assert tc.init_worker_cache(str(ETOPO_DIR), "ETOPO") is True
+ assert tc.get_worker_cache() is cache
+
+ # Functional check: retrieve topo for one cell through the worker-cache
+ # path; should match reader output (this is the same contract used by
+ # the wired do_cell).
+ c_idx = 1086
+ topo_ref, _, lat_extent, lon_extent = _load_via_reader(grid, params, c_idx)
+ lat, lon, topo_arr = cache.get_etopo_data(
+ lat_extent, lon_extent, etopo_cg=params.etopo_cg
+ )
+ np.testing.assert_array_equal(topo_arr, topo_ref.topo)
+
+ # Cleanup leaves get_worker_cache failing again.
+ tc.close_worker_cache()
+ with pytest.raises(RuntimeError):
+ tc.get_worker_cache()
+
+
+@pytest.mark.parametrize("c_idx,description", TEST_CELLS)
+def test_etopo_equivalence(grid, params, cache, c_idx, description):
+ """Cache output must match the reference reader byte-for-byte for every cell."""
+ topo_ref, split_EW_ref, lat_extent, lon_extent = _load_via_reader(
+ grid, params, c_idx
+ )
+ lat_cache, lon_cache, topo_cache = _load_via_cache(
+ cache, params, lat_extent, lon_extent
+ )
+
+ # The free-function dateline detector must agree with the reader's own
+ # internal flag for the same vertex set.
+ assert (
+ compute_split_EW(lon_extent) == split_EW_ref
+ ), f"cell {c_idx}: compute_split_EW disagrees with reader ({description})"
+
+ np.testing.assert_array_equal(
+ lat_cache,
+ topo_ref.lat,
+ err_msg=f"cell {c_idx}: lat arrays differ ({description})",
+ )
+ np.testing.assert_array_equal(
+ lon_cache,
+ topo_ref.lon,
+ err_msg=f"cell {c_idx}: lon arrays differ ({description})",
+ )
+ np.testing.assert_array_equal(
+ topo_cache,
+ topo_ref.topo,
+ err_msg=f"cell {c_idx}: topo arrays differ ({description})",
+ )
diff --git a/tests/unit/test_io_simple.py b/tests/unit/test_io_simple.py
new file mode 100644
index 0000000..21debe9
--- /dev/null
+++ b/tests/unit/test_io_simple.py
@@ -0,0 +1,178 @@
+"""
+Simplified unit tests for I/O routines.
+
+Tests basic NetCDF reading functionality for topographic data.
+"""
+
+import pytest
+import numpy as np
+from pathlib import Path
+from pycsa.core import io, var
+
+
+class TestNetCDFReader:
+ """Test NetCDF data reading functionality."""
+
+ @pytest.fixture
+ def data_dir(self):
+ """Return path to test data directory."""
+ return Path(__file__).parent.parent.parent / "data"
+
+ def test_ncdata_initialization(self):
+ """Test ncdata object initialization."""
+ reader = io.ncdata(padding=10, padding_tol=50)
+ assert reader.padding == 60
+ assert reader.read_merit == False
+
+ def test_read_grid_data(self, data_dir):
+ """Test reading grid data from NetCDF file."""
+ grid_path = data_dir / "icon_compact_alaska.nc"
+ if not grid_path.exists():
+ pytest.skip(f"Test data not found: {grid_path}")
+
+ grid = var.grid()
+ reader = io.ncdata()
+ reader.read_dat(str(grid_path), grid)
+
+ assert grid.clat is not None
+ assert grid.clon is not None
+ assert len(grid.clat) > 0
+
+ def test_read_topography_data(self, data_dir):
+ """Test reading topography data from NetCDF file."""
+ topo_path = data_dir / "topo_compact_alaska.nc"
+ if not topo_path.exists():
+ pytest.skip(f"Test data not found: {topo_path}")
+
+ topo = var.topo_cell()
+ reader = io.ncdata()
+ reader.read_dat(str(topo_path), topo)
+
+ assert topo.lat is not None
+ assert topo.lon is not None
+ assert topo.topo is not None
+ assert topo.topo.size > 0
+
+
+class TestETOPOLoader:
+ """Test ETOPO 2022 15 arc-second data loading."""
+
+ @pytest.fixture
+ def etopo_dir(self, project_root):
+ """Return path to ETOPO data directory."""
+ etopo_path = project_root / "data" / "etopo_15s"
+ if not etopo_path.exists():
+ pytest.skip(f"ETOPO data not found: {etopo_path}")
+ return etopo_path
+
+ @pytest.fixture
+ def test_params(self, etopo_dir):
+ """Create test parameters for ETOPO loading."""
+
+ class TestParams:
+ def __init__(self):
+ self.path_etopo = str(etopo_dir) + "/"
+ self.lat_extent = [35.0, 40.0]
+ self.lon_extent = [-120.0, -115.0]
+ self.etopo_cg = 4 # Use coarse-graining for faster testing
+
+ return TestParams()
+
+ def test_etopo_loader_initialization(self, test_params, etopo_dir):
+ """Test ETOPO loader initialization and basic loading."""
+ cell = var.topo_cell()
+
+ loader = io.ncdata.read_etopo_topo(cell, test_params, verbose=False)
+
+ # Check that data was loaded
+ assert cell.lat is not None, "Latitude not loaded"
+ assert cell.lon is not None, "Longitude not loaded"
+ assert cell.topo is not None, "Topography not loaded"
+
+ # Check dimensions
+ assert len(cell.lat) > 0, "Latitude array is empty"
+ assert len(cell.lon) > 0, "Longitude array is empty"
+ assert cell.topo.size > 0, "Topography array is empty"
+
+ # Check that loaded region matches requested extent (with small tolerance)
+ # Note: Due to coarse-graining, exact boundaries may not be matched
+ assert cell.lat.min() <= test_params.lat_extent[0] + 0.1
+ assert cell.lat.max() >= test_params.lat_extent[1] - 0.1
+ assert cell.lon.min() <= test_params.lon_extent[0] + 0.1
+ assert cell.lon.max() >= test_params.lon_extent[1] - 0.1
+
+ def test_etopo_data_values(self, test_params, etopo_dir):
+ """Test that loaded ETOPO data has reasonable values."""
+ cell = var.topo_cell()
+
+ loader = io.ncdata.read_etopo_topo(cell, test_params, verbose=False)
+
+ # Check for reasonable elevation values (California coast to Sierra Nevada)
+ # Should have values from below sea level to several thousand meters
+ assert (
+ cell.topo.min() >= -11000
+ ), "Topography minimum too low (deepest ocean ~11km)"
+ assert cell.topo.max() <= 9000, "Topography maximum too high (Mt Everest ~9km)"
+
+ # Check for fill values (should not be present after loading)
+ assert not np.any(cell.topo == -99999), "Fill values present in loaded data"
+
+ # Check that data is not all zeros
+ assert not np.all(cell.topo == 0), "Topography data is all zeros"
+
+ def test_etopo_coarse_graining(self, etopo_dir):
+ """Test that coarse-graining reduces data size as expected."""
+
+ class ParamsCG1:
+ def __init__(self):
+ self.path_etopo = str(etopo_dir) + "/"
+ self.lat_extent = [36.0, 37.0]
+ self.lon_extent = [-119.0, -118.0]
+ self.etopo_cg = 1
+
+ class ParamsCG4:
+ def __init__(self):
+ self.path_etopo = str(etopo_dir) + "/"
+ self.lat_extent = [36.0, 37.0]
+ self.lon_extent = [-119.0, -118.0]
+ self.etopo_cg = 4
+
+ # Load with no coarse-graining
+ cell1 = var.topo_cell()
+ loader1 = io.ncdata.read_etopo_topo(cell1, ParamsCG1(), verbose=False)
+
+ # Load with 4x coarse-graining
+ cell4 = var.topo_cell()
+ loader4 = io.ncdata.read_etopo_topo(cell4, ParamsCG4(), verbose=False)
+
+ # Check that coarse-graining reduces size
+ size_ratio = cell1.topo.size / cell4.topo.size
+
+ # Should be approximately 4x4 = 16 times reduction
+ assert (
+ size_ratio > 10
+ ), f"Coarse-graining didn't reduce size enough: {size_ratio}x"
+ assert size_ratio < 20, f"Coarse-graining reduced size too much: {size_ratio}x"
+
+ def test_etopo_grid_structure(self, test_params, etopo_dir):
+ """Test that loaded grid has correct structure."""
+ cell = var.topo_cell()
+
+ loader = io.ncdata.read_etopo_topo(cell, test_params, verbose=False)
+
+ # Check that lat/lon are 1D arrays
+ assert cell.lat.ndim == 1, "Latitude should be 1D"
+ assert cell.lon.ndim == 1, "Longitude should be 1D"
+
+ # Check that topo is 2D
+ assert cell.topo.ndim == 2, "Topography should be 2D"
+
+ # Check that dimensions match
+ assert cell.topo.shape == (
+ len(cell.lat),
+ len(cell.lon),
+ ), f"Topography shape {cell.topo.shape} doesn't match lat/lon ({len(cell.lat)}, {len(cell.lon)})"
+
+ # Check that lat/lon are sorted
+ assert np.all(np.diff(cell.lat) > 0), "Latitude should be sorted ascending"
+ assert np.all(np.diff(cell.lon) > 0), "Longitude should be sorted ascending"