foxes 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

foxes/input/states/field_data_nc.py

@@ -1,14 +1,55 @@
 import numpy as np
 import pandas as pd
 import xarray as xr
+from copy import copy
 from scipy.interpolate import interpn
 
-from foxes.core import States
+from foxes.core import States, get_engine
 from foxes.utils import wd2uv, uv2wd, import_module
 from foxes.data import STATES, StaticData
 import foxes.variables as FV
 import foxes.constants as FC
-from foxes.config import config, get_path
+from foxes.config import config, get_input_path
+
+
+def _read_nc_file(
+    fpath,
+    coords,
+    vars,
+    weight_var,
+    nc_engine,
+    sel,
+    isel,
+    minimal,
+):
+    """Helper function for nc file reading"""
+    data = xr.open_dataset(fpath, engine=nc_engine)
+    for c in coords:
+        if c is not None and c not in data:
+            raise KeyError(
+                f"Missing coordinate '{c}' in file {fpath}, got: {list(data.coords.keys())}"
+            )
+    if minimal:
+        weights = None
+        if weight_var is not None:
+            if weight_var not in data.data_vars:
+                raise KeyError(
+                    f"Missing weight var '{weight_var}' in file {fpath}, found: {list(data.data_vars.keys())}"
+                )
+            if data[weight_var].dims != (coords[0],):
+                raise ValueError(
+                    f"Wrong dimensions for variable '{weight_var}' in file {fpath}. Expecting {(coords[0],)}, got {data[weight_var].dims}"
+                )
+            weights = data[weight_var].to_numpy()
+        return data[coords[0]].to_numpy(), weights
+    else:
+        data = data[vars]
+        data.attrs = {}
+        if isel is not None and len(isel):
+            data = data.isel(**isel)
+        if sel is not None and len(sel):
+            data = data.sel(**sel)
+        return data
 
 
 class FieldDataNC(States):
@@ -37,9 +78,13 @@ class FieldDataNC(States):
         The y coordinate name in the data
     h_coord: str
         The height coordinate name in the data
-    pre_load: bool
-        Flag for loading all data into memory during
-        initialization
+    load_mode: str
+        The load mode, choices: preload, lazy, fly.
+        preload loads all data during initialization,
+        lazy lazy-loads the data using dask, and fly
+        reads only states index and weights during initialization
+        and then opens the relevant files again within
+        the chunk calculation
     weight_ncvar: str
         Name of the weight data variable in the nc file(s)
     bounds_error: bool
@@ -52,6 +97,9 @@ class FieldDataNC(States):
         Linearly interpolate nan values
     interpn_pars: dict, optional
         Additional parameters for scipy.interpolate.interpn
+    bounds_extra_space: float or str
+        The extra space, either float in m,
+        or str for units of D, e.g. '2.5D'
 
     :group: input.states
 
@@ -67,13 +115,13 @@ class FieldDataNC(States):
         x_coord="UTMX",
         y_coord="UTMY",
         h_coord="height",
-        pre_load=True,
+        load_mode="preload",
         weight_ncvar=None,
         time_format="%Y-%m-%d_%H:%M:%S",
         sel=None,
         isel=None,
         interp_nans=False,
