From 9de7998ca506a1ade31c56421839281ffce6ac03 Mon Sep 17 00:00:00 2001
From: Jiemin Li <jeremy20150118@gmail.com>
Date: Wed, 9 Oct 2024 12:19:35 -0400
Subject: [PATCH] update the condition in the function of
 perpendicular_momentum

---
 src/test_data/arpes.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/src/test_data/arpes.py b/src/test_data/arpes.py
index a7ea2fd..c3e47db 100644
--- a/src/test_data/arpes.py
+++ b/src/test_data/arpes.py
@@ -363,7 +363,24 @@ def perpendicular_momentum(photon_energy, parallel_momentum,
     $\k_{perpendicular}$ are the surface parallel and perpendicular momentum,
     $E_{ph}$ is the photon energy, $E_{b}$ is the electron binding energy and
     $\Phi$ is the work function.
+    Parameters
+    ----------
+    photon_energy : float or 1D numpy.ndarray,
+        The photon energy (or energies) in eV.
+    parallel_momentum : float or 1D numpy.ndarray,
+        The parallel momentum in the unit of the inverse of angstrom.
+    binding_energy : float or 1D numpy.ndarray,
+        The binding energy (or energies) in eV.
+
+    inner_potential: float or 1D numpy.ndarray
+        The inner potential in eV.
+    work_function: float or 1D numpy.ndarray
+        The work function in eV.
 
+    Returns
+    -------
+    $k_z$ : float or 1D numpy.ndarray,
+        The perpendicular momentum in the unit of the inverse of angstrom.
     """
     check_data_len = [len(value)
                       if (isinstance(value, (np.ndarray)) and
@@ -378,11 +395,10 @@ def perpendicular_momentum(photon_energy, parallel_momentum,
     if not np.all(check_data_len):
         raise ValueError(f'The input must be an int, float or 1D numpy array!')
     else:
-        check_data_size = [x for x in check_data_len if (x!=1)]
-        if len(check_data_size)>1:
-            if check_data_size[1:] != check_data_size[:-1]:
-                raise ValueError(f'The length of inputs must be the same if '
-                                 f'their length is larger than 1!')
+        check_data_len = list(filter(lambda a: a!=1,check_data_len))
+        if check_data_len[1:] != check_data_len[:-1]:
+            raise ValueError(f'The length of inputs must be the same if '
+                             f'their length is larger than 1!')
 
     # h_bar = sci_const.hbar  # in J.s
     # m_e = sci_const.m_e  # in Kg or J.s^2/m^2 (E=mc^2)