diff --git a/euphonic/force_constants.py b/euphonic/force_constants.py index 4e562ca3c..971534303 100644 --- a/euphonic/force_constants.py +++ b/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, diff --git a/euphonic/powder.py b/euphonic/powder.py index 3fa3eccb5..c20ebf455 100644 --- a/euphonic/powder.py +++ b/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 diff --git a/euphonic/qpoint_frequencies.py b/euphonic/qpoint_frequencies.py index ec9f5c632..7418f6b93 100644 --- a/euphonic/qpoint_frequencies.py +++ b/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 diff --git a/euphonic/qpoint_phonon_modes.py b/euphonic/qpoint_phonon_modes.py index e8996a2dc..36aed6b9e 100644 --- a/euphonic/qpoint_phonon_modes.py +++ b/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: