foxes 0.8.2__py3-none-any.whl → 1.1.0.2__py3-none-any.whl

foxes/core/data.py

@@ -25,7 +25,14 @@ class Data(Dict):
 
     """
 
-    def __init__(self, data, dims, loop_dims, name="data"):
+    def __init__(
+        self,
+        data,
+        dims,
+        loop_dims,
+        states_i0=None,
+        name="data",
+    ):
         """
         Constructor.
 
@@ -39,6 +46,8 @@ class Data(Dict):
         loop_dims: array_like of str
             List of the loop dimensions during xarray's
             `apply_ufunc` calculations
+        states_i0: int, optional
+            The index of the first state
         name: str
             The data container name
 
@@ -49,6 +58,8 @@ class Data(Dict):
         self.dims = dims
         self.loop_dims = loop_dims
 
+        self.__states_i0 = states_i0
+
         self.sizes = {}
         for v, d in data.items():
             self._run_entry_checks(v, d, dims[v])
@@ -81,7 +92,7 @@ class Data(Dict):
         """
         return self.sizes[FC.TURBINE] if FC.TURBINE in self.sizes else None
 
-    def states_i0(self, counter=False, algo=None):
+    def states_i0(self, counter=False):
         """
         Get the state counter for first state in chunk
 
@@ -89,8 +100,6 @@ class Data(Dict):
         ----------
         counter: bool
             Return the state counter instead of the index
-        algo: foxes.core.Algorithm, optional
-            The algorithm, required for state counter
 
         Returns
         -------
@@ -102,9 +111,9 @@ class Data(Dict):
         if FC.STATE not in self:
             return None
         elif counter:
-            if algo is None:
-                raise KeyError(f"{self.name}: Missing algo for deducing state counter")
-            return np.argwhere(algo.states.index() == self[FC.STATE][0])[0][0]
+            if self.__states_i0 is None:
+                raise KeyError(f"Data '{self.name}': states_i0 requested but not set")
+            return self.__states_i0
         else:
             return self[FC.STATE][0]
 
@@ -129,6 +138,15 @@ class Data(Dict):
 
             self.dims[FV.TXYH] = tuple(list(dims[FV.X]) + [FC.XYH])
 
+        allc = set()
+        for dms in self.dims.values():
+            if dms is not None:
+                allc.update(dms)
+        allc = allc.difference(set(data.keys()))
+        for c in allc.intersection(self.sizes.keys()):
+            data[c] = np.arange(self.sizes[c])
+            dims[c] = (c,)
+
     def _run_entry_checks(self, name, data, dims):
         """Run entry checks on new data"""
         # remove axes of size 1, added by dask for extra loop dimensions:
@@ -165,12 +183,14 @@ class Data(Dict):
         self._run_entry_checks(name, data, dims)
         self._auto_update()
 
-    def get_slice(self, s, dim_map={}, name=None, keep=True):
+    def get_slice(self, variables, s, dim_map={}, name=None):
         """
         Get a slice of data.
 
         Parameters
         ----------
+        variables: list of str
+            The variable list that corresponds to s
         s: slice
             The slice
         dim_map: dict
@@ -178,9 +198,6 @@ class Data(Dict):
             If not found, same dimensions are assumed.
         name: str, optional
             The name of the data object
-        keep: bool
-            Keep non-matching fields as they are, else
-            throw them out
 
         Returns
         -------
@@ -188,20 +205,103 @@ class Data(Dict):
             The new data object, containing slices
 
         """
+        if not isinstance(variables, (list, tuple, np.ndarray)):
+            variables = [variables]
+        if not isinstance(s, (list, tuple, np.ndarray)):
+            s = [s]
+
         data = {}
         dims = {}
         for v in self.keys():
-            try:
-                d = self.dims[v]
-                data[v] = self[v][s]
-                dims[v] = dim_map.get(d, d)
-            except IndexError:
-                if keep:
-                    data[v] = self[v]
-                    dims[v] = self.dims[v]
+            d = self.dims[v]
+            if d is not None:
+                hs = tuple(
+                    [s[variables.index(w)] if w in variables else np.s_[:] for w in d]
+                )
+                data[v] = self[v][hs]
+                dims[v] = (
+                    tuple([dim_map.get(dd, dd) for dd in d]) if len(dim_map) else d
+                )
         if name is None:
             name = self.name
