roms-tools 3.3.0__py3-none-any.whl → 3.4.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- roms_tools/__init__.py +1 -1
- roms_tools/analysis/cdr_ensemble.py +10 -13
- roms_tools/analysis/roms_output.py +5 -304
- roms_tools/{download.py → datasets/download.py} +1 -0
- roms_tools/{setup → datasets}/lat_lon_datasets.py +76 -64
- roms_tools/{setup → datasets}/river_datasets.py +9 -4
- roms_tools/datasets/roms_dataset.py +767 -0
- roms_tools/datasets/utils.py +475 -0
- roms_tools/{setup/fill.py → fill.py} +110 -13
- roms_tools/plot.py +4 -4
- roms_tools/setup/boundary_forcing.py +51 -43
- roms_tools/setup/cdr_release.py +2 -4
- roms_tools/setup/grid.py +29 -12
- roms_tools/setup/initial_conditions.py +19 -19
- roms_tools/setup/nesting.py +8 -4
- roms_tools/setup/river_forcing.py +4 -4
- roms_tools/setup/surface_forcing.py +14 -9
- roms_tools/setup/tides.py +1 -1
- roms_tools/setup/topography.py +10 -2
- roms_tools/setup/utils.py +72 -524
- roms_tools/tests/test_analysis/test_cdr_ensemble.py +4 -6
- roms_tools/tests/test_analysis/test_roms_output.py +1 -220
- roms_tools/tests/{test_setup → test_datasets}/test_lat_lon_datasets.py +4 -4
- roms_tools/tests/{test_setup → test_datasets}/test_river_datasets.py +1 -1
- roms_tools/tests/test_datasets/test_roms_dataset.py +539 -0
- roms_tools/tests/test_datasets/test_utils.py +527 -0
- roms_tools/tests/{test_setup/test_fill.py → test_fill.py} +72 -9
- roms_tools/tests/test_setup/test_boundary_forcing.py +57 -138
- roms_tools/tests/test_setup/test_cdr_release.py +4 -5
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zarr.json +293 -2021
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/zarr.json +294 -2022
- roms_tools/tests/test_setup/test_grid.py +42 -1
- roms_tools/tests/test_setup/test_initial_conditions.py +3 -94
- roms_tools/tests/test_setup/test_nesting.py +2 -1
- roms_tools/tests/test_setup/test_surface_forcing.py +1 -1
- roms_tools/tests/test_setup/test_tides.py +1 -1
- roms_tools/tests/test_setup/test_utils.py +100 -15
- roms_tools/tests/test_tiling/test_partition.py +63 -15
- roms_tools/tests/test_utils.py +78 -0
- roms_tools/tiling/partition.py +81 -211
- roms_tools/utils.py +193 -0
- {roms_tools-3.3.0.dist-info → roms_tools-3.4.0.dist-info}/METADATA +1 -1
- {roms_tools-3.3.0.dist-info → roms_tools-3.4.0.dist-info}/RECORD +46 -170
- {roms_tools-3.3.0.dist-info → roms_tools-3.4.0.dist-info}/WHEEL +1 -1
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_ALT_CO2_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_ALT_CO2_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/ALK_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_ALT_CO2_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_ALT_CO2_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DIC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOCr_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOCr_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DON_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DON_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DONr_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DONr_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOPr_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/DOPr_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Fe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Fe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Lig_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/Lig_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NH4_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NH4_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NO3_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/NO3_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/O2_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/O2_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/PO4_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/PO4_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/SiO3_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/SiO3_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatChl_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatChl_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatFe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatFe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatSi_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diatSi_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazChl_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazChl_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazFe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazFe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/diazP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spCaCO3_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spCaCO3_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spChl_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spChl_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spFe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spFe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/spP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zooC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_climatology.zarr/zooC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/ALK_ALT_CO2_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/ALK_ALT_CO2_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/ALK_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/ALK_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DIC_ALT_CO2_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DIC_ALT_CO2_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DIC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DIC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOCr_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOCr_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DON_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DON_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DONr_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DONr_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOPr_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/DOPr_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/Fe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/Fe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/Lig_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/Lig_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/NH4_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/NH4_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/NO3_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/NO3_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/O2_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/O2_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/PO4_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/PO4_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/SiO3_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/SiO3_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatChl_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatChl_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatFe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatFe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatSi_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diatSi_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazChl_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazChl_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazFe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazFe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/diazP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spC_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spCaCO3_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spCaCO3_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spChl_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spChl_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spFe_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spFe_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spP_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/spP_west/zarr.json +0 -54
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/zooC_west/c/0/0/0 +0 -0
- roms_tools/tests/test_setup/test_data/bgc_boundary_forcing_from_unified_climatology.zarr/zooC_west/zarr.json +0 -54
- {roms_tools-3.3.0.dist-info → roms_tools-3.4.0.dist-info}/licenses/LICENSE +0 -0
- {roms_tools-3.3.0.dist-info → roms_tools-3.4.0.dist-info}/top_level.txt +0 -0
roms_tools/tests/test_utils.py
CHANGED
|
@@ -3,9 +3,12 @@ from pathlib import Path
|
|
|
3
3
|
from unittest import mock
|
|
4
4
|
|
|
5
5
|
import numpy as np
|
|
6
|
+
import pandas as pd
|
|
6
7
|
import pytest
|
|
7
8
|
import xarray as xr
|
|
8
9
|
|
|
10
|
+
from roms_tools.datasets.download import download_test_data
|
|
11
|
+
from roms_tools.datasets.lat_lon_datasets import ERA5Correction
|
|
9
12
|
from roms_tools.utils import (
|
|
10
13
|
_path_list_from_input,
|
|
11
14
|
generate_focused_coordinate_range,
|
|
@@ -13,6 +16,8 @@ from roms_tools.utils import (
|
|
|
13
16
|
has_copernicus,
|
|
14
17
|
has_dask,
|
|
15
18
|
has_gcsfs,
|
|
19
|
+
interpolate_cyclic_time,
|
|
20
|
+
interpolate_from_climatology,
|
|
16
21
|
load_data,
|
|
17
22
|
)
|
|
18
23
|
|
|
@@ -241,3 +246,76 @@ def test_time_chunking_false_roms():
|
|
|
241
246
|
dim_names = {"time": "ocean_time"}
|
|
242
247
|
result = get_dask_chunks(dim_names, time_chunking=False)
|
|
243
248
|
assert "ocean_time" not in result
|
|
249
|
+
|
|
250
|
+
|
|
251
|
+
# test interpolate_from_climatology
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
@pytest.fixture
|
|
255
|
+
def climatology_data():
|
|
256
|
+
"""Create a simple annual cycle dataset with 12 time points."""
|
|
257
|
+
time_coord = np.arange(1, 13) # months as day_of_year approximation
|
|
258
|
+
da = xr.DataArray(np.arange(12), dims=("time",), coords={"time": time_coord})
|
|
259
|
+
ds = xr.Dataset({"var1": da, "var2": da * 2})
|
|
260
|
+
return da, ds, "time", "time"
|
|
261
|
+
|
|
262
|
+
|
|
263
|
+
def test_interpolate_dataarray_single_time(climatology_data):
|
|
264
|
+
da, _, time_dim, time_coord = climatology_data
|
|
265
|
+
target_time = pd.Timestamp("2000-03-15") # day_of_year ~ 75
|
|
266
|
+
interpolated = interpolate_from_climatology(da, time_dim, time_coord, target_time)
|
|
267
|
+
assert isinstance(interpolated, xr.DataArray)
|
|
268
|
+
assert interpolated.sizes[time_dim] == 1
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
def test_interpolate_dataset_multiple_times(climatology_data):
|
|
272
|
+
_, ds, time_dim, time_coord = climatology_data
|
|
273
|
+
target_times = pd.date_range("2000-01-01", periods=3, freq="ME")
|
|
274
|
+
interpolated = interpolate_from_climatology(ds, time_dim, time_coord, target_times)
|
|
275
|
+
assert isinstance(interpolated, xr.Dataset)
|
|
276
|
+
assert all(interpolated[var].sizes[time_dim] == 3 for var in interpolated.data_vars)
|
|
277
|
+
|
|
278
|
+
|
|
279
|
+
def test_interpolate_dataarray_time_dim_not_equal_time_coord():
|
|
280
|
+
time_values = np.arange(1, 13)
|
|
281
|
+
da = xr.DataArray(
|
|
282
|
+
np.arange(12),
|
|
283
|
+
dims=("time_dim",),
|
|
284
|
+
coords={"time_coord": ("time_dim", time_values)},
|
|
285
|
+
)
|
|
286
|
+
target_time = pd.Timestamp("2000-06-15")
|
|
287
|
+
interpolated = interpolate_from_climatology(
|
|
288
|
+
da, time_dim="time_dim", time_coord="time_coord", time=target_time
|
|
289
|
+
)
|
|
290
|
+
assert interpolated.sizes["time_dim"] == 1
|
|
291
|
+
assert np.issubdtype(interpolated.dtype, np.number)
|
|
292
|
+
|
|
293
|
+
|
|
294
|
+
def test_interpolate_cyclic_time_basic():
|
|
295
|
+
time_values = np.arange(1, 13)
|
|
296
|
+
da = xr.DataArray(np.arange(12), dims=("time",), coords={"time": time_values})
|
|
297
|
+
target_days = [0.5, 6.5, 12.5] # fractional days, include cyclic behavior
|
|
298
|
+
interpolated = interpolate_cyclic_time(
|
|
299
|
+
da, time_dim="time", time_coord="time", day_of_year=target_days
|
|
300
|
+
)
|
|
301
|
+
assert isinstance(interpolated, xr.DataArray)
|
|
302
|
+
assert interpolated.sizes["time"] == len(target_days)
|
|
303
|
+
|
|
304
|
+
|
|
305
|
+
def test_interpolate_from_climatology_invalid_input():
|
|
306
|
+
with pytest.raises(TypeError):
|
|
307
|
+
interpolate_from_climatology(
|
|
308
|
+
"not a dataset", "time", "time", pd.Timestamp("2000-01-01")
|
|
309
|
+
)
|
|
310
|
+
|
|
311
|
+
|
|
312
|
+
def test_interpolate_from_real_climatology(use_dask):
|
|
313
|
+
fname = download_test_data("ERA5_regional_test_data.nc")
|
|
314
|
+
era5_times = xr.open_dataset(fname).time
|
|
315
|
+
|
|
316
|
+
climatology = ERA5Correction(use_dask=use_dask)
|
|
317
|
+
field = climatology.ds["ssr_corr"]
|
|
318
|
+
field["time"] = field["time"].dt.days
|
|
319
|
+
|
|
320
|
+
interpolated_field = interpolate_from_climatology(field, "time", "time", era5_times)
|
|
321
|
+
assert len(interpolated_field.time) == len(era5_times)
|
roms_tools/tiling/partition.py
CHANGED
|
@@ -7,6 +7,61 @@ import xarray as xr
|
|
|
7
7
|
|
|
8
8
|
from roms_tools.utils import save_datasets
|
|
9
9
|
|
|
10
|
+
DIM_INFO = {
|
|
11
|
+
# eta-direction
|
|
12
|
+
"eta_rho": dict(axis="eta", ghost=2, edge="both"),
|
|
13
|
+
"eta_u": dict(axis="eta", ghost=2, edge="both"),
|
|
14
|
+
"eta_v": dict(axis="eta", ghost=1, edge="upper"),
|
|
15
|
+
"eta_psi": dict(axis="eta", ghost=3, edge="both"),
|
|
16
|
+
"eta_coarse": dict(axis="eta", ghost=2, edge="both"),
|
|
17
|
+
# xi-direction
|
|
18
|
+
"xi_rho": dict(axis="xi", ghost=2, edge="both"),
|
|
19
|
+
"xi_u": dict(axis="xi", ghost=1, edge="upper"),
|
|
20
|
+
"xi_v": dict(axis="xi", ghost=2, edge="both"),
|
|
21
|
+
"xi_psi": dict(axis="xi", ghost=3, edge="both"),
|
|
22
|
+
"xi_coarse": dict(axis="xi", ghost=2, edge="both"),
|
|
23
|
+
}
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _exact_division(size: int, nparts: int, dim: str) -> int:
|
|
27
|
+
if size % nparts != 0:
|
|
28
|
+
raise ValueError(
|
|
29
|
+
f"Dimension '{dim}' of size {size} cannot be evenly divided into {nparts} partitions."
|
|
30
|
+
)
|
|
31
|
+
return size // nparts
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _compute_partition_sizes(
|
|
35
|
+
total_size: int,
|
|
36
|
+
nparts: int,
|
|
37
|
+
ghost: int,
|
|
38
|
+
edge: str,
|
|
39
|
+
dim: str,
|
|
40
|
+
) -> list[int]:
|
|
41
|
+
"""Compute per-tile sizes including ghost cells."""
|
|
42
|
+
if nparts == 1:
|
|
43
|
+
return [total_size]
|
|
44
|
+
|
|
45
|
+
core = _exact_division(total_size - ghost, nparts, dim)
|
|
46
|
+
|
|
47
|
+
sizes = [core] * nparts
|
|
48
|
+
|
|
49
|
+
if edge == "both":
|
|
50
|
+
sizes[0] += ghost // 2
|
|
51
|
+
sizes[-1] += ghost - ghost // 2
|
|
52
|
+
elif edge == "upper":
|
|
53
|
+
sizes[-1] += ghost
|
|
54
|
+
else:
|
|
55
|
+
raise ValueError(f"Unknown edge rule '{edge}' for {dim}")
|
|
56
|
+
|
|
57
|
+
return sizes
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def _cumsum(sizes: list[int]) -> np.ndarray:
|
|
61
|
+
out = np.zeros(len(sizes) + 1, dtype=int)
|
|
62
|
+
out[1:] = np.cumsum(sizes)
|
|
63
|
+
return out
|
|
64
|
+
|
|
10
65
|
|
|
11
66
|
def partition(
|
|
12
67
|
ds: xr.Dataset, np_eta: int = 1, np_xi: int = 1, include_coarse_dims: bool = True
|
|
@@ -16,8 +71,8 @@ def partition(
|
|
|
16
71
|
|
|
17
72
|
This function divides the input dataset into `np_eta` by `np_xi` tiles, where each tile
|
|
18
73
|
represents a subdomain of the original dataset. The partitioning is performed along
|
|
19
|
-
the spatial dimensions `eta_rho`, `
|
|
20
|
-
depending on which dimensions are present in the dataset.
|
|
74
|
+
the spatial dimensions `eta_rho`, `eta_u`, `eta_v`, `eta_coarse`, `xi_rho`, `xi_u`, `xi_v`,
|
|
75
|
+
and `xi_coarse`, depending on which dimensions are present in the dataset.
|
|
21
76
|
|
|
22
77
|
Parameters
|
|
23
78
|
----------
|
|
@@ -71,227 +126,42 @@ def partition(
|
|
|
71
126
|
):
|
|
72
127
|
raise ValueError("np_eta and np_xi must be positive integers")
|
|
73
128
|
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
"
|
|
79
|
-
"eta_psi",
|
|
80
|
-
"xi_psi",
|
|
129
|
+
# Select applicable dimensions
|
|
130
|
+
dims_to_partition = [
|
|
131
|
+
d
|
|
132
|
+
for d in DIM_INFO
|
|
133
|
+
if d in ds.dims and (include_coarse_dims or "coarse" not in d)
|
|
81
134
|
]
|
|
82
|
-
coarse_dims = ["eta_coarse", "xi_coarse"]
|
|
83
|
-
partitionable_dims = base_dims + coarse_dims if include_coarse_dims else base_dims
|
|
84
|
-
|
|
85
|
-
dims_to_partition = [d for d in partitionable_dims if d in ds.dims]
|
|
86
|
-
|
|
87
|
-
# if eta is periodic there are no ghost cells along those dimensions
|
|
88
|
-
if "eta_v" in ds.sizes and ds.sizes["eta_rho"] == ds.sizes["eta_v"]:
|
|
89
|
-
# TODO how are we supposed to know if eta is periodic if eta_v doesn't appear? partit.F doesn't say...
|
|
90
|
-
n_eta_ghost_cells = 0
|
|
91
|
-
else:
|
|
92
|
-
n_eta_ghost_cells = 1
|
|
93
|
-
|
|
94
|
-
# if xi is periodic there are no ghost cells along those dimensions
|
|
95
|
-
if "xi_u" in ds.sizes and ds.sizes["xi_rho"] == ds.sizes["xi_u"]:
|
|
96
|
-
n_xi_ghost_cells = 0
|
|
97
|
-
else:
|
|
98
|
-
n_xi_ghost_cells = 1
|
|
99
|
-
|
|
100
|
-
def integer_division_or_raise(a: int, b: int, dimension: str) -> int:
|
|
101
|
-
"""Perform integer division and ensure that the division is exact.
|
|
102
|
-
|
|
103
|
-
Parameters
|
|
104
|
-
----------
|
|
105
|
-
a : int
|
|
106
|
-
The numerator for the division.
|
|
107
|
-
b : int
|
|
108
|
-
The denominator for the division.
|
|
109
|
-
dimension : str
|
|
110
|
-
The name of the dimension being partitioned, used for error reporting.
|
|
111
|
-
|
|
112
|
-
Returns
|
|
113
|
-
-------
|
|
114
|
-
int
|
|
115
|
-
The result of the integer division.
|
|
116
|
-
|
|
117
|
-
Raises
|
|
118
|
-
------
|
|
119
|
-
ValueError
|
|
120
|
-
If the division is not exact, indicating that the domain cannot be evenly divided
|
|
121
|
-
along the specified dimension.
|
|
122
|
-
"""
|
|
123
|
-
remainder = a % b
|
|
124
|
-
if remainder == 0:
|
|
125
|
-
return a // b
|
|
126
|
-
else:
|
|
127
|
-
raise ValueError(
|
|
128
|
-
f"Dimension '{dimension}' of size {a} cannot be evenly divided into {b} partitions."
|
|
129
|
-
)
|
|
130
|
-
|
|
131
|
-
if "eta_rho" in dims_to_partition:
|
|
132
|
-
eta_rho_domain_size = integer_division_or_raise(
|
|
133
|
-
ds.sizes["eta_rho"] - 2 * n_eta_ghost_cells, np_eta, "eta_rho"
|
|
134
|
-
)
|
|
135
135
|
|
|
136
|
-
|
|
137
|
-
xi_rho_domain_size = integer_division_or_raise(
|
|
138
|
-
ds.sizes["xi_rho"] - 2 * n_xi_ghost_cells, np_xi, "xi_rho"
|
|
139
|
-
)
|
|
136
|
+
partitioned_sizes = {}
|
|
140
137
|
|
|
141
|
-
|
|
142
|
-
|
|
143
|
-
|
|
144
|
-
)
|
|
145
|
-
|
|
146
|
-
if "xi_u" in dims_to_partition:
|
|
147
|
-
xi_u_domain_size = integer_division_or_raise(
|
|
148
|
-
ds.sizes["xi_u"] - 1 * n_xi_ghost_cells, np_xi, "xi_u"
|
|
149
|
-
)
|
|
150
|
-
|
|
151
|
-
if "eta_psi" in dims_to_partition:
|
|
152
|
-
eta_psi_domain_size = integer_division_or_raise(
|
|
153
|
-
ds.sizes["eta_psi"] - 3 * n_eta_ghost_cells, np_eta, "eta_psi"
|
|
154
|
-
)
|
|
155
|
-
|
|
156
|
-
if "xi_psi" in dims_to_partition:
|
|
157
|
-
xi_psi_domain_size = integer_division_or_raise(
|
|
158
|
-
ds.sizes["xi_psi"] - 3 * n_xi_ghost_cells, np_xi, "xi_psi"
|
|
159
|
-
)
|
|
138
|
+
for dim in dims_to_partition:
|
|
139
|
+
info = DIM_INFO[dim]
|
|
140
|
+
nparts = np_eta if info["axis"] == "eta" else np_xi
|
|
160
141
|
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
142
|
+
partitioned_sizes[dim] = _compute_partition_sizes(
|
|
143
|
+
total_size=ds.sizes[dim],
|
|
144
|
+
nparts=nparts,
|
|
145
|
+
ghost=info["ghost"],
|
|
146
|
+
edge=info["edge"],
|
|
147
|
+
dim=dim,
|
|
164
148
|
)
|
|
165
|
-
if "xi_coarse" in dims_to_partition:
|
|
166
|
-
xi_coarse_domain_size = integer_division_or_raise(
|
|
167
|
-
ds.sizes["xi_coarse"] - 2 * n_xi_ghost_cells, np_xi, "xi_coarse"
|
|
168
|
-
)
|
|
169
|
-
|
|
170
|
-
# unpartitioned dimensions should have sizes unchanged
|
|
171
|
-
partitioned_sizes = {
|
|
172
|
-
dim: [size] for dim, size in ds.sizes.items() if dim in dims_to_partition
|
|
173
|
-
}
|
|
174
|
-
|
|
175
|
-
# TODO refactor to use two functions for odd- and even-length dimensions
|
|
176
|
-
if "eta_v" in dims_to_partition:
|
|
177
|
-
partitioned_sizes["eta_v"] = [eta_v_domain_size] * (np_eta - 1) + [
|
|
178
|
-
eta_v_domain_size + n_eta_ghost_cells
|
|
179
|
-
]
|
|
180
|
-
if "xi_u" in dims_to_partition:
|
|
181
|
-
partitioned_sizes["xi_u"] = [xi_u_domain_size] * (np_xi - 1) + [
|
|
182
|
-
xi_u_domain_size + n_xi_ghost_cells
|
|
183
|
-
]
|
|
184
|
-
|
|
185
|
-
if np_eta > 1:
|
|
186
|
-
if "eta_rho" in dims_to_partition:
|
|
187
|
-
partitioned_sizes["eta_rho"] = (
|
|
188
|
-
[eta_rho_domain_size + n_eta_ghost_cells]
|
|
189
|
-
+ [eta_rho_domain_size] * (np_eta - 2)
|
|
190
|
-
+ [eta_rho_domain_size + n_eta_ghost_cells]
|
|
191
|
-
)
|
|
192
|
-
if "eta_psi" in dims_to_partition:
|
|
193
|
-
partitioned_sizes["eta_psi"] = (
|
|
194
|
-
[n_eta_ghost_cells + eta_psi_domain_size]
|
|
195
|
-
+ [eta_psi_domain_size] * (np_eta - 2)
|
|
196
|
-
+ [eta_psi_domain_size + 2 * n_eta_ghost_cells]
|
|
197
|
-
)
|
|
198
|
-
if "eta_coarse" in dims_to_partition:
|
|
199
|
-
partitioned_sizes["eta_coarse"] = (
|
|
200
|
-
[eta_coarse_domain_size + n_eta_ghost_cells]
|
|
201
|
-
+ [eta_coarse_domain_size] * (np_eta - 2)
|
|
202
|
-
+ [eta_coarse_domain_size + n_eta_ghost_cells]
|
|
203
|
-
)
|
|
204
|
-
|
|
205
|
-
if np_xi > 1:
|
|
206
|
-
if "xi_rho" in dims_to_partition:
|
|
207
|
-
partitioned_sizes["xi_rho"] = (
|
|
208
|
-
[xi_rho_domain_size + n_xi_ghost_cells]
|
|
209
|
-
+ [xi_rho_domain_size] * (np_xi - 2)
|
|
210
|
-
+ [xi_rho_domain_size + n_xi_ghost_cells]
|
|
211
|
-
)
|
|
212
|
-
if "xi_psi" in dims_to_partition:
|
|
213
|
-
partitioned_sizes["xi_psi"] = (
|
|
214
|
-
[n_xi_ghost_cells + xi_psi_domain_size]
|
|
215
|
-
+ [xi_psi_domain_size] * (np_xi - 2)
|
|
216
|
-
+ [xi_psi_domain_size + 2 * n_xi_ghost_cells]
|
|
217
|
-
)
|
|
218
|
-
if "xi_coarse" in dims_to_partition:
|
|
219
|
-
partitioned_sizes["xi_coarse"] = (
|
|
220
|
-
[xi_coarse_domain_size + n_xi_ghost_cells]
|
|
221
|
-
+ [xi_coarse_domain_size] * (np_xi - 2)
|
|
222
|
-
+ [xi_coarse_domain_size + n_xi_ghost_cells]
|
|
223
|
-
)
|
|
224
|
-
|
|
225
|
-
def cumsum(pmf):
|
|
226
|
-
"""Implementation of cumsum which ensures the result starts with zero."""
|
|
227
|
-
cdf = np.empty(len(pmf) + 1, dtype=int)
|
|
228
|
-
cdf[0] = 0
|
|
229
|
-
np.cumsum(pmf, out=cdf[1:])
|
|
230
|
-
return cdf
|
|
231
149
|
|
|
232
150
|
file_numbers = []
|
|
233
151
|
partitioned_datasets = []
|
|
152
|
+
|
|
234
153
|
for i in range(np_eta):
|
|
235
154
|
for j in range(np_xi):
|
|
236
|
-
|
|
237
|
-
file_numbers.append(file_number)
|
|
238
|
-
|
|
155
|
+
file_numbers.append(j + i * np_xi)
|
|
239
156
|
indexers = {}
|
|
240
157
|
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
xi_rho_partition_indices = cumsum(partitioned_sizes["xi_rho"])
|
|
249
|
-
indexers["xi_rho"] = slice(
|
|
250
|
-
int(xi_rho_partition_indices[j]),
|
|
251
|
-
int(xi_rho_partition_indices[j + 1]),
|
|
252
|
-
)
|
|
253
|
-
|
|
254
|
-
if "eta_v" in dims_to_partition:
|
|
255
|
-
eta_v_partition_indices = cumsum(partitioned_sizes["eta_v"])
|
|
256
|
-
indexers["eta_v"] = slice(
|
|
257
|
-
int(eta_v_partition_indices[i]),
|
|
258
|
-
int(eta_v_partition_indices[i + 1]),
|
|
259
|
-
)
|
|
260
|
-
if "xi_u" in dims_to_partition:
|
|
261
|
-
xi_u_partition_indices = cumsum(partitioned_sizes["xi_u"])
|
|
262
|
-
indexers["xi_u"] = slice(
|
|
263
|
-
int(xi_u_partition_indices[j]), int(xi_u_partition_indices[j + 1])
|
|
264
|
-
)
|
|
265
|
-
if "eta_psi" in dims_to_partition:
|
|
266
|
-
eta_psi_partition_indices = cumsum(partitioned_sizes["eta_psi"])
|
|
267
|
-
indexers["eta_psi"] = slice(
|
|
268
|
-
int(eta_psi_partition_indices[i]),
|
|
269
|
-
int(eta_psi_partition_indices[i + 1]),
|
|
270
|
-
)
|
|
271
|
-
if "xi_psi" in dims_to_partition:
|
|
272
|
-
xi_psi_partition_indices = cumsum(partitioned_sizes["xi_psi"])
|
|
273
|
-
indexers["xi_psi"] = slice(
|
|
274
|
-
int(xi_psi_partition_indices[j]),
|
|
275
|
-
int(xi_psi_partition_indices[j + 1]),
|
|
276
|
-
)
|
|
277
|
-
|
|
278
|
-
if "eta_coarse" in dims_to_partition:
|
|
279
|
-
eta_coarse_partition_indices = cumsum(partitioned_sizes["eta_coarse"])
|
|
280
|
-
indexers["eta_coarse"] = slice(
|
|
281
|
-
int(eta_coarse_partition_indices[i]),
|
|
282
|
-
int(eta_coarse_partition_indices[i + 1]),
|
|
283
|
-
)
|
|
284
|
-
|
|
285
|
-
if "xi_coarse" in dims_to_partition:
|
|
286
|
-
xi_coarse_partition_indices = cumsum(partitioned_sizes["xi_coarse"])
|
|
287
|
-
indexers["xi_coarse"] = slice(
|
|
288
|
-
int(xi_coarse_partition_indices[j]),
|
|
289
|
-
int(xi_coarse_partition_indices[j + 1]),
|
|
290
|
-
)
|
|
291
|
-
|
|
292
|
-
partitioned_ds = ds.isel(**indexers)
|
|
293
|
-
|
|
294
|
-
partitioned_datasets.append(partitioned_ds)
|
|
158
|
+
for dim, sizes in partitioned_sizes.items():
|
|
159
|
+
info = DIM_INFO[dim]
|
|
160
|
+
idx = i if info["axis"] == "eta" else j
|
|
161
|
+
bounds = _cumsum(sizes)
|
|
162
|
+
indexers[dim] = slice(bounds[idx], bounds[idx + 1])
|
|
163
|
+
|
|
164
|
+
partitioned_datasets.append(ds.isel(**indexers))
|
|
295
165
|
|
|
296
166
|
return file_numbers, partitioned_datasets
|
|
297
167
|
|
roms_tools/utils.py
CHANGED
|
@@ -10,6 +10,7 @@ from pathlib import Path
|
|
|
10
10
|
from typing import TypeAlias
|
|
11
11
|
|
|
12
12
|
import numpy as np
|
|
13
|
+
import pandas as pd
|
|
13
14
|
import xarray as xr
|
|
14
15
|
|
|
15
16
|
from roms_tools.constants import R_EARTH
|
|
@@ -1031,3 +1032,195 @@ def get_pkg_error_msg(purpose: str, package_name: str, option_name: str) -> str:
|
|
|
1031
1032
|
• `pip install {package_name}` or
|
|
1032
1033
|
• `conda install {package_name}`
|
|
1033
1034
|
Alternatively, install `roms-tools` with conda to include all dependencies.""")
|
|
1035
|
+
|
|
1036
|
+
|
|
1037
|
+
def wrap_longitudes(ds: xr.Dataset, straddle: bool) -> xr.Dataset:
|
|
1038
|
+
"""
|
|
1039
|
+
Safely adjust longitude variables for datasets that may or may not cross
|
|
1040
|
+
the dateline. Only modifies longitude-like variables that are present.
|
|
1041
|
+
|
|
1042
|
+
Parameters
|
|
1043
|
+
----------
|
|
1044
|
+
ds : xr.Dataset
|
|
1045
|
+
Dataset containing longitude variables (e.g., lon_rho, lon_u, lon_v).
|
|
1046
|
+
straddle : bool
|
|
1047
|
+
- True: force longitudes into [-180, 180]
|
|
1048
|
+
- False: force longitudes into [0, 360]
|
|
1049
|
+
|
|
1050
|
+
Returns
|
|
1051
|
+
-------
|
|
1052
|
+
xr.Dataset
|
|
1053
|
+
A new dataset with adjusted longitude values.
|
|
1054
|
+
"""
|
|
1055
|
+
lon_vars = ["lon_rho", "lon_u", "lon_v"]
|
|
1056
|
+
|
|
1057
|
+
for lon_dim in lon_vars:
|
|
1058
|
+
if lon_dim not in ds:
|
|
1059
|
+
continue # skip missing coordinate
|
|
1060
|
+
|
|
1061
|
+
lon = ds[lon_dim]
|
|
1062
|
+
|
|
1063
|
+
if straddle:
|
|
1064
|
+
# wrap into [-180, 180]
|
|
1065
|
+
lon_wrapped = xr.where(lon > 180, lon - 360, lon)
|
|
1066
|
+
else:
|
|
1067
|
+
# wrap into [0, 360]
|
|
1068
|
+
lon_wrapped = xr.where(lon < 0, lon + 360, lon)
|
|
1069
|
+
|
|
1070
|
+
# preserve attributes
|
|
1071
|
+
lon_wrapped.attrs.update(lon.attrs)
|
|
1072
|
+
ds[lon_dim] = lon_wrapped
|
|
1073
|
+
|
|
1074
|
+
return ds
|
|
1075
|
+
|
|
1076
|
+
|
|
1077
|
+
def interpolate_from_climatology(
|
|
1078
|
+
field: xr.DataArray | xr.Dataset,
|
|
1079
|
+
time_dim: str,
|
|
1080
|
+
time_coord: str,
|
|
1081
|
+
time: xr.DataArray | pd.DatetimeIndex,
|
|
1082
|
+
) -> xr.DataArray | xr.Dataset:
|
|
1083
|
+
"""Interpolates a climatological field to specified time points.
|
|
1084
|
+
|
|
1085
|
+
This function interpolates the input `field` based on `day_of_year` values
|
|
1086
|
+
extracted from the provided `time` points. If `field` is an `xarray.Dataset`,
|
|
1087
|
+
interpolation is applied to all its data variables individually.
|
|
1088
|
+
|
|
1089
|
+
Parameters
|
|
1090
|
+
----------
|
|
1091
|
+
field : xarray.DataArray or xarray.Dataset
|
|
1092
|
+
The input field to be interpolated.
|
|
1093
|
+
- If `field` is an `xarray.DataArray`, it must have a time dimension identified by `time_dim_name`.
|
|
1094
|
+
- If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along `time_dim_name`.
|
|
1095
|
+
The time dimension is assumed to represent `day_of_year` for climatological purposes.
|
|
1096
|
+
time_dim : str
|
|
1097
|
+
The name of the time dimension in `field`.
|
|
1098
|
+
time_coord : str
|
|
1099
|
+
The name of the time coordinate in `field`.
|
|
1100
|
+
time : xarray.DataArray or pandas.DatetimeIndex
|
|
1101
|
+
The target time points for interpolation. These are internally converted to `day_of_year`
|
|
1102
|
+
before performing interpolation.
|
|
1103
|
+
|
|
1104
|
+
Returns
|
|
1105
|
+
-------
|
|
1106
|
+
xarray.DataArray or xarray.Dataset
|
|
1107
|
+
The interpolated field, maintaining the same type (`xarray.DataArray` or `xarray.Dataset`)
|
|
1108
|
+
but aligned to the specified `time` values.
|
|
1109
|
+
|
|
1110
|
+
Notes
|
|
1111
|
+
-----
|
|
1112
|
+
- This function assumes that `field` represents a climatological dataset, where time is expressed as `day_of_year` (1-365).
|
|
1113
|
+
- The `time` input is automatically converted to `day_of_year`, so manual conversion is not required before calling this function.
|
|
1114
|
+
"""
|
|
1115
|
+
|
|
1116
|
+
def np_times_to_fractional_days(
|
|
1117
|
+
np_times: np.ndarray | pd.DatetimeIndex | np.datetime64 | pd.Timestamp,
|
|
1118
|
+
) -> np.ndarray:
|
|
1119
|
+
"""Convert datetime(s) to fractional day-of-year values."""
|
|
1120
|
+
pd_times = pd.to_datetime(np_times)
|
|
1121
|
+
|
|
1122
|
+
# scalar input -> make it a 1-element array
|
|
1123
|
+
if np.isscalar(pd_times):
|
|
1124
|
+
pd_times = np.array([pd_times])
|
|
1125
|
+
|
|
1126
|
+
fractional_days = pd_times.dayofyear + (
|
|
1127
|
+
pd_times.hour / 24 + pd_times.minute / 1440 + pd_times.second / 86400
|
|
1128
|
+
)
|
|
1129
|
+
return (
|
|
1130
|
+
fractional_days.values
|
|
1131
|
+
if hasattr(fractional_days, "values")
|
|
1132
|
+
else np.array(fractional_days)
|
|
1133
|
+
)
|
|
1134
|
+
|
|
1135
|
+
def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
|
|
1136
|
+
if isinstance(time, xr.DataArray):
|
|
1137
|
+
day_of_year = time.dt.dayofyear
|
|
1138
|
+
else:
|
|
1139
|
+
day_of_year = np_times_to_fractional_days(time)
|
|
1140
|
+
|
|
1141
|
+
data_array_interpolated = interpolate_cyclic_time(
|
|
1142
|
+
data_array, time_dim, time_coord, day_of_year
|
|
1143
|
+
)
|
|
1144
|
+
|
|
1145
|
+
# expand dims if single element
|
|
1146
|
+
if day_of_year.size == 1:
|
|
1147
|
+
data_array_interpolated = data_array_interpolated.expand_dims({time_dim: 1})
|
|
1148
|
+
return data_array_interpolated
|
|
1149
|
+
|
|
1150
|
+
if isinstance(field, xr.DataArray):
|
|
1151
|
+
return interpolate_single_field(field)
|
|
1152
|
+
elif isinstance(field, xr.Dataset):
|
|
1153
|
+
interpolated_data_vars = {
|
|
1154
|
+
var: interpolate_single_field(data_array)
|
|
1155
|
+
for var, data_array in field.data_vars.items()
|
|
1156
|
+
}
|
|
1157
|
+
return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
|
|
1158
|
+
|
|
1159
|
+
else:
|
|
1160
|
+
raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
|
|
1161
|
+
|
|
1162
|
+
|
|
1163
|
+
def interpolate_cyclic_time(
|
|
1164
|
+
data_array: xr.DataArray,
|
|
1165
|
+
time_dim: str,
|
|
1166
|
+
time_coord: str,
|
|
1167
|
+
day_of_year: int | float | np.ndarray | xr.DataArray | Sequence[int | float],
|
|
1168
|
+
) -> xr.DataArray:
|
|
1169
|
+
"""Interpolates a DataArray cyclically across the start and end of the year.
|
|
1170
|
+
|
|
1171
|
+
This function extends the data cyclically by appending the last time step
|
|
1172
|
+
(shifted back by one year) at the beginning and the first time step
|
|
1173
|
+
(shifted forward by one year) at the end. It then performs linear interpolation
|
|
1174
|
+
to match the specified `day_of_year` values.
|
|
1175
|
+
|
|
1176
|
+
Parameters
|
|
1177
|
+
----------
|
|
1178
|
+
data_array : xr.DataArray
|
|
1179
|
+
The input data array containing a time-like dimension.
|
|
1180
|
+
time_dim : str
|
|
1181
|
+
The name of the time dimension in the dataset.
|
|
1182
|
+
time_coord : str
|
|
1183
|
+
The name of the time coordinate in the dataset.
|
|
1184
|
+
day_of_year : Union[int, float, np.ndarray, xr.DataArray, Sequence[Union[int, float]]]
|
|
1185
|
+
The target day(s) of the year for interpolation. This can be:
|
|
1186
|
+
- A single integer or float representing the day of the year.
|
|
1187
|
+
- A NumPy array or xarray DataArray containing multiple days.
|
|
1188
|
+
- A list or tuple of integers or floats for multiple target days.
|
|
1189
|
+
|
|
1190
|
+
Returns
|
|
1191
|
+
-------
|
|
1192
|
+
xr.DataArray
|
|
1193
|
+
The interpolated DataArray, ensuring cyclic continuity across year boundaries.
|
|
1194
|
+
|
|
1195
|
+
Notes
|
|
1196
|
+
-----
|
|
1197
|
+
- This function is useful for interpolating climatological data, where the time axis
|
|
1198
|
+
represents a repeating annual cycle.
|
|
1199
|
+
- The `day_of_year` values should be within the range [1, 365] or [1, 366] for leap years.
|
|
1200
|
+
"""
|
|
1201
|
+
# Concatenate across the beginning and end of the year
|
|
1202
|
+
time_concat = xr.concat(
|
|
1203
|
+
[
|
|
1204
|
+
data_array[time_coord][-1] - 365.25, # Shift last time backward
|
|
1205
|
+
data_array[time_coord],
|
|
1206
|
+
data_array[time_coord][0] + 365.25, # Shift first time forward
|
|
1207
|
+
],
|
|
1208
|
+
dim=time_dim,
|
|
1209
|
+
)
|
|
1210
|
+
|
|
1211
|
+
data_array_concat = xr.concat(
|
|
1212
|
+
[
|
|
1213
|
+
data_array.isel(**{time_dim: -1}), # Append last value at the beginning
|
|
1214
|
+
data_array,
|
|
1215
|
+
data_array.isel(**{time_dim: 0}), # Append first value at the end
|
|
1216
|
+
],
|
|
1217
|
+
dim=time_dim,
|
|
1218
|
+
)
|
|
1219
|
+
data_array_concat[time_dim] = time_concat
|
|
1220
|
+
|
|
1221
|
+
# Interpolate to specified times
|
|
1222
|
+
data_array_interpolated = data_array_concat.interp(
|
|
1223
|
+
**{time_dim: day_of_year}, method="linear"
|
|
1224
|
+
)
|
|
1225
|
+
|
|
1226
|
+
return data_array_interpolated
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: roms-tools
|
|
3
|
-
Version: 3.
|
|
3
|
+
Version: 3.4.0
|
|
4
4
|
Summary: Tools for running and analysing UCLA-ROMS simulations
|
|
5
5
|
Author-email: Nora Loose <nora.loose@gmail.com>, Thomas Nicholas <tom@cworthy.org>, Scott Eilerman <scott.eilerman@cworthy.org>
|
|
6
6
|
License: Apache-2
|