foxes 1.2__py3-none-any.whl → 1.2.1__py3-none-any.whl

foxes/output/rose_plot.py

@@ -1,7 +1,10 @@
 import numpy as np
-import pandas as pd
+import matplotlib.pyplot as plt
+from xarray import Dataset
+from matplotlib.cm import ScalarMappable
+from matplotlib.projections.polar import PolarAxes
+from matplotlib.lines import Line2D
 
-from foxes.utils import wd2uv, uv2wd, TabWindroseAxes
 from foxes.algorithms import Downwind
 from foxes.core import WindFarm, Turbine
 from foxes.models import ModelBook
@@ -24,30 +27,38 @@ class RosePlotOutput(Output):
 
     """
 
-    def __init__(self, results, **kwargs):
+    def __init__(
+        self,
+        farm_results=None,
+        point_results=None,
+        use_points=False,
+        **kwargs,
+    ):
         """
         Constructor.
 
         Parameters
         ----------
-        results: xarray.Dataset
-            The calculation results (farm or points)
+        farm_results: xarray.Dataset, optional
+            The farm results
+        point_results: xarray.Dataset, optional
+            The point results
+        use_points: bool
+            Flag for using points in cases where both
+            farm and point results are given
         kwargs: dict, optional
             Additional parameters for the base class
 
         """
         super().__init__(**kwargs)
-        dims = list(results.sizes.keys())
-        if dims[1] == FC.TURBINE:
-            self._rtype = FC.TURBINE
-        elif dims[1] == FC.POINT:
+        if use_points or (farm_results is None and point_results is not None):
+            self.results = point_results
             self._rtype = FC.POINT
+        elif farm_results is not None:
+            self.results = farm_results
+            self._rtype = FC.TURBINE
         else:
-            raise KeyError(
-                f"Results dimension 1 is neither '{FC.TURBINE}' nor '{FC.POINT}': dims = {results.dims}"
-            )
-
-        self.results = results.to_dataframe()
+            raise KeyError(f"Require either farm_results or point_results")
 
     @classmethod
     def get_data_info(cls, dname):
@@ -104,224 +115,198 @@ class RosePlotOutput(Output):
 
     def get_data(
         self,
-        sectors,
-        var,
-        var_bins,
+        wd_sectors,
+        ws_var,
+        ws_bins,
         wd_var=FV.AMB_WD,
-        turbine=None,
-        point=None,
-        start0=False,
+        turbine=0,
+        point=0,
+        add_inf=False,
     ):
         """
-        Get pandas DataFrame with wind rose data.
+        Generates the plot data
 
         Parameters
         ----------
-        sectors: int
-            The number of wind direction sectors
-        var: str
-            The data variable name
-        var_bins: list of float
-            The variable bin separation values
-        wd_var: str, optional
-            The wind direction variable name
-        turbine: int, optional
-            Only relevant in case of farm results.
-            If None, mean over all turbines.
-            Else, data from a single turbine
-        point: int, optional
-            Only relevant in case of point results.
-            If None, mean over all points.
-            Else, data from a single point
-        start0: bool
-            Flag for starting the first sector at
-            zero degrees instead of minus half width
+        wd_sectors: int
+            The number of wind rose sectors
+        ws_var: str
+            The wind speed variable
+        ws_bins: list of float
+            The wind speed bins
+        wd_var: str
+            The wind direction variable
+        turbine: int
+            The turbine index, for weights and for
+            data if farm_results are given
+        point: int
+            The point index, for data if point_results
+            are given
+        add_inf: bool
+            Add an upper bin up to infinity
 
         Returns
         -------