-        return type(self)(data, dims, loop_dims=self.loop_dims, name=name)
+        if FC.STATE in variables and self.__states_i0 is not None:
+            i0 = self.states_i0(counter=True)
+            a = s[variables.index(FC.STATE)]
+            sts = np.arange(i0, i0 + self.n_states)[a]
+            if len(sts) == 1:
+                states_i0 = sts[0]
+            elif np.all(sts == np.arange(sts[0], sts[0] + len(sts))):
+                states_i0 = sts[0]
+            else:
+                raise ValueError(
+                    f"Cannot determine states_i0 for states slices {a}, leading to selection {list(sts)}"
+                )
+        else:
+            states_i0 = None
+        return type(self)(
+            data, dims, loop_dims=self.loop_dims, name=name, states_i0=states_i0
+        )
+
+    @classmethod
+    def from_dataset(cls, ds, *args, callback=None, s_states=None, copy=True, **kwargs):
+        """
+        Create Data object from a dataset
+
+        Parameters
+        ----------
+        ds: xarray.Dataset
+            The dataset
+        args: tuple, optional
+            Additional parameters for the constructor
+        callback: Function, optional
+            Function f(data, dims) that manipulates
+            the data and dims dicts before construction
+        s_states: slice, optional
+            Slice object for states
+        copy: bool
+            Flag for copying data
+        kwargs: dict, optional
+            Additional parameters for the constructor
+
+        Returns
+        -------
+        data: Data
+            The data object
+
+        """
+        data = {}
+        dims = {}
+
+        for c, d in ds.coords.items():
+            if c == FC.STATE:
+                s = np.s_[:] if s_states is None else s_states
+                data[c] = d.to_numpy()[s].copy() if copy else d.to_numpy()[s]
+            else:
+                data[c] = d.to_numpy().copy() if copy else d.to_numpy()
+            dims[c] = d.dims
+
+        n_states = None
+        for v, d in ds.data_vars.items():
+            if FC.STATE in d.dims:
+                if d.dims[0] != FC.STATE:
+                    raise ValueError(
+                        f"Expecting coordinate '{FC.STATE}' at position 0 for data variable '{v}', got {d.dims}"
+                    )
+                n_states = len(d.to_numpy())
+                s = np.s_[:] if s_states is None else s_states
+                data[v] = d.to_numpy()[s].copy() if copy else d.to_numpy()[s]
+            else:
+                data[v] = d.to_numpy().copy() if copy else d.to_numpy()
+            dims[v] = d.dims
+
+        if callback is not None:
+            callback(data, dims)
+
+        if FC.STATE not in data and s_states is not None and n_states is not None:
+            data[FC.STATE] = np.arange(n_states)[s_states]
+            dims[FC.STATE] = (FC.STATE,)
+
+        return cls(*args, data=data, dims=dims, **kwargs)
 
 
 class MData(Data):
@@ -284,6 +384,46 @@ class FData(Data):
                         f"FData '{self.name}': Missing '{x}' in sizes, got {sorted(list(self.sizes.keys()))}"
                     )
 
+    @classmethod
+    def from_dataset(cls, ds, *args, mdata=None, callback=None, **kwargs):
+        """
+        Create Data object from a dataset
+
+        Parameters
+        ----------
+        ds: xarray.Dataset
+            The dataset
+        args: tuple, optional
+            Additional parameters for the constructor
+        mdata: MData, optional
+            The mdata object
+        callback: Function, optional
+            Function f(data, dims) that manipulates
+            the data and dims dicts before construction
+        kwargs: dict, optional
+            Additional parameters for the constructor
+
+        Returns
+        -------
+        data: Data
+            The data object
+
+        """
+        if mdata is None:
+            return super().from_dataset(ds, *args, callback=callback, **kwargs)
+        else:
+
+            def cb(data, dims):
+                if FC.STATE not in data:
+                    data[FC.STATE] = mdata[FC.STATE]
+                    dims[FC.STATE] = mdata.dims[FC.STATE]
+                    data[FV.WEIGHT] = mdata[FV.WEIGHT]
+                    dims[FV.WEIGHT] = mdata.dims[FV.WEIGHT]
+                if callback is not None:
+                    callback(data, dims)
+
+            return super().from_dataset(ds, *args, callback=cb, **kwargs)
+
 
 class TData(Data):
     """
@@ -427,6 +567,21 @@ class TData(Data):
         """
         return np.einsum("stp...,p->st...", self[variable], self[FC.TWEIGHTS])
 
