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

foxes/engines/default.py

@@ -0,0 +1,75 @@
+import numpy as np
+
+from foxes.core import Engine
+import foxes.constants as FC
+
+
+class DefaultEngine(Engine):
+    """
+    The case size dependent default engine.
+
+    :group: engines
+
+    """
+
+    def run_calculation(
+        self,
+        algo,
+        model,
+        model_data=None,
+        farm_data=None,
+        point_data=None,
+        **kwargs,
+    ):
+        """
+        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, optional
+            The initial model data
+        farm_data: xarray.Dataset, optional
+            The initial farm data
+        point_data: xarray.Dataset, optional
+            The initial point data
+
+        Returns
+        -------
+        results: xarray.Dataset
+            The model results
+
+        """
+        max_n = np.sqrt(self.n_procs) * (500 / algo.n_turbines) ** 1.5
+
+        if (algo.n_states >= max_n) or (
+            point_data is not None
+            and self.chunk_size_points is not None
+            and point_data.sizes[FC.TARGET] > self.chunk_size_points
+        ):
+            ename = "process"
+        else:
+            ename = "single"
+
+        self.print(f"{type(self).__name__}: Selecting engine '{ename}'", level=1)
+
+        self.finalize()
+
+        with Engine.new(
+            ename,
+            n_procs=self.n_procs,
+            chunk_size_states=self.chunk_size_states,
+            chunk_size_points=self.chunk_size_points,
+            verbosity=self.verbosity,
+        ) as e:
+            results = e.run_calculation(
+                algo, model, model_data, farm_data, point_data=point_data, **kwargs
+            )
+
+        self.initialize()
+
+        return results

foxes/engines/futures.py

@@ -0,0 +1,72 @@
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
+
+from .pool import PoolEngine
+
+
+class ThreadsEngine(PoolEngine):
+    """
+    The threads engine for foxes calculations.
+
+    :group: engines
+
+    """
+
+    def _create_pool(self):
+        """Creates the pool"""
+        self._pool = ThreadPoolExecutor(max_workers=self.n_procs)
+
+    def _submit(self, f, *args, **kwargs):
+        """
+        Submits to the pool
+
+        Parameters
+        ----------
+        f: Callable
+            The function f(*args, **kwargs) to be
+            submitted
+        args: tuple, optional
+            Arguments for the function
+        kwargs: dict, optional
+            Arguments for the function
+
+        Returns
+        -------
+        future: object
+            The future object
+
+        """
+        return self._pool.submit(f, *args, **kwargs)
+
+    def _result(self, future):
+        """
+        Waits for result from a future
+
+        Parameters
+        ----------
+        future: object
+            The future
+
+        Returns
+        -------
+        result: object
+            The calculation result
+
+        """
+        return future.result()
+
+    def _shutdown_pool(self):
+        """Shuts down the pool"""
+        self._pool.shutdown()
+
+
+class ProcessEngine(ThreadsEngine):
+    """
+    The processes engine for foxes calculations.
+
+    :group: engines
+
+    """
+
+    def _create_pool(self):
+        """Creates the pool"""
+        self._pool = ProcessPoolExecutor(max_workers=self.n_procs)

foxes/engines/mpi.py

@@ -0,0 +1,38 @@
+from foxes.utils import import_module
+
+from .pool import PoolEngine
+from .futures import ProcessEngine
+
+
+def load_mpi():
+    """On-demand loading of the mpi4py package"""
+    global MPIPoolExecutor
+    MPIPoolExecutor = import_module(
+        "mpi4py.futures", hint="pip install mpi4py"
+    ).MPIPoolExecutor
+
+
+class MPIEngine(ProcessEngine):
+    """
+    The MPI engine for foxes calculations.
+
+    Examples
+    --------
+    Run command, e.g. for 12 processors and a script run.py:
+
+    >>> mpiexec -n 12 -m mpi4py.futures run.py
+
+    :group: engines
+
+    """
+
+    def initialize(self):
+        """
+        Initializes the engine.
+        """
+        load_mpi()
+        PoolEngine.initialize(self)
+
+    def _create_pool(self):
+        """Creates the pool"""
+        self._pool = MPIPoolExecutor(max_workers=self.n_procs)

foxes/engines/multiprocess.py

