Reformat using Black

Author: zichunhao

angular_momentum/__init__.py

@@ -1,2 +1,2 @@
 from .angular_momentum import *
-from .ladder_ops import *
+from .ladder_ops import *

angular_momentum/angular_momentum.py

@@ -1,22 +1,26 @@
 import numpy as np
 from .ladder_ops import J_plus, J_minus
 
+
 def J_x(j: int) -> np.ndarray:
     return (J_plus(j) + J_minus(j)) / 2
 
+
 def J_y(j: int) -> np.ndarray:
     return (J_plus(j) - J_minus(j)) / 2j
 
+
 def J_z(j: int) -> np.ndarray:
     return np.identity(int(2 * j + 1))
 
 
 def Jx(j: int) -> np.ndarray:
     return J_x(j)
 
+
 def Jy(j: int) -> np.ndarray:
     return J_y(j)
 
+
 def Jz(j: int) -> np.ndarray:
     return J_z(j)
-

angular_momentum/ladder_ops.py

@@ -1,9 +1,10 @@
 import numpy as np
 import math
 
+
 def J_plus(j: int) -> np.ndarray:
     if int(2 * j + 1) != 2 * j + 1:
-        raise ValueError(f'j must be a half integer. Found: {j}')
+        raise ValueError(f"j must be a half integer. Found: {j}")
 
     dim = int(2 * j + 1)
 
@@ -21,7 +22,7 @@ def J_plus(j: int) -> np.ndarray:
 
 def J_minus(j: int) -> np.ndarray:
     if int(2 * j + 1) != 2 * j + 1:
-        raise ValueError(f'j must be a half integer. Found: {j}')
+        raise ValueError(f"j must be a half integer. Found: {j}")
 
     dim = int(2 * j + 1)
 
@@ -37,46 +38,41 @@ def J_minus(j: int) -> np.ndarray:
     return mat
 
 
-def J_plus_component(
-    j_prime: int, 
-    m_prime: int, 
-    j: int, 
-    m: int
-) -> float:
-    '''
+def J_plus_component(j_prime: int, m_prime: int, j: int, m: int) -> float:
+    """
     Get the matrix element of the raising operator
     ..math::
         \langle j_\prime, m_\prime | J_+ | j, m \rangle
         \sqrt{(j - m) * (j + m + 1)} \delta_{j_\prime, j} \delta_{m_\prime, m + 1}.
-    '''
+    """
     if (j_prime != j) or (m_prime != m + 1):
         return 0
     return J_plus_coefficient(j, m)
 
 
 def J_minus_component(j_prime: int, m_prime: int, j: int, m: int) -> float:
-    '''
+    """
     Get the matrix element of the lowering operator
     ..math::
         \langle j_\prime, m_\prime | J_+ | j, m \rangle
         \sqrt{(j + m) * (j - m + 1)} \delta_{j_\prime, j} \delta_{m_\prime, m - 1}.
-    '''
+    """
     if (j_prime != j) or (m_prime != m - 1):
         return 0
     return J_minus_coefficient(j, m)
 
 
 def J_plus_coefficient(j: int, m: int) -> float:
-    '''
+    """
     Applies raising operator on the state :math:`|j, m \rangle`
     and returns the coefficient (:math:`\sqrt{(j + m) (j - m + 1)}`).
-    '''
+    """
     return math.sqrt((j - m) * (j + m + 1))
 
 
 def J_minus_coefficient(j: int, m: int) -> float:
-    '''
+    """
     Applies lowering operator on the state :math:`|j, m \rangle`
     and returns the coefficient (:math:`\sqrt{(j - m) (j + m + 1)}`).
-    '''
+    """
     return math.sqrt((j + m) * (j - m + 1))

cg/__init__.py

@@ -1,2 +1,2 @@
 from .clebsch_gordan_coefficients import *
-from .methods import *
+from .methods import *

cg/clebsch_gordan_coefficients.py

@@ -4,11 +4,12 @@
 from typing import Dict, Optional, Tuple, Union, List
 from .methods import *
 
+
 class ClebschGordanCoefficients:
     def __init__(self, jmax: int):
         if int(2 * jmax) != 2 * jmax:
             jmax = int(2 * jmax) / 2
-        
+
         self.__jmax = jmax
         jmin = 1 / 2
         self.__j_list = np.linspace(jmin, jmax, int(2 * (jmax - jmin) + 1))
@@ -23,133 +24,129 @@ def __getitem__(self, *args) -> Union[pd.DataFrame, Dict[float, pd.DataFrame]]:
             if j1 > self.jmax or j2 > self.jmax:
                 self.set_jmax(max(j1, j2))
             return self.__cg_dict.get((j1, j2))
-        
+
         elif len(*args) == 3:
             j1, j2, m = args[0]
             if j1 > self.jmax or j2 > self.jmax:
                 self.set_jmax(max(j1, j2))
             return self.__cg_dict.get((j1, j2)).get(m)
-        
+
         else:
-            raise IndexError(f'Invalid number of indices: {len(*args)}')
-    
+            raise IndexError(f"Invalid number of indices: {len(*args)}")
+
     def cg_matrix(
-        self, 
-        j1: int, 
-        j2: int, 
-        m: int, 
-        return_indices: bool = False
+        self, j1: int, j2: int, m: int, return_indices: bool = False
     ) -> Union[Tuple[np.ndarray, np.ndarray, np.ndarray], np.ndarray]:
         cg_dict = self.__cg_dict.get((j1, j2))
         if cg_dict is None:
             self.update(max(j1, j2))
-        
+
         mat = cg_dict.get(m)
-        
+
         if mat is None:
             return None
-        
+
         if return_indices:
-          return mat.to_numpy(), mat.index.to_numpy(), mat.columns.to_numpy()
-        
+            return mat.to_numpy(), mat.index.to_numpy(), mat.columns.to_numpy()
+
         return mat.to_numpy()
-    
-    def cg_matrices_all_m(self, j1: int, j2: int) -> Tuple[
-        List[np.ndarray], 
-        List[Tuple[
-            Tuple[float, float], # (j1, j2)
-            float # m
-        ]]
+
+    def cg_matrices_all_m(
+        self, j1: int, j2: int
+    ) -> Tuple[
+        List[np.ndarray], List[Tuple[Tuple[float, float], float]]  # (j1, j2)  # m
     ]:
-        '''
+        """
         Returns the Clebsch-Gordan coefficients for all m for the given j1, j2 in matrix forms
         along with the index labels of the matrix elements.
