Add axis kwarg to measurement functions

Fixes #118
mhostetter · Oct 14, 2023 · 246ad2c · 246ad2c
1 parent 76a62bc
commit 246ad2c
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 8 deletions.
diff --git a/src/sdr/_measurement/_energy.py b/src/sdr/_measurement/_energy.py
@@ -11,14 +11,20 @@
 
 
 @export
-def energy(x: npt.ArrayLike, db: bool = False) -> float:
+def energy(
+    x: npt.ArrayLike,
+    axis: int | tuple[int, ...] | None = None,
+    db: bool = False,
+) -> float:
     r"""
     Measures the energy of a time-domain signal $x[n]$.
 
     $$E = \sum_{n=0}^{N-1} \left| x[n] \right|^2$$
 
     Arguments:
         x: The time-domain signal $x[n]$ to measure.
+        axis: Axis or axes along which to compute the energy. The default is `None`, which computes the energy of
+            the entire array.
         db: Indicates whether to return the result in decibels (dB).
 
     Returns:
@@ -28,7 +34,7 @@ def energy(x: npt.ArrayLike, db: bool = False) -> float:
         measurement-energy
     """
     x = np.asarray(x)
-    E = np.sum(np.abs(x) ** 2)
+    E = np.sum(np.abs(x) ** 2, axis=axis)
     if db:
         E = to_db(E)
     return E
diff --git a/src/sdr/_measurement/_power.py b/src/sdr/_measurement/_power.py
@@ -11,14 +11,20 @@
 
 
 @export
-def peak_power(x: npt.ArrayLike, db: bool = False) -> float:
+def peak_power(
+    x: npt.ArrayLike,
+    axis: int | tuple[int, ...] | None = None,
+    db: bool = False,
+) -> float:
     r"""
     Measures the peak power of a time-domain signal $x[n]$.
 
     $$P_{\text{peak}} = \max \left| x[n] \right|^2$$
 
     Arguments:
         x: The time-domain signal $x[n]$ to measure.
+        axis: Axis or axes along which to compute the peak power. The default is `None`, which computes the peak power
+            of the entire array.
         db: Indicates whether to return the result in decibels (dB).
 
     Returns:
@@ -29,21 +35,27 @@ def peak_power(x: npt.ArrayLike, db: bool = False) -> float:
         measurement-power
     """
     x = np.asarray(x)
-    P_peak = np.max(np.abs(x) ** 2)
+    P_peak = np.max(np.abs(x) ** 2, axis=axis)
     if db:
         P_peak = to_db(P_peak, type="power")
     return P_peak
 
 
 @export
-def average_power(x: npt.ArrayLike, db: bool = False) -> float:
+def average_power(
+    x: npt.ArrayLike,
+    axis: int | tuple[int, ...] | None = None,
+    db: bool = False,
+) -> float:
     r"""
     Measures the average power of a time-domain signal $x[n]$.
 
     $$P_{\text{avg}} = \frac{E}{N} = \frac{1}{N} \sum_{n=0}^{N-1} \left| x[n] \right|^2$$
 
     Arguments:
         x: The time-domain signal $x[n]$ to measure.
+        axis: Axis or axes along which to compute the average power. The default is `None`, which computes the average
+            power of the entire array.
         db: Indicates whether to return the result in decibels (dB).
 
     Returns:
@@ -54,21 +66,26 @@ def average_power(x: npt.ArrayLike, db: bool = False) -> float:
         measurement-power
     """
     x = np.asarray(x)
-    P_avg = np.mean(np.abs(x) ** 2)
+    P_avg = np.mean(np.abs(x) ** 2, axis=axis)
     if db:
         P_avg = to_db(P_avg, type="power")
     return P_avg
 
 
 @export
-def papr(x: npt.ArrayLike) -> float:
+def papr(
+    x: npt.ArrayLike,
+    axis: int | tuple[int, ...] | None = None,
+) -> float:
     r"""
     Measures the peak-to-average power ratio (PAPR) of a time-domain signal $x[n]$.
 
     $$\text{PAPR} = 10 \log_{10} \frac{P_{\text{peak}}}{P_{\text{avg}}}$$
 
     Arguments:
         x: The time-domain signal $x[n]$ to measure.
+        axis: Axis or axes along which to compute the PAPR. The default is `None`, which computes the PAPR of the
+            entire array.
 
     Returns:
         The PAPR of $x[n]$ in dB.
@@ -83,4 +100,4 @@ def papr(x: npt.ArrayLike) -> float:
         measurement-power
     """
     x = np.asarray(x)
-    return to_db(peak_power(x) / average_power(x))
+    return to_db(peak_power(x, axis=axis) / average_power(x, axis=axis))