fix gradient calculation for JimenezWakeDeflection and GCLHillDeflection

py_wake/deflection_models/gcl_hill_vortex.py

@@ -2,6 +2,7 @@ from numpy import newaxis as na
 import numpy as np
 from py_wake.deflection_models import DeflectionModel
 from py_wake.utils.gradients import hypot
+from py_wake.utils import gradients
 
 
 class GCLHillDeflection(DeflectionModel):
@@ -54,7 +55,7 @@ class GCLHillDeflection(DeflectionModel):
         U_d_ijlkx = -0.4 * U_w_ijlx * np.sin(theta_ilk)[:, na, :, :, na]
         U_a_ijlkx = WS_eff_ilk[:, na, :, :, na] - 0.4 * U_w_ijlx * np.cos(theta_ilk)[:, na, :, :, na]
 
-        deflection_ijlk = np.trapz(U_d_ijlkx / U_a_ijlkx, dw_ijlkx, axis=4)
+        deflection_ijlk = gradients.trapz(U_d_ijlkx / U_a_ijlkx, dw_ijlkx, axis=4)
         self.hcw_ijlk = hcw_ijl[..., na] - deflection_ijlk * np.cos(theta_deflection_ilk[:, na])
         self.dh_ijlk = dh_ijl[..., na] + deflection_ijlk * np.sin(theta_deflection_ilk[:, na])
         return dw_ijl[..., na], self.hcw_ijlk, self.dh_ijlk

py_wake/deflection_models/jimenez.py

@@ -2,6 +2,7 @@ from numpy import newaxis as na
 import numpy as np
 from py_wake.deflection_models import DeflectionModel
 from py_wake.utils.gradients import hypot
+from py_wake.utils import gradients
 
 
 class JimenezWakeDeflection(DeflectionModel):
@@ -25,7 +26,7 @@ class JimenezWakeDeflection(DeflectionModel):
         denominator_ilk = np.cos(theta_ilk)**2 * np.sin(theta_ilk) * (ct_ilk / 2)
         nominator_ijxl = (1 + (self.beta / D_src_il)[:, na, na, :] * np.maximum(dw_ijxl, 0))**2
         alpha = denominator_ilk[:, na, na] / nominator_ijxl[..., na]
-        deflection_ijlk = np.trapz(np.sin(alpha), dw_ijxl[..., na], axis=2)
+        deflection_ijlk = gradients.trapz(np.sin(alpha), dw_ijxl[..., na], axis=2)
         self.hcw_ijlk = hcw_ijl[..., na] + deflection_ijlk * np.cos(theta_deflection_ilk[:, na])
         self.dh_ijlk = dh_ijl[..., na] + deflection_ijlk * np.sin(theta_deflection_ilk[:, na])
         return dw_ijl[..., na], self.hcw_ijlk, self.dh_ijlk

py_wake/tests/test_utils/test_gradients.py

@@ -31,6 +31,7 @@ from py_wake.site.distance import StraightDistance
 from py_wake.tests.check_speed import timeit
 from py_wake.wind_turbines.power_ct_functions import CubePowerSimpleCt
 from py_wake.examples.data.iea34_130rwt._iea34_130rwt import IEA34_130_1WT_Surrogate
+from py_wake.deflection_models.jimenez import JimenezWakeDeflection
 
 
 @pytest.mark.parametrize('obj', [_wind_turbines, WindTurbines, V80().power, _wind_turbines.__dict__])
@@ -363,6 +364,44 @@ def test_Interpolators(interpolator):
         plt.show()
 
 
+@pytest.mark.parametrize('y,x,axis', [([2, 3, 7, 9], [1, 2, 4, 8], 0),
+                                      (lambda x:-x**2 + 9, np.linspace(-3, 3, 10), 0)])
+def test_trapz(y, x, axis):
+    if callable(y):
+        y = y(x)
+
+    npt.assert_array_equal(np.trapz(y, x, axis), gradients.trapz(y, x, axis))
+    dtrapz_dy_lst = [method(gradients.trapz, True)(y, x, axis) for method in [fd, cs, autograd]]
+    npt.assert_array_almost_equal(dtrapz_dy_lst[0], dtrapz_dy_lst[1])
+    npt.assert_array_equal(dtrapz_dy_lst[1], dtrapz_dy_lst[2])
+
+    if x is not None:
+        dtrapz_dx_lst = [method(gradients.trapz, True, argnum=1)(y, x, axis) for method in [fd, cs, autograd]]
+        npt.assert_array_almost_equal(dtrapz_dx_lst[0], dtrapz_dx_lst[1])
+        npt.assert_array_almost_equal(dtrapz_dx_lst[1], dtrapz_dx_lst[2], 14)
+
+
+@pytest.mark.parametrize('test', [
+    lambda y, x: gradients.trapz(np.reshape(y, (2, 4)), np.reshape(x, (2, 4)), axis=1),
+    lambda y, x: gradients.trapz(np.reshape(y, (2, 4)).T, np.reshape(x, (2, 4)).T, axis=0),
+    lambda y, x: gradients.trapz(np.reshape(y, (1, 2, 4, 1)), np.reshape(x, (1, 2, 4, 1)), axis=2)
+])
+def test_trapz_axis(test):
+    y, x = [2, 3, 7, 9] * 2, [1, 2, 4, 8] * 2
+
+    autograd(test, True, argnum=1)(y, x)
+    autograd(test, True)(y, x)
+    dtrapz_dy_lst = [method(test, True)(y, x) for method in [fd, cs, autograd]]
+    npt.assert_array_almost_equal(dtrapz_dy_lst[0], dtrapz_dy_lst[1])
+    npt.assert_array_equal(dtrapz_dy_lst[1], dtrapz_dy_lst[2])
+
+    if x is not None:
+
+        dtrapz_dx_lst = [method(test, True, argnum=1)(y, x) for method in [fd, cs, autograd]]
+        npt.assert_array_almost_equal(dtrapz_dx_lst[0], dtrapz_dx_lst[1])
+        npt.assert_array_almost_equal(dtrapz_dx_lst[1], dtrapz_dx_lst[2], 14)
+
+
 def test_manual_vs_autograd_speed():
     cubePowerSimpleCt = CubePowerSimpleCt()
     x = np.random.random(10000) * 30

