Name commonly-used Literal groups

pace-neutrons · Sep 6, 2023 · 3cfc18a · 3cfc18a
1 parent 5b66185
commit 3cfc18a
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 16 deletions.
diff --git a/euphonic/broadening.py b/euphonic/broadening.py
@@ -14,6 +14,9 @@
 from euphonic import ureg
 
 
+ErrorFit = Literal['cheby-log', 'cubic']
+
+
 def variable_width_broadening(
     bins: Quantity,
     x: Quantity,
@@ -23,7 +26,7 @@ def variable_width_broadening(
     width_convention: Literal['FWHM', 'STD'] = 'FWHM',
     adaptive_error: float = 1e-2,
     shape: Literal['gauss', 'lorentz'] = 'gauss',
-    fit: Literal['cheby-log', 'cubic'] = 'cheby-log'
+    fit: ErrorFit = 'cheby-log'
     ) -> Quantity:
     r"""Apply x-dependent Gaussian broadening to 1-D data series
 
@@ -104,7 +107,7 @@ def width_interpolated_broadening(
     weights: np.ndarray,
     adaptive_error: float,
     shape: Literal['gauss', 'lorentz'] = 'gauss',
-    fit: Literal['cheby-log', 'cubic'] = 'cheby-log'
+    fit: ErrorFit = 'cheby-log'
     ) -> Quantity:
     """
     Uses a fast, approximate method to broaden a spectrum
@@ -161,7 +164,7 @@ def _lorentzian(x: np.ndarray, gamma: np.ndarray) -> np.ndarray:
 
 def _get_spacing(error,
                  shape: Literal['gauss', 'lorentz'] = 'gauss',
-                 fit: Literal['cheby-log', 'cubic'] = 'cheby-log'):
+                 fit: ErrorFit = 'cheby-log'):
     """
     Determine suitable spacing value for mode_width given accepted error level
 
@@ -207,7 +210,7 @@ def _width_interpolated_broadening(
     weights: np.ndarray,
     adaptive_error: float,
     shape: Literal['gauss', 'lorentz'] = 'gauss',
-    fit: Literal['cheby-log', 'cubic'] = 'cheby-log') -> np.ndarray:
+    fit: ErrorFit = 'cheby-log') -> np.ndarray:
     """
     Broadens a spectrum using a variable-width kernel, taking the
     same arguments as `variable_width` but expects arrays with

diff --git a/euphonic/qpoint_frequencies.py b/euphonic/qpoint_frequencies.py
@@ -10,7 +10,10 @@
 from euphonic.util import (_calc_abscissa, get_qpoint_labels)
 from euphonic import (ureg, Crystal, Quantity, Spectrum1D,
                       Spectrum1DCollection, Spectrum2D)
-from euphonic.broadening import _width_interpolated_broadening
+from euphonic.broadening import ErrorFit, _width_interpolated_broadening
+
+
+AdaptiveMethod = Literal['reference', 'fast']
 
 
 class QpointFrequencies:
@@ -93,9 +96,9 @@ def calculate_dos(
         dos_bins: Quantity,
         mode_widths: Optional[Quantity] = None,
         mode_widths_min: Quantity = Quantity(0.01, 'meV'),
-        adaptive_method: Literal['reference', 'fast'] = 'reference',
+        adaptive_method: AdaptiveMethod = 'reference',
         adaptive_error: float = 0.01,
-        adaptive_error_fit: Literal['cheby-log', 'cubic'] = 'cubic'
+        adaptive_error_fit: ErrorFit = 'cubic'
         ) -> Spectrum1D:
         """
         Calculates a density of states, in units of modes per atom per
@@ -160,9 +163,9 @@ def _calculate_dos(
         mode_widths: Optional[Quantity] = None,
         mode_widths_min: Quantity = Quantity(0.01, 'meV'),
         mode_weights: Optional[np.ndarray] = None,
-        adaptive_method: Literal['reference', 'fast'] = 'reference',
+        adaptive_method: AdaptiveMethod = 'reference',
         adaptive_error: float = 0.01,
-        adaptive_error_fit: Literal['cheby-log', 'cubic'] = 'cubic',
+        adaptive_error_fit: ErrorFit = 'cubic',
         q_idx: Optional[int] = None
         ) -> Quantity:
         """

