diff --git a/blobmodel/blob_shape.py b/blobmodel/blob_shape.py index 4bb2cf1..fa1f1e8 100644 --- a/blobmodel/blob_shape.py +++ b/blobmodel/blob_shape.py @@ -20,6 +20,137 @@ def get_blob_shape_perp(self, theta: np.ndarray, **kwargs) -> np.ndarray: raise NotImplementedError +def _get_exponential_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the exponential pulse shape. + + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + + Returns + ------- + np.ndarray + Array representing the exponential pulse shape. + """ + return np.exp(theta) * np.heaviside(-1.0 * theta, 1) + + +def _get_lorentz_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the Lorentzian pulse shape. + + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + + Returns + ------- + np.ndarray + Array representing the Lorentzian pulse shape. + """ + return 1 / (np.pi * (1 + theta**2)) + + +def _get_double_exponential_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the double-exponential pulse shape. + + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + lam : float + Asymmetry parameter controlling the shape. + + Returns + ------- + np.ndarray + Array representing the double-exponential pulse shape. + """ + lam = kwargs["lam"] + assert (lam > 0.0) & (lam < 1.0) + kern = np.zeros(shape=np.shape(theta)) + kern[theta < 0] = np.exp(theta[theta < 0] / lam) + kern[theta >= 0] = np.exp(-theta[theta >= 0] / (1 - lam)) + return kern + + +def _get_gaussian_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the Gaussian pulse shape. + + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + + Returns + ------- + np.ndarray + Array representing the Gaussian pulse shape. + """ + return 1 / np.sqrt(2 * np.pi) * np.exp(-(theta**2) / 2) + + +def _get_rectangle_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the hard ellipse pulse shape. + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + Returns + ------- + np.ndarray + Array representing the rectangle pulse shape. + """ + return np.abs(theta) < 0.5 + + +def _get_secant_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the secant pulse shape. + + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + + Returns + ------- + np.ndarray + Array representing the secant pulse shape. + """ + return 2 / np.pi / (np.exp(theta) + np.exp(-theta)) + + +def _get_dipole_shape(theta: np.ndarray, **kwargs) -> np.ndarray: + """Compute the diople pulse shape as a derivative of a gaussian pulse shape. + + Parameters + ---------- + theta : np.ndarray + Array of theta values. + kwargs + Additional keyword arguments. + + Returns + ------- + np.ndarray + Array representing the dipole pulse shape. + """ + return -2 * theta / np.sqrt(2 * np.pi) * np.exp(-(theta**2) / 2) + + class BlobShapeImpl(AbstractBlobShape): """Implementation of the AbstractPulseShape class.""" @@ -77,132 +208,6 @@ def get_blob_shape_perp(self, theta: np.ndarray, **kwargs) -> np.ndarray: """ raise NotImplementedError - def _get_exponential_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the exponential pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - - Returns - ------- - np.ndarray - Array representing the exponential pulse shape. - """ - return np.exp(theta) * np.heaviside(-1.0 * theta, 1) - - def _get_lorentz_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the Lorentzian pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - - Returns - ------- - np.ndarray - Array representing the Lorentzian pulse shape. - """ - return 1 / (np.pi * (1 + theta**2)) - - def _get_double_exponential_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the double-exponential pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - lam : float - Asymmetry parameter controlling the shape. - - Returns - ------- - np.ndarray - Array representing the double-exponential pulse shape. - """ - lam = kwargs["lam"] - assert (lam > 0.0) & (lam < 1.0) - kern = np.zeros(shape=np.shape(theta)) - kern[theta < 0] = np.exp(theta[theta < 0] / lam) - kern[theta >= 0] = np.exp(-theta[theta >= 0] / (1 - lam)) - return kern - - def _get_gaussian_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the Gaussian pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - - Returns - ------- - np.ndarray - Array representing the Gaussian pulse shape. - """ - return 1 / np.sqrt(2 * np.pi) * np.exp(-(theta**2) / 2) - - def _get_secant_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the secant pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - - Returns - ------- - np.ndarray - Array representing the secant pulse shape. - """ - return 2 / np.pi / (np.exp(theta) + np.exp(-theta)) - - def _get_dipole_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the diople pulse shape as a derivative of a gaussian pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - - Returns - ------- - np.ndarray - Array representing the dipole pulse shape. - """ - return -2 * theta / np.sqrt(2 * np.pi) * np.exp(-(theta**2) / 2) - - def _get_rectangle_shape(theta: np.ndarray, **kwargs) -> np.ndarray: - """Compute the hard ellipse pulse shape. - - Parameters - ---------- - theta : np.ndarray - Array of theta values. - kwargs - Additional keyword arguments. - - Returns - ------- - np.ndarray - Array representing the rectangle pulse shape. - """ - return np.abs(theta) < 0.5 - __GENERATORS = { "exp": _get_exponential_shape, "gauss": _get_gaussian_shape, diff --git a/examples/blob_tilting.py b/examples/blob_tilting.py index 32773e3..e696e25 100644 --- a/examples/blob_tilting.py +++ b/examples/blob_tilting.py @@ -2,8 +2,8 @@ import numpy as np # velocities -vx = 0 -vy = 1 +vx = 1 +vy = 0 # sizes wx = 3 diff --git a/tests/test_pulse_shape.py b/tests/test_pulse_shape.py index 688bdcc..1d077ab 100644 --- a/tests/test_pulse_shape.py +++ b/tests/test_pulse_shape.py @@ -1,5 +1,5 @@ import pytest -from blobmodel import BlobShapeImpl, AbstractBlobShape, BlobShapeImpl +from blobmodel import AbstractBlobShape, BlobShapeImpl import numpy as np @@ -47,24 +47,33 @@ def test__get_double_exponential_shape(): theta = np.array([-1, 0, 1]) lam = 0.5 expected_result = np.array([0.13533528, 1.0, 0.13533528]) - assert np.allclose( - BlobShapeImpl._get_double_exponential_shape(theta, lam=lam), expected_result - ) + ps = BlobShapeImpl("2-exp", "2-exp") + values = ps.get_blob_shape_perp(theta, lam=lam) + assert np.max(np.abs(values - expected_result)) < 1e-5, "Wrong shape" def test__get_secant_shape(): theta = np.array([1, 2, 3]) expected_result = np.array([0.20628208, 0.08460748, 0.03161706]) - assert np.allclose(BlobShapeImpl._get_secant_shape(theta), expected_result) + + ps = BlobShapeImpl("secant", "secant") + values = ps.get_blob_shape_perp(theta) + assert np.max(np.abs(values - expected_result)) < 1e-5, "Wrong shape" def test__get_lorentz_shape(): theta = np.array([1, 2, 3]) expected_result = np.array([0.15915494, 0.06366198, 0.03183099]) - assert np.allclose(BlobShapeImpl._get_lorentz_shape(theta), expected_result) + + ps = BlobShapeImpl("lorentz", "lorentz") + values = ps.get_blob_shape_perp(theta) + assert np.max(np.abs(values - expected_result)) < 1e-5, "Wrong shape" def test__get_dipole_shape(): theta = np.array([1, 2, 3]) expected_result = np.array([-0.48394145, -0.21596387, -0.02659109]) - assert np.allclose(BlobShapeImpl._get_dipole_shape(theta), expected_result) + + ps = BlobShapeImpl("dipole", "dipole") + values = ps.get_blob_shape_perp(theta) + assert np.max(np.abs(values - expected_result)) < 1e-5, "Wrong shape"