-        pd.DataFrame:
-            The wind rose data
+        data: xarray.Dataset
+            The plot data
 
         """
-
-        dwd = 360.0 / sectors
-        wds = np.arange(0.0, 360.0, dwd)
-        wdb = np.append(wds, 360) if start0 else np.arange(-dwd / 2, 360.0, dwd)
-        lgd = f"interval_{var}"
-
-        data = self.results[[wd_var, FV.WEIGHT]].copy()
-        data[lgd] = self.results[var]
-        uv = wd2uv(data[wd_var].to_numpy())
-        data["u"] = uv[:, 0]
-        data["v"] = uv[:, 1]
-
-        data[FV.WEIGHT] *= 100
-        data = data.rename(columns={FV.WEIGHT: "frequency"})
-
-        el = turbine if self._rtype == FC.TURBINE else point
-        if el is None:
-            data = data.groupby(level=0).mean()
-        else:
-            sname = data.index.names[0]
-            grp = data.reset_index().groupby(self._rtype)
-            data = grp.get_group(el).set_index(sname)
-
-        data["wd"] = uv2wd(data[["u", "v"]].to_numpy())
-        data.drop(["u", "v"], axis=1, inplace=True)
-        if not start0:
-            data.loc[data["wd"] > 360.0 - dwd / 2, "wd"] -= 360.0
-
-        data[wd_var] = pd.cut(data["wd"], wdb, labels=wds)
-        data[lgd] = pd.cut(data[lgd], var_bins, right=False, include_lowest=True)
-
-        grp = data[[wd_var, lgd, "frequency"]].groupby([wd_var, lgd], observed=False)
-        data = grp.sum().reset_index()
-
-        data[wd_var] = data[wd_var].astype(np.float64)
-        data[lgd] = list(data[lgd])
-        if start0:
-            data[wd_var] += dwd / 2
-
-        ii = pd.IntervalIndex(data[lgd])
-        data[var] = ii.mid
-        data[f"bin_min_{var}"] = ii.left
-        data[f"bin_max_{var}"] = ii.right
-        data[f"bin_min_{wd_var}"] = data[wd_var] - dwd / 2
-        data[f"bin_max_{wd_var}"] = data[wd_var] + dwd / 2
-        data["sector"] = (data[wd_var] / dwd).astype(int)
-
-        data = data[
-            [
-                wd_var,
-                var,
-                "sector",
-                f"bin_min_{wd_var}",
-                f"bin_max_{wd_var}",
-                f"bin_min_{var}",
-                f"bin_max_{var}",
-                lgd,
-                "frequency",
-            ]
-        ]
-        data.index.name = "bin"
+        if add_inf:
+            ws_bins = list(ws_bins) + [np.inf]
+        w = self.results[FV.WEIGHT].to_numpy()[:, turbine]
+        t = turbine if self._rtype == FC.TURBINE else point
+        ws = self.results[ws_var].to_numpy()[:, t]
+        wd = self.results[wd_var].to_numpy()[:, t].copy()
+        wd_delta = 360 / wd_sectors
+        wd[wd >= 360 - wd_delta / 2] -= 360
+        wd_bins = np.arange(-wd_delta / 2, 360, wd_delta)
+        ws_bins = np.asarray(ws_bins, dtype=ws.dtype)
+
+        freq = 100 * np.histogram2d(wd, ws, (wd_bins, ws_bins), weights=w)[0]
+
+        data = Dataset(
+            coords={
+                wd_var: np.arange(0, 360, wd_delta),
+                ws_var: 0.5 * (ws_bins[:-1] + ws_bins[1:]),
+            },
+            data_vars={
+                f"bin_min_{wd_var}": (wd_var, wd_bins[:-1]),
+                f"bin_max_{wd_var}": (wd_var, wd_bins[1:]),
+                f"bin_min_{ws_var}": (ws_var, ws_bins[:-1]),
+                f"bin_max_{ws_var}": (ws_var, ws_bins[1:]),
+                "frequency": ((wd_var, ws_var), freq),
+            },
+            attrs={
+                f"{wd_var}_bounds": wd_bins,
+                f"{ws_var}_bounds": ws_bins,
+            },
+        )
 
         return data
 
     def get_figure(
         self,
-        sectors,
-        var,
-        var_bins,
+        wd_sectors,
+        ws_var,
+        ws_bins,
         wd_var=FV.AMB_WD,
-        turbine=None,
-        point=None,
-        title=None,
-        legend=None,
-        design="bar",
-        start0=False,
         fig=None,
+        ax=None,
         figsize=None,
-        rect=None,
+        freq_delta=3,
+        cmap="summer",
+        title=None,
+        legend_pars=None,
         ret_data=False,
         **kwargs,
     ):
         """
-        Creates figure object
+        Creates the figure
 
         Parameters
         ----------
-        sectors: int
-            The number of wind direction sectors
-        var: str
-            The data variable name
-        var_bins: list of float
-            The variable bin separation values
-        wd_var: str, optional
-            The wind direction variable name
-        turbine: int, optional
-            Only relevant in case of farm results.
-            If None, mean over all turbines.
-            Else, data from a single turbine
-        point: int, optional
-            Only relevant in case of point results.
-            If None, mean over all points.
-            Else, data from a single point
-        title. str, optional
-            The title
-        legend: str, optional
-            The data legend string
-        design: str
-            The wind rose design: bar, contour, ...
-        start0: bool
-            Flag for starting the first sector at
-            zero degrees instead of minus half width
-        fig: matplotlib.Figure
-            The figure to which to add an axis
+        wd_sectors: int
+            The number of wind rose sectors
+        ws_var: str
+            The wind speed variable
+        ws_bins: list of float
+            The wind speed bins
+        wd_var: str
+            The wind direction variable
+        fig: pyplot.Figure, optional
+            The figure object
+        ax: pyplot.Axes, optional
+            The axes object
         figsize: tuple, optional
-            The figsize of the newly created figure
-        rect: list, optional
-            The rectangle of the figure which to fill,
-            e.g. [0.1, 0.1, 0.8, 0.8]
+            The figsize argument for plt.subplots
+        freq_delta: int
+            The frequency delta for the label
+            in percent
+        cmap: str
+            The color map
+        title: str, optional
+            The title
+        legend_pars: dict, optional
+            Parameters for the legend
         ret_data: bool
             Flag for returning wind rose data
         kwargs: dict, optional
-            Additional arguments for TabWindroseAxes
-            plot function
+            Additional parameters for get_data
 
         Returns
         -------
-        fig: matplotlib.pyplot.Figure
-            The rose plot figure
-        data: pd.DataFrame, optional
-            The wind rose data
+        ax: pyplot.Axes
+            The axes object
+        data: xarray.Dataset, optional
+            The plot data
 
         """
-        lg = legend
-        if title is None or legend is None:
-            ttl, lg = self.get_data_info(var)
-        if title is not None:
-            ttl = title
-
-        wrdata = self.get_data(
-            sectors=sectors,
-            var=var,
-            var_bins=var_bins,
-            wd_var=wd_var,
-            turbine=turbine,
-            point=point,
-            start0=start0,
-        )
+        data = self.get_data(wd_sectors, ws_var, ws_bins, wd_var, **kwargs)
+
+        n_wsb = data.sizes[ws_var]
+        n_wdb = data.sizes[wd_var]
+        ws_bins = np.asarray(data.attrs[f"{ws_var}_bounds"])
+        wd_cent = np.mod(90 - data[wd_var].to_numpy(), 360)
+        wd_cent = np.radians(wd_cent)
+        wd_delta = 360 / n_wdb
+        wd_width = np.radians(0.9 * wd_delta)
+        freq = data["frequency"].to_numpy()
+
+        if ax is not None:
+            if not isinstance(ax, PolarAxes):
+                raise TypeError(
+                    f"Require axes of type '{PolarAxes.__name__}' for '{type(self).__name__}', got '{type(ax).__name__}'"
+                )
+        else:
+            fig, ax = plt.subplots(figsize=figsize, subplot_kw={"projection": "polar"})
+
+        bcmap = plt.get_cmap(cmap, n_wsb)
+        color_list = bcmap(np.linspace(0, 1, n_wsb))
+
+        bottom = np.zeros(n_wdb)
+        for wsi in range(n_wsb):
+            ax.bar(
+                wd_cent,
+                freq[:, wsi],
+                bottom=bottom,
+                width=wd_width,
+                color=color_list[wsi],
+            )
+            bottom += freq[:, wsi]
+
+        fmax = np.max(np.sum(freq, axis=1))
+        freq_delta = int(freq_delta)
+        freq_ticks = np.arange(0, fmax + freq_delta / 2, freq_delta, dtype=np.int32)[1:]
+
+        tksl = np.arange(0, 360, max(wd_delta, 30))
+        tks = np.radians(np.mod(90 - tksl, 360))
+        ax.set_xticks(tks, [f"{int(d)}°" for d in tksl])
+        ax.set_yticks(freq_ticks, [f"{f}%" for f in freq_ticks])
+        ax.set_title(title)
 
-        ax = TabWindroseAxes.from_ax(fig=fig, rect=rect, figsize=figsize)
-        fig = ax.get_figure()
-
-        plfun = getattr(ax, design)
-        plfun(
-            direction=wrdata[wd_var].to_numpy(),
-            var=wrdata[var].to_numpy(),
-            weights=wrdata["frequency"].to_numpy(),
-            bin_min_dir=np.sort(wrdata[f"bin_min_{wd_var}"].unique()),
-            bin_min_var=np.sort(wrdata[f"bin_min_{var}"].unique()),
-            bin_max_var=np.sort(wrdata[f"bin_max_{var}"].unique()),
-            **kwargs,
+        llines = [Line2D([0], [0], color=c, lw=10) for c in np.flip(color_list, axis=0)]
+        lleg = [
+            f"[{ws_bins[i]:.1f}, {ws_bins[i+1]:.1f})" for i in range(n_wsb - 1, -1, -1)
+        ]
+        lpars = dict(
+            loc="upper left",
+            bbox_to_anchor=(0.8, 0.5),
+            title=f"{ws_var}",
         )
-        ax.set_legend(title=lg)
-        ax.set_title(ttl)
+        wsl = [FV.WS, FV.REWS, FV.REWS2, FV.REWS3]
+        wsl += [FV.var2amb[v] for v in wsl]
+        if ws_var in wsl:
+            lpars["title"] += " [m/s]"
+        if legend_pars is not None:
+            lpars.update(legend_pars)
+        ax.legend(llines, lleg, **lpars)
 
         if ret_data:
-            return fig, wrdata
+            return ax, data
         else:
-            return fig
+            return ax
 
-    def write_figure(
-        self,
-        file_name,
-        sectors,
-        var,
-        var_bins,
-        ret_data=False,
-        **kwargs,
-    ):
+    def write_figure(self, file_name, *args, ret_data=False, **kwargs):
         """
         Write rose plot to file
 
@@ -329,16 +314,12 @@ class RosePlotOutput(Output):
         ----------
         file_name: str
             Name of the output file
-        sectors: int
-            The number of wind direction sectors
-        var: str
-            The data variable name
-        var_bins: list of float
-            The variable bin separation values
+        args: tuple, optional
+            Additional parameters for get_figure
         ret_data: bool
             Flag for returning wind rose data
         kwargs: dict, optional
-            Additional parameters for get_figure()
+            Additional parameters for get_figure
 
         Returns
         -------
