foxes 0.5.0.2__py3-none-any.whl → 0.5.2__py3-none-any.whl

foxes/output/calc_points.py

@@ -0,0 +1,143 @@
+import numpy as np
+from xarray import Dataset
+
+import foxes.constants as FC
+import foxes.variables as FV
+from foxes.utils import write_nc
+
+from .output import Output
+
+class PointCalculator(Output):
+    """
+    Computes results at given points
+    
+    Attributes
+    ----------
+    algo: foxes.Algorithm
+        The algorithm for point calculation
+    farm_results: xarray.Dataset
+        The farm results
+    
+    :group: output
+
+    """
+
+    def __init__(self, algo, farm_results, **kwargs):
+        """
+        Constructor.
+        
+        Parameters
+        ----------
+        algo: foxes.Algorithm
+            The algorithm for point calculation
+        farm_results: xarray.Dataset
+            The farm results
+        kwargs: dict, optional
+            Additional parameters for the base class
+
+        """
+        super().__init__(**kwargs)
+        self.algo = algo
+        self.farm_results = farm_results
+
+    def calculate(
+            self, 
+            points, 
+            *args, 
+            states_mean=False,
+            weight_turbine=0,
+            to_file=None,
+            write_vars=None,
+            write_pars={},
+            **kwargs
+        ):
+        """
+        Calculate point results
+        
+        Parameters
+        ----------
+        points: numpy.ndarray
+            The points, shape: (n_points, 3) 
+            or (n_states, n_points, 3)
+        args: tuple, optional
+            Additional arguments for algo.calc_points
+        states_mean: bool
+            Flag for taking the mean over states
+        weight_turbine: int, optional
+            Index of the turbine from which to take the weight
+        to_file: str, optional
+            Path to the output netCDF file
+        write_vars: list of str
+            The variables to be written to file, or None
+            for all
+        write_pars: dict, optional
+            Additional parameters for write_nc
+        kwargs: tuple, optional
+            Additional arguments for algo.calc_points
+
+        Returns
+        -------
+        point_results: xarray.Dataset
+            The point results. The calculated variables have
+            dimensions (state, point)
+
+        """
+        if points.shape[-1] == 3 and len(points.shape) == 3:
+            pts = points
+            p_has_s = True
+        elif points.shape[-1] == 3 and len(points.shape) == 2:
+            pts = np.zeros([self.algo.n_states] + list(points.shape), dtype=FC.DTYPE)
+            pts[:] = points[None, :]
+            p_has_s = False
+        else:
+            raise ValueError(f"Expecting point shape (n_states, n_points, 3) or (n_points, 3), got {points.shape}")
+        
+        pres = self.algo.calc_points(self.farm_results, pts, *args, **kwargs)
+
+        if states_mean:
+            weights = self.farm_results[FV.WEIGHT].to_numpy()[:, weight_turbine]
+            pres = Dataset(
+                data_vars={
+                    v: np.einsum('s,sp->p', weights, pres[v].to_numpy()) 
+                    for v in pres.data_vars.keys()
+                }
+            )
+
+        vrs = list(pres.data_vars.keys()) if write_vars is None else write_vars
+        if to_file is not None:
+            if states_mean:
+                if p_has_s:
+                    points = np.einsum('s,spd->pd', weights, points)
+                dvars = {
+                    "x": ((FC.POINT,), points[..., 0]),
+                    "y": ((FC.POINT,), points[..., 1]),
+                    "z": ((FC.POINT,), points[..., 2])
+                }   
+                dvars.update({v: ((FC.POINT,), pres[v].to_numpy()) 
+                            for v in vrs})
+                ds = Dataset(data_vars=dvars)
+            else:
+                if p_has_s:
+                    dvars = {
+                        "x": ((FC.STATE, FC.POINT), points[..., 0]),
+                        "y": ((FC.STATE, FC.POINT), points[..., 1]),
+                        "z": ((FC.STATE, FC.POINT), points[..., 2])
+                    }
+                else:
+                    dvars = {
+                        "x": ((FC.POINT,), points[..., 0]),
+                        "y": ((FC.POINT,), points[..., 1]),
+                        "z": ((FC.POINT,), points[..., 2])
+                    }   
+                dvars.update({v: ((FC.STATE, FC.POINT), pres[v].to_numpy()) 
+                            for v in vrs})
+                ds = Dataset(
+                    coords={FC.STATE: pres[FC.STATE].to_numpy()},
+                    data_vars=dvars
+                )
+
+            fpath = self.get_fpath(to_file)
+            write_nc(ds, fpath, **write_pars)
+
+        return pres
+    

