More Literal typing

euphonic/force_constants.py

@@ -2,10 +2,12 @@
 import os
 import sys
 import warnings
-from typing import Optional, Tuple, Union, TypeVar, Sequence, Dict, Any, Type
+from typing import (Literal, Optional, Tuple, Union,
+                    TypeVar, Sequence, Dict, Any, Type)
 from multiprocessing import cpu_count
 
 import numpy as np
+from pint import Quantity
 from scipy.special import erfc
 from threadpoolctl import threadpool_limits
 
@@ -19,7 +21,7 @@
 from euphonic.util import (is_gamma, get_all_origins,
                            mode_gradients_to_widths,
                            _get_supercell_relative_idx)
-from euphonic import (ureg, Quantity, Crystal, QpointPhononModes,
+from euphonic import (ureg, Crystal, QpointPhononModes,
                       QpointFrequencies)
 
 
@@ -165,7 +167,7 @@ def calculate_qpoint_phonon_modes(
             self,
             qpts: np.ndarray,
             weights: Optional[np.ndarray] = None,
-            asr: Optional[str] = None,
+            asr: Optional[Literal['realspace', 'reciprocal']] = None,
             dipole: bool = True,
             dipole_parameter: float = 1.0,
             splitting: bool = True,
@@ -190,10 +192,10 @@ def calculate_qpoint_phonon_modes(
             If not given, equal weights are
             applied
         asr
-            One of {'realspace', 'reciprocal'}. Which acoustic sum rule
-            correction to apply. 'realspace' applies the correction to the
-            force constant matrix in real space. 'reciprocal' applies the
-            correction to the dynamical matrix at every q-point
+            Which acoustic sum rule correction to apply. 'realspace' applies
+            the correction to the force constant matrix in real
+            space. 'reciprocal' applies the correction to the dynamical matrix
+            at every q-point
         dipole
             Whether to calculate the dipole tail correction to the dynamical
             matrix at each q-point using the Ewald sum, if the Born
@@ -396,7 +398,7 @@ def calculate_qpoint_frequencies(
             self,
             qpts: np.ndarray,
             weights: Optional[np.ndarray] = None,
-            asr: Optional[str] = None,
+            asr: Optional[Literal['realspace', 'reciprocal']] = None,
             dipole: bool = True,
             dipole_parameter: float = 1.0,
             splitting: bool = True,
@@ -427,7 +429,7 @@ def _calculate_phonons_at_qpts(
             self,
             qpts: np.ndarray,
             weights: Optional[np.ndarray],
-            asr: Optional[str],
+            asr: Optional[Literal['realspace', 'reciprocal']],
             dipole: bool,
             dipole_parameter: float,
             splitting: bool,

euphonic/powder.py

@@ -1,7 +1,7 @@
 """Functions for averaging spectra in spherical q bins"""
 
 import numpy as np
-from typing import Optional, Union, Dict
+from typing import Literal, Optional, Union, Dict
 
 from euphonic import (Crystal, DebyeWaller, ForceConstants,
                       QpointFrequencies, QpointPhononModes, Spectrum1D,
@@ -10,9 +10,16 @@
 from euphonic.util import mp_grid, get_reference_data
 
 
+SphericalSamplingOptions = Literal['golden',
+                                   'sphere-projected-grid',
+                                   'spherical-polar-grid',
+                                   'spherical-polar-improved',
+                                   'random-sphere']
+
+
 def sample_sphere_dos(fc: ForceConstants,
                       mod_q: Quantity,
-                      sampling: str = 'golden',
+                      sampling: SphericalSamplingOptions = 'golden',
                       npts: int = 1000, jitter: bool = False,
                       energy_bins: Quantity = None,
                       **calc_modes_args
@@ -92,7 +99,7 @@ def sample_sphere_dos(fc: ForceConstants,
 def sample_sphere_pdos(
         fc: ForceConstants,
         mod_q: Quantity,
-        sampling: str = 'golden',
+        sampling: SphericalSamplingOptions = 'golden',
         npts: int = 1000, jitter: bool = False,
         energy_bins: Quantity = None,
         weighting: Optional[str] = None,
@@ -199,7 +206,7 @@ def sample_sphere_structure_factor(
     dw: DebyeWaller = None,
     dw_spacing: Quantity = 0.025 * ureg('1/angstrom'),
     temperature: Optional[Quantity] = 273. * ureg('K'),
-    sampling: str = 'golden',
+    sampling: SphericalSamplingOptions = 'golden',
     npts: int = 1000, jitter: bool = False,
     energy_bins: Quantity = None,
     scattering_lengths: Union[str, Dict[str, Quantity]] = 'Sears1992',
@@ -342,7 +349,7 @@ def _qpts_cart_to_frac(qpts: Quantity,
 
 
 def _get_qpts_sphere(npts: int,
-                     sampling: str = 'golden',
+                     sampling: SphericalSamplingOptions = 'golden',
                      jitter: bool = False) -> np.ndarray:
     """Get q-point coordinates according to specified sampling scheme
 

euphonic/qpoint_frequencies.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Any, Dict, List, Optional, Tuple, TypeVar, Type
+from typing import Any, Dict, List, Literal, Optional, Tuple, TypeVar, Type
 
 import numpy as np
 
@@ -88,13 +88,15 @@ def __setattr__(self, name: str, value: Any) -> None:
                                ['frequencies_unit'])
         super(QpointFrequencies, self).__setattr__(name, value)
 
-    def calculate_dos(self, dos_bins: Quantity,
-                      mode_widths: Optional[Quantity] = None,
-                      mode_widths_min: Quantity = Quantity(0.01, 'meV'),
-                      adaptive_method: str = 'reference',
-                      adaptive_error: float = 0.01,
-                      adaptive_error_fit: str = 'cubic'
-                      ) -> Spectrum1D:
+    def calculate_dos(
+        self,
+        dos_bins: Quantity,
+        mode_widths: Optional[Quantity] = None,
+        mode_widths_min: Quantity = Quantity(0.01, 'meV'),
+        adaptive_method: Literal['reference', 'fast'] = 'reference',
+        adaptive_error: float = 0.01,
+        adaptive_error_fit: Literal['cheby-log', 'cubic'] = 'cubic'
+        ) -> Spectrum1D:
         """
         Calculates a density of states, in units of modes per atom per
         energy unit, such that the integrated area is equal to 3.
@@ -114,8 +116,7 @@ def calculate_dos(self, dos_bins: Quantity,
             infinitely sharp peaks
         adaptive_method
             String. Specifies whether to use slow, reference adaptive method or
-            faster, approximate method. Allowed options are 'reference'
-            or 'fast', default is 'reference'.
+            faster, approximate method.
         adaptive_error
             Scalar float. Acceptable error for gaussian approximations
             when using the fast adaptive method, defined as the absolute
@@ -153,15 +154,17 @@ def calculate_dos(self, dos_bins: Quantity,
                                   adaptive_error_fit=adaptive_error_fit)
         return Spectrum1D(dos_bins, dos)
 
-    def _calculate_dos(self, dos_bins: Quantity,
-                       mode_widths: Optional[Quantity] = None,
-                       mode_widths_min: Quantity = Quantity(0.01, 'meV'),
-                       mode_weights: Optional[np.ndarray] = None,
-                       adaptive_method: str = 'reference',
-                       adaptive_error: float = 0.01,
-                       adaptive_error_fit: str = 'cubic',
-                       q_idx: Optional[int] = None
-                       ) -> Quantity:
+    def _calculate_dos(
+        self,
+        dos_bins: Quantity,
+        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_error: float = 0.01,
+        adaptive_error_fit: Literal['cheby-log', 'cubic'] = 'cubic',
+        q_idx: Optional[int] = None
+        ) -> Quantity:
         """
         Calculates a density of states (same arg defs as calculate_dos),
         but returns just a Quantity containing the y_data, rather

euphonic/qpoint_phonon_modes.py

@@ -1,5 +1,5 @@
 import math
-from typing import Dict, Optional, Union, TypeVar, Any, Type
+from typing import Dict, Literal, Optional, Union, TypeVar, Any, Type
 from collections.abc import Mapping
 
 import numpy as np
@@ -424,9 +424,9 @@ def calculate_pdos(
             self, dos_bins: Quantity,
             mode_widths: Optional[Quantity] = None,
             mode_widths_min: Quantity = Quantity(0.01, 'meV'),
-            adaptive_method: Optional[str] = 'reference',
+            adaptive_method: Literal['reference', 'fast'] = 'reference',
             adaptive_error: Optional[float] = 0.01,
-            adaptive_error_fit: str = 'cubic',
+            adaptive_error_fit: Literal['cheby-log', 'cubic'] = 'cubic',
             weighting: Optional[str] = None,
             cross_sections: Union[str, Dict[str, Quantity]] = 'BlueBook',
             ) -> Spectrum1DCollection: