roms-tools 2.6.1__py3-none-any.whl → 2.7.0__py3-none-any.whl

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_long_name/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        2
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<U21",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_long_name/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "ntracers"
+    ],
+    "long_name": "Tracer long name"
+}

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_unit/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        2
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<U15",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        2
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_no_climatology.zarr/tracer_unit/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "ntracers"
+    ],
+    "long_name": "Tracer unit"
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/.zmetadata

@@ -86,7 +86,7 @@
             ],
             "long_name": "River ID (1-based Fortran indexing)"
         },
-        "river_flux/.zarray": {
+        "river_fraction/.zarray": {
             "chunks": [
                 20,
                 20
@@ -108,12 +108,42 @@
             ],
             "zarr_format": 2
         },
-        "river_flux/.zattrs": {
+        "river_fraction/.zattrs": {
             "_ARRAY_DIMENSIONS": [
                 "eta_rho",
                 "xi_rho"
             ],
-            "long_name": "River ID plus local volume fraction",
+            "long_name": "River volume fraction",
+            "units": "none"
+        },
+        "river_index/.zarray": {
+            "chunks": [
+                20,
+                20
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<f8",
+            "fill_value": "NaN",
+            "filters": null,
+            "order": "C",
+            "shape": [
+                20,
+                20
+            ],
+            "zarr_format": 2
+        },
+        "river_index/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "eta_rho",
+                "xi_rho"
+            ],
+            "long_name": "River ID",
             "units": "none"
         },
         "river_name/.zarray": {
@@ -204,7 +234,7 @@
                 "ntracers",
                 "nriver"
             ],
-            "coordinates": "abs_time month river_name tracer_name",
+            "coordinates": "abs_time month river_name tracer_long_name tracer_name tracer_unit",
             "long_name": "River tracer data"
         },
         "river_volume/.zarray": {
@@ -238,6 +268,32 @@
             "long_name": "River volume flux",
             "units": "m^3/s"
         },
+        "tracer_long_name/.zarray": {
+            "chunks": [
+                34
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<U43",
+            "fill_value": null,
+            "filters": null,
+            "order": "C",
+            "shape": [
+                34
+            ],
+            "zarr_format": 2
+        },
+        "tracer_long_name/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "ntracers"
+            ],
+            "long_name": "Tracer long name"
+        },
         "tracer_name/.zarray": {
             "chunks": [
                 34
@@ -263,6 +319,32 @@
                 "ntracers"
             ],
             "long_name": "Tracer name"
+        },
+        "tracer_unit/.zarray": {
+            "chunks": [
+                34
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<U15",
+            "fill_value": null,
+            "filters": null,
+            "order": "C",
+            "shape": [
+                34
+            ],
+            "zarr_format": 2
+        },
+        "tracer_unit/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "ntracers"
+            ],
+            "long_name": "Tracer unit"
         }
     },
     "zarr_consolidated_format": 1

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_fraction/.zarray

@@ -0,0 +1,22 @@
+{
+    "chunks": [
+        20,
+        20
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f8",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        20,
+        20
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_fraction/.zattrs

@@ -0,0 +1,8 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "eta_rho",
+        "xi_rho"
+    ],
+    "long_name": "River volume fraction",
+    "units": "none"
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_index/.zarray

@@ -0,0 +1,22 @@
+{
+    "chunks": [
+        20,
+        20
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<f8",
+    "fill_value": "NaN",
+    "filters": null,
+    "order": "C",
+    "shape": [
+        20,
+        20
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_index/.zattrs

@@ -0,0 +1,8 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "eta_rho",
+        "xi_rho"
+    ],
+    "long_name": "River ID",
+    "units": "none"
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/river_tracer/.zattrs

@@ -4,6 +4,6 @@
         "ntracers",
         "nriver"
     ],
-    "coordinates": "abs_time month river_name tracer_name",
+    "coordinates": "abs_time month river_name tracer_long_name tracer_name tracer_unit",
     "long_name": "River tracer data"
 }

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/tracer_long_name/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        34
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<U43",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        34
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/tracer_long_name/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "ntracers"
+    ],
+    "long_name": "Tracer long name"
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/tracer_unit/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        34
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<U15",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        34
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/river_forcing_with_bgc.zarr/tracer_unit/.zattrs