-        '''
+        """
         cg_mats = []
         cg_indices = []
-        
+
         if j1 > self.jmax or j2 > self.jmax:
             self.set_jmax(max(j1, j2))
-        
+
         m_list = self.__cg_dict[(j1, j2)].keys()
         for m in m_list:
             mat, row, col = self.cg_matrix(j1, j2, m, return_indices=True)
             cg_mats.append(mat)
             cg_indices.append((row, col))
         return cg_mats, cg_indices
-    
-    
+
     def set_jmax(self, jmax: int):
-        '''Updates jmax and appends larger CG coefficients'''
+        """Updates jmax and appends larger CG coefficients"""
         if jmax <= self.__jmax:
             return
-        
+
         jmax_old = self.__jmax
-        j_new_list = np.linspace(jmax_old, jmax, int(2 * (jmax - jmax_old) + 1))[1:]  # exclude jmax_old
-        
+        j_new_list = np.linspace(jmax_old, jmax, int(2 * (jmax - jmax_old) + 1))[
+            1:
+        ]  # exclude jmax_old
+
         new_dict = dict()
-        new_dict.update({
-            (j1, j2): cg_tables_all_m(j1, j2)
-            for j1, j2 in itertools.product(self.__j_list, j_new_list)
-        })
-        new_dict.update({
-            (j1, j2): cg_tables_all_m(j1, j2)
-            for j1, j2 in itertools.product(j_new_list, self.__j_list)
-        })
-        new_dict.update({
-            (j1, j2): cg_tables_all_m(j1, j2)
-            for j1, j2 in itertools.product(j_new_list, j_new_list)
-        })
-        
+        new_dict.update(
+            {
+                (j1, j2): cg_tables_all_m(j1, j2)
+                for j1, j2 in itertools.product(self.__j_list, j_new_list)
+            }
+        )
+        new_dict.update(
+            {
+                (j1, j2): cg_tables_all_m(j1, j2)
+                for j1, j2 in itertools.product(j_new_list, self.__j_list)
+            }
+        )
+        new_dict.update(
+            {
+                (j1, j2): cg_tables_all_m(j1, j2)
+                for j1, j2 in itertools.product(j_new_list, j_new_list)
+            }
+        )
+
         self.__cg_dict.update(new_dict)
         self.__j_list = np.concatenate((self.__j_list, j_new_list))
         self.__jmax = jmax
-        
+
         return True
-        
+
     def __repr__(self) -> str:
-        return f'ClebschGordanCoefficients(jmax={self.__jmax})'
-    
+        return f"ClebschGordanCoefficients(jmax={self.__jmax})"
+
     def __str__(self) -> str:
         return self.__repr__()
-    
+
     def display_all_tables(
-        self, 
-        jmin: Optional[int] = None, 
-        jmax: Optional[int] = None
+        self, jmin: Optional[int] = None, jmax: Optional[int] = None
     ) -> None:
-        '''
+        """
         Displays all CG coefficients stored in the object
-            If the bounds jmin and/or jmax are/is specified, display within the given bounds 
-        '''
+            If the bounds jmin and/or jmax are/is specified, display within the given bounds
+        """
         if jmax is not None and jmax > self.jmax:
             self.set_jmax(jmax)
-        
+
         for key, tables in self.__cg_dict.items():
             j1, j2 = key
             if (jmax is not None and jmax > 0) and (jmax < j1 or jmax < j2):
                 continue
-            
+
             if (jmin is not None and jmin > 0) and (jmin > j1 or jmin > j2):
                 continue
-            
-            print(f'{j1 = }, {j2 = }')
+
+            print(f"{j1 = }, {j2 = }")
             for m, table in tables.items():
-                print(f'{m = }:')
+                print(f"{m = }:")
                 display(table)
         return
-    
+
     @property
-    def dicts(self) -> Dict[
-        Tuple[float, float],
-        Dict[float, pd.DataFrame]
-    ]:
+    def dicts(self) -> Dict[Tuple[float, float], Dict[float, pd.DataFrame]]:
         return self.__cg_dict
-    
+
     @property
     def jmax(self):
         return self.__jmax
-    
+
     @property
     def keys(self):
         return self.__cg_dict.keys()