-        verbosity=1,
+        bounds_extra_space=1000,
         **interpn_pars,
     ):
         """
@@ -98,11 +146,15 @@ class FieldDataNC(States):
             The x coordinate name in the data
         y_coord: str
             The y coordinate name in the data
-        h_coord: str
+        h_coord: str, optional
             The height coordinate name in the data
-        pre_load: bool
-            Flag for loading all data into memory during
-            initialization
+        load_mode: str
+            The load mode, choices: preload, lazy, fly.
+            preload loads all data during initialization,
+            lazy lazy-loads the data using dask, and fly
+            reads only states index and weights during initialization
+            and then opens the relevant files again within
+            the chunk calculation
         weight_ncvar: str, optional
             Name of the weight data variable in the nc file(s)
         time_format: str
@@ -113,8 +165,9 @@ class FieldDataNC(States):
             Subset selection via xr.Dataset.isel()
         interp_nans: bool
             Linearly interpolate nan values
-        verbosity: int
-            Verbosity level for pre_load file reading
+        bounds_extra_space: float or str, optional
+            The extra space, either float in m,
+            or str for units of D, e.g. '2.5D'
         interpn_pars: dict, optional
             Additional parameters for scipy.interpolate.interpn
 
@@ -128,12 +181,13 @@ class FieldDataNC(States):
         self.y_coord = y_coord
         self.h_coord = h_coord
         self.weight_ncvar = weight_ncvar
-        self.pre_load = pre_load
+        self.load_mode = load_mode
         self.time_format = time_format
-        self.sel = sel
-        self.isel = isel
+        self.sel = sel if sel is not None else {}
+        self.isel = isel if isel is not None else {}
         self.interpn_pars = interpn_pars
         self.interp_nans = interp_nans
+        self.bounds_extra_space = bounds_extra_space
 
         self.var2ncvar = {
             v: var2ncvar.get(v, v) for v in output_vars if v not in fixed_vars
@@ -145,57 +199,6 @@ class FieldDataNC(States):
         self.__weights = None
         self.__inds = None
 
-        # pre-load file reading:
-        if not isinstance(self.data_source, xr.Dataset):
-            if "*" in str(self.data_source):
-                pass
-            else:
-                self.__data_source = get_path(self.data_source)
-                if not self.data_source.is_file():
-                    self.__data_source = StaticData().get_file_path(
-                        STATES, self.data_source.name, check_raw=False
-                    )
-            if verbosity:
-                if pre_load:
-                    print(
-                        f"States '{self.name}': Reading data from '{self.data_source}'"
-                    )
-                else:
-                    print(
-                        f"States '{self.name}': Reading index from '{self.data_source}'"
-                    )
-
-            def _read_ds():
-                fpath = get_path(self.data_source)
-                if fpath.is_file():
-                    return xr.open_dataset(fpath)
-                else:
-                    # try to read multiple files, needs dask:
-                    try:
-                        return xr.open_mfdataset(
-                            str(fpath),
-                            parallel=False,
-                            concat_dim=self.states_coord,
-                            combine="nested",
-                            data_vars="minimal",
-                            coords="minimal",
-                            compat="override",
-                        )
-                    except ValueError as e:
-                        import_module("dask", hint="pip install dask")
-                        raise e
-
-            self.__data_source = _read_ds()
-
-        if sel is not None:
-            self.__data_source = self.data_source.sel(self.sel)
-        if isel is not None:
-            self.__data_source = self.data_source.isel(self.isel)
-        if pre_load:
-            self.__data_source.load()
-
-        self._get_inds(self.data_source)
-
     @property
     def data_source(self):
         """
@@ -207,56 +210,34 @@ class FieldDataNC(States):
             The data source
 
         """
-        if self.pre_load and self.running:
+        if self.load_mode in ["preload", "fly"] and self.running:
             raise ValueError(
-                f"States '{self.name}': Cannot access data_source while running"
+                f"States '{self.name}': Cannot access data_source while running for load mode '{self.load_mode}'"
             )
         return self.__data_source
 
-    def _get_inds(self, ds):
-        """
-        Helper function for index and weights
-        reading
-        """
-        for c in [self.states_coord, self.x_coord, self.y_coord, self.h_coord]:
-            if not c in ds:
-                raise KeyError(
-                    f"States '{self.name}': Missing coordinate '{c}' in data"
-                )
-
-        self.__inds = ds[self.states_coord].to_numpy()
-        if self.time_format is not None:
-            self.__inds = pd.to_datetime(
-                self.__inds, format=self.time_format
-            ).to_numpy()
-        self._N = len(self.__inds)
-
-        if self.weight_ncvar is not None:
-            self.__weights = ds[self.weight_ncvar].to_numpy()
-
-        for v in self.ovars:
-            if v in self.var2ncvar:
-                ncv = self.var2ncvar[v]
-                if not ncv in ds:
-                    raise KeyError(
-                        f"States '{self.name}': nc variable '{ncv}' not found in data, found: {sorted(list(ds.keys()))}"
-                    )
-            elif v not in self.fixed_vars:
-                raise ValueError(
-                    f"States '{self.name}': Variable '{v}' neither found in var2ncvar not in fixed_vars"
-                )
-
-    def _get_data(self, ds, verbosity):
+    def _get_data(self, ds, coords, verbosity):
         """
         Helper function for data extraction
         """
-        x = ds[self.x_coord].to_numpy()
-        y = ds[self.y_coord].to_numpy()
-        h = ds[self.h_coord].to_numpy()
-        n_x = len(x)
-        n_y = len(y)
-        n_h = len(h)
-        n_sts = ds.sizes[self.states_coord]
+        for ci, c in enumerate(coords):
+            found = False
+            if c is not None:
+                for v in ds.data_vars.values():
+                    if c in v.dims:
+                        found = True
+                        break
+            if not found:
+                coords[ci] = None
+
+        dlst = []
+        for c in coords:
+            if c is not None:
+                dlst.append(np.atleast_1d(ds[c].to_numpy()))
+            else:
+                dlst.append(np.array([0], dtype=config.dtype_double))
+        sts, h, y, x = dlst
+        n_sts, n_h, n_y, n_x = [len(u) for u in dlst]
 
         cor_shxy = (self.states_coord, self.h_coord, self.x_coord, self.y_coord)
         cor_shyx = (self.states_coord, self.h_coord, self.y_coord, self.x_coord)
@@ -278,8 +259,13 @@ class FieldDataNC(States):
             elif ds[ncv].dims == cor_s:
                 vars_s.append(v)
             else:
+                expc = [
+                    c
+                    for c in [cor_shxy, cor_shyx, cor_sxy, cor_syx, cor_sh, cor_s]
+                    if None not in c
+                ]
                 raise ValueError(
-                    f"States '{self.name}': Wrong coordinate order for variable '{ncv}': Found {ds[ncv].dims}, expecting {cor_shxy}, {cor_shyx}, {cor_sxy}, {cor_syx}, {cor_sh} or {cor_s}"
+                    f"States '{self.name}': Wrong coordinates for variable '{ncv}': Found {ds[ncv].dims}, expecting one of {expc}"
                 )
 
         data = np.zeros(
@@ -321,7 +307,7 @@ class FieldDataNC(States):
                     f"  {v}: {np.nanmin(d)} --> {np.nanmax(d)}, nans: {nn} ({100*nn/len(d.flat):.2f}%)"
                 )
 
-        return data
+        return sts, h, y, x, data
 
     def output_point_vars(self, algo):
         """
