Merge pull request #160 from roxyboy/axisnum

Update `axis_num` when `detrend='linear'` is taken
xgcm · Jul 2, 2021 · 49497c7 · 49497c7
2 parents 4086c40 + 2d66615
commit 49497c7
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/xrft/detrend.py b/xrft/detrend.py
@@ -9,7 +9,8 @@
 
 
 def detrend(da, dim, detrend_type="constant"):
-    """Detrend a DataArray
+    """
+    Detrend a DataArray
 
     Parameters
     ----------

diff --git a/xrft/xrft.py b/xrft/xrft.py
@@ -411,8 +411,12 @@ def dft(
         delta_x.append(delta)
         lag_x.append(lag)
 
-    if detrend:
-        da = _detrend(da, dim, detrend_type=detrend)
+    if detrend is not None:
+        if detrend == "linear":
+            orig_dims = da.dims
+            da = _detrend(da, dim, detrend_type=detrend).transpose(*orig_dims)
+        else:
+            da = _detrend(da, dim, detrend_type=detrend)
 
     if window is not None:
         _, da = _apply_window(da, dim, window_type=window)
@@ -566,7 +570,7 @@ def idft(
     if chunks_to_segments:
         daft = _stack_chunks(daft, dim)
 
-    rawdims = daft.dims  # take care of segmented dimesions, if any
+    rawdims = daft.dims  # take care of segmented dimensions, if any
 
     if real_dim is not None:
         daft = daft.transpose(