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

foxes/core/engine.py

@@ -0,0 +1,640 @@
+import os
+import numpy as np
+from abc import ABC, abstractmethod
+from tqdm import tqdm
+from xarray import Dataset
+
+from foxes.core import MData, FData, TData
+from foxes.utils import all_subclasses
+import foxes.constants as FC
+
+__global_engine_data__ = dict(
+    engine=None,
+)
+
+
+class Engine(ABC):
+    """
+    Abstract base clas for foxes calculation engines
+
+    Attributes
+    ----------
+    chunk_size_states: int
+        The size of a states chunk
+    chunk_size_points: int
+        The size of a points chunk
+    n_procs: int, optional
+        The number of processes to be used,
+        or None for automatic
+    verbosity: int
+        The verbosity level, 0 = silent
+
+    :group: core
+
+    """
+
+    def __init__(
+        self,
+        chunk_size_states=None,
+        chunk_size_points=None,
+        n_procs=None,
+        verbosity=1,
+    ):
+        """
+        Constructor.
+
+        Parameters
+        ----------
+        chunk_size_states: int, optional
+            The size of a states chunk
+        chunk_size_points: int, optional
+            The size of a points chunk
+        n_procs: int, optional
+            The number of processes to be used,
+            or None for automatic
+        verbosity: int
+            The verbosity level, 0 = silent
+
+        """
+        self.chunk_size_states = chunk_size_states
+        self.chunk_size_points = chunk_size_points
+        try:
+            self.n_procs = n_procs if n_procs is not None else os.process_cpu_count()
+        except AttributeError:
+            self.n_procs = os.cpu_count()
+        self.verbosity = verbosity
+        self.__initialized = False
+        self.__entered = False
+
+    def __repr__(self):
+        s = f"n_procs={self.n_procs}, chunk_size_states={self.chunk_size_states}, chunk_size_points={self.chunk_size_points}"
+        return f"{type(self).__name__}({s})"
+
+    def __enter__(self):
+        if self.__entered:
+            raise ValueError(f"Enter called for already entered engine")
+        self.__entered = True
+        if not self.initialized:
+            self.initialize()
+        return self
+
+    def __exit__(self, *exit_args):
+        if not self.__entered:
+            raise ValueError(f"Exit called for not entered engine")
+        self.__entered = False
+        if self.initialized:
+            self.finalize(*exit_args)
+
+    def __del__(self):
+        if self.initialized:
+            self.finalize()
+
+    @property
+    def entered(self):
+        """
+        Flag that this model has been entered.
+
+        Returns
+        -------
+        flag: bool
+            True if the model has been entered.
+
+        """
+        return self.__entered
+
+    @property
+    def initialized(self):
+        """
+        Initialization flag.
+
+        Returns
+        -------
+        flag: bool
+            True if the model has been initialized.
+
+        """
+        return self.__initialized
+
+    def initialize(self):
+        """
+        Initializes the engine.
+        """
+        if not self.entered:
+            self.__enter__()
+        elif not self.initialized:
+            if get_engine(error=False, default=False) is not None:
+                raise ValueError(
+                    f"Cannot initialize engine '{type(self).__name__}', since engine already set to '{type(get_engine()).__name__}'"
+                )
+            global __global_engine_data__
+            __global_engine_data__["engine"] = self
+            self.__initialized = True
+
+    def finalize(self, type=None, value=None, traceback=None):
+        """
+        Finalizes the engine.
+
+        Parameters
+        ----------
+        type: object, optional
+            Dummy argument for the exit function
+        value: object, optional
+            Dummy argument for the exit function
+        traceback: object, optional
+            Dummy argument for the exit function
+
+        """
+        if self.entered:
+            self.__exit__(type, value, traceback)
+        elif self.initialized:
+            global __global_engine_data__
+            __global_engine_data__["engine"] = None
+            self.__initialized = False
+
+    def print(self, *args, level=1, **kwargs):
+        """Prints based on verbosity"""
+        if self.verbosity >= level:
+            print(*args, **kwargs)
+
+    @property
+    def loop_dims(self):
+        """
+        Gets the loop dimensions (possibly chunked)
+
+        Returns
+        -------
+        dims: list of str
+            The loop dimensions (possibly chunked)
+
+        """
+        if self.chunk_size_states is None and self.chunk_size_states is None:
+            return []
+        elif self.chunk_size_states is None:
+            return [FC.TARGET]
+        elif self.chunk_size_points is None:
+            return [FC.STATE]
+        else:
+            return [FC.STATE, FC.TARGET]
+
+    def select_subsets(self, *datasets, sel=None, isel=None):
+        """
+        Takes subsets of datasets
+
+        Parameters
+        ----------
+        datasets: tuple
+            The xarray.Dataset or xarray.Dataarray objects
+        sel: dict, optional
+            The selection dictionary
+        isel: dict, optional
+            The index selection dictionary
+
+        Returns
+        -------
+        subsets: list
+            The subsets of the input data
+
+        """
+        if sel is not None:
+            new_datasets = []
+            for data in datasets:
+                if data is not None:
+                    s = {c: u for c, u in sel.items() if c in data.coords}
+                    new_datasets.append(data.sel(s) if len(s) else data)
+                else:
+                    new_datasets.append(data)
+            datasets = new_datasets
+
+        if isel is not None:
+            new_datasets = []
+            for data in datasets:
+                if data is not None:
+                    s = {c: u for c, u in isel.items() if c in data.coords}
+                    new_datasets.append(data.isel(s) if len(s) else data)
+                else:
+                    new_datasets.append(data)
+            datasets = new_datasets
+
+        return datasets
+
+    def calc_chunk_sizes(self, n_states, n_targets=1):
+        """
+        Computes the sizes of states and points chunks
+
+        Parameters
+        ----------
+        n_states: int
+            The number of states
+        n_targets: int
+            The number of point targets
+
+        Returns
+        -------
+        chunk_sizes_states: numpy.ndarray
+            The sizes of all states chunks, shape: (n_chunks_states,)
+        chunk_sizes_targets: numpy.ndarray
+            The sizes of all targets chunks, shape: (n_chunks_targets,)
+
+        """
+        # determine states chunks:
+        if self.chunk_size_states is None:
+            n_chunks_states = min(self.n_procs, n_states)
+            chunk_size_states = max(int(n_states / self.n_procs), 1)
+        else:
+            chunk_size_states = min(n_states, self.chunk_size_states)
+            n_chunks_states = max(int(n_states / chunk_size_states), 1)
+            if int(n_states / n_chunks_states) > chunk_size_states:
+                n_chunks_states += 1
+                chunk_size_states = int(n_states / n_chunks_states)
+        chunk_sizes_states = np.full(n_chunks_states, chunk_size_states)
+        extra = n_states - n_chunks_states * chunk_size_states
+        if extra > 0:
+            chunk_sizes_states[-extra:] += 1
+
+        s = np.sum(chunk_sizes_states)
+        assert (
+            s == n_states
+        ), f"States count mismatch: Expecting {n_states}, chunks sum is {s}. Chunks: {[int(c) for c in chunk_sizes_states]}"
+
+        # determine points chunks:
+        chunk_sizes_targets = [n_targets]
+        if n_targets > 1:
+            if self.chunk_size_points is None:
+                if n_chunks_states == 1:
+                    n_chunks_targets = min(self.n_procs, n_targets)
+                    chunk_size_targets = max(int(n_targets / self.n_procs), 1)
+                else:
+                    chunk_size_targets = n_targets
+                    n_chunks_targets = 1
+            else:
+                chunk_size_targets = min(n_targets, self.chunk_size_points)
+                n_chunks_targets = max(int(n_targets / chunk_size_targets), 1)
+            if int(n_targets / n_chunks_targets) > chunk_size_targets:
+                n_chunks_targets += 1
+                chunk_size_targets = int(n_targets / n_chunks_targets)
+            chunk_sizes_targets = np.full(n_chunks_targets, chunk_size_targets)
+            extra = n_targets - n_chunks_targets * chunk_size_targets
+            if extra > 0:
+                chunk_sizes_targets[-extra:] += 1
+
+            s = np.sum(chunk_sizes_targets)
+            assert (
+                s == n_targets
+            ), f"Targets count mismatch: Expecting {n_targets}, chunks sum is {s}. Chunks: {[int(c) for c in chunk_sizes_targets]}"
+
+        return chunk_sizes_states, chunk_sizes_targets
+
+    def get_chunk_input_data(
+        self,
+        algo,
+        model_data,
+        farm_data,
+        point_data,
+        states_i0_i1,
+        targets_i0_i1,
+        out_vars,
+    ):
+        """
+        Extracts the data for a single chunk calculation
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The algorithm object
+        model_data: xarray.Dataset
+            The initial model data
+        farm_data: xarray.Dataset
+            The initial farm data
+        point_data: xarray.Dataset
+            The initial point data
+        states_i0_i1: tuple
+            The (start, end) values of the states
+        targets_i0_i1: tuple
+            The (start, end) values of the targets
+        out_vars: list of str
+            Names of the output variables
+
+        Returns
+        -------
+        data: list of foxes.core.Data
+            Either [mdata, fdata] or [mdata, fdata, tdata]
+
+        """
+        # prepare:
+        i0_states, i1_states = states_i0_i1
+        i0_targets, i1_targets = targets_i0_i1
+        s_states = np.s_[i0_states:i1_states]
+        s_targets = np.s_[i0_targets:i1_targets]
+
+        # create mdata:
+        mdata = MData.from_dataset(
+            model_data,
+            s_states=s_states,
+            loop_dims=[FC.STATE],
+            states_i0=i0_states,
+            copy=True,
+        )
+
+        # create fdata:
+        if point_data is None:
+
+            def cb(data, dims):
+                n_states = i1_states - i0_states
+                for o in set(out_vars).difference(data.keys()):
+                    data[o] = np.full(
+                        (n_states, algo.n_turbines), np.nan, dtype=FC.DTYPE
+                    )
+                    dims[o] = (FC.STATE, FC.TURBINE)
+
+        else:
+            cb = None
+        fdata = FData.from_dataset(
+            farm_data,
+            mdata=mdata,
+            s_states=s_states,
+            callback=cb,
+            loop_dims=[FC.STATE],
+            states_i0=i0_states,
+            copy=True,
+        )
+
+        # create tdata:
+        tdata = None
+        if point_data is not None:
+
+            def cb(data, dims):
+                n_states = i1_states - i0_states
+                n_targets = i1_targets - i0_targets
+                for o in set(out_vars).difference(data.keys()):
+                    data[o] = np.full((n_states, n_targets, 1), np.nan, dtype=FC.DTYPE)
+                    dims[o] = (FC.STATE, FC.TARGET, FC.TPOINT)
+
+            tdata = TData.from_dataset(
+                point_data,
+                mdata=mdata,
+                s_states=s_states,
+                s_targets=s_targets,
+                callback=cb,
+                loop_dims=[FC.STATE, FC.TARGET],
+                states_i0=i0_states,
+                copy=True,
+            )
+
+        return [d for d in [mdata, fdata, tdata] if d is not None]
+
+    def combine_results(
+        self,
+        algo,
+        results,
+        model_data,
+        out_vars,
+        out_coords,
+        n_chunks_states,
+        n_chunks_targets,
+        goal_data,
+        iterative,
+    ):
+        """
+        Combines chunk results into final Dataset
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The algorithm object
+        results: dict
+            The results from the chunk calculations,
+            key: (chunki_states, chunki_targets),
+            value: dict with numpy.ndarray values
+        model_data: xarray.Dataset
+            The initial model data
+        out_vars: list of str
+            Names of the output variables
+        out_coords: list of str
+            Names of the output coordinates
+        n_chunks_states: int
+            The number of states chunks
+        n_chunks_targets: int
+            The number of targets chunks
+        goal_data: foxes.core.Data
+            Either fdata or tdata
+        iterative: bool
+            Flag for use within the iterative algorithm
+
+        Returns
+        -------
+        ds: xarray.Dataset
+            The final results dataset
+
+        """
+        self.print(f"{type(self).__name__}: Combining results", level=2)
+        pbar = tqdm(total=len(out_vars)) if self.verbosity > 1 else None
+        data_vars = {}
+        for v in out_vars:
+            if v in results[(0, 0)][0]:
+                data_vars[v] = [out_coords, []]
+
+                if n_chunks_targets == 1:
+                    alls = 0
+                    for chunki_states in range(n_chunks_states):
+                        r, cstore = results[(chunki_states, 0)]
+                        data_vars[v][1].append(r[v])
+                        alls += data_vars[v][1][-1].shape[0]
+                        if iterative:
+                            for k, c in cstore.items():
+                                if k in algo.chunk_store:
+                                    algo.chunk_store[k].update(c)
+                                else:
+                                    algo.chunk_store[k] = c
+                else:
+                    for chunki_states in range(n_chunks_states):
+                        tres = []
+                        for chunki_points in range(n_chunks_targets):
+                            r, cstore = results[(chunki_states, chunki_points)]
+                            tres.append(r[v])
+                            if iterative:
+                                for k, c in cstore.items():
+                                    if k in algo.chunk_store:
+                                        algo.chunk_store[k].update(c)
+                                    else:
+                                        algo.chunk_store[k] = c
+                        data_vars[v][1].append(np.concatenate(tres, axis=1))
+                    del tres
+                del r, cstore
+                data_vars[v][1] = np.concatenate(data_vars[v][1], axis=0)
+            else:
+                data_vars[v] = (goal_data[v].dims, goal_data[v].to_numpy())
+
+            if pbar is not None:
+                pbar.update()
+        del results
+        if pbar is not None:
+            pbar.close()
+
+        # if not iterative or algo.final_iteration:
+        #    algo.reset_chunk_store()
+
+        coords = {}
+        if FC.STATE in out_coords and FC.STATE in model_data.coords:
+            coords[FC.STATE] = model_data[FC.STATE].to_numpy()
+
+        return Dataset(
+            coords=coords,
+            data_vars={v: tuple(d) for v, d in data_vars.items()},
+        )
+
+    @abstractmethod
+    def run_calculation(self, algo, model, model_data, farm_data, point_data=None):
+        """
+        Runs the model calculation
+
+        Parameters
+        ----------
+        algo: foxes.core.Algorithm
+            The algorithm object
+        model: foxes.core.DataCalcModel
+            The model that whose calculate function
+            should be run
+        model_data: xarray.Dataset
+            The initial model data
+        farm_data: xarray.Dataset
+            The initial farm data
+        point_data: xarray.Dataset, optional
+            The initial point data
+
+        Returns
+        -------
+        results: xarray.Dataset
+            The model results
+
+        """
+        n_states = model_data.sizes[FC.STATE]
+        if point_data is None:
+            self.print(
+                f"{type(self).__name__}: Calculating {n_states} states for {algo.n_turbines} turbines"
+            )
+        else:
+            self.print(
+                f"{type(self).__name__}: Calculating data at {point_data.sizes[FC.TARGET]} points for {n_states} states"
+            )
+        if not self.initialized:
+            raise ValueError(f"Engine '{type(self).__name__}' not initialized")
+        if not model.initialized:
+            raise ValueError(f"Model '{model.name}' not initialized")
+
+    @classmethod
+    def new(cls, engine_type, *args, **kwargs):
+        """
+        Run-time engine factory.
+
+        Parameters
+        ----------
+        engine_type: str
+            The selected derived class name
+        args: tuple, optional
+            Additional parameters for constructor
+        kwargs: dict, optional
+            Additional parameters for constructor
+
+        """
+
+        if engine_type is None:
+            return None
+        else:
+            engine_type = dict(
+                default="DefaultEngine",
+                threads="ThreadsEngine",
+                process="ProcessEngine",
+                xarray="XArrayEngine",
+                dask="DaskEngine",
+                multiprocess="MultiprocessEngine",
+                local_cluster="LocalClusterEngine",
+                slurm_cluster="SlurmClusterEngine",
+                mpi="MPIEngine",
+                ray="RayEngine",
+                numpy="NumpyEngine",
+                single="SingleChunkEngine",
+            ).get(engine_type, engine_type)
+
+        allc = all_subclasses(cls)
+        found = engine_type in [scls.__name__ for scls in allc]
+
+        if found:
+            for scls in allc:
+                if scls.__name__ == engine_type:
+                    return scls(*args, **kwargs)
+
+        else:
+            estr = "engine type '{}' is not defined, available types are \n {}".format(
+                engine_type, sorted([i.__name__ for i in allc])
+            )
+            raise KeyError(estr)
+
+
+def get_engine(error=True, default=True):
+    """
+    Gets the global calculation engine
+
+    Parameters
+    ----------
+    error: bool
+        Flag for raising ValueError if no
+        engine is found
+    default: bool or dict or Engine
+        Set default engine if not set yet
+
+    Returns
+    -------
+    engine: foxes.core.Engine
+        The foxes calculation engine
+
+    :group: core
+
+    """
+    engine = __global_engine_data__["engine"]
+    if engine is None:
+        if isinstance(default, dict):
+            engine = Engine.new(**default)
+            print(f"Selecting default engine '{engine}'")
+            engine.initialize()
+            return engine
+        elif isinstance(default, Engine):
+            print(f"Selecting default engine '{default}'")
+            default.initialize()
+            return default
+        elif isinstance(default, bool) and default:
+            engine = Engine.new(
+                engine_type="DefaultEngine", chunk_size_points=20000, verbosity=0
+            )
+            print(f"Selecting '{engine}'")
+            engine.initialize()
+            return engine
+        elif error:
+            raise ValueError("Engine not found.")
+    return engine
+
+
+def has_engine():
+    """
+    Flag that checks if engine has been set
+
+    Returns
+    -------
+    flag: bool
+        True if engine has been set
+
+    :group: core
+
+    """
+    return __global_engine_data__["engine"] is not None
+
+
+def reset_engine():
+    """
+    Resets the global calculation engine
+
+    :group: core
+
+    """
+    engine = get_engine(error=False, default=False)
+    if engine is not None:
+        engine.finalize(type=None, value=None, traceback=None)