wavespectra 4.6.0__tar.gz → 4.7.0__tar.gz

{wavespectra-4.6.0/wavespectra.egg-info → wavespectra-4.7.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wavespectra
-Version: 4.6.0
+Version: 4.7.0
 Summary: Library for ocean wave spectra
 Author-email: Wavespectra Developers <r.guedes@oceanum.science>
 Maintainer-email: Rafael Guedes <r.guedes@oceanum.science>, Tom Durrant <t.durrant@oceanum.science>

{wavespectra-4.6.0 → wavespectra-4.7.0}/pyproject.toml

@@ -173,7 +173,7 @@ extend-select = ["B015", "B018"]
 github_url = "https://github.com/wavespectra/wavespectra/"
 
 [tool.tbump.version]
-current = "4.6.0"
+current = "4.7.0"
 regex = '''
   (?P<major>\d+)
   \.(?P<minor>\d+)

{wavespectra-4.6.0 → wavespectra-4.7.0}/tests/test_partition.py

@@ -2,6 +2,7 @@ import pytest
 from pathlib import Path
 import numpy as np
 
+import wavespectra
 from wavespectra import read_ww3
 from wavespectra.partition.partition import (
     Partition,
@@ -18,6 +19,13 @@ from wavespectra.partition.tracking import track_partitions
 HERE = Path(__file__).parent
 
 
+@pytest.fixture(autouse=True)
+def dataset_transforms():
+    """Run under the future dataset transforms behaviour in this module."""
+    with wavespectra.set_options(dataset_transforms=True):
+        yield
+
+
 class BasePTM:
     def setup_class(self):
         self.dset = read_ww3(HERE / "sample_files/ww3file.nc")
@@ -437,11 +445,18 @@ class TestPartitionAndTrack(BasePTM):
         self._track("hp01", wspd=self.dset.wspd, wdir=self.dset.wdir, dpt=self.dset.dpt)
 
     def test_track_hp01_no_wind(self):
+        # Wind variables are resolved from the dataset so partitioning from
+        # the DataArray is required to exercise the no-wind path
+        pt = Partition(self.dset.efth)
         with pytest.warns(UserWarning):
-            dspart = self.pt.track(method="hp01", swells=2)
+            dspart = pt.track(method="hp01", swells=2)
         dspart = dspart.load()
         assert (dspart.ntracks.values > 0).all()
 
+    def test_track_hp01_wind_from_dataset(self):
+        dspart = self.pt.track(method="hp01", swells=2).load()
+        assert (dspart.ntracks.values > 0).all()
+
     def test_track_invalid_method(self):
         with pytest.raises(ValueError):
             self.pt.track(method="ptm4")

wavespectra-4.7.0/tests/test_specdataset.py

@@ -0,0 +1,192 @@
+"""Testing the dataset-aware behaviour of the SpecDataset accessor."""
+
+import os
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from wavespectra import read_ww3, set_options, get_options
+from wavespectra.partition.partition import Partition
+
+
+FILES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "sample_files")
+
+
+@pytest.fixture(autouse=True)
+def dataset_transforms():
+    """Opt in to the future dataset transforms behaviour in this module."""
+    with set_options(dataset_transforms=True):
+        yield
+
+
+TRANSFORMS = {
+    "oned": {},
+    "split": {"fmin": 0.05, "fmax": 0.4},
+    "smooth": {},
+    "rotate": {"angle": 30},
+    "interp": {"freq": np.arange(0.04, 0.4, 0.02)},
+    "to_energy": {},
+    "scale_by_hs": {"expr": "2 * hs"},
+}
+
+
+@pytest.fixture(scope="module")
+def dset():
+    yield read_ww3(os.path.join(FILES_DIR, "ww3file.nc"))
+
+
+@pytest.mark.parametrize("method, kwargs", TRANSFORMS.items())
+def test_transform_returns_dataset_with_variables(dset, method, kwargs):
+    """Dataset accessor transforms return Dataset with non-spectral variables."""
+    dsout = getattr(dset.spec, method)(**kwargs)
+    assert isinstance(dsout, xr.Dataset)
+    for name in ["wspd", "wdir", "dpt", "lat", "lon"]:
+        assert name in dsout.data_vars
+
+
+@pytest.mark.parametrize("method, kwargs", TRANSFORMS.items())
+def test_transform_dataarray_accessor_unchanged(dset, method, kwargs):
+    """DataArray accessor transforms still return a DataArray with same values."""
+    darr = getattr(dset.efth.spec, method)(**kwargs)
+    assert isinstance(darr, xr.DataArray)
+    dsout = getattr(dset.spec, method)(**kwargs)
+    specname = "energy" if method == "to_energy" else "efth"
+    assert dsout[specname].values == pytest.approx(darr.values, nan_ok=True)
+
+
+def test_interp_like_returns_dataset(dset):
+    dsout = dset.spec.interp_like(dset.efth)
+    assert isinstance(dsout, xr.Dataset)
+    assert "wspd" in dsout.data_vars
+    assert isinstance(dset.efth.spec.interp_like(dset.efth), xr.DataArray)
+
+
+def test_transform_drops_spectral_variables(dset):
+    """Extra variables that depend on spectral dims are dropped."""
+    ds = dset.copy()
+    ds["extra"] = ds.efth * 2
+    dsout = ds.spec.smooth()
+    assert "extra" not in dsout
+    assert "wspd" in dsout.data_vars
+
+
+def test_transform_keeps_dataset_attrs(dset):
+    ds = dset.copy()
+    ds.attrs["title"] = "test title"
+    assert ds.spec.oned().attrs["title"] == "test title"
+
+
+def test_scalar_params_still_dataarray(dset):
+    assert isinstance(dset.spec.hs(), xr.DataArray)
+    assert isinstance(dset.spec.stats(["hs", "tp"]), xr.Dataset)
+
+
+def test_transform_chaining(dset):
+    dsout = dset.spec.split(fmax=0.2).spec.oned()
+    assert isinstance(dsout, xr.Dataset)
+    assert "wspd" in dsout.data_vars
+
+
+class TestPartitionDatasetAware:
+    def test_partition_returns_dataset(self, dset):
+        dspart = dset.spec.partition.ptm1(swells=2)
+        assert isinstance(dspart, xr.Dataset)
+        for name in ["wspd", "wdir", "dpt"]:
+            assert name in dspart.data_vars
+        assert "part" in dspart.efth.dims
+
+    def test_partition_wind_defaults_from_dataset(self, dset):
+        implicit = dset.spec.partition.ptm1(swells=2)
+        explicit = dset.spec.partition.ptm1(
+            wspd=dset.wspd, wdir=dset.wdir, dpt=dset.dpt, swells=2
+        )
+        xr.testing.assert_allclose(implicit, explicit)
+
+    @pytest.mark.parametrize("method", ["ptm2", "ptm4"])
+    def test_partition_wind_defaults_other_methods(self, dset, method):
+        dspart = getattr(dset.spec.partition, method)()
+        assert isinstance(dspart, xr.Dataset)
+
+    def test_partition_dataarray_returns_dataarray(self, dset):
+        dspart = dset.efth.spec.partition.ptm1(
+            wspd=dset.wspd, wdir=dset.wdir, dpt=dset.dpt, swells=2
+        )
+        assert isinstance(dspart, xr.DataArray)
+
+    def test_partition_wind_required_without_dataset(self, dset):
+        with pytest.raises(ValueError):
+            dset.efth.spec.partition.ptm1()
+
+    def test_partition_no_wind_dims(self, dset):
+        dspart = dset.spec.partition.ptm5(fcut=0.1)
+        assert isinstance(dspart, xr.Dataset)
+        assert "wspd" in dspart.data_vars
+
+    def test_track_returns_dataset_with_variables(self, dset):
+        dsout = dset.spec.partition.track(swells=2)
+        assert isinstance(dsout, xr.Dataset)
+        for name in ["efth", "track_id", "ntracks", "wspd", "wdir", "dpt"]:
+            assert name in dsout.data_vars
+
+    def test_track_dataarray_no_extra_variables(self, dset):
+        dsout = Partition(dset.efth).track(
+            wspd=dset.wspd, wdir=dset.wdir, dpt=dset.dpt, swells=2
+        )
+        assert isinstance(dsout, xr.Dataset)
+        assert "wspd" not in dsout.data_vars
+
+
+class TestDefaultBehaviour:
+    """Without the dataset_transforms option the old behaviour is preserved."""
+
+    def test_transform_warns_and_returns_dataarray(self, dset):
+        with set_options(dataset_transforms=False):
+            with pytest.warns(FutureWarning, match="oned"):
+                dsout = dset.spec.oned()
+        assert isinstance(dsout, xr.DataArray)
+
+    def test_transform_values_unchanged(self, dset):
+        with set_options(dataset_transforms=False):
+            with pytest.warns(FutureWarning):
+                darr = dset.spec.smooth()
+        xr.testing.assert_allclose(darr, dset.efth.spec.smooth())
+
+    def test_dataarray_accessor_does_not_warn(self, dset):
+        import warnings
+
+        with set_options(dataset_transforms=False):
+            with warnings.catch_warnings():
+                warnings.simplefilter("error", FutureWarning)
+                darr = dset.efth.spec.oned()
+        assert isinstance(darr, xr.DataArray)
+
+    def test_partition_warns_and_returns_dataarray(self, dset):
+        with set_options(dataset_transforms=False):
+            with pytest.warns(FutureWarning, match="Partitioning"):
+                dspart = dset.spec.partition.ptm1(
+                    wspd=dset.wspd, wdir=dset.wdir, dpt=dset.dpt, swells=2
+                )
+        assert isinstance(dspart, xr.DataArray)
+
+    def test_hp01_does_not_default_wind_from_dataset(self, dset):
+        with set_options(dataset_transforms=False):
+            with pytest.warns((FutureWarning, UserWarning)):
+                dspart = dset.spec.partition.hp01(swells=2)
+        assert isinstance(dspart, xr.DataArray)
+        # No wind classification, partition 0 is null
+        assert float(dspart.isel(part=0).sum()) == 0
+
+
+class TestSetOptions:
+    def test_option_toggles_and_restores(self, dset):
+        with set_options(dataset_transforms=False):
+            assert get_options()["dataset_transforms"] is False
+            with set_options(dataset_transforms=True):
+                assert get_options()["dataset_transforms"] is True
+                assert isinstance(dset.spec.oned(), xr.Dataset)
+            assert get_options()["dataset_transforms"] is False
+
+    def test_invalid_option_raises(self):
+        with pytest.raises(ValueError):
+            set_options(not_an_option=True)

{wavespectra-4.6.0 → wavespectra-4.7.0}/wavespectra/__init__.py

@@ -7,6 +7,8 @@
 
 import warnings
 
+from wavespectra.core.options import set_options, get_options
+
 try:
     from wavespectra.specdataset import SpecDataset
     from wavespectra.specarray import SpecArray
@@ -14,7 +16,7 @@ except ImportError as exc:
     warnings.warn(f"Cannot import accessors at the main module level:\n{exc}")
 
 
-__version__ = "4.6.0"
+__version__ = "4.7.0"
 
 
 def _import_functions(pkgname="input", prefix="read"):

{wavespectra-4.6.0 → wavespectra-4.7.0}/wavespectra/construct/__init__.py

@@ -145,7 +145,7 @@ def partition_and_reconstruct(
             if v not in dset.data_vars:
                 raise ValueError(f"Missing variable '{v}' in dset for partitioning.")
         kw = dict(wspd=dset.wspd, wdir=dset.wdir, dpt=dset.dpt, swells=parts - 1)
-    dspart = getattr(dset.spec.partition, partition_method)(**kw)
+    dspart = getattr(dset[attrs.SPECNAME].spec.partition, partition_method)(**kw)
 
     # Calculating parameters
     stats = list(set(STATS) - set(use_defaults))

wavespectra-4.7.0/wavespectra/core/options.py

@@ -0,0 +1,57 @@
+"""Package-level options."""
+
+DATASET_TRANSFORMS = "dataset_transforms"
+
+OPTIONS = {
+    DATASET_TRANSFORMS: False,
+}
+
+
+class set_options:
+    """Set package-level options for wavespectra.
+
+    Can be used either as a function to set options globally or as a context
+    manager to set them within a controlled block, restoring the previous
+    values when exiting the block.
+
+    Currently supported options:
+
+    - ``dataset_transforms``: If True, methods that transform the spectral
+      variable such as ``interp``, ``smooth``, ``split``, ``oned`` and the
+      partitioning methods return a Dataset preserving the non-spectral
+      variables when called from the Dataset accessor, and the ``wspd``,
+      ``wdir`` and ``dpt`` arguments of the ``hp01`` and ``track``
+      partitioning methods default to the dataset variables with those names.
+      If False (default), those methods return a bare spectral DataArray and
+      emit a FutureWarning. This will become the default behaviour in
+      wavespectra 5.0.
+
+    Examples:
+        >>> import wavespectra
+        >>> wavespectra.set_options(dataset_transforms=True)  # global
+        >>> with wavespectra.set_options(dataset_transforms=True):
+        ...     dsout = dset.spec.oned()  # within context only
+
+    """
+
+    def __init__(self, **kwargs):
+        self.old = {}
+        for key, value in kwargs.items():
+            if key not in OPTIONS:
+                raise ValueError(
+                    f"Argument '{key}' is not in the set of valid options "
+                    f"{set(OPTIONS)}"
+                )
+            self.old[key] = OPTIONS[key]
+        OPTIONS.update(kwargs)
+
+    def __enter__(self):
+        return
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        OPTIONS.update(self.old)
+
+
+def get_options():
+    """Return a dictionary with the current package-level options."""
+    return dict(OPTIONS)

{wavespectra-4.6.0 → wavespectra-4.7.0}/wavespectra/core/utils.py

@@ -280,6 +280,36 @@ def check_same_coordinates(*args):
             )
 
 
+def dataset_from_transform(transformed, dset):
+    """Merge transformed spectral variable back into its source dataset.
+
+    Args:
+        - transformed (DataArray, Dataset): Transformed spectral variable(s)
+          calculated from the spectral variable in dset.
+        - dset (Dataset): Source dataset the transform was calculated from.
+
+    Returns:
+        - dsout (Dataset): Dataset with the transformed spectral variable(s)
+          and the non-spectral variables from dset.
+
+    Note:
+        - Variables in dset that depend on the spectral dims (freq, dir) are
+          dropped since their coordinates may no longer be consistent with
+          those in the transformed spectra.
+
+    """
+    if isinstance(transformed, xr.DataArray):
+        transformed = transformed.to_dataset(name=transformed.name or attrs.SPECNAME)
+    drop_dims = [d for d in (attrs.FREQNAME, attrs.DIRNAME) if d in dset.dims]
+    others = dset.drop_dims(drop_dims)
+    others = others.drop_vars(
+        [v for v in others.data_vars if v in transformed.data_vars]
+    )
+    dsout = xr.merge([transformed, others])
+    dsout.attrs = dset.attrs
+    return dsout
+
+
 def load_function(module_name, func_name, prefix=None):
     """Returns a function object from string.
 

{wavespectra-4.6.0 → wavespectra-4.7.0}/wavespectra/input/swan.py

@@ -176,8 +176,11 @@ def read_swans(
 
     Args:
         - fileglob (str, list): glob pattern specifying files to read.
-        - ndays (float): number of days to keep from each file, choose None to
-          keep entire period.
+        - ndays (float): number of days to keep from the start of each file,
+          choose None to keep the entire period. Useful for constructing
+          continuous timeseries from files of overlapping forecast cycles.
+          Reading stops once the period is read and an exception is raised
+          if a file does not span the requested period.
         - int_freq (ndarray, bool): frequency array for interpolating onto:
 
             - ndarray: 1d array specifying frequencies to interpolate onto.
@@ -192,8 +195,11 @@ def read_swans(
             - False: No interpolation performed in direction space.
 
         - dirorder (bool): if True ensures directions are sorted.
-        - ntimes (int): use it to read only specific number of times, useful
-          for checking headers only.
+        - ntimes (int): number of times to read from the start of each file,
+          choose None to read all times. Reading stops early once the records
+          are read which makes it useful for quickly inspecting the first few
+          records of large files. Can be combined with ndays, whichever limit
+          is reached first stops the reading.
 
     Returns:
         - dset (SpecDataset): spectra dataset object read from file with
@@ -207,6 +213,8 @@ def read_swans(
           have same number of sites.
         - Either all or none of the spectra in fileglob must have tabfile
           associated to provide wind/depth data.
+        - Winds and depth from tab files are sliced at the times actually
+          read when the ndays or ntimes options are used.
         - Concatenation is done with numpy arrays for efficiency.
 
     """
@@ -249,16 +257,40 @@ def read_swans(
         freqs = swanfile.freqs
         dirs = swanfile.dirs
 
-        if ntimes is None:
-            spec_list = [s for s in swanfile.readall()]
-        else:
-            spec_list = [swanfile.read() for itime in range(ntimes)]
+        # Read spectra records, stopping early once ntimes records or ndays
+        # worth of times have been read so the rest of the file is not parsed
+        spec_list = []
+        tend = None
+        truncated = False
+        while ntimes is None or len(spec_list) < ntimes:
+            sset = swanfile.read()
+            if not sset:
+                break
+            if ndays is not None:
+                if tend is None:
+                    tend = times[0] + datetime.timedelta(days=ndays)
+                if times[-1] > tend:
+                    # Ignore the time parsed beyond the period and stop reading
+                    times.pop()
+                    truncated = True
+                    break
+            spec_list.append(sset)
+
+        # Ensure file spans the requested period unless ntimes stopped first
+        ntimes_reached = ntimes is not None and len(spec_list) == ntimes
+        if tend is not None and not truncated and not ntimes_reached:
+            if times[-1] < tend:
+                raise OSError(
+                    "Times in %s does not extend for %0.2f days" % (filename, ndays)
+                )
 
-        # Read tab files for winds / depth
+        # Read tab files for winds / depth, slicing at the times actually read
+        tab = pd.DataFrame()
         if swanfile.is_tab:
             try:
                 tab = read_tab(swanfile.tabfile).rename(columns={"dep": attrs.DEPNAME})
-                if len(swanfile.times) == tab.index.size:
+                if isinstance(times, list) and set(times).issubset(tab.index):
+                    tab = tab.loc[times]
                     if "X-wsp" in tab and "Y-wsp" in tab:
                         tab[attrs.WSPDNAME], tab[attrs.WDIRNAME] = uv_to_spddir(
                             tab["X-wsp"], tab["Y-wsp"], coming_from=True
@@ -277,23 +309,10 @@ def read_swans(
                     )
                 ]
             except Exception as exc:
+                tab = pd.DataFrame()
                 warnings.warn(
                     f"Cannot parse depth and winds from {swanfile.tabfile}:\n{exc}"
                 )
-        else:
-            tab = pd.DataFrame()
-
-        # Shrinking times
-        if ndays is not None:
-            tend = times[0] + datetime.timedelta(days=ndays)
-            if tend > times[-1]:
-                raise OSError(
-                    "Times in %s does not extend for %0.2f days" % (filename, ndays)
-                )
-            iend = times.index(min(times, key=lambda d: abs(d - tend)))
-            times = times[0 : iend + 1]
-            spec_list = spec_list[0 : iend + 1]
-            tab = tab.loc[times[0] : tend] if tab is not None else tab
 
         spec_list = flatten_list(spec_list, [])
 

{wavespectra-4.6.0 → wavespectra-4.7.0}/wavespectra/input/triaxys.py

@@ -10,7 +10,7 @@ import xarray as xr
 
 from wavespectra.specdataset import SpecDataset
 from wavespectra.core.attributes import attrs, set_spec_attributes
-from wavespectra.core.utils import interp_spec
+from wavespectra.core.utils import interp_spec, dataset_from_transform
 
 
 def read_triaxys(
@@ -47,7 +47,9 @@ def read_triaxys(
         dset = dset.assign_coords({"dir": dirs + magnetic_variation})
         # Regrid
         if regrid_dir:
-            dset = dset.spec.interp(dir=dirs)
+            dset = dataset_from_transform(
+                dset[attrs.SPECNAME].spec.interp(dir=dirs), dset
+            )
     return dset
 
 

{wavespectra-4.6.0 → wavespectra-4.7.0}/wavespectra/output/octopus.py

@@ -107,7 +107,7 @@ def to_octopus(
                 dset.spec.momf(mom=1).rename("momf1"),
                 dset.spec.momf(mom=2).rename("momf2"),
                 dset.spec.momd(mom=0)[0].rename("momd"),
-                dset.spec.to_energy(),
+                dset[attrs.SPECNAME].spec.to_energy(),
                 (dset.efth.spec.df * dset.spec.momd(mom=0)[0])
                 .transpose(attrs.TIMENAME, attrs.SITENAME, attrs.FREQNAME)
                 .rename("fSpec"),