py_wake/utils/gradients.py

@@ -6,7 +6,7 @@ import autograd.numpy as anp
 from autograd.numpy.numpy_boxes import ArrayBox
 import inspect
 from autograd.core import defvjp, primitive
-from autograd.differential_operators import grad, jacobian, elementwise_grad
+from autograd.differential_operators import jacobian, elementwise_grad
 import matplotlib.pyplot as plt
 import sys
 import os
@@ -128,8 +128,9 @@ def set_vjp(df):
             df_lst = [df_lst]
 
         pf = primitive(f)
-        first_arg = int(inspect.getfullargspec(f).args[0] == 'self')
-        defvjp(pf, *[lambda ans, *args: lambda g: g * df(*args) for df in df_lst], argnums=count(first_arg))
+        first_arg = int(len(inspect.getfullargspec(f).args) > 0 and inspect.getfullargspec(f).args[0] == 'self')
+        defvjp(pf, *[lambda ans, *args, df=df, **kwargs: lambda g: g * df(*args, **kwargs)
+                     for df in df_lst], argnums=count(first_arg))
         return pf
     return get_func
 
@@ -243,26 +244,29 @@ def cabs(a):
         return np.where(a < 0, -a, a)
 
 
-def dinterp(xp, x, y):
+def dinterp_dxp(xp, x, y):
     if len(x) > 1:
         return np.interp(xp, np.repeat(x, 2)[1:-1], np.repeat(np.diff(y) / np.diff(x), 2))
     else:
         return np.ones_like(xp)
 
 
-@primitive
+@set_vjp([dinterp_dxp])
 def interp(xp, x, y, *args, **kwargs):
-    if np.isrealobj(xp):
+    if all([np.isrealobj(v) for v in [xp, x, y]]):
         return np.interp(xp, x, y, *args, **kwargs)
     else:
+        # yp = np.interp(xp.real, x.real, y.real, *args, **kwargs)
+        # dyp_dxp = dinterp_dxp(xp.real, x.real, y.real)
+        # dyp_dx = fd(np.interp, True, 1)(xp.real, x.real, y.real)
+        # dyp_dy = fd(np.interp, True, 2)(xp.real, x.real, y.real)
+        # return yp + 1j * (xp.imag * dyp_dxp + x.imag * dyp_dx + y.imag * dyp_dy)
+
         yp = np.interp(xp.real, x, y, *args, **kwargs)
-        dyp = dinterp(xp.real, x, y)
+        dyp = dinterp_dxp(xp.real, x, y)
         return yp + xp.imag * 1j * dyp
 
 
-defvjp(interp, lambda ans, xp, x, y: lambda g: g * dinterp(xp, x, y))
-
-
 def logaddexp(x, y):
     if isinstance(x, ArrayBox) or isinstance(y, ArrayBox):
         return anp.logaddexp(x, y)
@@ -300,5 +304,21 @@ class UnivariateSpline(scipy_UnivariateSpline):
         return y
 
 
-# def get_dtype(arg_lst):
-#     return (float, np.complex128)[any([np.iscomplexobj(v) for v in arg_lst])]
+def trapz(y, x, axis=-1):
+    if isinstance(y, ArrayBox) or isinstance(x, ArrayBox):
+        x, y = asarray(x), asarray(y)
+        axis = np.arange(len(np.shape(y)))[axis]
+        # Silly implementation but np.take, np.diff and np.trapz did not seem to work with autograd
+        # I tried to implement gradients of np.trapz manually but failed to make it work for arbitrary axis
+        if axis == 0:
+            return ((y[:-1] + y[1:]) / 2 * (x[1:] - x[:-1])).sum(0)
+        elif axis == 1:
+            return ((y[:, :-1] + y[:, 1:]) / 2 * (x[:, 1:] - x[:, :-1])).sum(1)
+        elif axis == 2:
+            return ((y[:, :, :-1] + y[:, :, 1:]) / 2 * (x[:, :, 1:] - x[:, :, :-1])).sum(2)
+        elif axis == 4:
+            return ((y[:, :, :, :, :-1] + y[:, :, :, :, 1:]) / 2 * (x[:, :, :, :, 1:] - x[:, :, :, :, :-1])).sum(4)
+        else:   # pragma: no cover
+            raise NotImplementedError()
+    else:
+        return np.trapz(y, x, axis=axis)