@@ -364,13 +350,129 @@ class FieldDataNC(States):
 
         """
 
-        if (FV.WS in self.ovars and FV.WD not in self.ovars) or (
-            FV.WS not in self.ovars and FV.WD in self.ovars
-        ):
-            raise KeyError(
-                f"States '{self.name}': Missing '{FV.WS}' or '{FV.WD}' in output variables {self.ovars}"
+        # pre-load file reading:
+        coords = [self.states_coord, self.h_coord, self.y_coord, self.x_coord]
+        if not isinstance(self.data_source, xr.Dataset):
+
+            # check variables:
+            for v in self.ovars:
+                if v not in self.var2ncvar and v not in self.fixed_vars:
+                    raise ValueError(
+                        f"States '{self.name}': Variable '{v}' neither found in var2ncvar not in fixed_vars"
+                    )
+            if (FV.WS in self.ovars and FV.WD not in self.ovars) or (
+                FV.WS not in self.ovars and FV.WD in self.ovars
+            ):
+                raise KeyError(
+                    f"States '{self.name}': Missing '{FV.WS}' or '{FV.WD}' in output variables {self.ovars}"
+                )
+
+            # check static data:
+            fpath = get_input_path(self.data_source)
+            if "*" not in str(self.data_source):
+                if not fpath.is_file():
+                    fpath = StaticData().get_file_path(
+                        STATES, fpath.name, check_raw=False
+                    )
+
+            # find bounds:
+            if self.x_coord is not None and self.x_coord not in self.sel:
+                xy_min, xy_max = algo.farm.get_xy_bounds(
+                    extra_space=self.bounds_extra_space, algo=algo
+                )
+                if verbosity > 0:
+                    print(
+                        f"States '{self.name}': Restricting to bounds {xy_min} - {xy_max}"
+                    )
+                self.sel.update(
+                    {
+                        self.x_coord: slice(xy_min[0], xy_max[1]),
+                        self.y_coord: slice(xy_min[1], xy_max[1]),
+                    }
+                )
+
+            # read files:
+            if verbosity > 0:
+                if self.load_mode == "preload":
+                    print(
+                        f"States '{self.name}': Reading data from '{self.data_source}'"
+                    )
+                elif self.load_mode == "lazy":
+                    print(
+                        f"States '{self.name}': Reading header from '{self.data_source}'"
+                    )
+                else:
+                    print(
+                        f"States '{self.name}': Reading states from '{self.data_source}'"
+                    )
+            files = sorted(list(fpath.resolve().parent.glob(fpath.name)))
+            vars = list(self.var2ncvar.values())
+            if self.weight_ncvar is not None:
+                vars += [self.weight_ncvar]
+            self.__data_source = get_engine().map(
+                _read_nc_file,
+                files,
+                coords=coords,
+                weight_var=self.weight_ncvar,
+                vars=vars,
+                nc_engine=config.nc_engine,
+                isel=self.isel,
+                sel=self.sel,
+                minimal=self.load_mode == "fly",
             )
 
+            if self.load_mode in ["preload", "lazy"]:
+
+                if self.load_mode == "lazy":
+                    try:
+                        self.__data_source = [ds.chunk() for ds in self.__data_source]
+                    except (ModuleNotFoundError, ValueError) as e:
+                        import_module("dask")
+                        raise e
+                self.__data_source = xr.concat(
+                    self.__data_source,
+                    dim=self.states_coord,
+                    coords="minimal",
+                    data_vars="minimal",
+                    compat="equals",
+                    join="exact",
+                    combine_attrs="drop",
+                )
+                if self.load_mode == "preload":
+                    self.__data_source.load()
+                if self.weight_ncvar is not None:
+                    self.__weights = self.__data_source[self.weight_ncvar].to_numpy()
+                self.__inds = self.__data_source[self.states_coord].to_numpy()
+                self._N = len(self.__inds)
+
+            elif self.load_mode == "fly":
+                self.__inds, weights = zip(*self.__data_source)
+                self.__data_source = fpath
+                self._files_maxi = {f: len(inds) for f, inds in zip(files, self.__inds)}
+                self.__inds = np.concatenate(self.__inds, axis=0)
+                self._N = len(self.__inds)
+                if weights[0] is not None:
+                    self.__weights = np.zeros(
+                        (self._N, algo.n_turbines), dtype=config.dtype_double
+                    )
+                    self.__weights[:] = np.concatenate(weights, axis=0)[:, None]
+                    del weights
+
+            else:
+                raise KeyError(
+                    f"States '{self.name}': Unknown load_mode '{self.load_mode}', choices: preload, lazy, fly"
+                )
+
+            if self.time_format is not None:
+                self.__inds = pd.to_datetime(
+                    self.__inds, format=self.time_format
+                ).to_numpy()
+
+            if self.__weights is None:
+                self.__weights = np.full(
+                    (self._N, algo.n_turbines), 1.0 / self._N, dtype=config.dtype_double
+                )
+
         # ensure WD and WS get the first two slots of data:
         self._dkys = {}
         if FV.WS in self.ovars:
@@ -382,34 +484,25 @@ class FieldDataNC(States):
                 self._dkys[v] = len(self._dkys)
         self._n_dvars = len(self._dkys)
 
-        if self.__weights is None:
-            self.__weights = np.full(
-                (self._N, algo.n_turbines), 1.0 / self._N, dtype=config.dtype_double
-            )
-
         idata = super().load_data(algo, verbosity)
 
-        if self.pre_load:
+        if self.load_mode == "preload":
             self.X = self.var(FV.X)
             self.Y = self.var(FV.Y)
             self.H = self.var(FV.H)
             self.VARS = self.var("vars")
             self.DATA = self.var("data")
 
-            ds = self.data_source
+            __, h, y, x, data = self._get_data(self.data_source, coords, verbosity)
+            self._prl_coords = coords
 
-            h = ds[self.h_coord].to_numpy()
-            y = ds[self.y_coord].to_numpy()
-            x = ds[self.x_coord].to_numpy()
-            v = list(self._dkys.keys())
             coos = (FC.STATE, self.H, self.Y, self.X, self.VARS)
-            data = self._get_data(ds, verbosity)
             data = (coos, data)
 
             idata["coords"][self.H] = h
             idata["coords"][self.Y] = y
             idata["coords"][self.X] = x
-            idata["coords"][self.VARS] = v
+            idata["coords"][self.VARS] = list(self._dkys.keys())
             idata["data_vars"][self.DATA] = data
 
         return idata
@@ -452,7 +545,7 @@ class FieldDataNC(States):
         )
         del self.__weights, self.__inds
 
-        if self.pre_load:
+        if self.load_mode == "preload":
             data_stash[self.name]["data_source"] = self.__data_source
             del self.__data_source
 
@@ -489,7 +582,7 @@ class FieldDataNC(States):
         self.__weights = data.pop("weights")
         self.__inds = data.pop("inds")
 
-        if self.pre_load:
+        if self.load_mode == "preload":
             self.__data_source = data.pop("data_source")
 
     def size(self):
@@ -589,26 +682,73 @@ class FieldDataNC(States):
         points = tdata[FC.TARGETS].reshape(n_states, n_targets * n_tpoints, 3)
         n_pts = points.shape[1]
         n_states = fdata.n_states
+        coords = [self.states_coord, self.h_coord, self.y_coord, self.x_coord]
 
-        # pick pre-loaded data:
-        if self.pre_load:
+        # case preload:
+        if self.load_mode == "preload":
             x = mdata[self.X]
             y = mdata[self.Y]
             h = mdata[self.H]
             data = mdata[self.DATA].copy()
+            coords = self._prl_coords
 
-        # read data for this chunk:
-        else:
+        # case lazy:
+        elif self.load_mode == "lazy":
             i0 = mdata.states_i0(counter=True)
             s = slice(i0, i0 + n_states)
             ds = self.data_source.isel({self.states_coord: s}).load()
+            __, h, y, x, data = self._get_data(ds, coords, verbosity=0)
+            del ds
 
-            x = ds[self.x_coord].to_numpy()
-            y = ds[self.y_coord].to_numpy()
-            h = ds[self.h_coord].to_numpy()
-            data = self._get_data(ds, verbosity=0)
+        # case fly:
+        elif self.load_mode == "fly":
+            vars = list(self.var2ncvar.values())
+            if self.weight_ncvar is not None:
+                vars += [self.weight_ncvar]
+
+            i0 = mdata.states_i0(counter=True)
+            i1 = i0 + n_states
+            j0 = 0
+            data = []
+            for fpath, n in self._files_maxi.items():
+                if i0 < j0:
+                    break
+                else:
+                    j1 = j0 + n
+                    if i0 < j1:
+                        a = i0 - j0
+                        b = min(i1, j1) - j0
+                        isel = copy(self.isel)
+                        isel[self.states_coord] = slice(a, b)
+
+                        data.append(
+                            _read_nc_file(
+                                fpath,
+                                coords=coords,
+                                weight_var=self.weight_ncvar,
+                                vars=vars,
+                                nc_engine=config.nc_engine,
+                                isel=isel,
+                                sel=self.sel,
+                                minimal=False,
+                            )
+                        )
+
+                        i0 += b - a
+                    j0 = j1
+
+            assert (
+                i0 == i1
+            ), f"States '{self.name}': Missing states for load_mode '{self.load_mode}': (i0, i1) = {(i0, i1)}"
+
+            data = xr.concat(data, dim=self.states_coord)
+            __, h, y, x, data = self._get_data(data, coords, verbosity=0)
+
+        else:
+            raise KeyError(
+                f"States '{self.name}': Unknown load_mode '{self.load_mode}', choices: preload, lazy, fly"
+            )
 
-            del ds
         n_h = len(h)
         n_y = len(y)
         n_x = len(x)
@@ -635,6 +775,11 @@ class FieldDataNC(States):
         pts = np.flip(pts, axis=1)
         gvars = (sts, h, y, x)
 
+        # reset None coordinate data, since that should not be interpolated:
+        for i, (c, g) in enumerate(zip(coords, gvars)):
+            if c is None:
+                pts[..., i] = g[0]
+
         # interpolate nan values:
         if self.interp_nans and np.any(np.isnan(data)):
             df = pd.DataFrame(

foxes/input/states/multi_height.py

@@ -6,7 +6,7 @@ from scipy.interpolate import interp1d
 from foxes.core import States
 from foxes.utils import PandasFileHelper
 from foxes.data import STATES
-from foxes.config import config, get_path
+from foxes.config import config, get_input_path
 from foxes.utils import wd2uv, uv2wd
 import foxes.variables as FV
 import foxes.constants as FC
@@ -198,7 +198,7 @@ class MultiHeightStates(States):
 
         """
         if not isinstance(self.data_source, pd.DataFrame):
-            self._data_source = get_path(self.data_source)
+            self._data_source = get_input_path(self.data_source)
             if not self.data_source.is_file():
                 if verbosity:
                     print(
@@ -625,7 +625,7 @@ class MultiHeightNCStates(MultiHeightStates):
 
         """
         if not isinstance(self.data_source, Dataset):
-            self._data_source = get_path(self.data_source)
+            self._data_source = get_input_path(self.data_source)
             if not self.data_source.is_file():
                 if verbosity:
                     print(

foxes/input/states/states_table.py

@@ -6,7 +6,7 @@ from pathlib import Path
 from foxes.core import States, VerticalProfile
 from foxes.utils import PandasFileHelper, read_tab_file
 from foxes.data import STATES
-from foxes.config import config, get_path
+from foxes.config import config, get_input_path
 import foxes.variables as FV
 import foxes.constants as FC
 
@@ -215,7 +215,7 @@ class StatesTable(States):
             data = self.data_source
             isorg = True
         else:
-            self._data_source = get_path(self.data_source)
+            self._data_source = get_input_path(self.data_source)
             if not self.data_source.is_file():
                 if verbosity:
                     print(
@@ -550,7 +550,7 @@ class TabStates(StatesTable):
         """
         if self.data_source is None:
             if self.__tab_data is None:
-                self.__tab_source = get_path(self.__tab_source)
+                self.__tab_source = get_input_path(self.__tab_source)
                 if not self.__tab_source.is_file():
                     if verbosity:
                         print(