diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 31bd8371c..8c37ad3eb 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -1,5 +1,5 @@ `Unreleased `_ ----------- +------------------------------------------------------------------------------- - Requirements @@ -49,7 +49,7 @@ `v1.2.1 `_ ------- +----------------------------------------------------------------------------- - Improvements @@ -82,7 +82,7 @@ - A deprecation in Numpy 1.25, which indirectly caused a test failure, has been addressed. `v1.2.0 `_ ------- +----------------------------------------------------------------------------- - Improvements: @@ -105,7 +105,7 @@ and ``--brille-npts-density`` arguments. `v1.1.0 `_ ----------- +----------------------------------------------------------------------------- - New features: @@ -121,7 +121,7 @@ interactions, see `#239 `_. `v1.0.0 `_ ----------- +----------------------------------------------------------------------------- - Changes: @@ -163,7 +163,7 @@ by default. To recover previous behaviour set this to ``False``. `v0.6.5 `_ ------- +----------------------------------------------------------------------------- - New Features: @@ -212,7 +212,7 @@ Previously this was interpreted as an astronomical unit by Pint. `v0.6.4 `_ ------- +----------------------------------------------------------------------------- - Improvements: @@ -243,7 +243,7 @@ which will just emit a warning. `v0.6.3 `_ ------- +----------------------------------------------------------------------------- - New Features: @@ -273,7 +273,7 @@ they refer to, and to match other arguments like ``vmin``, ``vmax``. `v0.6.2 `_ ------- +----------------------------------------------------------------------------- - Improvements: @@ -289,7 +289,7 @@ factor is not present `v0.6.1 `_ ------- +----------------------------------------------------------------------------- - Bug fixes: @@ -302,7 +302,7 @@ can now be correctly recovered by multiplying S(Q,w) by the energy bin width. `v0.6.0 `_ ------- +---------------------------------------------------------------------------- - Euphonic can now calculate neutron-weighted partial density of states, and has new ``Spectra`` features to handle PDOS data: @@ -361,7 +361,7 @@ ``euphonic-optimise-dipole-parameter``. `v0.5.2 `_ ------- +----------------------------------------------------------------------------- - Improvements: @@ -386,7 +386,7 @@ ``ebins - 1`` bins were generated `v0.5.1 `_ ----------- +----------------------------------------------------------------------------- - New Features: @@ -417,7 +417,7 @@ implementation: moved ``1/2`` factor and added explicit q-point weights `v0.5.0 `_ ----------- +----------------------------------------------------------------------------- - New Features: @@ -436,7 +436,7 @@ ``Spectrum1DCollection.from_castep_phonon_dos`` - **Adaptive broadening** is now available for DOS, which can obtain a more representative DOS than standard fixed-width broadening. See - `the docs `_ + `the docs `__ for details - Adaptive broadening can be used in the ``euphonic-dos`` tool with the ``--adaptive`` argument @@ -494,7 +494,7 @@ in implementation `v0.4.0 `_ ----------- +----------------------------------------------------------------------------- - There have been some major changes and improvements to spectra, plotting and command line tools, including: @@ -549,7 +549,7 @@ `#108 `_ `v0.3.2 `_ ----------- +----------------------------------------------------------------------------- - New Features: @@ -594,7 +594,7 @@ there is only one q-point (which is gamma) and ``splitting=True`` `v0.3.1 `_ ----------- +----------------------------------------------------------------------------- - New Features: @@ -625,7 +625,7 @@ `#77 `_ `v0.3.0 `_ ----------- +----------------------------------------------------------------------------- - Breaking Changes: @@ -665,14 +665,14 @@ in C, `#45 `_ `v0.2.2 `_ ------- +----------------------------------------------------------------------------- - Bug fixes: - Add MANIFEST.in for PyPI distribution `v0.2.1 `_ ------- +----------------------------------------------------------------------------- - Bug fixes: @@ -680,7 +680,7 @@ file, so refactor C files to avoid this `v0.2.0 `_ ------- +-------------------------------------------------------------------------------- - There are several breaking changes: diff --git a/doc/source/python-api.rst b/doc/source/python-api.rst index 74f2896ca..00d9504fb 100644 --- a/doc/source/python-api.rst +++ b/doc/source/python-api.rst @@ -9,6 +9,7 @@ Python API Force Constants Phonon Frequencies and Eigenvectors Phonon Frequencies Only + Crystal Structure Density of States Structure Factors Scattering Intensities diff --git a/doc/source/spectra.rst b/doc/source/spectra.rst index c5faf53ec..aa7c5cee5 100644 --- a/doc/source/spectra.rst +++ b/doc/source/spectra.rst @@ -86,7 +86,7 @@ apply this resolution function to a Spectrum1D we package the polynomial into a def tosca_resolution(energy): poly_in_wavenumber = Polynomial([2.5, 0.005, 1e-7]) return poly_in_wavenumber(energy.to('1/cm').magnitude) * ureg('1/cm') - dos_tosca = dos.broaden(tosca_resolution, shape='gauss', width_convention='STD') + dos_tosca = dos.broaden(tosca_resolution, shape='gauss', width_convention='std') Plotting -------- diff --git a/euphonic/broadening.py b/euphonic/broadening.py index 7d675fce5..8e0f6339d 100644 --- a/euphonic/broadening.py +++ b/euphonic/broadening.py @@ -1,27 +1,34 @@ """ Functions for broadening spectra """ -from typing import Callable, Union +from typing import Callable, Literal, Union import warnings import numpy as np from numpy.polynomial import Chebyshev +from pint import Quantity from scipy.optimize import nnls from scipy.stats import norm from scipy.signal import convolve -from euphonic import ureg, Quantity +from euphonic import ureg -def variable_width_broadening(bins: Quantity, - x: Quantity, - width_function: Callable[[Quantity], Quantity], - weights: Union[np.ndarray, Quantity], - width_lower_limit: Quantity = None, - width_convention: str = 'fwhm', - adaptive_error: float = 1e-2, - shape: str = 'gauss', - fit: str = 'cheby-log') -> Quantity: +ErrorFit = Literal['cheby-log', 'cubic'] +KernelShape = Literal['gauss', 'lorentz'] + + +def variable_width_broadening( + bins: Quantity, + x: Quantity, + width_function: Callable[[Quantity], Quantity], + weights: Union[np.ndarray, Quantity], + width_lower_limit: Quantity = None, + width_convention: Literal['fwhm', 'std'] = 'fwhm', + adaptive_error: float = 1e-2, + shape: KernelShape = 'gauss', + fit: ErrorFit = 'cheby-log' + ) -> Quantity: r"""Apply x-dependent Gaussian broadening to 1-D data series Typically this is an energy-dependent instrumental resolution function. @@ -49,14 +56,14 @@ def variable_width_broadening(bins: Quantity, (default) the bin width will be used. To disable any lower limit, set to 0 or lower. width_convention - Either 'std' or 'fwhm', to indicate if polynomial function yields - standard deviation (sigma) or full-width half-maximum. + Indicate if polynomial function yields standard deviation (sigma) or + full-width half-maximum. adaptive_error Acceptable error for gaussian approximations, defined as the absolute difference between the areas of the true and approximate gaussians. shape - Select 'gauss' or 'lorentz' kernel. + Select broadening kernel function. fit Select parametrisation of kernel width spacing to adaptive_error. 'cheby-log' is recommended: for shape 'gauss', 'cubic' is also @@ -76,15 +83,10 @@ def variable_width_broadening(bins: Quantity, widths = sigma_function(x) - # With newer versions of Numpy/Pint we could dispense with most of the unit - # and magnitude shuffling as the numpy functions are handled more cleanly. - if width_lower_limit is None: - width_lower_limit = np.diff(bins.magnitude).max() * bins.units + width_lower_limit = np.diff(bins).max() - widths = np.maximum(widths.magnitude, - width_lower_limit.to(widths.units).magnitude - ) * widths.units + widths = np.maximum(widths, width_lower_limit) if isinstance(weights, np.ndarray): weights = weights * ureg('dimensionless') @@ -99,13 +101,15 @@ def variable_width_broadening(bins: Quantity, fit=fit) * weights_unit -def width_interpolated_broadening(bins: Quantity, - x: Quantity, - widths: Quantity, - weights: np.ndarray, - adaptive_error: float, - shape: str = 'gauss', - fit: str = 'cheby-log') -> Quantity: +def width_interpolated_broadening( + bins: Quantity, + x: Quantity, + widths: Quantity, + weights: np.ndarray, + adaptive_error: float, + shape: KernelShape = 'gauss', + fit: ErrorFit = 'cheby-log' + ) -> Quantity: """ Uses a fast, approximate method to broaden a spectrum with a variable-width kernel. Exact Gaussians are calculated @@ -131,8 +135,8 @@ def width_interpolated_broadening(bins: Quantity, as the absolute difference between the areas of the true and approximate gaussians. shape - Select 'gauss' or 'lorentz' kernel. Widths will correspond to sigma or - gamma parameters respectively. + Select kernel shape. Widths will correspond to sigma or gamma + parameters respectively. fit Select parametrisation of kernel width spacing to adaptive_error. 'cheby-log' is recommended: for shape 'gauss', 'cubic' is also @@ -159,7 +163,9 @@ def _lorentzian(x: np.ndarray, gamma: np.ndarray) -> np.ndarray: return gamma / (2 * np.pi * (x**2 + (gamma / 2)**2)) -def _get_spacing(error, shape='gauss', fit='cheby-log'): +def _get_spacing(error, + shape: KernelShape = 'gauss', + fit: ErrorFit = 'cheby-log'): """ Determine suitable spacing value for mode_width given accepted error level @@ -198,13 +204,14 @@ def _get_spacing(error, shape='gauss', fit='cheby-log'): f'"Lorentz" shapes.') -def _width_interpolated_broadening(bins: np.ndarray, - x: np.ndarray, - widths: np.ndarray, - weights: np.ndarray, - adaptive_error: float, - shape='gauss', - fit='cheby-log') -> np.ndarray: +def _width_interpolated_broadening( + bins: np.ndarray, + x: np.ndarray, + widths: np.ndarray, + weights: np.ndarray, + adaptive_error: float, + shape: KernelShape = 'gauss', + 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 @@ -277,7 +284,8 @@ def _width_interpolated_broadening(bins: np.ndarray, return spectrum -def find_coeffs(spacing: float, shape='gauss') -> np.ndarray: +def find_coeffs(spacing: float, + 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/crystal.py b/euphonic/crystal.py index 4abd44c04..57191a54a 100644 --- a/euphonic/crystal.py +++ b/euphonic/crystal.py @@ -113,19 +113,15 @@ def reciprocal_cell(self) -> Quantity: Shape (3, 3) float Quantity in 1/length units, the reciprocal lattice vectors """ - cv = self._cell_vectors + cv = self.cell_vectors bxc = np.cross(cv[1], cv[2]) cxa = np.cross(cv[2], cv[0]) axb = np.cross(cv[0], cv[1]) - vol = self._cell_volume() + vol = self.cell_volume() norm = 2*np.pi/vol - recip = np.array([norm*bxc, - norm*cxa, - norm*axb])*(1/ureg.bohr) - - return recip.to(1/ureg(self.cell_vectors_unit)) + return np.vstack((norm * bxc, norm * cxa, norm * axb)) def cell_volume(self) -> Quantity: """ @@ -136,11 +132,7 @@ def cell_volume(self) -> Quantity: volume Scalar float Quantity in length**3 units. The cell volume """ - vol = self._cell_volume()*ureg.bohr**3 - return vol.to(ureg(self.cell_vectors_unit)**3) - - def _cell_volume(self) -> float: - return _cell_vectors_to_volume(self._cell_vectors) + return _cell_vectors_to_volume(self.cell_vectors) def get_mp_grid_spec(self, spacing: Quantity = 0.1 * ureg('1/angstrom') diff --git a/euphonic/force_constants.py b/euphonic/force_constants.py index 4e562ca3c..b1e29a91a 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 @@ -370,7 +372,8 @@ def calculate_qpoint_phonon_modes( .. math:: - \\frac{d\\omega_{\\mathbf{q}\\nu}}{d\\mathbf{q}} = \\frac{1}{2\\omega_{\\mathbf{q}\\nu}}{(E^{-1}\\frac{dD}{d\\mathbf{q}}E)} + \\frac{d\\omega_{\\mathbf{q}\\nu}}{d\\mathbf{q}} = \\frac{1}{2\\omega_{\\mathbf{q}\\nu}}{(E^{-1}\\frac{dD}{d\\mathbf{q}}E)} + :math:`\\frac{dD}{d\\mathbf{q}}` can be obtained by differentiating the Fourier equation above: .. math:: @@ -396,7 +399,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 +430,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, @@ -1097,7 +1100,7 @@ def _dipole_correction_init(crystal: Crystal, recip_q0 += recip_q0_tmp else: break - vol = crystal._cell_volume() + vol = crystal.cell_volume().to('bohr^3').magnitude recip_q0 *= math.pi/(vol*lambda_2) # Fill in remaining entries by symmetry @@ -1207,7 +1210,7 @@ def _calculate_dipole_correction( /(gvec_phases[:, i:]*q_phases[i:])) recip_dipole[i, i:] = np.einsum( 'ikl,ij->jkl', recip_exp, phase_exp) - cell_volume = crystal._cell_volume() + cell_volume = crystal.cell_volume().to('bohr^3').magnitude recip_dipole *= math.pi/(cell_volume*lambda_2) # Fill in remaining entries by symmetry @@ -1300,7 +1303,7 @@ def _calculate_gamma_correction(self, q_dir: np.ndarray) -> np.ndarray: recip_cell = self.crystal.reciprocal_cell().to('1/bohr').magnitude q_dir_cart = np.einsum('ij,i->j', recip_cell, q_dir) - cell_volume = self.crystal._cell_volume() + cell_volume = self.crystal.cell_volume().to('bohr^3').magnitude denominator = np.einsum('ij,i,j', dielectric, q_dir_cart, q_dir_cart) factor = 4*math.pi/(cell_volume*denominator) diff --git a/euphonic/plot.py b/euphonic/plot.py index 25a3a0705..7b4272425 100644 --- a/euphonic/plot.py +++ b/euphonic/plot.py @@ -61,8 +61,7 @@ def plot_1d_to_axis(spectra: Union[Spectrum1D, Spectrum1DCollection], f"same as the number of lines to plot (got {len(spectra)})") # Find where there are two identical x_data points in a row - breakpoints = (np.where(spectra.x_data.magnitude[:-1] - == spectra.x_data.magnitude[1:])[0] + breakpoints = (np.where(spectra.x_data[:-1] == spectra.x_data[1:])[0] + 1).tolist() breakpoints = [0] + breakpoints + [None] diff --git a/euphonic/powder.py b/euphonic/powder.py index 141eeb6d1..b51105e07 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', @@ -292,11 +299,9 @@ def sample_sphere_structure_factor( **calc_modes_args) dw = dw_phonons.calculate_debye_waller(temperature ) # type: DebyeWaller - else: - if not np.isclose(dw.temperature.to('K').magnitude, - temperature.to('K').magnitude): - raise ValueError('Temperature argument is not consistent with ' - 'temperature stored in DebyeWaller object.') + elif not np.isclose(dw.temperature, temperature): + raise ValueError('Temperature argument is not consistent with ' + 'temperature stored in DebyeWaller object.') qpts_cart = _get_qpts_sphere(npts, sampling=sampling, jitter=jitter ) * mod_q @@ -319,7 +324,7 @@ def _get_default_bins(phonons: Union[QpointPhononModes, QpointFrequencies], nbins: int = 1000) -> Quantity: """Get a default set of energy bin edges for set of phonon modes""" max_energy = np.max(phonons.frequencies) * 1.05 - return np.linspace(0, max_energy.magnitude, (nbins + 1)) * max_energy.units + return np.linspace(0, max_energy, (nbins + 1)) def _qpts_cart_to_frac(qpts: Quantity, @@ -340,13 +345,12 @@ def _qpts_cart_to_frac(qpts: Quantity, """ lattice = crystal.reciprocal_cell() - return np.linalg.solve(lattice.to(ureg('1/bohr')).magnitude.T, - qpts.to(ureg('1/bohr')).magnitude.T - ).T + # Cast 'dimensionless' to ensure length units are cancelled properly + return np.linalg.solve(lattice.T, qpts.T).T.to('dimensionless').magnitude 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..63310b2d2 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 @@ -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: @@ -88,13 +91,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: AdaptiveMethod = 'reference', + adaptive_error: float = 0.01, + adaptive_error_fit: ErrorFit = '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 +119,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 +157,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: AdaptiveMethod = 'reference', + adaptive_error: float = 0.01, + adaptive_error_fit: ErrorFit = '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 3b2b44b3b..2f63253a6 100644 --- 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: Optional[str] = 'reference', + adaptive_method: AdaptiveMethod = 'reference', adaptive_error: Optional[float] = 0.01, - adaptive_error_fit: str = 'cubic', + adaptive_error_fit: ErrorFit = 'cubic', weighting: Optional[str] = None, cross_sections: Union[str, Dict[str, Quantity]] = 'BlueBook', ) -> Spectrum1DCollection: @@ -567,7 +569,7 @@ def calculate_pdos( raise ValueError(( f'Unexpected dimensionality in cross_sections units, ' f'expected {ex_units}, got {str(cs[0].dimensionality)}')) - cs = [x.to('mbarn').magnitude for x in cs]*ureg('mbarn') + cs = [x.to('mbarn') for x in cs] else: cs = None diff --git a/euphonic/readers/castep.py b/euphonic/readers/castep.py index 387ac64ce..65455a0d9 100644 --- a/euphonic/readers/castep.py +++ b/euphonic/readers/castep.py @@ -1,7 +1,7 @@ from collections import defaultdict import re import struct -from typing import (Any, BinaryIO, Dict, List, NamedTuple, +from typing import (Any, BinaryIO, Dict, Iterator, List, NamedTuple, Optional, TextIO, Tuple, Union) import numpy as np @@ -53,15 +53,11 @@ def read_phonon_dos_data( weights = np.empty((0,)) freqs = np.empty((0, n_branches)) mode_grads = np.empty((0, n_branches)) - while True: - frequency_block = _read_frequency_block( + for frequency_block in _read_frequency_blocks( f, n_branches, extra_columns=[0], - terminator=' END GRADIENTS\n') - if frequency_block is None: - # We've reached 'END GRADIENTS' line - break + terminator=' END GRADIENTS\n'): qmode_grad = frequency_block.extra @@ -73,17 +69,11 @@ def read_phonon_dos_data( if 'BEGIN DOS' not in line: raise RuntimeError( f'Expected "BEGIN DOS" in {filename}, got {line}') - # max_rows arg not available until Numpy 1.16.0 - try: - dos_data = np.loadtxt(f, max_rows=n_bins) - except TypeError: - data = f.readlines() - dos_data = np.array([[float(elem) for elem in line.split()] - for line in data[:n_bins]]) + + dos_data = np.loadtxt(f, max_rows=n_bins) data_dict: Dict[str, Any] = {} - data_dict['crystal'] = {} - cry_dict = data_dict['crystal'] + cry_dict = data_dict['crystal'] = {} cry_dict['n_atoms'] = n_atoms cry_dict['cell_vectors'] = (cell_vectors*ureg('angstrom').to( cell_vectors_unit)).magnitude @@ -195,12 +185,7 @@ def read_phonon_data( repeated_qpt_ids = defaultdict(set) loto_split_indices = set() - while True: - frequency_block = _read_frequency_block(f, n_branches) - if frequency_block is None: - # Reached empty line, this should be end of file - break - + for frequency_block in _read_frequency_blocks(f, n_branches): qpt_id = frequency_block.qpt_id if prefer_non_loto and frequency_block.direction is not None: @@ -243,9 +228,8 @@ def read_phonon_data( # Multiple qpts with same CASTEP q-pt index: correct weights for qpt_id, indices in repeated_qpt_ids.items(): - intersection = indices & loto_split_indices if (prefer_non_loto - and intersection + and (intersection := indices & loto_split_indices) and len(indices) > len(intersection)): # Repeated q-point has both split and un-split variations; # set weights of split points to zero @@ -396,11 +380,9 @@ def _read_frequency_block( qpt = np.array([float(x) for x in floats[:3]]) qweight = float(floats[3]) - direction: Optional[np.ndarray] + direction: Optional[np.ndarray] = None if len(floats) >= 6: direction = floats[4:7] - else: - direction = None freq_lines = [f.readline().split() for i in range(n_branches)] @@ -419,6 +401,15 @@ def _read_frequency_block( return _FrequencyBlock(i_qpt, qpt, direction, qweight, qfreq, extra) +def _read_frequency_blocks(*args, **kwargs) -> Iterator[_FrequencyBlock]: + """Iterate over frequency blocks""" + while True: + frequency_block = _read_frequency_block(*args, **kwargs) + if frequency_block is None: + break + yield frequency_block + + def read_interpolation_data( filename: str, cell_vectors_unit: str = 'angstrom', @@ -458,11 +449,9 @@ def read_interpolation_data( with open(filename, 'rb') as f: int_type = '>i4' float_type = '>f8' - header = '' first_cell_read = True - while header.strip() != b'END': - header = _read_entry(f) - if header.strip() == b'BEGIN_UNIT_CELL': + while (header := _read_entry(f).strip()) != b'END': + if header == b'BEGIN_UNIT_CELL': # CASTEP writes the cell twice: the first is the # geometry optimised cell, the second is the original # cell. We only want the geometry optimised cell. @@ -470,7 +459,7 @@ def read_interpolation_data( (n_atoms, cell_vectors, atom_r, atom_mass, atom_type) = _read_cell(f, int_type, float_type) first_cell_read = False - elif header.strip() == b'FORCE_CON': + elif header == b'FORCE_CON': sc_matrix = np.transpose(np.reshape( _read_entry(f, int_type), (3, 3))) n_cells_in_sc = int(np.rint(np.absolute( @@ -483,16 +472,15 @@ def read_interpolation_data( cell_origins = np.reshape( _read_entry(f, int_type), (n_cells_in_sc, 3)) _ = _read_entry(f, int_type) # FC row not used - elif header.strip() == b'BORN_CHGS': + elif header == b'BORN_CHGS': born = np.reshape( _read_entry(f, float_type), (n_atoms, 3, 3)) - elif header.strip() == b'DIELECTRIC': + elif header == b'DIELECTRIC': dielectric = np.transpose(np.reshape( _read_entry(f, float_type), (3, 3))) data_dict: Dict[str, Any] = {} - data_dict['crystal'] = {} - cry_dict = data_dict['crystal'] + cry_dict = data_dict['crystal'] = {} cry_dict['n_atoms'] = n_atoms cry_dict['cell_vectors'] = cell_vectors*ureg( 'bohr').to(cell_vectors_unit).magnitude @@ -563,29 +551,27 @@ def _read_cell(file_obj: BinaryIO, int_type: str, float_type: str Shape (n_atoms,) string ndarray. The chemical symbols of each atom in the unit cell """ - header = '' - while header.strip() != b'END_UNIT_CELL': - header = _read_entry(file_obj) - if header.strip() == b'CELL%NUM_IONS': + while (header := _read_entry(file_obj).strip()) != b'END_UNIT_CELL': + if header == b'CELL%NUM_IONS': n_atoms = _read_entry(file_obj, int_type) - elif header.strip() == b'CELL%REAL_LATTICE': + elif header == b'CELL%REAL_LATTICE': cell_vectors = np.transpose(np.reshape( _read_entry(file_obj, float_type), (3, 3))) - elif header.strip() == b'CELL%NUM_SPECIES': + elif header == b'CELL%NUM_SPECIES': n_species = _read_entry(file_obj, int_type) - elif header.strip() == b'CELL%NUM_IONS_IN_SPECIES': + elif header == b'CELL%NUM_IONS_IN_SPECIES': n_atoms_in_species = _read_entry(file_obj, int_type) if n_species == 1: n_atoms_in_species = np.array([n_atoms_in_species]) - elif header.strip() == b'CELL%IONIC_POSITIONS': + elif header == b'CELL%IONIC_POSITIONS': max_atoms_in_species = max(n_atoms_in_species) atom_r_tmp = np.reshape(_read_entry(file_obj, float_type), (n_species, max_atoms_in_species, 3)) - elif header.strip() == b'CELL%SPECIES_MASS': + elif header == b'CELL%SPECIES_MASS': atom_mass_tmp = _read_entry(file_obj, float_type) if n_species == 1: atom_mass_tmp = np.array([atom_mass_tmp]) - elif header.strip() == b'CELL%SPECIES_SYMBOL': + elif header == b'CELL%SPECIES_SYMBOL': # Need to decode binary string for Python 3 compatibility if n_species == 1: atom_type_tmp = [_read_entry(file_obj, 'S8') diff --git a/euphonic/readers/phonopy.py b/euphonic/readers/phonopy.py index ce4f48f92..b1b0bc980 100644 --- a/euphonic/readers/phonopy.py +++ b/euphonic/readers/phonopy.py @@ -331,8 +331,8 @@ def convert_eigenvector_phases(phonon_dict: Dict[str, np.ndarray] eigvecs = np.reshape(phonon_dict['eigenvectors'], (n_qpts, n_atoms, 3, n_atoms, 3)) - na = np.newaxis - rk_diff = atom_r[:, na, :] - atom_r[na, :, :] + + rk_diff = atom_r[:, None, :] - atom_r[None, :, :] conversion = np.exp(-2j*np.pi*np.einsum('il,jkl->ijk', qpts, rk_diff)) eigvecs = np.einsum('ijklm,ijl->ijklm', eigvecs, conversion) return np.reshape(eigvecs, (n_qpts, 3*n_atoms, n_atoms, 3)) @@ -750,8 +750,7 @@ def read_interpolation_data( ufc = summary_dict['ufc'] data_dict: Dict[str, Any] = {} - data_dict['crystal'] = {} - cry_dict = data_dict['crystal'] + cry_dict = data_dict['crystal'] = {} cry_dict['cell_vectors'] = summary_dict['cell_vectors']*ureg( ulength).to(cell_vectors_unit).magnitude cry_dict['cell_vectors_unit'] = cell_vectors_unit diff --git a/euphonic/spectra.py b/euphonic/spectra.py index f6e14f369..3626cd29d 100644 --- a/euphonic/spectra.py +++ b/euphonic/spectra.py @@ -5,16 +5,17 @@ import math import json from numbers import Integral, Real -from typing import (Any, Callable, Dict, List, Optional, overload, +from typing import (Any, Callable, Dict, List, Literal, Optional, overload, Sequence, Tuple, TypeVar, Union, Type) import warnings -from pint import DimensionalityError +from pint import DimensionalityError, Quantity import numpy as np from scipy.ndimage import correlate1d, gaussian_filter -from euphonic import ureg, Quantity, __version__ -from euphonic.broadening import variable_width_broadening +from euphonic import ureg, __version__ +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 @@ -143,7 +144,7 @@ def _split_by_indices(self: T, indices: Union[Sequence[int], np.ndarray] def _split_by_tol(self: T, btol: float = 10.0) -> List[T]: """Split data along x-axis at detected breakpoints""" - diff = np.diff(self.x_data.magnitude) + diff = np.diff(self.x_data) median = np.median(diff) breakpoints = np.where((diff / median) > btol)[0] + 1 return self._split_by_indices(breakpoints) @@ -212,8 +213,9 @@ 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: str = 'gauss', - method: Optional[str] = None, + widths: Sequence[float], + shape: KernelShape = 'gauss', + method: Optional[Literal['convolve']] = None, ) -> np.ndarray: """ Broaden data along each axis by widths, returning the broadened @@ -295,18 +297,18 @@ def _gfwhm_to_sigma(fwhm: float, ax_bin_centres: np.ndarray) -> float: return sigma_bin @staticmethod - def _bin_edges_to_centres(bin_edges: np.ndarray) -> np.ndarray: + def _bin_edges_to_centres(bin_edges: Quantity) -> Quantity: return bin_edges[:-1] + 0.5*np.diff(bin_edges) @staticmethod - def _bin_centres_to_edges(bin_centres: np.ndarray) -> np.ndarray: + def _bin_centres_to_edges(bin_centres: Quantity) -> Quantity: return np.concatenate(( [bin_centres[0]], (bin_centres[1:] + bin_centres[:-1])/2, [bin_centres[-1]])) @staticmethod - def _is_bin_edge(data_length, bin_length) -> bool: + def _is_bin_edge(data_length: int, bin_length: int) -> bool: """Determine if axis data are bin edges or centres""" if bin_length == data_length + 1: return True @@ -333,8 +335,7 @@ def get_bin_edges(self) -> Quantity: if self._is_bin_edge(self.y_data.shape[-1], self.x_data.shape[0]): return self.x_data else: - return self._bin_centres_to_edges( - self.x_data.magnitude)*self.x_data.units + return self._bin_centres_to_edges(self.x_data) def get_bin_centres(self) -> Quantity: """ @@ -348,8 +349,7 @@ def get_bin_centres(self) -> Quantity: # Need to use -1 index for y_data so it also works for # Spectrum1DCollection which has y_data shape (n_spectra, bins) if self._is_bin_edge(self.y_data.shape[-1], self.x_data.shape[0]): - return self._bin_edges_to_centres( - self.x_data.magnitude)*self.x_data.units + return self._bin_edges_to_centres(self.x_data) else: return self.x_data @@ -357,8 +357,7 @@ def get_bin_widths(self) -> Quantity: """ Get x-axis bin widths """ - bins = self.get_bin_edges() - return np.diff(bins.magnitude) * bins.units + return np.diff(self.get_bin_edges()) def assert_regular_bins(self, message: str = '', rtol: float = 1e-5, @@ -379,8 +378,8 @@ def assert_regular_bins(self, message: str = '', """ bin_widths = self.get_bin_widths() - if not np.all(np.isclose(bin_widths.magnitude, - bin_widths.magnitude[0], + # Need to cast to magnitude to use isclose() with atol before Pint 0.21 + if not np.all(np.isclose(bin_widths.magnitude, bin_widths.magnitude[0], rtol=rtol, atol=atol)): raise AssertionError("Not all x-axis bins are the same width. " + message) @@ -568,26 +567,29 @@ def from_castep_phonon_dos(cls: Type[T], filename: str, @overload def broaden(self: T, x_width: Quantity, - shape: str = 'gauss', method: Optional[str] = None) -> T: + shape: KernelShape = 'gauss', + method: Optional[Literal['convolve']] = None + ) -> T: ... @overload def broaden(self: T, x_width: CallableQuantity, - shape: str = 'gauss', method: Optional[str] = None, + shape: KernelShape = 'gauss', + method: Optional[Literal['convolve']] = None, width_lower_limit: Optional[Quantity] = None, - width_convention: str = 'FWHM', + width_convention: Literal['fwhm', 'std'] = 'fwhm', width_interpolation_error: float = 0.01, - width_fit: str = 'cheby-log' + width_fit: ErrorFit = 'cheby-log' ) -> T: # noqa: F811 ... - def broaden(self: T, x_width: Union[Quantity, CallableQuantity], - shape: str = 'gauss', - method: Optional[str] = None, - width_lower_limit: Quantity = None, - width_convention: str = 'FWHM', - width_interpolation_error: float = 0.01, - width_fit: str = 'cheby-log' + def broaden(self: T, x_width, + shape='gauss', + method=None, + width_lower_limit=None, + width_convention='fwhm', + width_interpolation_error=0.01, + width_fit='cheby-log' ) -> T: # noqa: F811 """ Broaden y_data and return a new broadened spectrum object @@ -611,8 +613,8 @@ def broaden(self: T, x_width: Union[Quantity, CallableQuantity], Set a lower bound to width obtained calling x_width function. By default, this is equal to bin width. To disable, set to -Inf. width_convention - By default ('FWHM'), x_width is interpreted as full-width - half-maximum. Set to 'STD' to instead define standard deviation. + By default ('fwhm'), x_width is interpreted as full-width + half-maximum. Set to 'std' to instead define standard deviation. width_interpolation_error When x_width is a callable function, variable-width broadening is implemented by an approximate kernel-interpolation scheme. This @@ -781,7 +783,7 @@ def _split_by_indices(self, for x0, x1 in ranges] def __len__(self): - return self.y_data.magnitude.shape[0] + return self.y_data.shape[0] @overload def __getitem__(self, item: int) -> Spectrum1D: @@ -839,7 +841,7 @@ def _type_check(spectrum): _type_check(spectra[0]) x_data = spectra[0].x_data x_tick_labels = spectra[0].x_tick_labels - y_data_length = len(spectra[0].y_data.magnitude) + y_data_length = len(spectra[0].y_data) y_data_magnitude = np.empty((len(spectra), y_data_length)) y_data_magnitude[0, :] = spectra[0].y_data.magnitude y_data_units = spectra[0].y_data.units @@ -847,7 +849,7 @@ def _type_check(spectrum): for i, spectrum in enumerate(spectra[1:]): _type_check(spectrum) assert spectrum.y_data.units == y_data_units - assert np.allclose(spectrum.x_data.magnitude, x_data.magnitude) + assert np.allclose(spectrum.x_data, x_data) assert spectrum.x_data.units == x_data.units assert spectrum.x_tick_labels == x_tick_labels y_data_magnitude[i + 1, :] = spectrum.y_data.magnitude @@ -1034,27 +1036,30 @@ def from_castep_phonon_dos(cls: Type[T], filename: str) -> T: @overload def broaden(self: T, x_width: Quantity, - shape: str = 'gauss', method: Optional[str] = None) -> T: + shape: KernelShape = 'gauss', + method: Optional[Literal['convolve']] = None + ) -> T: ... @overload def broaden(self: T, x_width: CallableQuantity, - shape: str = 'gauss', method: Optional[str] = None, - width_lower_limit: Quantity = None, - width_convention: str = 'FWHM', + shape: KernelShape = 'gauss', + method: Optional[Literal['convolve']] = None, + width_lower_limit: Optional[Quantity] = None, + width_convention: Literal['fwhm', 'std'] = 'fwhm', width_interpolation_error: float = 0.01, - width_fit: str = 'cheby-log' + width_fit: ErrorFit = 'cheby-log' ) -> T: # noqa: F811 ... def broaden(self: T, - x_width: Union[Quantity, CallableQuantity], - shape: str = 'gauss', - method: Optional[str] = None, - width_lower_limit: Quantity = None, - width_convention: str = 'FWHM', - width_interpolation_error: float = 0.01, - width_fit: str = 'cheby-log' + x_width, + shape='gauss', + method=None, + width_lower_limit=None, + width_convention='fwhm', + width_interpolation_error=0.01, + width_fit='cheby-log' ) -> T: # noqa: F811 """ Individually broaden each line in y_data, returning a new @@ -1079,8 +1084,8 @@ def broaden(self: T, Set a lower bound to width obtained calling x_width function. By default, this is equal to bin width. To disable, set to -Inf. width_convention - By default ('FWHM'), x_width is interpreted as full-width - half-maximum. Set to 'STD' to instead define standard deviation. + By default ('fwhm'), x_width is interpreted as full-width + half-maximum. Set to 'std' to instead define standard deviation. width_interpolation_error When x_width is a callable function, variable-width broadening is implemented by an approximate kernel-interpolation scheme. This @@ -1346,13 +1351,13 @@ def _split_by_indices(self, def broaden(self: T, x_width: Optional[Union[Quantity, CallableQuantity]] = None, y_width: Optional[Union[Quantity, CallableQuantity]] = None, - shape: str = 'gauss', - method: Optional[str] = None, + shape: KernelShape = 'gauss', + method: Optional[Literal['convolve']] = None, x_width_lower_limit: Quantity = None, y_width_lower_limit: Quantity = None, - width_convention: str = 'FWHM', + width_convention: Literal['fwhm', 'std'] = 'fwhm', width_interpolation_error: float = 0.01, - width_fit: str = 'cheby-log' + width_fit: ErrorFit = 'cheby-log' ) -> T: """ Broaden z_data and return a new broadened Spectrum2D object @@ -1386,8 +1391,8 @@ def broaden(self: T, Set a lower bound to width obtained calling y_width function. By default, this is equal to y bin width. To disable, set to -Inf. width_convention - By default ('FWHM'), widths are interpreted as full-width - half-maximum. Set to 'STD' to instead define standard deviation. + By default ('fwhm'), widths are interpreted as full-width + half-maximum. Set to 'std' to instead define standard deviation. width_interpolation_error When x_width is a callable function, variable-width broadening is implemented by an approximate kernel-interpolation scheme. This @@ -1431,12 +1436,11 @@ def broaden(self: T, widths_in_bin_units, shape=shape, method=method) - spectrum = Spectrum2D( - np.copy(self.x_data.magnitude)*ureg(self.x_data_unit), - np.copy(self.y_data.magnitude)*ureg(self.y_data_unit), - z_broadened*ureg(self.z_data_unit), - copy.copy(self.x_tick_labels), - copy.deepcopy(self.metadata)) + spectrum = Spectrum2D(np.copy(self.x_data), + np.copy(self.y_data), + z_broadened*ureg(self.z_data_unit), + copy.copy(self.x_tick_labels), + copy.deepcopy(self.metadata)) else: spectrum = self @@ -1457,22 +1461,22 @@ def broaden(self: T, def _broaden_spectrum2d_with_function( spectrum: 'Spectrum2D', width_function: Callable[[Quantity], Quantity], - axis: str = 'y', + axis: Literal['x', 'y'] = 'y', width_lower_limit: Quantity = None, - width_convention: str = 'fwhm', + width_convention: Literal['fwhm', 'std'] = 'fwhm', width_interpolation_error: float = 1e-2, - shape: str = 'gauss', - width_fit: str = 'cheby-log') -> 'Spectrum2D': + shape: KernelShape = 'gauss', + width_fit: ErrorFit = 'cheby-log' + ) -> 'Spectrum2D': """ Apply value-dependent Gaussian broadening to one axis of Spectrum2D """ assert axis in ('x', 'y') bins = spectrum.get_bin_edges(bin_ax=axis) - bin_widths = np.diff(bins.magnitude) * bins.units + bin_widths = np.diff(bins) - if not np.all(np.isclose(bin_widths.magnitude, - bin_widths.magnitude[0])): + if not np.all(np.isclose(bin_widths, bin_widths[0])): bin_width = bin_widths.mean() else: bin_width = bin_widths[0] @@ -1484,7 +1488,7 @@ def _broaden_spectrum2d_with_function( z_data = z_data.T # Output data: matches input units - z_broadened = np.empty_like(z_data.magnitude) * spectrum.z_data.units + z_broadened = np.empty_like(z_data) * spectrum.z_data.units for i, row in enumerate(z_data): z_broadened[i] = variable_width_broadening( @@ -1501,12 +1505,11 @@ def _broaden_spectrum2d_with_function( if axis == 'x': z_broadened = z_broadened.T - return Spectrum2D( - np.copy(spectrum.x_data.magnitude) * ureg(spectrum.x_data_unit), - np.copy(spectrum.y_data.magnitude) * ureg(spectrum.y_data_unit), - z_broadened, - copy.copy(spectrum.x_tick_labels), - copy.copy(spectrum.metadata)) + return Spectrum2D(np.copy(spectrum.x_data), + np.copy(spectrum.y_data), + z_broadened, + copy.copy(spectrum.x_tick_labels), + copy.copy(spectrum.metadata)) def copy(self: T) -> T: """Get an independent copy of spectrum""" @@ -1516,7 +1519,7 @@ def copy(self: T) -> T: copy.copy(self.x_tick_labels), copy.deepcopy(self.metadata)) - def get_bin_edges(self, bin_ax: str = 'x') -> Quantity: + def get_bin_edges(self, bin_ax: Literal['x', 'y'] = 'x') -> Quantity: """ Get bin edges for the axis specified by bin_ax. If the size of bin_ax is one element larger than z_data along that axis, bin_ax @@ -1539,10 +1542,9 @@ def get_bin_edges(self, bin_ax: str = 'x') -> Quantity: if self._is_bin_edge(data_ax_len, bin_data.shape[0]): return bin_data else: - return self._bin_centres_to_edges(bin_data.magnitude - )*bin_data.units + return self._bin_centres_to_edges(bin_data) - def get_bin_centres(self, bin_ax: str = 'x') -> Quantity: + def get_bin_centres(self, bin_ax: Literal['x', 'y'] = 'x') -> Quantity: """ Get bin centres for the axis specified by bin_ax. If the size of bin_ax is the same size as z_data along that axis, bin_ax is @@ -1561,12 +1563,11 @@ def get_bin_centres(self, bin_ax: str = 'x') -> Quantity: bin_data = getattr(self, f'{bin_ax}_data') data_ax_len = self.z_data.shape[enum[bin_ax]] if self._is_bin_edge(data_ax_len, bin_data.shape[0]): - return self._bin_edges_to_centres( - bin_data.magnitude)*bin_data.units + return self._bin_edges_to_centres(bin_data) else: return bin_data - def get_bin_widths(self, bin_ax: str = 'x') -> Quantity: + def get_bin_widths(self, bin_ax: Literal['x', 'y'] = 'x') -> Quantity: """ Get x-bin widths along specified axis @@ -1576,9 +1577,9 @@ def get_bin_widths(self, bin_ax: str = 'x') -> Quantity: Axis of interest, 'x' or 'y' """ bins = self.get_bin_edges(bin_ax) - return np.diff(bins.magnitude) * bins.units + return np.diff(bins) - def assert_regular_bins(self, bin_ax: str = 'x', + def assert_regular_bins(self, bin_ax: Literal['x', 'y'], message: str = '', rtol: float = 1e-5, atol: float = 0.) -> None: @@ -1601,8 +1602,7 @@ def assert_regular_bins(self, bin_ax: str = 'x', """ bin_widths = self.get_bin_widths(bin_ax) - if not np.all(np.isclose(bin_widths.magnitude, - bin_widths.magnitude[0], + if not np.all(np.isclose(bin_widths, bin_widths[0], rtol=rtol, atol=atol)): raise AssertionError( f"Not all {bin_ax}-axis bins are the same width. " + message) @@ -1768,7 +1768,9 @@ def _get_cos_range(angle_range: Tuple[float]) -> Tuple[float]: return max(limiting_values), min(limiting_values) -def _distribution_1d(xbins: np.ndarray, xwidth: float, shape: str = 'lorentz' +def _distribution_1d(xbins: np.ndarray, + xwidth: float, + shape: Literal['lorentz'] = 'lorentz', ) -> np.ndarray: x = _get_dist_bins(xbins) if shape == 'lorentz': diff --git a/euphonic/structure_factor.py b/euphonic/structure_factor.py index 7fd83c5ea..97d2ddbd1 100644 --- a/euphonic/structure_factor.py +++ b/euphonic/structure_factor.py @@ -158,8 +158,7 @@ def calculate_1d_average(self, sqw_map = self._bose_corrected_structure_factor( e_bins, calc_bose=calc_bose, temperature=temperature) - spectrum = np.average( - sqw_map.magnitude, axis=0, weights=weights)*sqw_map.units + spectrum = np.average(sqw_map, axis=0, weights=weights) return Spectrum1D(e_bins, spectrum) def calculate_sqw_map(self, @@ -199,7 +198,7 @@ def calculate_sqw_map(self, StructureFactor.structure_factors is defined as the mode-resolved :math:`S(Q, \\omega_{\\mathbf{q}\\nu})` per atom of sample. To create an :math:`S(Q,\\omega)` map, it is binned in :math:`\\omega` - and the Bose population factor is applied [1]: + and the Bose population factor is applied [1]_: .. math:: @@ -318,8 +317,7 @@ def _bose_factor(self, temperature: Optional[Quantity] = None """ if self.temperature is not None: if (temperature is not None - and not np.isclose(temperature.to('K').magnitude, - self.temperature.to('K').magnitude)): + and not np.isclose(temperature, self.temperature)): raise ValueError(( 'Temperature provided to calculate the Bose factor ' '({:~P}) is not consistent with the temperature ' diff --git a/euphonic/util.py b/euphonic/util.py index ec71f9640..efe84ea66 100644 --- a/euphonic/util.py +++ b/euphonic/util.py @@ -49,9 +49,7 @@ def direction_changed(qpts: np.ndarray, tolerance: float = 5e-6 modq = np.linalg.norm(delta, axis=1) # Determine how much the direction has changed (dot) relative to the # vector sizes (modq) - direction_changed = (np.abs(np.abs(dot) - modq[1:]*modq[:-1]) > tolerance) - - return direction_changed + return np.abs(np.abs(dot) - modq[1:]*modq[:-1]) > tolerance def is_gamma(qpt: np.ndarray) -> Union[bool, np.ndarray]: @@ -253,10 +251,10 @@ def custom_decode(dct): try: unit = ureg(unit_str) - except UndefinedUnitError: + except UndefinedUnitError as exc: raise ValueError( f'Units "{unit_str}" from data file "{filename}" ' - 'are not supported by the Euphonic unit register.') + 'are not supported by the Euphonic unit register.') from exc return {key: value * unit for key, value in data.items() @@ -292,10 +290,11 @@ def mode_gradients_to_widths(mode_gradients: Quantity, cell_vectors: Quantity raise ValueError( f'Unexpected shape for mode_gradients {modg.shape}, ' f'expected (n_qpts, n_modes) or (n_qpts, n_modes, 3)') - cell_volume = (_cell_vectors_to_volume( - cell_vectors.magnitude)*cell_vectors.units**3).to('bohr**3').magnitude - q_spacing = 2/(np.cbrt(len(mode_gradients)*cell_volume)) - mode_widths = q_spacing*modg + + cell_volume = _cell_vectors_to_volume(cell_vectors) # type: Quantity + q_spacing = 2 / (np.cbrt(len(mode_gradients) + * cell_volume.to('bohr**3').magnitude)) + mode_widths = q_spacing * modg return mode_widths*ureg('hartree').to( mode_gradients.units/cell_vectors.units) @@ -408,9 +407,9 @@ def convert_fc_phases(force_constants: np.ndarray, atom_r: np.ndarray, co_idx = j break else: - raise Exception(( - 'Couldn\'t determine cell origins for ' - 'force constants matrix')) + raise ValueError( + "Couldn't determine cell origins for " + "force constants matrix") cell_origins_map[i] = co_idx if force_constants.shape[0] == force_constants.shape[1]: @@ -435,10 +434,12 @@ def convert_fc_phases(force_constants: np.ndarray, atom_r: np.ndarray, return fc_converted, cell_origins -def _cell_vectors_to_volume(cell_vectors: np.ndarray) -> float: +def _cell_vectors_to_volume(cell_vectors: Quantity) -> Quantity: """Convert 3x3 cell vectors to volume""" - return np.dot(cell_vectors[0], - np.cross(cell_vectors[1], cell_vectors[2])) + volume = np.dot(cell_vectors[0], + np.cross(cell_vectors[1], cell_vectors[2])) + assert isinstance(volume, Quantity) + return volume def _get_unique_elems_and_idx( @@ -711,7 +712,7 @@ def _get_supercell_relative_idx(cell_origins: np.ndarray, if np.all(dist_min <= 16*sys.float_info.epsilon): break if np.any(dist_min > 16*sys.float_info.epsilon): - raise Exception('Couldn\'t find supercell relative index') + raise ValueError('Couldn\'t find supercell relative index') return sc_relative_index