Fix parallel dAEP with timeseries site

py_wake/tests/test_wind_farm_models/test_parallel_daep.py

+import numpy.testing as npt
+from py_wake.utils.gradients import cs
+from py_wake.examples.data.hornsrev1 import V80
+from py_wake.literature.turbopark import Nygaard_2022
+from py_wake.site.xrsite import XRSite
+from py_wake.utils.gradients import autograd
+import numpy as np
+import xarray as xr
+import numpy as np  # fmt: skip
+np.random.seed(42)
+
+
+def get_xrsite_ds():
+    x_dim, y_dim, h_dim, time_dim = 100, 40, 8, 48
+    h = np.array([30.0, 50.0, 100.0, 150.0, 200.0, 250.0, 300.0, 400.0])
+    x = np.linspace(6.947e5, 8.932e5, x_dim)
+    y = np.linspace(8.449e5, 9.225e5, y_dim)
+    time = np.arange(48)
+    WS = np.random.uniform(3.0, 15.0, size=(y_dim, x_dim, h_dim, time_dim))
+    WD = np.random.uniform(250.0, 300.0, size=(y_dim, x_dim, h_dim, time_dim))
+    P = 1.0
+    TI = 0.06
+    ds = xr.Dataset(
+        {
+            "WS": (["y", "x", "h", "time"], WS),
+            "WD": (["y", "x", "h", "time"], WD),
+            "P": ([], P),
+            "TI": ([], TI),
+        },
+        coords={
+            "h": (["h"], h),
+            "x": (["x"], x),
+            "y": (["y"], y),
+            "time": (["time"], time),
+        },
+    )
+    return ds
+
+
+def test_parallel_daep_with_time_site():
+    ds = get_xrsite_ds()
+    ds = ds.transpose("x", "y", "h", "time")
+    T = np.arange(len(ds.time))
+    site = XRSite(ds)
+    wfm = Nygaard_2022(site, V80())
+
+    test_x, test_y = [
+        ds.x.mean().item(),
+        ds.x.mean().item(),
+        ds.x.mean().item(),
+    ], [
+        ds.y.mean().item(),
+        ds.y.mean().item() + 250,
+        ds.y.mean().item() + 500,
+    ]
+
+    def aep_func(x, y, n_cpu=1):  # full=False, **kwargs
+        _f = np.zeros_like(T)
+        sim_res = wfm(x, y, time=T, ws=_f, wd=_f, n_cpu=n_cpu)
+        return sim_res.aep().sum().values
+
+    def aep_jac(x, y, n_cpu=1, **kwargs):
+        _f = np.zeros_like(T)
+        jx, jy = wfm.aep_gradients(
+            gradient_method=autograd,
+            wrt_arg=["x", "y"],
+            x=x,
+            y=y,
+            ws=_f,
+            wd=_f,
+            time=T,
+            n_cpu=n_cpu,
+        )
+        return np.array([np.atleast_2d(jx), np.atleast_2d(jy)])
+
+    aep1 = aep_func(test_x, test_y, n_cpu=1)
+    aep2 = aep_func(test_x, test_y, n_cpu=2)
+    npt.assert_almost_equal(
+        aep1, aep2, decimal=6, err_msg="AEP is not the same for n_cpu=1 and n_cpu=2"
+    )
+
+    cs_grad1 = cs(aep_func, True, argnum=[0, 1])(test_x, test_y, n_cpu=1)
+    ad_grad1 = aep_jac(test_x, test_y, n_cpu=1)
+    npt.assert_almost_equal(
+        np.stack(cs_grad1, axis=0).reshape(ad_grad1.shape),
+        ad_grad1,
+        decimal=6,
+        err_msg="Gradient is not the same for n_cpu=1 compared to complex step",
+    )
+
+    cs_grad2 = cs(aep_func, True, argnum=[0, 1])(test_x, test_y, n_cpu=2)
+    ad_grad2 = aep_jac(test_x, test_y, n_cpu=2)
+    npt.assert_almost_equal(
+        cs_grad1,
+        cs_grad2,
+        decimal=6,
+    )
+    npt.assert_almost_equal(
+        ad_grad1,
+        ad_grad2,
+        decimal=6,
+        err_msg="[AD] Gradient is not the same for n_cpu=1 and n_cpu=2",
+    )
+    npt.assert_almost_equal(
+        np.stack(cs_grad2, axis=0).reshape(ad_grad2.shape),
+        ad_grad2,
+        decimal=6,
+        err_msg="[CS vs AD] Gradient is not the same for n_cpu=1 and n_cpu=2",
+    )

