Merge pull request #21 from roxyboy/coords

More robust coordinate setting for the wavelet

Author: roxyboy

xwavelet/tests/test_wavelet.py

@@ -120,12 +120,12 @@ def test_isotropic_ps_slope(chunk, N=128, dL=1.0, amp=1e0, slope=-3.0, xo=50):
         dL,
         amp,
         slope,
-        other_dim_sizes=[20],
+        other_dim_sizes=[30],
         dim_order=True,
     )
 
     if chunk:
-        theta = theta.chunk({"d0": 30, "y": 64, "x": 64})
+        theta = theta.chunk({"d0": 5, "y": 64, "x": 64})
 
     s = xr.DataArray(
         np.linspace(0.1, 1.0, 20),

xwavelet/wavelet.py

@@ -79,11 +79,11 @@ def _morlet(xo, ntheta, a, s, y, x, dim):
     th = xr.DataArray(th, dims=["angle"], coords={"angle": th})
 
     # rotated positions
-    yp = np.sin(th) * s**-1 * y
-    xp = np.cos(th) * s**-1 * x
+    yp = np.sin(th) * s**-1 * (y - y.mean())
+    xp = np.cos(th) * s**-1 * (x - x.mean())
 
     arg1 = 2j * np.pi * ko * (yp - xp)
-    arg2 = -(x**2 + y**2) / 2 / s**2 / xo**2
+    arg2 = -((x - x.mean()) ** 2 + (y - y.mean()) ** 2) / 2 / s**2 / xo**2
     m = a * np.exp(arg1) * np.exp(arg2)
 
     return m, th