diff --git a/euphonic/qpoint_phonon_modes.py b/euphonic/qpoint_phonon_modes.py
@@ -5,7 +5,9 @@
 import numpy as np
 
 from euphonic.validate import _check_constructor_inputs
+from euphonic.broadening import ErrorFit
 from euphonic.io import _obj_from_json_file, _obj_to_dict, _process_dict
+from euphonic.qpoint_frequencies import AdaptiveMethod
 from euphonic.readers import castep, phonopy
 from euphonic.util import (direction_changed, is_gamma, get_reference_data)
 from euphonic import (ureg, Quantity, Crystal, DebyeWaller, QpointFrequencies,
@@ -424,9 +426,9 @@ def calculate_pdos(
             self, dos_bins: Quantity,
             mode_widths: Optional[Quantity] = None,
             mode_widths_min: Quantity = Quantity(0.01, 'meV'),
-            adaptive_method: Literal['reference', 'fast'] = 'reference',
+            adaptive_method: AdaptiveMethod = 'reference',
             adaptive_error: Optional[float] = 0.01,
-            adaptive_error_fit: Literal['cheby-log', 'cubic'] = 'cubic',
+            adaptive_error_fit: ErrorFit = 'cubic',
             weighting: Optional[str] = None,
             cross_sections: Union[str, Dict[str, Quantity]] = 'BlueBook',
             ) -> Spectrum1DCollection:

diff --git a/euphonic/spectra.py b/euphonic/spectra.py
@@ -14,7 +14,7 @@
 from scipy.ndimage import correlate1d, gaussian_filter
 
 from euphonic import ureg, __version__
-from euphonic.broadening import variable_width_broadening
+from euphonic.broadening import ErrorFit, 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
@@ -578,7 +578,7 @@ def broaden(self: T, x_width: CallableQuantity,
                 width_lower_limit: Optional[Quantity] = None,
                 width_convention: Literal['FWHM', 'STD'] = 'FWHM',
                 width_interpolation_error: float = 0.01,
-                width_fit: Literal['cheby-log', 'cubic'] = 'cheby-log'
+                width_fit: ErrorFit = 'cheby-log'
                 ) -> T:  # noqa: F811
         ...
 
@@ -1047,7 +1047,7 @@ def broaden(self: T, x_width: CallableQuantity,
                 width_lower_limit: Optional[Quantity] = None,
                 width_convention: Literal['FWHM', 'STD'] = 'FWHM',
                 width_interpolation_error: float = 0.01,
-                width_fit: Literal['cheby-log', 'cubic'] = 'cheby-log'
+                width_fit: ErrorFit = 'cheby-log'
                 ) -> T:  # noqa: F811
         ...
 
@@ -1356,7 +1356,7 @@ def broaden(self: T,
                 y_width_lower_limit: Quantity = None,
                 width_convention: Literal['FWHM', 'STD'] = 'FWHM',
                 width_interpolation_error: float = 0.01,
-                width_fit: Literal['cheby-log', 'cubic'] = 'cheby-log'
+                width_fit: ErrorFit = 'cheby-log'
                 ) -> T:
         """
         Broaden z_data and return a new broadened Spectrum2D object
@@ -1465,7 +1465,7 @@ def _broaden_spectrum2d_with_function(
             width_convention: Literal['FWHM', 'STD'] = 'FWHM',
             width_interpolation_error: float = 1e-2,
             shape: Literal['gauss', 'lorentz'] = 'gauss',
-            width_fit: Literal['cheby-log', 'cubic'] = 'cheby-log'
+            width_fit: ErrorFit = 'cheby-log'
         ) -> 'Spectrum2D':
         """
         Apply value-dependent Gaussian broadening to one axis of Spectrum2D