@@ -0,0 +1,6 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "ntracers"
+    ],
+    "long_name": "Tracer unit"
+}

roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/.zattrs

@@ -1,7 +1,6 @@
 {
-    "allan_factor": 2.0,
     "model_reference_date": "2000-01-01 00:00:00",
-    "roms_tools_version": "0.1.dev157+dirty",
+    "roms_tools_version": "0.1.dev177",
     "source": "TPXO",
     "title": "ROMS tidal forcing created by ROMS-Tools"
 }

roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/.zmetadata

@@ -1,15 +1,39 @@
 {
     "metadata": {
         ".zattrs": {
-            "allan_factor": 2.0,
             "model_reference_date": "2000-01-01 00:00:00",
-            "roms_tools_version": "0.1.dev157+dirty",
+            "roms_tools_version": "0.1.dev177",
             "source": "TPXO",
             "title": "ROMS tidal forcing created by ROMS-Tools"
         },
         ".zgroup": {
             "zarr_format": 2
         },
+        "ntides/.zarray": {
+            "chunks": [
+                1
+            ],
+            "compressor": {
+                "blocksize": 0,
+                "clevel": 5,
+                "cname": "lz4",
+                "id": "blosc",
+                "shuffle": 1
+            },
+            "dtype": "<U3",
+            "fill_value": null,
+            "filters": null,
+            "order": "C",
+            "shape": [
+                1
+            ],
+            "zarr_format": 2
+        },
+        "ntides/.zattrs": {
+            "_ARRAY_DIMENSIONS": [
+                "ntides"
+            ]
+        },
         "omega/.zarray": {
             "chunks": [
                 1
@@ -33,9 +57,7 @@
         "omega/.zattrs": {
             "_ARRAY_DIMENSIONS": [
                 "ntides"
-            ],
-            "long_name": "Tidal frequencies",
-            "units": "1/s"
+            ]
         },
         "pot_Im/.zarray": {
             "chunks": [

roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/ntides/.zarray

@@ -0,0 +1,20 @@
+{
+    "chunks": [
+        1
+    ],
+    "compressor": {
+        "blocksize": 0,
+        "clevel": 5,
+        "cname": "lz4",
+        "id": "blosc",
+        "shuffle": 1
+    },
+    "dtype": "<U3",
+    "fill_value": null,
+    "filters": null,
+    "order": "C",
+    "shape": [
+        1
+    ],
+    "zarr_format": 2
+}

roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/ntides/.zattrs

@@ -0,0 +1,5 @@
+{
+    "_ARRAY_DIMENSIONS": [
+        "ntides"
+    ]
+}

roms_tools/tests/test_setup/test_data/tidal_forcing.zarr/omega/.zattrs

@@ -1,7 +1,5 @@
 {
     "_ARRAY_DIMENSIONS": [
         "ntides"
-    ],
-    "long_name": "Tidal frequencies",
-    "units": "1/s"
+    ]
 }

roms_tools/tests/test_setup/test_datasets.py

@@ -8,9 +8,11 @@ from roms_tools.setup.datasets import (
     GLORYSDataset,
     ERA5Correction,
     CESMBGCDataset,
+    TPXODataset,
 )
 from roms_tools.download import download_test_data
 from pathlib import Path
+from collections import OrderedDict
 
 
 @pytest.fixture
@@ -529,12 +531,114 @@ def test_climatology_error(use_dask):
     [
         ("era5_data", 0.25),
         ("glorys_data", 1 / 12),
-        ("tpxo_data", 1 / 6),
+        # ("tpxo_data", 1 / 6),
         ("cesm_bgc_data", 1.0),
         ("cesm_surface_bgc_data", 1.0),
+        ("unified_bgc_data", 2.0),
+        ("unified_surface_bgc_data", 2.0),
     ],
 )
 def test_horizontal_resolution(data_fixture, expected_resolution, request):
 
     data = request.getfixturevalue(data_fixture)
     assert np.isclose(data.resolution, expected_resolution)
+
+
+class TestTPXODataset:
+    @pytest.fixture
+    def regional_tpxo_dataset(self, use_dask):
+        """TPXO dataset with regional coverage and 25 constituents: M2, S2, N2, K2, K1, O1, P1, Q1, MM, Mf, MSF, M4, Mn4, Ms4, ..."""
+        fname_grid = Path(download_test_data("regional_grid_tpxo10v2.nc"))
+        fname_h = Path(download_test_data("regional_h_tpxo10v2.nc"))
+
+        return TPXODataset(
+            filename=fname_h,
+            grid_filename=fname_grid,
+            location="h",
+            var_names={"ssh_Re": "hRe", "ssh_Im": "hIm"},
+            use_dask=use_dask,
+        )
+
+    @pytest.fixture
+    def global_tpxo_dataset(self, use_dask):
+        """TPXO dataset with global coverage and 1 constituent: M2"""
+        fname_grid = Path(download_test_data("global_grid_tpxo10.v2.nc"))
+        fname_h = Path(download_test_data("global_h_tpxo10.v2.nc"))
+
+        return TPXODataset(
+            filename=fname_h,
+            grid_filename=fname_grid,
+            location="h",
+            var_names={"ssh_Re": "hRe", "ssh_Im": "hIm"},
+            use_dask=use_dask,
+        )
+
+    @pytest.fixture
+    def omega(self):
+        return {
+            "m2": 1.405189e-04,
+            "s2": 1.454441e-04,
+            "n2": 1.378797e-04,
+            "k2": 1.458423e-04,
+            "k1": 7.292117e-05,
+            "o1": 6.759774e-05,
+            "p1": 7.252295e-05,
+            "q1": 6.495854e-05,
+            "mm": 0.026392e-04,
+            "mf": 0.053234e-04,
+            "m4": 2.810377e-04,
+            "mn4": 2.783984e-04,
+            "ms4": 2.859630e-04,
+            "2n2": 1.352405e-04,
+            "s1": 7.2722e-05,
+        }
+
+    @pytest.mark.parametrize(
+        "tpxo_dataset_fixture",
+        [
+            "regional_tpxo_dataset",
+            "global_tpxo_dataset",
+        ],
+    )
+    def test_initialization_and_clean_up(self, tpxo_dataset_fixture, request):
+
+        tpxo_dataset = request.getfixturevalue(tpxo_dataset_fixture)
+        assert tpxo_dataset.location == "h"
+
+        assert "hRe" in tpxo_dataset.ds.data_vars
+        assert "hIm" in tpxo_dataset.ds.data_vars
+        assert "mask" in tpxo_dataset.ds.data_vars
+        assert "longitude" in tpxo_dataset.ds.dims
+        assert "latitude" in tpxo_dataset.ds.dims
+        assert "longitude" in tpxo_dataset.ds.variables
+        assert "latitude" in tpxo_dataset.ds.variables
+
+    def test_select_fewer_constituents(self, regional_tpxo_dataset, omega):
+
+        regional_tpxo_dataset.select_constituents(2, omega)
+        assert (
+            regional_tpxo_dataset.ds["ntides"].values.astype("U3") == ["m2", "s2"]
+        ).all()
+
+    def test_select_constituents_with_reordering(self, regional_tpxo_dataset, omega):
+
+        regional_tpxo_dataset.select_constituents(11, omega)
+
+        assert len(regional_tpxo_dataset.ds["ntides"]) == 11
+        # check that m4 has been moved from 12th to 11th position to follow TPXO9 order
+        assert (
+            regional_tpxo_dataset.ds["ntides"].isel(ntides=10).item().decode("utf-8")
+            == "m4"
+        )
+
+    def test_select_constituents_omega_mismatch(self, regional_tpxo_dataset, omega):
+
+        omega = OrderedDict(
+            list(omega.items())[:3] + [("fake", 6.495854e-05)] + list(omega.items())[3:]
+        )
+        with pytest.raises(ValueError, match="The dataset contains tidal constituents"):
+            regional_tpxo_dataset.select_constituents(11, omega)
+
+    def test_select_constituents_too_few(self, global_tpxo_dataset, omega):
+        with pytest.raises(ValueError, match="The dataset contains tidal constituents"):
+            global_tpxo_dataset.select_constituents(11, omega)