Name the gauss/lorentz literal type

pace-neutrons · Sep 6, 2023 · ecc5a3f · ecc5a3f
1 parent b4f3a9d
commit ecc5a3f
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 13 deletions.
diff --git a/euphonic/broadening.py b/euphonic/broadening.py
@@ -15,6 +15,7 @@
 
 
 ErrorFit = Literal['cheby-log', 'cubic']
+KernelShape = Literal['gauss', 'lorentz']
 
 
 def variable_width_broadening(
@@ -25,7 +26,7 @@ def variable_width_broadening(
     width_lower_limit: Quantity = None,
     width_convention: Literal['FWHM', 'STD'] = 'FWHM',
     adaptive_error: float = 1e-2,
-    shape: Literal['gauss', 'lorentz'] = 'gauss',
+    shape: KernelShape = 'gauss',
     fit: ErrorFit = 'cheby-log'
     ) -> Quantity:
     r"""Apply x-dependent Gaussian broadening to 1-D data series
@@ -106,7 +107,7 @@ def width_interpolated_broadening(
     widths: Quantity,
     weights: np.ndarray,
     adaptive_error: float,
-    shape: Literal['gauss', 'lorentz'] = 'gauss',
+    shape: KernelShape = 'gauss',
     fit: ErrorFit = 'cheby-log'
     ) -> Quantity:
     """
@@ -163,7 +164,7 @@ def _lorentzian(x: np.ndarray, gamma: np.ndarray) -> np.ndarray:
 
 
 def _get_spacing(error,
-                 shape: Literal['gauss', 'lorentz'] = 'gauss',
+                 shape: KernelShape = 'gauss',
                  fit: ErrorFit = 'cheby-log'):
     """
     Determine suitable spacing value for mode_width given accepted error level
@@ -209,7 +210,7 @@ def _width_interpolated_broadening(
     widths: np.ndarray,
     weights: np.ndarray,
     adaptive_error: float,
-    shape: Literal['gauss', 'lorentz'] = 'gauss',
+    shape: KernelShape = 'gauss',
     fit: ErrorFit = 'cheby-log') -> np.ndarray:
     """
     Broadens a spectrum using a variable-width kernel, taking the
@@ -284,7 +285,7 @@ def _width_interpolated_broadening(
 
 
 def find_coeffs(spacing: float,
-                shape: Literal['gauss', 'lorentz'] = 'gauss') -> np.ndarray:
+                shape: KernelShape = 'gauss') -> np.ndarray:
     """"
     Function that, for a given spacing value, gives the coefficients of the
     polynomial which describes the relationship between kernel width and the

diff --git a/euphonic/spectra.py b/euphonic/spectra.py
@@ -14,7 +14,8 @@
 from scipy.ndimage import correlate1d, gaussian_filter
 
 from euphonic import ureg, __version__
-from euphonic.broadening import ErrorFit, variable_width_broadening
+from euphonic.broadening import (ErrorFit, KernelShape,
+                                 variable_width_broadening)
 from euphonic.io import (_obj_to_json_file, _obj_from_json_file,
                          _obj_to_dict, _process_dict)
 from euphonic.readers.castep import read_phonon_dos_data
@@ -213,7 +214,7 @@ def split(self: T, indices: Union[Sequence[int], np.ndarray] = None,
     @staticmethod
     def _broaden_data(data: np.ndarray, bin_centres: Sequence[np.ndarray],
                       widths: Sequence[float],
-                      shape: Literal['gauss', 'lorentz'] = 'gauss',
+                      shape: KernelShape = 'gauss',
                       method: Optional[Literal['convolve']] = None,
                       ) -> np.ndarray:
         """
@@ -566,14 +567,14 @@ def from_castep_phonon_dos(cls: Type[T], filename: str,
 
     @overload
     def broaden(self: T, x_width: Quantity,
-                shape: Literal['gauss', 'lorentz'] = 'gauss',
+                shape: KernelShape = 'gauss',
                 method: Optional[Literal['convolve']] = None
                 ) -> T:
         ...
 
     @overload
     def broaden(self: T, x_width: CallableQuantity,
-                shape: Literal['gauss', 'lorentz'] = 'gauss',
+                shape: KernelShape = 'gauss',
                 method: Optional[Literal['convolve']] = None,
                 width_lower_limit: Optional[Quantity] = None,
                 width_convention: Literal['FWHM', 'STD'] = 'FWHM',
@@ -1035,14 +1036,14 @@ def from_castep_phonon_dos(cls: Type[T], filename: str) -> T:
 
     @overload
     def broaden(self: T, x_width: Quantity,
-                shape: Literal['gauss', 'lorentz'] = 'gauss',
+                shape: KernelShape = 'gauss',
                 method: Optional[Literal['convolve']] = None
                 ) -> T:
         ...
 
     @overload
     def broaden(self: T, x_width: CallableQuantity,
-                shape: Literal['gauss', 'lorentz'] = 'gauss',
+                shape: KernelShape = 'gauss',
                 method: Optional[Literal['convolve']] = None,
                 width_lower_limit: Optional[Quantity] = None,
                 width_convention: Literal['FWHM', 'STD'] = 'FWHM',
@@ -1350,7 +1351,7 @@ def _split_by_indices(self,
     def broaden(self: T,
                 x_width: Optional[Union[Quantity, CallableQuantity]] = None,
                 y_width: Optional[Union[Quantity, CallableQuantity]] = None,
-                shape: Literal['gauss', 'lorentz'] = 'gauss',
+                shape: KernelShape = 'gauss',
                 method: Optional[Literal['convolve']] = None,
                 x_width_lower_limit: Quantity = None,
                 y_width_lower_limit: Quantity = None,
@@ -1464,7 +1465,7 @@ def _broaden_spectrum2d_with_function(
             width_lower_limit: Quantity = None,
             width_convention: Literal['FWHM', 'STD'] = 'FWHM',
             width_interpolation_error: float = 1e-2,
-            shape: Literal['gauss', 'lorentz'] = 'gauss',
+            shape: KernelShape = 'gauss',
             width_fit: ErrorFit = 'cheby-log'
         ) -> 'Spectrum2D':
         """