foxes/output/flow_plots_2d/__init__.py

@@ -1,2 +1,3 @@
 from .flow_plots import FlowPlots2D
 from .seq_flow_ani_plugin import SeqFlowAnimationPlugin
+from .get_fig import get_fig

foxes/output/flow_plots_2d/common.py

@@ -1,6 +1,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from mpl_toolkits.axes_grid1 import make_axes_locatable
+from xarray import Dataset
 
 from foxes.utils import wd2uv
 import foxes.variables as FV
@@ -213,7 +214,6 @@ def calc_point_results(
     **kwargs,
 ):
     """
-
     Helper function that calculates results at grid points.
 
     Parameters
@@ -593,3 +593,106 @@ def get_grid_yz(
         z_pos,
         g_pts.reshape(n_states, n_pts, 3),
     )
+
+round_defaults = {v: 4 for v in FV.__dict__.keys() if isinstance(v, str)}
+round_defaults[FV.WD] = 3
+round_defaults[FV.TI] = 6
+round_defaults[FV.RHO] = 6
+round_defaults[FC.XYH] = 3
+round_defaults.update({FV.var2amb[v]: round_defaults[v] for v in FV.var2amb.keys()})
+
+def data2xr(
+        x_pos, 
+        y_pos, 
+        z_pos, 
+        point_results, 
+        vars=None,
+        round="auto", 
+        to_file=None,
+        complevel=5,
+        verbosity=1,
+    ):
+    """
+    Converts the image data to xarray data
+    
+    Parameter
+    ---------
+    x_pos: numpy.ndarray or float
+        The x grid positions, shape: (n_x, 3)
+    y_pos: numpy.ndarray or float
+        The y grid positions, shape: (n_y, 3)
+    z_pos: numpy.ndarray or float
+        The z grid positions, shape: (n_z, 3)
+    point_results: xarray.Dataset
+        Results of calc_points
+    vars: list of str, optional
+        Variable selection, or None for all
+    round: dict, optional
+        Round variables to given digits, or 'auto'
+        for default
+    to_file: str, optional
+        Write to nc file
+    complevel: int
+        The compression level
+    verbosity: int
+        The verbostiy level, 0 = silent
+
+    Returns
+    -------
+    ds: xarray.Dataset
+        The xarray data object
+
+    """
+    if round == "auto":
+        round = round_defaults
+    if round is not None:
+        x_pos = np.round(x_pos, round[FC.XYH])
+        y_pos = np.round(y_pos, round[FC.XYH])
+        z_pos = np.round(z_pos, round[FC.XYH])
+
+    if vars is None:
+        vars = list(point_results.data_vars.keys())
+    data = {}
+    for v in vars:
+        data[v] = point_results[v].to_numpy()
+        if v in round:
+            data[v] = np.round(data[v], round[v])
+
+    allc = [x_pos, y_pos, z_pos]
+    allcn = ["x", "y", "z"]
+    ci = [i for i, x in enumerate(allc) if isinstance(x, np.ndarray)]
+    cj = [i for i in range(3) if i not in ci][0]
+    cl = list(reversed([len(allc[i]) for i in ci]))
+    cn = list(reversed([allcn[i] for i in ci]))
+
+    coords = {}
+    attrs = {allcn[cj]: allc[cj].to_numpy()}
+    if FC.STATE in point_results.coords:
+        if point_results.dims[FC.STATE] > 1:
+            coords[FC.STATE] = point_results[FC.STATE].to_numpy()
+        else:
+            attrs[FC.STATE] = str(point_results[FC.STATE][0].to_numpy())
+    coords.update({
+        allcn[i]: allc[i] for i in reversed(ci)
+    })
+
+    dvars = {}
+    for v, d in data.items():
+        if len(d.shape) == 1:
+            dvars[v] = (cn, np.swapaxes(d.reshape(*cl), 0, 1))
+        else:
+            dvars[v] = ([FC.STATE] + cn, np.swapaxes(d.reshape(d.shape[0], *cl), 1, 2))
+
+    ds = Dataset(
+        coords=coords,
+        data_vars=dvars,
+        attrs=attrs
+    )
+
+    if to_file is not None:
+        if verbosity > 0:
+            print("Writing file", to_file)
+        enc = {k: {"zlib": True, "complevel": complevel} for k in ds.data_vars}
+        ds.to_netcdf(to_file, encoding=enc)
+    
+    return ds