+    def targets_i0(self):
+        """
+        Get the target counter for first target in chunk
+
+        Returns
+        -------
+        int:
+            The target index for first target in chunk
+
+        """
+        if FC.TARGET not in self:
+            return None
+        else:
+            return self[FC.TARGET][0]
+
     @classmethod
     def from_points(
         cls,
@@ -515,3 +670,80 @@ class TData(Data):
         data[FC.TWEIGHTS] = tweights
         dims[FC.TWEIGHTS] = (FC.TPOINT,)
         return cls(data, dims, [FC.STATE], name=name, **kwargs)
+
+    @classmethod
+    def from_dataset(
+        cls,
+        ds,
+        *args,
+        s_targets=None,
+        mdata=None,
+        callback=None,
+        **kwargs,
+    ):
+        """
+        Create Data object from a dataset
+
+        Parameters
+        ----------
+        ds: xarray.Dataset
+            The dataset
+        args: tuple, optional
+            Additional parameters for the constructor
+        s_targets: slice, optional
+            Slice object for targets
+        mdata: MData, optional
+            The mdata object
+        callback: Function, optional
+            Function f(data, dims) that manipulates
+            the data and dims dicts before construction
+        kwargs: dict, optional
+            Additional parameters for the constructor
+
+        Returns
+        -------
+        data: Data
+            The data object
+
+        """
+        if mdata is None:
+            cb0 = callback
+        else:
+
+            def cb_mdata(data, dims):
+                if FC.STATE not in data:
+                    data[FC.STATE] = mdata[FC.STATE]
+                    dims[FC.STATE] = mdata.dims[FC.STATE]
+                if callback is not None:
+                    callback(data, dims)
+
+            cb0 = cb_mdata
+
+        if s_targets is None:
+            cb1 = cb0
+        else:
+
+            def cb_targets(data, dims):
+                if FC.TARGET not in data:
+                    data[FC.TARGET] = np.arange(ds.sizes[FC.TARGET])
+                    dims[FC.TARGET] = (FC.TARGET,)
+                for v, d in data.items():
+                    if FC.TARGET in dims[v]:
+                        if dims[v] == (FC.TARGET,):
+                            data[v] = d[s_targets].copy()
+                        elif len(dims[v]) < 3 or dims[v][:3] != (
+                            FC.STATE,
+                            FC.TARGET,
+                            FC.TPOINT,
+                        ):
+                            raise ValueError(
+                                f"Expecting coordinates '{ (FC.STATE, FC.TARGET, FC.TPOINT)}' at positions 0-2 for data variable '{v}', got {dims[v]}"
+                            )
+                        else:
+                            data[v] = d[:, s_targets]
+                if cb0 is not None:
+                    cb0(data, dims)
+
+            cb1 = cb_targets
+
+        return super().from_dataset(ds, *args, callback=cb1, **kwargs)

foxes/core/data_calc_model.py

@@ -1,19 +1,11 @@
-import numpy as np
-import xarray as xr
 from abc import abstractmethod
-from dask.distributed import progress
-from dask.diagnostics import ProgressBar
 
 from .model import Model
-from .data import MData, FData, TData
-from foxes.utils.runners import DaskRunner
-import foxes.constants as FC
-import foxes.variables as FV
 
 
 class DataCalcModel(Model):
     """
-    Abstract base class for models with
+    Abstract base class for models
     that run calculation on xarray Dataset
     data.
 
@@ -28,9 +20,22 @@ class DataCalcModel(Model):
 
     """
 
+    @abstractmethod
+    def output_coords(self):
+        """
+        Gets the coordinates of all output arrays
+
+        Returns
+        -------
+        dims: tuple of str
+            The coordinates of all output arrays
+
+        """
+        pass
+
     @abstractmethod
     def calculate(self, algo, *data, **parameters):
-        """ "
+        """
         The main model calculation.
 
         This function is executed on a single chunk of data,
@@ -53,284 +58,3 @@ class DataCalcModel(Model):
 
         """
         pass
-
-    def _wrap_calc(
-        self,
-        *ldata,
-        algo,
-        dvars,
-        lvars,
-        ldims,
-        evars,
-        edims,
-        edata,
-        loop_dims,
-        out_vars,
-        out_dims,
-        calc_pars,
-        init_vars,
-    ):
-        """
-        Wrapper that mitigates between apply_ufunc and `calculate`.
-        """
-        n_prev = len(init_vars)
-        if n_prev:
-            prev = ldata[:n_prev]
-            ldata = ldata[n_prev:]
-
-        # reconstruct original data:
-        data = []
-        for i, hvars in enumerate(dvars):
-            v2l = {v: lvars.index(v) for v in hvars if v in lvars}
-            v2e = {v: evars.index(v) for v in hvars if v in evars}
-
-            hdata = {v: ldata[v2l[v]] if v in v2l else edata[v2e[v]] for v in hvars}
-            hdims = {v: ldims[v2l[v]] if v in v2l else edims[v2e[v]] for v in hvars}
-
-            if i == 0:
-                data.append(MData(data=hdata, dims=hdims, loop_dims=loop_dims))
-            elif i == 1:
-                data.append(FData(data=hdata, dims=hdims, loop_dims=loop_dims))
-            elif i == 2:
-                data.append(TData(data=hdata, dims=hdims, loop_dims=loop_dims))
-            else:
-                raise NotImplementedError(
-                    f"Not more than 3 data sets implemented, found {len(dvars)}"
-                )
-
-            del hdata, hdims, v2l, v2e
-
-        # deduce output shape:
-        oshape = []
-        for li, l in enumerate(out_dims):
-            for i, dims in enumerate(ldims):
-                if l in dims:
-                    oshape.append(ldata[i].shape[dims.index(l)])
-                    break
-            if len(oshape) != li + 1:
-                raise ValueError(f"Model '{self.name}': Failed to find loop dimension")
-
-        # add zero output data arrays:
-        odims = {v: tuple(out_dims) for v in out_vars}
-        odata = {
-            v: (
-                np.full(oshape, np.nan, dtype=FC.DTYPE)
-                if v not in init_vars
-                else prev[init_vars.index(v)].copy()
-            )
-            for v in out_vars
-            if v not in data[-1]
-        }
-
-        if len(data) == 1:
-            data.append(FData(odata, odims, loop_dims))
-        else:
-            odata.update(data[-1])
-            odims.update(data[-1].dims)
-            if len(data) == 2:
-                data[-1] = FData(odata, odims, loop_dims)
-            else:
-                data[-1] = TData(odata, odims, loop_dims)
-        del odims, odata
-
-        # link chunk state indices from mdata to fdata and tdata:
-        if FC.STATE in data[0]:
-            for d in data[1:]:
-                d[FC.STATE] = data[0][FC.STATE]
-
-        # link weights from mdata to fdata:
-        if FV.WEIGHT in data[0]:
-            data[1][FV.WEIGHT] = data[0][FV.WEIGHT]
-            data[1].dims[FV.WEIGHT] = data[0].dims[FV.WEIGHT]
-
-        # run model calculation:
-        self.ensure_variables(algo, *data)
-        results = self.calculate(algo, *data, **calc_pars)
-
-        # replace missing results by first input data with matching shape:
-        missing = set(out_vars).difference(results.keys())
-        if len(missing):
-            found = set()
-            for v in missing:
-                for dta in data:
-                    if v in dta and dta[v].shape == tuple(oshape):
-                        results[v] = dta[v]
-                        found.add(v)
-                        break
-            missing -= found
-            if len(missing):
-                raise ValueError(
-                    f"Model '{self.name}': Missing results {list(missing)}, expected shape {oshape}"
-                )
-        del data
-
-        # create output:
-        n_vars = len(out_vars)
-        data = np.zeros(oshape + [n_vars], dtype=FC.DTYPE)
-        for v in out_vars:
-            data[..., out_vars.index(v)] = results[v]
-
-        return data
-
-    def run_calculation(
-        self,
-        algo,
-        *data,
-        out_vars,
-        loop_dims,
-        out_core_vars,
-        initial_results=None,
-        sel=None,
-        isel=None,
-        **calc_pars,
-    ):
-        """
-        Starts the model calculation in parallel, via
-        xarray's `apply_ufunc`.
-
-        Typically this function is called by algorithms.
-
-        Parameters
-        ----------
-        algo: foxes.core.Algorithm
-            The calculation algorithm
-        data: tuple of xarray.Dataset
-            The input data
-        out_vars: list of str
-            The calculation output variables
-        loop_dims: array_like of str
-            List of the loop dimensions during xarray's
-            `apply_ufunc` calculations
-        out_core_vars: list of str
-            The core dimensions of the output data, use
-            `FC.VARS` for variables dimension (required)
-        initial_results: xarray.Dataset, optional
-            Initial results
-        sel: dict, optional
-            Selection of loop_dim variable subset values
-        isel: dict, optional
-            Selection of loop_dim variable subset index values
-        calc_pars: dict, optional
-            Additional arguments for the `calculate` function
-
-        Returns
-        -------
-        results: xarray.Dataset
-            The calculation results
-
-        """
-        # check:
-        if not self.initialized:
-            raise ValueError(
-                f"DataCalcModel '{self.name}': run_calculation called for uninitialized model"
-            )
-
-        # prepare:
-        loopd = set(loop_dims)
-
-        # extract loop-var dependent and independent data:
-        ldata = []
-        lvars = []
-        ldims = []
-        edata = []
-        evars = []
-        edims = []
-        dvars = []
-        ivars = []
-        idims = []
-        if initial_results is not None:
-            ds = initial_results
-            hvarsl = [v for v, d in ds.items() if len(loopd.intersection(d.dims))]
-            ldata += [ds[v] for v in hvarsl]
-            idims += [ds[v].dims for v in hvarsl]
-            ivars += hvarsl
-
-        for ds in data:
-            hvarsl = [v for v, d in ds.items() if len(loopd.intersection(d.dims))]
-            ldata += [ds[v] for v in hvarsl]
-            ldims += [ds[v].dims for v in hvarsl]
-            lvars += hvarsl
-
-            hvarse = [v for v in ds.keys() if v not in hvarsl]
-            edata += [ds[v].values for v in hvarse]
-            edims += [ds[v].dims for v in hvarse]
-            evars += hvarse
-
-            for c, d in ds.coords.items():
-                if c in loopd:
-                    ldata.append(xr.DataArray(data=d.values, coords={c: d}, dims=[c]))
-                    ldims.append((c,))
-                    lvars.append(c)
-                else:
-                    edata.append(d.values)
-                    edims.append((c,))
-                    evars.append(c)
-
-            dvars.append(list(ds.keys()) + list(ds.coords.keys()))
-
-        # subset selection:
-        if sel is not None:
-            nldata = []
-            for ds in ldata:
-                s = {k: v for k, v in sel.items() if k in ds.coords}
-                if len(s):
-                    nldata.append(ds.sel(s))
-            ldata = nldata
-            del nldata
-        if isel is not None:
-            nldata = []
-            for ds in ldata:
-                s = {k: v for k, v in isel.items() if k in ds.coords}
-                if len(s):
-                    nldata.append(ds.isel(s))
-            ldata = nldata
-            del nldata
-
-        # setup dask options:
-        dargs = dict(output_sizes={FC.VARS: len(out_vars)})
-        if FC.TURBINE in loopd and FC.TURBINE not in ldims.values():
-            dargs["output_sizes"][FC.TURBINE] = algo.n_turbines
-        if FC.VARS not in out_core_vars:
-            raise ValueError(
-                f"Model '{self.name}': Expecting '{FC.VARS}' in out_core_vars, got {out_core_vars}"
-            )
-
-        # setup arguments for wrapper function:
-        out_dims = loop_dims + list(set(out_core_vars).difference([FC.VARS]))
-        wargs = dict(
-            algo=algo,
-            dvars=dvars,
-            lvars=lvars,
-            ldims=ldims,
-            evars=evars,
-            edims=edims,
-            edata=edata,
-            loop_dims=loop_dims,
-            out_vars=out_vars,
-            out_dims=out_dims,
-            calc_pars=calc_pars,
-            init_vars=ivars,
-        )
-
-        # run parallel computation:
-        iidims = [[c for c in d if c not in loopd] for d in idims]
-        icdims = [[c for c in d if c not in loopd] for d in ldims]
-        results = xr.apply_ufunc(
-            self._wrap_calc,
-            *ldata,
-            input_core_dims=iidims + icdims,
-            output_core_dims=[out_core_vars],
-            output_dtypes=[FC.DTYPE],
-            dask="parallelized",
-            dask_gufunc_kwargs=dargs,
-            kwargs=wargs,
-        )
-
-        # reorganize results Dataset:
-        results = results.assign_coords({FC.VARS: out_vars}).to_dataset(dim=FC.VARS)
-
-        if DaskRunner.is_distributed() and len(ProgressBar.active):
-            progress(results.persist())
-
-        # update data by calculation results:
-        return results.compute()