From 929504a485c644e1df04fb8b7914eab9e1734fbd Mon Sep 17 00:00:00 2001
From: Takaya Uchida <tuchida@277Takajupiter2.lan>
Date: Mon, 28 Aug 2023 11:46:04 -0400
Subject: [PATCH] Normalization for 1D

---
 xwavelet/wavelet.py | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/xwavelet/wavelet.py b/xwavelet/wavelet.py
index 77470ab..8d7913b 100644
--- a/xwavelet/wavelet.py
+++ b/xwavelet/wavelet.py
@@ -189,7 +189,7 @@ def cwvlt(
         dimensions will be transformed. If the inputs are dask arrays, the
         arrays must not be chunked along these dimensions.
     t0 : float
-        Length scale.
+        Time scale.
     a : float
         Amplitude of wavelet.
     wtype : str
@@ -440,7 +440,10 @@ def wvlt_power_spectrum(
     else:
         C = 1.0
 
-    return np.abs(dawt) ** 2 * (dawt[s.dims[0]]) ** -1 * x0**2 / C
+    if len(dim) == 1:
+        return np.abs(dawt) ** 2 * x0 / C
+    elif len(dim) == 2:
+        return np.abs(dawt) ** 2 * (dawt[s.dims[0]]) ** -1 * x0**2 / C
 
 
 def wvlt_cross_spectrum(
@@ -575,4 +578,7 @@ def wvlt_cross_spectrum(
     else:
         C = 1.0
 
-    return (dawt * np.conj(dawt1)).real * (dawt[s.dims[0]]) ** -1 * x0**2 / C
+    if len(dim) == 1:
+        return (dawt * np.conj(dawt1)).real * x0 / C
+    elif len(dim) == 2:
+        return (dawt * np.conj(dawt1)).real * (dawt[s.dims[0]]) ** -1 * x0**2 / C