@@ -347,19 +328,13 @@ class RosePlotOutput(Output):
 
         """
 
-        r = self.get_figure(
-            sectors=sectors,
-            var=var,
-            var_bins=var_bins,
-            ret_data=ret_data,
-            **kwargs,
-        )
+        r = self.get_figure(*args, ret_data=ret_data, **kwargs)
         fpath = self.get_fpath(file_name)
         if ret_data:
-            r[0].write_image(fpath)
+            r[0].get_figure().savefig(fpath, bbox_inches="tight")
             return r[1]
         else:
-            r.write_image(fpath)
+            r.get_figure().savefig(fpath, bbox_inches="tight")
 
 
 class StatesRosePlotOutput(RosePlotOutput):
@@ -409,3 +384,234 @@ class StatesRosePlotOutput(RosePlotOutput):
         results = algo.calc_farm(ambient=True).rename_vars({ws_var: FV.AMB_WS})
 
         super().__init__(results, **kwargs)
+
+
+class WindRoseBinPlot(Output):
+    """
+    Plots mean data in wind rose bins
+
+    Attributes
+    ----------
+    farm_results: xarray.Dataset
+        The wind farm results
+
+    :group: output
+
+    """
+
+    def __init__(self, farm_results, **kwargs):
+        """
+        Constructor
+
+        Parameters
+        ----------
+        farm_results: xarray.Dataset
+            The wind farm results
+        kwargs: dict, optional
+            Parameters for the base class
+
+        """
+        super().__init__(**kwargs)
+        self.farm_results = farm_results
+
+    def get_data(
+        self,
+        variable,
+        ws_bins,
+        wd_sectors=12,
+        wd_var=FV.AMB_WD,
+        ws_var=FV.AMB_REWS,
+        turbine=0,
+        contraction="weights",
+    ):
+        """
+        Generates the plot data
+
+        Parameters
+        ----------
+        variable: str
+            The variable name
+        ws_bins: list of float
+            The wind speed bins
+        wd_var: str
+            The wind direction variable
+        ws_var: str
+            The wind speed variable
+        turbine: int
+            The turbine index
+        contraction: str
+            The contraction method for states:
+            weights, mean_no_weights, sum_no_weights
+
+        Returns
+        -------
+        data: xarray.Dataset
+            The plot data
+
+        """
+        var = self.farm_results[variable].to_numpy()[:, turbine]
+        w = self.farm_results[FV.WEIGHT].to_numpy()[:, turbine]
+        ws = self.farm_results[ws_var].to_numpy()[:, turbine]
+        wd = self.farm_results[wd_var].to_numpy()[:, turbine].copy()
+        wd_delta = 360 / wd_sectors
+        wd[wd >= 360 - wd_delta / 2] -= 360
+        wd_bins = np.arange(-wd_delta / 2, 360, wd_delta)
+        ws_bins = np.asarray(ws_bins)
+
+        if contraction == "weights":
+            z = np.histogram2d(wd, ws, (wd_bins, ws_bins), weights=w)[0]
+            z[z < 1e-13] = np.nan
+            z = np.histogram2d(wd, ws, (wd_bins, ws_bins), weights=w * var)[0] / z
+        elif contraction == "mean_no_weights":
+            z = np.histogram2d(wd, ws, (wd_bins, ws_bins))[0].astype(w.dtype)
+            z[z < 1] = np.nan
+            z = np.histogram2d(wd, ws, (wd_bins, ws_bins), weights=var)[0] / z
+        elif contraction == "sum_no_weights":
+            z = np.histogram2d(wd, ws, (wd_bins, ws_bins), weights=var)[0]
+        else:
+            raise KeyError(
+                f"Contraction '{contraction}' not supported. Choices: weights, mean_no_weights, sum_no_weights"
+            )
+
+        data = Dataset(
+            coords={
+                wd_var: 0.5 * (wd_bins[:-1] + wd_bins[1:]),
+                ws_var: 0.5 * (ws_bins[:-1] + ws_bins[1:]),
+            },
+            data_vars={
+                variable: ((wd_var, ws_var), z),
+            },
+            attrs={
+                f"{wd_var}_bounds": wd_bins,
+                f"{ws_var}_bounds": ws_bins,
+            },
+        )
+
+        return data
+
+    def get_figure(
+        self,
+        variable,
+        ws_bins,
+        wd_sectors=12,
+        wd_var=FV.AMB_WD,
+        ws_var=FV.AMB_REWS,
+        turbine=0,
+        contraction="weights",
+        fig=None,
+        ax=None,
+        title=None,
+        figsize=None,
+        ret_data=False,
+        **kwargs,
+    ):
+        """
+        Creates the figure
+
+        Parameters
+        ----------
+        variable: str
+            The variable name
+        ws_bins: list of float
+            The wind speed bins
+        wd_var: str
+            The wind direction variable
+        ws_var: str
+            The wind speed variable
+        turbine: int
+            The turbine index
+        contraction: str
+            The contraction method for states:
+            weights, mean_no_weights, sum_no_weights
+        fig: pyplot.Figure, optional
+            The figure object
+        ax: pyplot.Axes, optional
+            The axes object
+        title: str, optional
+            The title
+        figsize: tuple, optional
+            The figsize argument for plt.subplots
+        ret_data: bool
+            Flag for returning wind rose data
+        kwargs: dict, optional
+            Additional parameters for plt.pcolormesh
+
+        Returns
+        -------
+        ax: pyplot.Axes
+            The axes object
+
+        """
+        data = self.get_data(
+            variable=variable,
+            ws_bins=ws_bins,
+            wd_sectors=wd_sectors,
+            wd_var=wd_var,
+            ws_var=ws_var,
+            turbine=turbine,
+            contraction=contraction,
+        )
+
+        wd_delta = 360 / data.sizes[wd_var]
+        wd_bins = np.mod(90 - data.attrs[f"{wd_var}_bounds"], 360)
+        wd_bins = np.radians(wd_bins)
+        ws_bins = data.attrs[f"{ws_var}_bounds"]
+
+        if ax is not None:
+            if not isinstance(ax, PolarAxes):
+                raise TypeError(
+                    f"Require axes of type '{PolarAxes.__name__}' for '{type(self).__name__}', got '{type(ax).__name__}'"
+                )
+        else:
+            fig, ax = plt.subplots(figsize=figsize, subplot_kw={"projection": "polar"})
+
+        y, x = np.meshgrid(ws_bins, wd_bins)
+        z = data[variable].to_numpy()
+
+        prgs = {"shading": "flat"}
+        prgs.update(kwargs)
+
+        img = ax.pcolormesh(x, y, z, **prgs)
+
+        tksl = np.arange(0, 360, max(wd_delta, 30))
+        tks = np.radians(np.mod(90 - tksl, 360))
+        ax.set_xticks(tks, [f"{d}°" for d in tksl])
+        ax.set_yticks(ws_bins)
+        ax.set_title(title)
+        cbar = fig.colorbar(img, ax=ax, pad=0.12)
+        cbar.ax.set_title(variable)
+
+        if ret_data:
+            return ax, data
+        else:
+            return ax
+
+    def write_figure(self, file_name, *args, ret_data=False, **kwargs):
+        """
+        Write rose plot to file
+
+        Parameters
+        ----------
+        file_name: str
+            Name of the output file
+        args: tuple, optional
+            Additional parameters for get_figure
+        ret_data: bool
+            Flag for returning wind rose data
+        kwargs: dict, optional
+            Additional parameters for get_figure
+
+        Returns
+        -------
+        data: pd.DataFrame, optional
+            The wind rose data
+
+        """
+
+        r = self.get_figure(*args, ret_data=ret_data, **kwargs)
+        fpath = self.get_fpath(file_name)
+        if ret_data:
+            r[0].get_figure().savefig(fpath, bbox_inches="tight")
+            return r[1]
+        else:
+            r.get_figure().savefig(fpath, bbox_inches="tight")