py_wake/wind_farm_models/wind_farm_model.py

@@ -322,7 +322,7 @@ class WindFarmModel(ABC):
             else:  # pragma: no cover
                 raise ValueError(f'Shape, {s}, of argument {k} is invalid')
 
-        arg_lst = [{'wd': wd[wd_slice], 'ws': ws[ws_slice], 'time':get_subtask_arg('time', time, wd_slice, ws_slice),
+        arg_lst = [{'wd': wd[wd_slice], 'ws': ws[ws_slice], 'time': get_subtask_arg('time', time, wd_slice, ws_slice),
                     ** {k: get_subtask_arg(k, v, wd_slice, ws_slice) for k, v in kwargs.items()}} for wd_slice, ws_slice in slice_lst]
 
         return map_func, arg_lst, wd_chunks, ws_chunks
@@ -330,16 +330,21 @@ class WindFarmModel(ABC):
     def _aep_chunk_wrapper(self, aep_function,
                            x, y, h=None, type=0, wd=None, ws=None,   # @ReservedAssignment
                            normalize_probabilities=False, with_wake_loss=True,
-                           n_cpu=1, wd_chunks=None, ws_chunks=None, **kwargs):
+                           n_cpu=1, wd_chunks=None, ws_chunks=None, time=False, **kwargs):
         wd, ws = self.site.get_defaults(wd, ws)
         wd_bin_size = self.site.wd_bin_size(wd)
 
         map_func, kwargs_lst, wd_chunks, ws_chunks = self._multiprocessing_chunks(
-            wd=wd, ws=ws, time=False, n_cpu=n_cpu, wd_chunks=wd_chunks, ws_chunks=ws_chunks,
+            wd=wd, ws=ws, time=time, n_cpu=n_cpu, wd_chunks=wd_chunks, ws_chunks=ws_chunks,
             x=x, y=y, h=h, type=type, **kwargs)
 
-        return np.sum([np.array(aep) / self.site.wd_bin_size(args['wd']) * wd_bin_size
-                       for args, aep in zip(kwargs_lst, map_func(aep_function, kwargs_lst))], 0)
+        if time is False:
+            return np.sum([np.array(aep) / self.site.wd_bin_size(args['wd']) * wd_bin_size
+                           for args, aep in zip(kwargs_lst, map_func(aep_function, kwargs_lst))], 0)
+
+        return np.mean(
+            [np.array(aep) for aep in map_func(aep_function, kwargs_lst)], axis=0
+        )
 
     def aep_gradients(self, gradient_method=autograd, wrt_arg=['x', 'y'], gradient_method_kwargs={},
                       n_cpu=1, wd_chunks=None, ws_chunks=None, **kwargs):
@@ -506,13 +511,13 @@ class SimulationResult(xr.Dataset):
 
         if linear_power_segments:
             s = "The linear_power_segments method "
-            assert all([n in self for n in ['Weibull_A', 'Weibull_k', 'Sector_frequency']]),\
+            assert all([n in self for n in ['Weibull_A', 'Weibull_k', 'Sector_frequency']]), \
                 s + "requires a site with weibull information"
             assert normalize_probabilities is False, \
                 s + "cannot be combined with normalize_probabilities"
-            assert np.all(self.Power.isel(ws=0) == 0) and np.all(self.Power.isel(ws=-1) == 0),\
+            assert np.all(self.Power.isel(ws=0) == 0) and np.all(self.Power.isel(ws=-1) == 0), \
                 s + "requires first wind speed to have no power (just below cut-in)"
-            assert np.all(self.Power.isel(ws=-1) == 0),\
+            assert np.all(self.Power.isel(ws=-1) == 0), \
                 s + "requires last wind speed to have no power (just above cut-out)"
             weighted_power = weibull.WeightedPower(
                 self.ws.values,