fix issue when computing AEP gradients with respect to ws and kwargs, e.g. yaw, tilt etc

py_wake/tests/test_wind_farm_models/test_wind_farm_model.py

@@ -110,10 +110,8 @@ def test_aep_chunks_input_dims(rho):
 
 @pytest.mark.parametrize('wrt_arg', ['x', 'y', 'h',
                                      ['x', 'y'], ['x', 'h'],
-                                     # 'wd', 'ws'
-                                     # 'yaw'
-                                     ])
-def test_aep_gradients_function(wrt_arg):
+                                     'wd', 'ws'])
+def test_aep_gradients_function1(wrt_arg):
     wfm = IEA37CaseStudy1(16, deflectionModel=JimenezWakeDeflection())
     x, y = wfm.site.initial_position[np.array([0, 2, 5, 8, 14])].T
     kwargs = {'x': x, 'y': y, 'h': x * 0 + wfm.windTurbines.hub_height(),
@@ -127,36 +125,43 @@ def test_aep_gradients_function(wrt_arg):
         ax1, ax2 = plt.subplots(1, 2)[1]
         wfm(**kwargs).flow_map(XYGrid(resolution=100)).plot_wake_map(ax=ax1)
         ax2.set_title(wrt_arg)
-        ax2.plot(dAEP_autograd.flatten(), '.')
+        ax2.plot(dAEP_autograd.flatten(), '.', label='autograd')
+        ax2.plot(dAEP_cs.flatten(), '.', label='cs')
+        ax2.plot(dAEP_fd.flatten(), '.', label='fd')
+        plt.legend()
         plt.show()
 
-    npt.assert_array_almost_equal(dAEP_autograd, dAEP_cs, 15)
-    npt.assert_array_almost_equal(dAEP_autograd, dAEP_fd, 6)
+    npt.assert_array_almost_equal(dAEP_autograd, dAEP_cs, 14)
+    npt.assert_array_almost_equal(dAEP_autograd, dAEP_fd, 5)
     npt.assert_array_equal(dAEP_autograd, wfm.aep_gradients(gradient_method=autograd, wrt_arg=wrt_arg, **kwargs))
 
 
-def test_aep_gradients_wrt_kwargs():
+@pytest.mark.parametrize('wrt_arg', ['x', 'y', 'h',
+                                     ['x', 'y'], ['x', 'h'],
+                                     'wd', 'ws', 'yaw'])
+def test_aep_gradients_function2(wrt_arg):
     wfm = IEA37CaseStudy1(16, deflectionModel=JimenezWakeDeflection())
     x, y = wfm.site.initial_position[np.array([0, 2, 5, 8, 14])].T
     kwargs = {'x': x, 'y': y, 'h': x * 0 + wfm.windTurbines.hub_height(),
               'wd': [0], 'ws': 9.8, 'yaw': np.arange(1, 6).reshape((5, 1, 1)) * 5, 'tilt': 0}
 
-    def aep(yaw, **kwargs):
-        return wfm.aep(yaw=yaw, **kwargs)
-
-    dAEP_autograd = autograd(aep)(**kwargs)
-    dAEP_cs = cs(aep)(**kwargs)
-    dAEP_fd = fd(aep)(**kwargs)
+    dAEP_autograd = wfm.aep_gradients(gradient_method=autograd, wrt_arg=wrt_arg, **kwargs)
+    dAEP_cs = wfm.aep_gradients(gradient_method=cs, wrt_arg=wrt_arg, **kwargs)
+    dAEP_fd = wfm.aep_gradients(gradient_method=fd, wrt_arg=wrt_arg, **kwargs)
 
     if 0:
         ax1, ax2 = plt.subplots(1, 2)[1]
         wfm(**kwargs).flow_map(XYGrid(resolution=100)).plot_wake_map(ax=ax1)
-
-        ax2.plot(dAEP_autograd.flatten(), '.')
+        ax2.set_title(wrt_arg)
+        ax2.plot(dAEP_autograd.flatten(), '.', label='autograd')
+        ax2.plot(dAEP_cs.flatten(), '.', label='cs')
+        ax2.plot(dAEP_fd.flatten(), '.', label='fd')
+        plt.legend()
         plt.show()
 
-    npt.assert_array_almost_equal(dAEP_autograd, dAEP_cs, 15)
-    npt.assert_array_almost_equal(dAEP_autograd, dAEP_fd, 6)
+    npt.assert_array_almost_equal(dAEP_autograd, dAEP_cs, 14)
+    npt.assert_array_almost_equal(dAEP_autograd, dAEP_fd, 5)
+    npt.assert_array_equal(dAEP_autograd, wfm.aep_gradients(gradient_method=autograd, wrt_arg=wrt_arg, **kwargs))
 
 
 def test_aep_gradients_parallel():

py_wake/wind_farm_models/engineering_models.py

@@ -663,8 +663,9 @@ class All2AllIterative(EngineeringWindFarmModel):
                              WS_eff_ilk, TI_eff_ilk,
                              D_i, time,
                              I, L, K, **kwargs):
-        if any([np.iscomplexobj(v) for v in [kwargs.get(k, 0)
-                                             for k in ['x_ilk', 'y_ilk', 'h_ilk', 'D_i', 'yaw_ilk', 'tilt_ilk']]]):
+        if any([np.iscomplexobj(v) for v in ([kwargs.get(k, 0)
+                                              for k in ['x_ilk', 'y_ilk', 'h_ilk', 'D_i', 'yaw_ilk', 'tilt_ilk']] +
+                                             [ws, wd])]):
             dtype = np.complex128
         else:
             dtype = float

py_wake/wind_farm_models/wind_farm_model.py

@@ -347,8 +347,11 @@ class WindFarmModel(ABC):
         compute the gradients of the aep:
         gradient_function = wfm.aep_gradietns(autograd, ['x','y'])
         gradients = gradient_function(x,y)
+        This behaiour only works when wrt_arg is one or more of ['x','y','h','wd', 'ws']
+
         2) With additional key-word arguments, kwargs, the method returns the gradients of the aep:
         gradients = wfm.aep_gradients(autograd,['x','y'],x=x,y=y)
+        This behaviour also works when wrt_arg is a keyword argument, e.g. yaw
 
         Parameters
         ----------
@@ -375,12 +378,18 @@ class WindFarmModel(ABC):
                 gradient_method_kwargs=gradient_method_kwargs,
                 n_cpu=n_cpu, wd_chunks=wd_chunks, ws_chunks=ws_chunks, **kwargs)
 
-        argnum = [['x', 'y', 'h', 'type', 'wd', 'ws', 'yaw', 'tilt'].index(a) for a in np.atleast_1d(wrt_arg)]
-        f = gradient_method(self.aep, True, argnum, **gradient_method_kwargs)
-
         if kwargs:
-            return f(**kwargs)
+            wrt_arg = np.atleast_1d(wrt_arg)
+
+            def wrap_aep(*args, **kwargs):
+                kwargs.update({n: v for n, v in zip(wrt_arg, args)})
+                return self.aep(**kwargs)
+
+            f = gradient_method(wrap_aep, True, tuple(range(len(wrt_arg))), **gradient_method_kwargs)
+            return np.array(f(*[kwargs.pop(n) for n in wrt_arg], **kwargs))
         else:
+            argnum = [['x', 'y', 'h', 'type', 'wd', 'ws'].index(a) for a in np.atleast_1d(wrt_arg)]
+            f = gradient_method(self.aep, True, argnum, **gradient_method_kwargs)
             return f
 
     def _aep_gradients_kwargs(self, kwargs):