diff --git a/euphonic/broadening.py b/euphonic/broadening.py index c0c905a1f..c21e5944b 100644 --- a/euphonic/broadening.py +++ b/euphonic/broadening.py @@ -1,11 +1,11 @@ """ Functions for broadening spectra """ -from typing import List +from typing import Tuple, Union import warnings import numpy as np -from pint import Unit +from numpy.polynomial import Polynomial from scipy.optimize import nnls from scipy.stats import norm from scipy.signal import convolve @@ -15,33 +15,36 @@ def polynomial_broadening(bins: Quantity, x: Quantity, - width_polynomial: List[float], - weights: np.ndarray, - width_unit: Unit = None, - width_lower_limit: float = None, + width_polynomial: Tuple[Polynomial, Quantity], + weights: Union[np.ndarray, Quantity], + width_lower_limit: Quantity = None, adaptive_error: float = 1e-2) -> Quantity: - """Use fast approximate method to apply x-dependent Gaussian broadening + r"""Use fast approximate method to apply x-dependent Gaussian broadening Typically this is an energy-dependent instrumental resolution function. + Data is binned and broadened to output array with reciprocal units Parameters ---------- bins - Energy bins for output spectrum + Data bins for output spectrum x Data positions (to be binned) width_polynomial - Coefficients of polynomial fit to energy/width relationship. Order - should be consistent with numpy.polyfit / numpy.polyval (i.e. from - largest exponent to smallest.) Width is expressed as standard deviation - (not FHWM); the conversion factor between these conventions is + A numpy Polynomial object encodes broadening width as a function of + binning axis (typically energy). This is paired in the input tuple with + a scale factor Quantity; x values will be divided by this to obtain the + dimensionless function input, and the function output values are + multiplied by this Quantity to obtain appropriately dimensioned width + values. + + Width is expressed as standard deviation (not FHWM); the conversion + factor between these conventions is :math:`FWHM = \sqrt{8 \ln 2} \sigma` weights Weight for each data point corresponding to x. Note that these should be "counts" rather than binned spectral weights; this function will bin the data and apply bin-width weighting. - width_unit - x-axis units for width_polynomial. (By default, use same as bins.) width_lower_limit A lower bound is set for broadening width in WIDTH_UNIT. If set to None (default) the bin width will be used. To disable any lower limit, set @@ -50,18 +53,32 @@ def polynomial_broadening(bins: Quantity, Acceptable error for gaussian approximations, defined as the absolute difference between the areas of the true and approximate gaussians. - + """ - if width_unit is None: - width_unit = bins.units + + if isinstance(weights, np.ndarray): + weights = weights * ureg('dimensionless') + assert isinstance(weights, Quantity) + + width_poly, width_unit = width_polynomial + widths = width_poly((x / width_unit).magnitude) * width_unit + + # With newer versions of Pint we could dispense with most of this unit + # and magnitude shuffling and the numpy functions are handled more cleanly. + if width_lower_limit is None: - width_lower_limit = (bins[1] - bins[0]).to(width_unit).magnitude + width_lower_limit = np.diff(bins.magnitude).max() * bins.units - widths = np.polyval(width_polynomial, x.to(width_unit).magnitude) - widths = np.maximum(widths, width_lower_limit) * width_unit - - return width_interpolated_broadening(bins, x, widths, weights, - adaptive_error=adaptive_error) + widths = np.maximum(widths.magnitude, + width_lower_limit.to(widths.units).magnitude + ) * widths.units + + weights_unit = weights.units + + return width_interpolated_broadening(bins, x, widths, + weights.magnitude, + adaptive_error=adaptive_error + ) * weights_unit def width_interpolated_broadening(bins: Quantity, @@ -108,6 +125,7 @@ def width_interpolated_broadening(bins: Quantity, weights, adaptive_error)/conv + def _width_interpolated_broadening( bins: np.ndarray, x: np.ndarray, @@ -181,6 +199,7 @@ def _width_interpolated_broadening( return spectrum + def find_coeffs(spacing: float) -> np.ndarray: """" Function that, for a given spacing value, gives the coefficients of the diff --git a/tests_and_analysis/test/euphonic_test/test_broadening.py b/tests_and_analysis/test/euphonic_test/test_broadening.py index 3cc7e4518..aaaa21224 100644 --- a/tests_and_analysis/test/euphonic_test/test_broadening.py +++ b/tests_and_analysis/test/euphonic_test/test_broadening.py @@ -4,7 +4,8 @@ from scipy.integrate import simps from euphonic.broadening import (find_coeffs, - width_interpolated_broadening) + width_interpolated_broadening, + polynomial_broadening) from euphonic import ureg from tests_and_analysis.test.euphonic_test.test_force_constants\ import get_fc_path @@ -17,9 +18,9 @@ def test_polynomial_close_to_exact(): """Check polynomial broadening agrees with exact for trivial case""" - from euphonic.broadening import polynomial_broadening from numpy.random import RandomState from scipy.ndimage import gaussian_filter + from numpy.polynomial import Polynomial rng = RandomState(123) @@ -35,9 +36,8 @@ def test_polynomial_close_to_exact(): poly_broadened = polynomial_broadening( bins=(bins * ureg('meV')), x=(x * ureg('meV')), - width_polynomial=[0., 0., sigma], + width_polynomial=(Polynomial([sigma, 0., 0.]), ureg('meV')), weights=(y * bin_width), # Convert from spectrum heights to counts - width_unit=ureg('meV').units, adaptive_error=1e-5) npt.assert_allclose(exact, poly_broadened.to('1/meV').magnitude, @@ -73,6 +73,7 @@ def test_area_unchanged_for_broadened_dos(material, qpt_freqs_json, ebins_centres) assert adaptively_broadened_dos_area == pytest.approx(dos_area, rel=0.01) + @pytest.mark.parametrize( ('material, qpt_freqs_json, mode_widths_json,' 'expected_dos_json, ebins'), [ @@ -96,6 +97,7 @@ def test_lower_bound_widths_broadened(material, qpt_freqs_json, expected_dos = get_expected_spectrum1d(expected_dos_json) npt.assert_allclose(expected_dos.y_data.magnitude, dos.magnitude) + def test_uneven_bin_width_raises_warning(): qpt_freqs = get_qpt_freqs('quartz', 'quartz_554_full_qpoint_frequencies.json')