@@ -0,0 +1,71 @@
+from foxes.utils import import_module
+
+from .pool import PoolEngine
+
+Pool = None
+
+
+def load_multiprocess():
+    """On-demand loading of the multiprocess package"""
+    global Pool
+    if Pool is None:
+        Pool = import_module("multiprocess", hint="pip install multiprocess").Pool
+
+
+class MultiprocessEngine(PoolEngine):
+    """
+    The multiprocessing engine for foxes calculations.
+
+    :group: engines
+
+    """
+
+    def _create_pool(self):
+        """Creates the pool"""
+        load_multiprocess()
+        self._pool = Pool(processes=self.n_procs)
+
+    def _submit(self, f, *args, **kwargs):
+        """
+        Submits to the pool
+
+        Parameters
+        ----------
+        f: Callable
+            The function f(*args, **kwargs) to be
+            submitted
+        args: tuple, optional
+            Arguments for the function
+        kwargs: dict, optional
+            Arguments for the function
+
+        Returns
+        -------
+        future: object
+            The future object
+
+        """
+        return self._pool.apply_async(f, args=args, kwds=kwargs)
+
+    def _result(self, future):
+        """
+        Waits for result from a future
+
+        Parameters
+        ----------
+        future: object
+            The future
+
+        Returns
+        -------
+        result: object
+            The calculation result
+
+        """
+        return future.get()
+
+    def _shutdown_pool(self):
+        """Shuts down the pool"""
+        self._pool.close()
+        self._pool.terminate()
+        self._pool.join()

foxes/engines/numpy.py

@@ -0,0 +1,167 @@
+from tqdm import tqdm
+from xarray import Dataset
+from tqdm import tqdm
+
+from foxes.core import Engine
+import foxes.constants as FC
+
+from .pool import _run
+
+
+class NumpyEngine(Engine):
+    """
+    The numpy engine for foxes calculations.
+
+    :group: engines
+
+    """
+
+    def run_calculation(
+        self,
+        algo,
+        model,
+        model_data=None,
+        farm_data=None,
+        point_data=None,
+        out_vars=[],
+        chunk_store={},
+        sel=None,
+        isel=None,
+        iterative=False,
+        **calc_pars,
+    ):
+        """
+        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
+            The initial point data
+        out_vars: list of str, optional
+            Names of the output variables
+        chunk_store: foxes.utils.Dict
+            The chunk store
+        sel: dict, optional
+            Selection of coordinate subsets
+        isel: dict, optional
+            Selection of coordinate subsets index values
+        iterative: bool
+            Flag for use within the iterative algorithm
+        calc_pars: dict, optional
+            Additional parameters for the model.calculate()
+
+        Returns
+        -------
+        results: xarray.Dataset
+            The model results
+
+        """
+        # subset selection:
+        model_data, farm_data, point_data = self.select_subsets(
+            model_data, farm_data, point_data, sel=sel, isel=isel
+        )
+
+        # basic checks:
+        super().run_calculation(algo, model, model_data, farm_data, point_data)
+
+        # prepare:
+        n_states = model_data.sizes[FC.STATE]
+        out_coords = model.output_coords()
+        coords = {}
+        if FC.STATE in out_coords and FC.STATE in model_data.coords:
+            coords[FC.STATE] = model_data[FC.STATE].to_numpy()
+        if farm_data is None:
+            farm_data = Dataset()
+        goal_data = farm_data if point_data is None else point_data
+
+        # DEBUG objec mem sizes:
+        # from foxes.utils import print_mem
+        # for m in [algo] + model.models:
+        #    print_mem(m, pre_str="MULTIP CHECKING LARGE DATA", min_csize=9999)
+
+        # calculate chunk sizes:
+        n_targets = point_data.sizes[FC.TARGET] if point_data is not None else 0
+        chunk_sizes_states, chunk_sizes_targets = self.calc_chunk_sizes(
+            n_states, n_targets
+        )
+        n_chunks_states = len(chunk_sizes_states)
+        n_chunks_targets = len(chunk_sizes_targets)
+        self.print(
+            f"{type(self).__name__}: Selecting n_chunks_states = {n_chunks_states}, n_chunks_targets = {n_chunks_targets}",
+            level=2,
+        )
+
+        # prepare and submit chunks:
+        n_chunks_all = n_chunks_states * n_chunks_targets
+        self.print(f"{type(self).__name__}: Looping over {n_chunks_all} chunks")
+        pbar = (
+            tqdm(total=n_chunks_all)
+            if self.verbosity > 0 and n_chunks_all > 1
+            else None
+        )
+        results = {}
+        i0_states = 0
+        for chunki_states in range(n_chunks_states):
+            i1_states = i0_states + chunk_sizes_states[chunki_states]
+            i0_targets = 0
+            for chunki_points in range(n_chunks_targets):
+                i1_targets = i0_targets + chunk_sizes_targets[chunki_points]
+
+                i = chunki_states * n_chunks_targets + chunki_points
+
+                # get this chunk's data:
+                data = self.get_chunk_input_data(
+                    algo=algo,
+                    model_data=model_data,
+                    farm_data=farm_data,
+                    point_data=point_data,
+                    states_i0_i1=(i0_states, i1_states),
+                    targets_i0_i1=(i0_targets, i1_targets),
+                    out_vars=out_vars,
+                )
+
+                # submit model calculation:
+                key = (chunki_states, chunki_points)
+                results[key] = _run(
+                    algo,
+                    model,
+                    data,
+                    iterative,
+                    chunk_store,
+                    (i0_states, i0_targets),
+                    **calc_pars,
+                )
+                chunk_store.update(results[key][1])
+                del data
+
+                i0_targets = i1_targets
+
+                if pbar is not None:
+                    pbar.update()
+
+            i0_states = i1_states
+
+        del calc_pars, farm_data, point_data
+        if pbar is not None:
+            pbar.close()
+
+        return self.combine_results(
+            algo=algo,
+            results=results,
+            model_data=model_data,
+            out_vars=out_vars,
+            out_coords=out_coords,
+            n_chunks_states=n_chunks_states,
+            n_chunks_targets=n_chunks_targets,
+            goal_data=goal_data,
+            iterative=iterative,
+        )

