diff --git a/news/test_squeeze.rst b/news/test_squeeze.rst new file mode 100644 index 0000000..750fbcf --- /dev/null +++ b/news/test_squeeze.rst @@ -0,0 +1,23 @@ +**Added:** + +* Polynomial squeeze of x-axis of morphed data + +**Changed:** + +* + +**Deprecated:** + +* + +**Removed:** + +* + +**Fixed:** + +* + +**Security:** + +* diff --git a/src/diffpy/morph/morphs/morphsqueeze.py b/src/diffpy/morph/morphs/morphsqueeze.py new file mode 100644 index 0000000..9a615b0 --- /dev/null +++ b/src/diffpy/morph/morphs/morphsqueeze.py @@ -0,0 +1,52 @@ +import numpy as np +from numpy.polynomial import Polynomial +from scipy.interpolate import CubicSpline + +from diffpy.morph.morphs.morph import LABEL_GR, LABEL_RA, Morph + + +class MorphSqueeze(Morph): + """Squeeze the morph function. + + This applies a polynomial to squeeze the morph non-linearly. The resulting + squeezed morph is interpolated to the (trimmed) target grid. + Only the overlapping region between the squeezed morph and the target + grid is used. The target is trimmed (or not) accordingly, and the final + outputs (morph and target) are returned on the same grid, defined by this + trimmed target range. + + Configuration Variables + ----------------------- + squeeze + list or array-like + Polynomial coefficients [a0, a1, ..., an] for the squeeze function. + """ + + # Define input output types + summary = "Squeeze morph by polynomial shift" + xinlabel = LABEL_RA + yinlabel = LABEL_GR + xoutlabel = LABEL_RA + youtlabel = LABEL_GR + parnames = ["squeeze"] + + def morph(self, x_morph, y_morph, x_target, y_target): + Morph.morph(self, x_morph, y_morph, x_target, y_target) + if self.squeeze is None: + self.x_morph_out = self.x_morph_in + self.y_morph_out = self.y_morph_in + return self.xyallout + + squeeze_polynomial = Polynomial(self.squeeze) + x_squeezed = self.x_morph_in + squeeze_polynomial(self.x_morph_in) + x_min = max(float(self.x_target_in[0]), float(x_squeezed[0])) + x_max = min(float(self.x_target_in[-1]), float(x_squeezed[-1])) + min_index = np.where(self.x_target_in >= x_min)[0][0] + max_index = np.where(self.x_target_in <= x_max)[0][-1] + self.x_target_out = self.x_target_in[min_index : max_index + 1] + self.y_target_out = self.y_target_in[min_index : max_index + 1] + self.y_morph_out = CubicSpline(x_squeezed, self.y_morph_in)( + self.x_target_out + ) + self.x_morph_out = self.x_target_out + return self.xyallout diff --git a/tests/test_morphsqueeze.py b/tests/test_morphsqueeze.py new file mode 100644 index 0000000..b3d1929 --- /dev/null +++ b/tests/test_morphsqueeze.py @@ -0,0 +1,101 @@ +import numpy as np +import pytest +from numpy.polynomial import Polynomial + +from diffpy.morph.morphs.morphsqueeze import MorphSqueeze + + +@pytest.mark.parametrize( + "squeeze_coeffs", + [ + # The order of coefficients is [a0, a1, a2, ..., an] + # Negative cubic squeeze coefficients + [-0.2, -0.01, -0.001, -0.001], + # Positive cubic squeeze coefficients + [0.2, 0.01, 0.001, 0.001], + # Positive and negative cubic squeeze coefficients + [0.2, -0.01, 0.002, -0.001], + # Quadratic squeeze coefficients + [-0.2, 0.005, -0.007], + # Linear squeeze coefficients + [0.1, 0.3], + # 4th order squeeze coefficients + [0.2, -0.01, 0.001, -0.001, 0.0004], + # Zeros and non-zeros, the full polynomial is applied + [0, 0.03, 0, -0.001], + # Testing zeros, expect no squeezing + [0, 0, 0, 0, 0, 0], + ], +) +def test_morphsqueeze_target_extends_beyond_morph(squeeze_coeffs): + # Target data extends beyond morph and different grids + x_target = np.linspace(-3, 25, 401) + y_target = np.sin(x_target) + x_morph = np.linspace(0, 10, 301) + squeeze_polynomial = Polynomial(squeeze_coeffs) + x_squeezed = x_morph + squeeze_polynomial(x_morph) + y_morph = np.sin(x_squeezed) + # Trim target data to the region overlapping with the squeezed morph + x_min = max(float(x_target[0]), float(x_squeezed[0])) + x_max = min(float(x_target[-1]), float(x_squeezed[-1])) + min_index = np.where(x_target >= x_min)[0][0] + max_index = np.where(x_target <= x_max)[0][-1] + x_morph_expected = x_target[min_index : max_index + 1] + y_morph_expected = y_target[min_index : max_index + 1] + morph = MorphSqueeze() + morph.squeeze = squeeze_coeffs + x_morph_actual, y_morph_actual, x_target_actual, y_target_actual = morph( + x_morph, y_morph, x_target, y_target + ) + assert np.allclose(y_morph_actual, y_morph_expected) + assert np.allclose(x_morph_actual, x_morph_expected) + assert np.allclose(x_target_actual, x_morph_expected) + assert np.allclose(y_target_actual, y_morph_expected) + + +@pytest.mark.parametrize( + "squeeze_coeffs", + [ + # The order of coefficients is [a0, a1, a2, ..., an] + # Negative cubic squeeze coefficients + [-0.2, -0.01, -0.001, -0.001], + # Positive cubic squeeze coefficients + [0.2, 0.01, 0.001, 0.001], + # Positive and negative cubic squeeze coefficients + [0.2, -0.01, 0.002, -0.001], + # Quadratic squeeze coefficients + [-0.2, 0.005, -0.007], + # Linear squeeze coefficients + [0.1, 0.3], + # 4th order squeeze coefficients + [0.2, -0.01, 0.001, -0.001, 0.0004], + # Zeros and non-zeros, the full polynomial is applied + [0, 0.03, 0, -0.001], + # Testing zeros, expect no squeezing + [0, 0, 0, 0, 0, 0], + ], +) +def test_morphsqueeze_morph_extends_beyond_target(squeeze_coeffs): + # Different grid for morph and target data to test different grids + x_target = np.linspace(0, 10, 101) + y_target = np.sin(x_target) + x_morph = np.linspace(-3, 15, 301) + squeeze_polynomial = Polynomial(squeeze_coeffs) + x_squeezed = x_morph + squeeze_polynomial(x_morph) + y_morph = np.sin(x_squeezed) + # Trim target data to the region overlapping with the squeezed morph + x_min = max(float(x_target[0]), float(x_squeezed[0])) + x_max = min(float(x_target[-1]), float(x_squeezed[-1])) + min_index = np.where(x_target >= x_min)[0][0] + max_index = np.where(x_target <= x_max)[0][-1] + x_morph_expected = x_target[min_index : max_index + 1] + y_morph_expected = y_target[min_index : max_index + 1] + morph = MorphSqueeze() + morph.squeeze = squeeze_coeffs + x_morph_actual, y_morph_actual, x_target_actual, y_target_actual = morph( + x_morph, y_morph, x_target, y_target + ) + assert np.allclose(y_morph_actual, y_morph_expected) + assert np.allclose(x_morph_actual, x_morph_expected) + assert np.allclose(x_target_actual, x_morph_expected) + assert np.allclose(y_target_actual, y_morph_expected)