foxes 1.3__py3-none-any.whl → 1.5__py3-none-any.whl

foxes/input/states/dataset_states.py

@@ -0,0 +1,710 @@
+import numpy as np
+import pandas as pd
+import xarray as xr
+from copy import copy, deepcopy
+
+from foxes.core import States, get_engine
+from foxes.utils import import_module
+from foxes.data import STATES, StaticData
+import foxes.variables as FV
+import foxes.constants as FC
+from foxes.config import config, get_input_path
+
+
+def _read_nc_file(
+    fpath,
+    coords,
+    vars,
+    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.sizes:
+            raise KeyError(
+                f"Missing coordinate '{c}' in file {fpath}, got: {list(data.sizes.keys())}"
+            )
+    if minimal:
+        return data[coords[0]].to_numpy()
+    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)
+        assert min(data.sizes.values()) > 0, (
+            f"States: No data in file {fpath}, isel={isel}, sel={sel}, resulting sizes={data.sizes}"
+        )
+        return data
+
+
+class DatasetStates(States):
+    """
+    Abstract base class for heterogeneous ambient states that
+    are based on data from NetCDF files or an xarray Dataset.
+
+    Attributes
+    ----------
+    data_source: str or xarray.Dataset
+        The data or the file search pattern, should end with
+        suffix '.nc'. One or many files.
+    ovars: list of str
+        The output variables
+    var2ncvar: dict
+        Mapping from variable names to variable names
+        in the nc file
+    fixed_vars: dict
+        Uniform values for output variables, instead
+        of reading from data
+    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
+    time_format: str
+        The datetime parsing format string
+    sel: dict, optional
+        Subset selection via xr.Dataset.sel()
+    isel: dict, optional
+        Subset selection via xr.Dataset.isel()
+    weight_factor: float
+        The factor to multiply the weights with
+
+    :group: input.states
+
+    """
+
+    def __init__(
+        self,
+        data_source,
+        output_vars,
+        var2ncvar={},
+        fixed_vars={},
+        load_mode="preload",
+        time_format="%Y-%m-%d_%H:%M:%S",
+        sel=None,
+        isel=None,
+        weight_factor=None,
+        **kwargs,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        data_source: str or xarray.Dataset
+            The data or the file search pattern, should end with
+            suffix '.nc'. One or many files.
+        output_vars: list of str
+            The output variables
+        var2ncvar: dict, optional
+            Mapping from variable names to variable names
+            in the nc file
+        fixed_vars: dict, optional
+            Uniform values for output variables, instead
+            of reading from data
+        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
+        time_format: str
+            The datetime parsing format string
+        sel: dict, optional
+            Subset selection via xr.Dataset.sel()
+        isel: dict, optional
+            Subset selection via xr.Dataset.isel()
+        weight_factor: float, optional
+            The factor to multiply the weights with
+        kwargs: dict, optional
+            Additional arguments for the base class
+
+        """
+        super().__init__(**kwargs)
+
+        self.ovars = list(output_vars)
+        self.fixed_vars = fixed_vars
+        self.load_mode = load_mode
+        self.var2ncvar = var2ncvar
+        self.time_format = time_format
+        self.sel = sel
+        self.isel = isel
+        self.weight_factor = weight_factor
+
+        self._N = None
+        self._inds = None
+        self.__data_source = data_source
+
+    @property
+    def data_source(self):
+        """
+        The data source
+
+        Returns
+        -------
+        s: object
+            The data source
+
+        """
+        if self.load_mode in ["preload", "fly"] and self.running:
+            raise ValueError(
+                f"States '{self.name}': Cannot access data_source while running for load mode '{self.load_mode}'"
+            )
+        return self.__data_source
+
+    def _read_ds(self, ds, cmap, variables, verbosity=0):
+        """
+        Helper function for _get_data, extracts data from the original Dataset.
+
+        Parameters
+        ----------
+        ds: xarray.Dataset
+            The Dataset to read data from
+        cmap: dict
+            A mapping from foxes variable names to Dataset dimension names
+        variables: list of str
+            The variables to extract from the Dataset
+        verbosity: int
+            The verbosity level, 0 = silent
+
+        Returns
+        -------
+        coords: dict
+            keys: Foxes variable names, values: 1D coordinate value arrays
+        data: dict
+            The extracted data, keys are variable names,
+            values are tuples (dims, data_array)
+            where dims is a tuple of dimension names and
+            data_array is a numpy.ndarray with the data values
+
+        """
+        data = {}
+        for v in variables:
+            w = self.var2ncvar.get(v, v)
+            if w in ds.data_vars:
+                d = ds[w]
+                i = [d.dims.index(c) for c in cmap.values() if c in d.dims]
+                assert len(i) == len(d.dims), (
+                    f"States '{self.name}': Variable '{w}' has dimensions {d.dims}, but not all of them are in the coordinate mapping {cmap}"
+                )
+                dms = tuple([v for v, c in cmap.items() if c in d.dims])
+                j = list(range(len(i)))
+                if i == j:
+                    data[v] = (dms, d.to_numpy())
+                elif len(i) == 2:
+                    data[v] = (dms, np.swapaxes(d.to_numpy(), 0, 1))
+                else:
+                    data[v] = (dms, np.moveaxis(d.to_numpy(), i, j))
+            else:
+                raise KeyError(
+                    f"States '{self.name}': Variable '{w}' not found in data source '{self.data_source}', available variables: {list(ds.data_vars)}"
+                )
+
+        coords = {v: ds[c].to_numpy() for v, c in cmap.items() if c in ds.coords}
+
+        if verbosity > 1:
+            if len(coords):
+                print(f"\n{self.name}: Coordinate ranges")
+                for c, d in coords.items():
+                    print(f"  {c}: {np.min(d)} --> {np.max(d)}")
+            print(f"\n{self.name}: Data ranges")
+            for v, d in data.items():
+                nn = np.sum(np.isnan(d))
+                print(
+                    f"  {v}: {np.nanmin(d)} --> {np.nanmax(d)}, nans: {nn} ({100 * nn / len(d.flat):.2f}%)"
+                )
+
+        return coords, data
+
+    def _get_data(self, ds, cmap, variables, verbosity=0):
+        """
+        Gets the data from the Dataset and prepares it for calculations.
+
+        Parameters
+        ----------
+        ds: xarray.Dataset
+            The Dataset to read data from
+        cmap: dict
+            A mapping from foxes variable names to Dataset dimension names
+        variables: list of str
+            The variables to extract from the Dataset
+        verbosity: int
+            The verbosity level, 0 = silent
+
+        Returns
+        -------
+        coords: dict
+            keys: Foxes variable names, values: 1D coordinate value arrays
+        data: dict
+            The extracted data, keys are dimension tuples,
+            values are tuples (DATA key, variables, data_array)
+            where DATA key is the name in the mdata object,
+            variables is a list of variable names, and
+            data_array is a numpy.ndarray with the data values,
+            the last dimension corresponds to the variables
+        weights: numpy.ndarray or None
+            The weights array, if only state dependent, otherwise
+            weights are among data. Shape: (n_states,)
+
+        """
+        coords, data0 = self._read_ds(ds, cmap, variables, verbosity=verbosity)
+
+        weights = None
+        if FV.WEIGHT in variables:
+            assert FV.WEIGHT in data0, (
+                f"States '{self.name}': Missing weights variable '{FV.WEIGHT}' in data, found {sorted(list(data0.keys()))}"
+            )
+            if self.weight_factor is not None:
+                data0[FV.WEIGHT][1] *= self.weight_factor
+            if data0[FV.WEIGHT][0] == (FC.STATE,):
+                weights = data0.pop(FV.WEIGHT)[1]
+
+        data = {}  # dim: [DATA key, variables, data array]
+        for v, (dims, d) in data0.items():
+            if dims not in data:
+                i = len(data)
+                data[dims] = [self.var(f"data{i}"), [], []]
+            data[dims][1].append(v)
+            data[dims][2].append(d)
+        for dims in data.keys():
+            data[dims][2] = np.stack(data[dims][2], axis=-1)
+        data = {
+            tuple(list(dims) + [f"vars{i}"]): d
+            for i, (dims, d) in enumerate(data.items())
+        }
+        return coords, data, weights
+
+    def _preload(self, algo, cmap, bounds_extra_space, verbosity=0):
+        """Helper function for preloading data."""
+
+        assert FC.STATE in cmap, (
+            f"States '{self.name}': States coordinate '{FC.STATE}' not in cmap {cmap}"
+        )
+        states_coord = cmap[FC.STATE]
+
+        if not isinstance(self.data_source, xr.Dataset):
+            # 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 bounds_extra_space is not None:
+                assert FV.X in cmap, (
+                    f"States '{self.name}': x coordinate '{FV.X}' not in cmap {cmap}"
+                )
+                assert FV.Y in cmap, (
+                    f"States '{self.name}': y coordinate '{FV.Y}' not in cmap {cmap}"
+                )
+
+                # if bounds and 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=bounds_extra_space, algo=algo
+                )
+                if verbosity > 0:
+                    print(
+                        f"States '{self.name}': Restricting to bounds {xy_min} - {xy_max}"
+                    )
+                if self.sel is None:
+                    self.sel = {}
+                self.sel.update(
+                    {
+                        cmap[FV.X]: slice(xy_min[0], xy_max[1]),
+                        cmap[FV.Y]: 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)))
+            coords = list(cmap.values())
+            vars = [self.var2ncvar.get(v, v) for v in self.variables]
+            self.__data_source = get_engine().map(
+                _read_nc_file,
+                files,
+                coords=coords,
+                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
+                if len(self.__data_source) == 1:
+                    self.__data_source = self.__data_source[0]
+                else:
+                    self.__data_source = xr.concat(
+                        self.__data_source,
+                        dim=states_coord,
+                        coords="minimal",
+                        data_vars="minimal",
+                        compat="equals",
+                        join="exact",
+                        combine_attrs="drop",
+                    )
+                if self.load_mode == "preload":
+                    self.__data_source.load()
+                self._inds = self.__data_source[states_coord].to_numpy()
+                self._N = len(self._inds)
+
+            elif self.load_mode == "fly":
+                self._inds = 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)
+
+            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()
+
+        # given data is already Dataset:
+        else:
+            self._inds = self.data_source[states_coord].to_numpy()
+            self._N = len(self._inds)
+
+        return self.__data_source
+
+    def load_data(
+        self,
+        algo,
+        cmap,
+        variables,
+        bounds_extra_space=None,
+        verbosity=0,
+    ):
+        """
+        Load and/or create all model data that is subject to chunking.
+
+        Such data should not be stored under self, for memory reasons. The
+        data returned here will automatically be chunked and then provided
+        as part of the mdata object during calculations.
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The calculation algorithm
+        cmap: dict
+            A mapping from foxes variable names to Dataset dimension names
+        variables: list of str
+            The variables to extract from the Dataset
+        bounds_extra_space: float, optional
+            The extra space in meters to add to the horizontal wind farm bounds
+        verbosity: int
+            The verbosity level, 0 = silent
+
+        Returns
+        -------
+        idata: dict
+            The dict has exactly two entries: `data_vars`,
+            a dict with entries `name_str -> (dim_tuple, data_ndarray)`;
+            and `coords`, a dict with entries `dim_name_str -> dim_array`
+
+        """
+        # preload data:
+        self._preload(algo, cmap, bounds_extra_space, verbosity=verbosity)
+
+        idata = super().load_data(algo, verbosity)
+
+        if self.load_mode == "preload":
+            self._coords, data, w = self._get_data(
+                self.data_source, cmap, variables, verbosity
+            )
+
+            if FC.STATE in self._coords:
+                idata["coords"][FC.STATE] = self._coords.pop(FC.STATE)
+            else:
+                del idata["coords"][FC.STATE]
+            if w is not None:
+                idata["data_vars"][FV.WEIGHT] = ((FC.STATE,), w)
+
+            vmap = {FC.STATE: FC.STATE}
+            self._data_state_keys = []
+            self._data_nostate = {}
+            for dims, d in data.items():
+                dms = tuple([vmap.get(c, self.var(c)) for c in dims])
+                if FC.STATE in dims:
+                    self._data_state_keys.append(d[0])
+                    idata["coords"][dms[-1]] = d[1]
+                    idata["data_vars"][d[0]] = (dms, d[2])
+                else:
+                    self._data_nostate[dims] = (d[1], d[2])
+            del data
+
+        return idata
+
+    def set_running(
+        self,
+        algo,
+        data_stash,
+        sel=None,
+        isel=None,
+        verbosity=0,
+    ):
+        """
+        Sets this model status to running, and moves
+        all large data to stash.
+
+        The stashed data will be returned by the
+        unset_running() function after running calculations.
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The calculation algorithm
+        data_stash: dict
+            Large data stash, this function adds data here.
+            Key: model name. Value: dict, large model data
+        sel: dict, optional
+            The subset selection dictionary
+        isel: dict, optional
+            The index subset selection dictionary
+        verbosity: int
+            The verbosity level, 0 = silent
+
+        """
+        super().set_running(algo, data_stash, sel, isel, verbosity)
+
+        data_stash[self.name] = dict(
+            inds=self._inds,
+        )
+        del self._inds
+
+        if self.load_mode == "preload":
+            data_stash[self.name]["data_source"] = self.__data_source
+            del self.__data_source
+
+    def unset_running(
+        self,
+        algo,
+        data_stash,
+        sel=None,
+        isel=None,
+        verbosity=0,
+    ):
+        """
+        Sets this model status to not running, recovering large data
+        from stash
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The calculation algorithm
+        data_stash: dict
+            Large data stash, this function adds data here.
+            Key: model name. Value: dict, large model data
+        sel: dict, optional
+            The subset selection dictionary
+        isel: dict, optional
+            The index subset selection dictionary
+        verbosity: int
+            The verbosity level, 0 = silent
+
+        """
+        super().unset_running(algo, data_stash, sel, isel, verbosity)
+
+        data = data_stash[self.name]
+        self._inds = data.pop("inds")
+
+        if self.load_mode == "preload":
+            self.__data_source = data.pop("data_source")
+
+    def output_point_vars(self, algo):
+        """
+        The variables which are being modified by the model.
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The calculation algorithm
+
+        Returns
+        -------
+        output_vars: list of str
+            The output variable names
+
+        """
+        return self.ovars
+
+    def size(self):
+        """
+        The total number of states.
+
+        Returns
+        -------
+        int:
+            The total number of states
+
+        """
+        return self._N
+
+    def index(self):
+        """
+        The index list
+
+        Returns
+        -------
+        indices: array_like
+            The index labels of states, or None for default integers
+
+        """
+        if self.running:
+            raise ValueError(f"States '{self.name}': Cannot access index while running")
+        return self._inds
+
+    def get_calc_data(self, mdata, cmap, variables):
+        """
+        Gathers data for calculations.
+
+        Call this function from the calculate function of the
+        derived class.
+
+        Parameters
+        ----------
+        mdata: foxes.core.MData
+            The mdata object
+        cmap: dict
+            A mapping from foxes variable names to Dataset dimension names
+        variables: list of str
+            The variables to extract from the Dataset
+
+        Returns
+        -------
+        coords: dict
+            keys: Foxes variable names, values: 1D coordinate value arrays
+        data: dict
+            The extracted data, keys are dimension tuples,
+            values are tuples (DATA key, variables, data_array)
+            where DATA key is the name in the mdata object,
+            variables is a list of variable names, and
+            data_array is a numpy.ndarray with the data values,
+            the last dimension corresponds to the variables
+        weights: numpy.ndarray or None
+            The weights array, if only state dependent, otherwise
+            weights are among data. Shape: (n_states,)
+
+        """
+        # prepare
+        assert FC.STATE in cmap, (
+            f"States '{self.name}': States coordinate '{FC.STATE}' not in cmap {cmap}"
+        )
+        states_coord = cmap[FC.STATE]
+        n_states = mdata.n_states
+
+        # case preload
+        if self.load_mode == "preload":
+            coords = self._coords
+            weights = mdata[FV.WEIGHT] if FV.WEIGHT in mdata else None
+            data = deepcopy(self._data_nostate)
+            for DATA in self._data_state_keys:
+                dims = mdata.dims[DATA]
+                vrs = mdata[dims[-1]].tolist()
+                dms = tuple(
+                    [self.unvar(c) if c != FC.STATE else FC.STATE for c in dims[:-1]]
+                    + [dims[-1]]
+                )
+                data[dms] = (vrs, mdata[DATA].copy())
+
+        # case lazy
+        elif self.load_mode == "lazy":
+            i0 = mdata.states_i0(counter=True)
+            s = slice(i0, i0 + n_states)
+            ds = self.data_source.isel({states_coord: s}).load()
+            coords, data, weights = self._get_data(ds, cmap, variables, verbosity=0)
+            data = {dims: (d[1], d[2]) for dims, d in data.items()}
+            del ds
+
+        # case fly
+        elif self.load_mode == "fly":
+            vars = [self.var2ncvar.get(v, v) for v in variables]
+            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) if self.isel is not None else {}
+                        isel[states_coord] = slice(a, b)
+
+                        data.append(
+                            _read_nc_file(
+                                fpath,
+                                coords=list(cmap.values()),
+                                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)}"
+            )
+            if len(data) == 1:
+                data = data[0]
+            else:
+                data = xr.concat(
+                    data,
+                    dim=states_coord,
+                    data_vars="minimal",
+                    coords="minimal",
+                    compat="override",
+                    join="exact",
+                    combine_attrs="drop",
+                )
+            coords, data, weights = self._get_data(data, cmap, variables, verbosity=0)
+            data = {dims: (d[1], d[2]) for dims, d in data.items()}
+
+        else:
+            raise KeyError(
+                f"States '{self.name}': Unknown load_mode '{self.load_mode}', choices: preload, lazy, fly"
+            )
+
+        return coords, data, weights