foxes/engines/pool.py

@@ -0,0 +1,249 @@
+import xarray as xr
+from abc import abstractmethod
+from tqdm import tqdm
+
+from foxes.core import Engine
+import foxes.constants as FC
+
+
+def _run(algo, model, data, iterative, chunk_store, i0_t0, **cpars):
+    """Helper function for running in a single process"""
+    algo.reset_chunk_store(chunk_store)
+    results = model.calculate(algo, *data, **cpars)
+    chunk_store = algo.reset_chunk_store() if iterative else {}
+    cstore = {i0_t0: chunk_store[i0_t0]} if i0_t0 in chunk_store else {}
+    return results, cstore
+
+
+class PoolEngine(Engine):
+    """
+    Abstract engine for pool type parallelizations.
+
+    :group: engines
+
+    """
+
+    @abstractmethod
+    def _create_pool(self):
+        """Creates the pool"""
+        pass
+
+    @abstractmethod
+    def _submit(self, f, *args, **kwargs):
+        """
+        Submits to the pool
+
+        Parameters
+        ----------
+        f: Callable
+            The function f(*args, **kwargs) to be
+            submitted
+        args: tuple, optional
+            Arguments for the function
+        kwargs: dict, optional
+            Arguments for the function
+
+        Returns
+        -------
+        future: object
+            The future object
+
+        """
+        pass
+
+    @abstractmethod
+    def _result(self, future):
+        """
+        Waits for result from a future
+
+        Parameters
+        ----------
+        future: object
+            The future
+
+        Returns
+        -------
+        result: object
+            The calculation result
+
+        """
+        pass
+
+    @abstractmethod
+    def _shutdown_pool(self):
+        """Shuts down the pool"""
+        pass
+
+    def __enter__(self):
+        self._create_pool()
+        return super().__enter__()
+
+    def __exit__(self, *exit_args):
+        self._shutdown_pool()
+        super().__exit__(*exit_args)
+
+    def run_calculation(
+        self,
+        algo,
+        model,
+        model_data=None,
+        farm_data=None,
+        point_data=None,
+        out_vars=[],
+        chunk_store={},
+        sel=None,
+        isel=None,
+        iterative=False,
+        **calc_pars,
+    ):
+        """
+        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
+            The initial point data
+        out_vars: list of str, optional
+            Names of the output variables
+        chunk_store: foxes.utils.Dict
+            The chunk store
+        sel: dict, optional
+            Selection of coordinate subsets
+        isel: dict, optional
+            Selection of coordinate subsets index values
+        iterative: bool
+            Flag for use within the iterative algorithm
+        calc_pars: dict, optional
+            Additional parameters for the model.calculate()
+
+        Returns
+        -------
+        results: xarray.Dataset
+            The model results
+
+        """
+        # subset selection:
+        model_data, farm_data, point_data = self.select_subsets(
+            model_data, farm_data, point_data, sel=sel, isel=isel
+        )
+
+        # basic checks:
+        super().run_calculation(algo, model, model_data, farm_data, point_data)
+
+        # prepare:
+        n_states = model_data.sizes[FC.STATE]
+        out_coords = model.output_coords()
+        coords = {}
+        if FC.STATE in out_coords and FC.STATE in model_data.coords:
+            coords[FC.STATE] = model_data[FC.STATE].to_numpy()
+        if farm_data is None:
+            farm_data = xr.Dataset()
+        goal_data = farm_data if point_data is None else point_data
+
+        # DEBUG objec mem sizes:
+        # from foxes.utils import print_mem
+        # for m in [algo] + model.models:
+        #    print_mem(m, pre_str="MULTIP CHECKING LARGE DATA", min_csize=9999)
+
+        # calculate chunk sizes:
+        n_targets = point_data.sizes[FC.TARGET] if point_data is not None else 0
+        chunk_sizes_states, chunk_sizes_targets = self.calc_chunk_sizes(
+            n_states, n_targets
+        )
+        n_chunks_states = len(chunk_sizes_states)
+        n_chunks_targets = len(chunk_sizes_targets)
+        self.print(
+            f"{type(self).__name__}: Selecting n_chunks_states = {n_chunks_states}, n_chunks_targets = {n_chunks_targets}",
+            level=2,
+        )
+
+        # prepare and submit chunks:
+        n_chunks_all = n_chunks_states * n_chunks_targets
+        self.print(
+            f"{type(self).__name__}: Submitting {n_chunks_all} chunks to {self.n_procs} processes",
+            level=2,
+        )
+        pbar = tqdm(total=n_chunks_all) if self.verbosity > 1 else None
+        jobs = {}
+        i0_states = 0
+        for chunki_states in range(n_chunks_states):
+            i1_states = i0_states + chunk_sizes_states[chunki_states]
+            i0_targets = 0
+            for chunki_points in range(n_chunks_targets):
+                i1_targets = i0_targets + chunk_sizes_targets[chunki_points]
+
+                # get this chunk's data:
+                data = self.get_chunk_input_data(
+                    algo=algo,
+                    model_data=model_data,
+                    farm_data=farm_data,
+                    point_data=point_data,
+                    states_i0_i1=(i0_states, i1_states),
+                    targets_i0_i1=(i0_targets, i1_targets),
+                    out_vars=out_vars,
+                )
+
+                # submit model calculation:
+                jobs[(chunki_states, chunki_points)] = self._submit(
+                    _run,
+                    algo,
+                    model,
+                    data,
+                    iterative,
+                    chunk_store,
+                    (i0_states, i0_targets),
+                    **calc_pars,
+                )
+                del data
+
+                i0_targets = i1_targets
+
+                if pbar is not None:
+                    pbar.update()
+
+            i0_states = i1_states
+
+        del calc_pars, farm_data, point_data
+        if pbar is not None:
+            pbar.close()
+
+        # wait for results:
+        if n_chunks_all > 1 or self.verbosity > 1:
+            self.print(
+                f"{type(self).__name__}: Computing {n_chunks_all} chunks using {self.n_procs} processes"
+            )
+        pbar = (
+            tqdm(total=n_chunks_all)
+            if n_chunks_all > 1 and self.verbosity > 0
+            else None
+        )
+        results = {}
+        for chunki_states in range(n_chunks_states):
+            for chunki_points in range(n_chunks_targets):
+                key = (chunki_states, chunki_points)
+                results[key] = self._result(jobs.pop((chunki_states, chunki_points)))
+                if pbar is not None:
+                    pbar.update()
+        if pbar is not None:
+            pbar.close()
+
+        return self.combine_results(
+            algo=algo,
+            results=results,
+            model_data=model_data,
+            out_vars=out_vars,
+            out_coords=out_coords,
+            n_chunks_states=n_chunks_states,
+            n_chunks_targets=n_chunks_targets,
+            goal_data=goal_data,
+            iterative=iterative,
+        )