diff --git a/blobmodel/blobs.py b/blobmodel/blobs.py index f5a67b5..5152ef2 100644 --- a/blobmodel/blobs.py +++ b/blobmodel/blobs.py @@ -3,6 +3,7 @@ from nptyping import NDArray import numpy as np from .pulse_shape import AbstractBlobShape +import cmath class Blob: @@ -23,6 +24,7 @@ def __init__( t_drain: Union[float, NDArray], prop_shape_parameters: dict = None, perp_shape_parameters: dict = None, + blob_alignment: bool = True, ) -> None: self.int = int self.blob_id = blob_id @@ -42,10 +44,8 @@ def __init__( self.perp_shape_parameters = ( {} if perp_shape_parameters is None else perp_shape_parameters ) - if self.v_x != 0: - self.theta = np.arctan(self.v_y / self.v_x) - else: - self.theta = np.pi / 2 * np.sign(self.v_y) + self.blob_alignment = blob_alignment + self.theta = cmath.phase(self.v_x + self.v_y * 1j) if blob_alignment else 0.0 def discretize_blob( self, @@ -87,7 +87,7 @@ def discretize_blob( prop_dir = ( self._prop_dir_blob_position(t) if type(t) in [int, float] # t has dimensionality = 0, used for testing - else self._prop_dir_blob_position(t)[0, 0] + else self._prop_dir_blob_position(t[0, 0]) ) number_of_y_propagations = ( prop_dir + adjusted_Ly - x_border @@ -147,6 +147,7 @@ def _single_blob( if one_dimensional else self._perpendicular_direction_shape( y_perp + y_offset, + t, Ly, periodic_y, number_of_y_propagations=number_of_y_propagations, @@ -183,6 +184,7 @@ def _propagation_direction_shape( def _perpendicular_direction_shape( self, y: NDArray, + t: NDArray, Ly: float, periodic_y: bool, number_of_y_propagations: NDArray, @@ -190,21 +192,29 @@ def _perpendicular_direction_shape( if periodic_y: y_diffs = ( y - - self._perp_dir_blob_position() + - self._perp_dir_blob_position(t) + number_of_y_propagations * Ly * np.cos(self.theta) ) else: - y_diffs = y - self._perp_dir_blob_position() + y_diffs = y - self._perp_dir_blob_position(t) theta_y = y_diffs / self.width_perp return self.blob_shape.get_pulse_shape_perp( theta_y, **self.perp_shape_parameters ) def _prop_dir_blob_position(self, t: NDArray) -> NDArray: - return self.pos_x + (self.v_x**2 + self.v_y**2) ** 0.5 * (t - self.t_init) + return ( + self.pos_x + (self.v_x**2 + self.v_y**2) ** 0.5 * (t - self.t_init) + if self.blob_alignment + else self.pos_x + self.v_x * (t - self.t_init) + ) - def _perp_dir_blob_position(self) -> float: - return self.pos_y + def _perp_dir_blob_position(self, t: NDArray) -> float: + return ( + self.pos_y + if self.blob_alignment + else self.pos_y + self.v_y * (t - self.t_init) + ) def _rotate( self, origin: Tuple[float, float], x: NDArray, y: NDArray, angle: float diff --git a/blobmodel/model.py b/blobmodel/model.py index d80fbdf..8e9fd03 100644 --- a/blobmodel/model.py +++ b/blobmodel/model.py @@ -227,11 +227,11 @@ def _compute_start_stop(self, blob: Blob, speed_up: bool, error: float): 0, int( ( - blob.t_init * blob.v_x + np.abs(blob.t_init * blob.v_x) + blob.width_prop * np.log(error * np.sqrt(np.pi)) + blob.pos_x ) - / (self._geometry.dt * blob.v_x) + / np.abs(self._geometry.dt * blob.v_x) ), ) # ignores t_drain when calculating stop time @@ -244,7 +244,7 @@ def _compute_start_stop(self, blob: Blob, speed_up: bool, error: float): + self._geometry.Lx - blob.pos_x ) - / (blob.v_x * self._geometry.dt) + / np.abs(blob.v_x * self._geometry.dt) ), ) return start, stop diff --git a/blobmodel/stochasticality.py b/blobmodel/stochasticality.py index 0e2a80a..ff37cfa 100644 --- a/blobmodel/stochasticality.py +++ b/blobmodel/stochasticality.py @@ -52,6 +52,7 @@ def __init__( vy_parameter: float = 1.0, shape_param_x_parameter: float = 0.5, shape_param_y_parameter: float = 0.5, + blob_alignment: bool = True, ) -> None: """The following distributions are implemented: @@ -62,6 +63,10 @@ def __init__( ray: rayleight distribution with mean 1 deg: degenerate distribution at `free_parameter` zeros: array of zeros + + blob_alignment: bool = True, optional, If True, the blobs are aligned with their velocity. + If blob_aligment == True, the blob shapes are rotated in the propagation direction of the blob. + If blob_aligment == False, the blob shapes are independent of the propagation direction. """ self.amplitude_dist = A_dist self.width_x_dist = wx_dist @@ -77,6 +82,7 @@ def __init__( self.velocity_y_parameter = vy_parameter self.shape_param_x_parameter = shape_param_x_parameter self.shape_param_y_parameter = shape_param_y_parameter + self.blob_alignment = blob_alignment def _draw_random_variables( self, @@ -161,6 +167,7 @@ def sample_blobs( t_drain=t_drain, prop_shape_parameters=spxs[i], perp_shape_parameters=spys[i], + blob_alignment=self.blob_alignment, ) for i in range(num_blobs) ] diff --git a/tests/test_blob.py b/tests/test_blob.py index 3263baa..594658f 100644 --- a/tests/test_blob.py +++ b/tests/test_blob.py @@ -1,5 +1,6 @@ from blobmodel import Blob, BlobShapeImpl import numpy as np +from unittest.mock import MagicMock blob = Blob( blob_id=0, @@ -29,6 +30,33 @@ def test_initial_blob(): assert error < 1e-10, "Numerical error too big" +def test_blob_non_alignment(): + # Dead pixels have already been preprocessed and have an array of nans at their site + blob = Blob( + blob_id=0, + blob_shape=BlobShapeImpl("exp", "exp"), + amplitude=1, + width_prop=1, + width_perp=1, + velocity_x=1, + velocity_y=1, + pos_x=0, + pos_y=0, + t_init=0, + t_drain=10**10, + blob_alignment=False, + ) + x = np.arange(0, 10, 0.1) + y = np.array([0, 1]) + + mesh_x, mesh_y = np.meshgrid(x, y) + mock = MagicMock(return_value=mesh_x) + blob.blob_shape.get_pulse_shape_prop = mock + blob.discretize_blob(x=mesh_x, y=mesh_y, t=0, periodic_y=False, Ly=10) + + np.testing.assert_array_equal(mesh_x, mock.call_args[0][0]) + + def test_periodicity(): x = np.arange(0, 10, 0.1) y = np.arange(0, 10, 1) @@ -75,6 +103,40 @@ def test_single_point(): assert error < 1e-10, "Numerical error too big" +def test_negative_radial_velocity(): + vx = -1 + blob_sp = Blob( + blob_id=0, + blob_shape=BlobShapeImpl("gauss"), + amplitude=1, + width_prop=1, + width_perp=1, + velocity_x=vx, + velocity_y=1, + pos_x=0, + pos_y=6, + t_init=0, + t_drain=10**100, + ) + + x = np.arange(-5, 5, 0.1) + y = 0 + t = 2 + + mesh_x, mesh_y = np.meshgrid(x, y) + blob_values = blob_sp.discretize_blob( + x=mesh_x, y=mesh_y, t=t, periodic_y=True, Ly=10 + ) + + # The exact analytical expression for the expected values is a bit cumbersome, thus we just check + # that the shape is correct + maxx = np.max(blob_values) + expected_values = maxx * np.exp(-((mesh_x - vx * t) ** 2)) + error = np.max(abs(expected_values - blob_values)) + + assert error < 1e-10, "Numerical error too big" + + def test_theta_0(): blob_sp = Blob( blob_id=0, @@ -137,8 +199,3 @@ def test_kwargs(): blob_sp.discretize_blob(x=mesh_x, y=mesh_y, t=0, periodic_y=True, Ly=10) mock_ps.get_pulse_shape_prop.assert_called_with([[0]], lam=0.2) - - -test_initial_blob() -test_periodicity() -test_single_point()