Apply pulse parameters

blobmodel/blobs.py

@@ -21,6 +21,8 @@ def __init__(
         pos_y: float,
         t_init: float,
         t_drain: Union[float, NDArray],
+        prop_shape_parameters: dict = None,
+        perp_shape_parameters: dict = None,
     ) -> None:
         self.int = int
         self.blob_id = blob_id
@@ -34,6 +36,12 @@ def __init__(
         self.pos_y = pos_y
         self.t_init = t_init
         self.t_drain = t_drain
+        self.prop_shape_parameters = (
+            {} if prop_shape_parameters is None else prop_shape_parameters
+        )
+        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:
@@ -168,7 +176,9 @@ def _propagation_direction_shape(
         else:
             x_diffs = x - self._prop_dir_blob_position(t)
         theta_x = x_diffs / self.width_prop
-        return self.blob_shape.get_pulse_shape_prop(theta_x, ...)
+        return self.blob_shape.get_pulse_shape_prop(
+            theta_x, **self.prop_shape_parameters
+        )
 
     def _perpendicular_direction_shape(
         self,
@@ -186,7 +196,9 @@ def _perpendicular_direction_shape(
         else:
             y_diffs = y - self._perp_dir_blob_position()
         theta_y = y_diffs / self.width_perp
-        return self.blob_shape.get_pulse_shape_perp(theta_y, ...)
+        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)

blobmodel/pulse_shape.py

@@ -1,6 +1,5 @@
 from abc import ABC, abstractmethod
 import numpy as np
-from nptyping import NDArray
 
 
 class AbstractBlobShape(ABC):
@@ -11,64 +10,63 @@ class AbstractBlobShape(ABC):
     """
 
     @abstractmethod
-    def get_pulse_shape_prop(self, theta_prop: NDArray, kwargs):
+    def get_pulse_shape_prop(self, theta: np.ndarray, **kwargs) -> np.ndarray:
         raise NotImplementedError
 
     @abstractmethod
-    def get_pulse_shape_perp(self, theta_perp: NDArray, kwargs):
+    def get_pulse_shape_perp(self, theta: np.ndarray, **kwargs) -> np.ndarray:
         raise NotImplementedError
 
 
 class BlobShapeImpl(AbstractBlobShape):
     """Implementation of the AbstractPulseShape class."""
 
-    __SHAPE_NAMES__ = {"exp", "gauss", "2-exp", "lorentz", "secant"}
-
     def __init__(self, pulse_shape_prop="gauss", pulse_shape_perp="gauss"):
-        self.__GENERATORS__ = {
-            "exp": BlobShapeImpl._get_exponential_shape,
-            "gauss": BlobShapeImpl._get_gaussian_shape,
-            "2-exp": BlobShapeImpl._get_double_exponential_shape,
-            "lorentz": BlobShapeImpl._get_lorentz_shape,
-            "secant": BlobShapeImpl._get_secant_shape,
-        }
-        self.pulse_shape_prop = pulse_shape_prop
-        self.pulse_shape_perp = pulse_shape_perp
         if (
-            pulse_shape_prop not in BlobShapeImpl.__SHAPE_NAMES__
-            or pulse_shape_perp not in BlobShapeImpl.__SHAPE_NAMES__
+            pulse_shape_prop not in BlobShapeImpl.__GENERATORS__.keys()
+            or pulse_shape_perp not in BlobShapeImpl.__GENERATORS__.keys()
         ):
             raise NotImplementedError(
                 f"{self.__class__.__name__}.blob_shape not implemented"
             )
+        self.get_pulse_shape_prop = BlobShapeImpl.__GENERATORS__.get(pulse_shape_prop)
+        self.get_pulse_shape_perp = BlobShapeImpl.__GENERATORS__.get(pulse_shape_perp)
 
-    def get_pulse_shape_prop(self, theta: np.ndarray, kwargs) -> np.ndarray:
-        return self.__GENERATORS__.get(self.pulse_shape_prop)(theta, kwargs)
+    def get_pulse_shape_prop(self, theta: np.ndarray, **kwargs) -> np.ndarray:
+        raise NotImplementedError
 
-    def get_pulse_shape_perp(self, theta: np.ndarray, kwargs) -> np.ndarray:
-        return self.__GENERATORS__.get(self.pulse_shape_perp)(theta, kwargs)
+    def get_pulse_shape_perp(self, theta: np.ndarray, **kwargs) -> np.ndarray:
+        raise NotImplementedError
 
     @staticmethod
-    def _get_exponential_shape(theta: np.ndarray, kwargs) -> np.ndarray:
+    def _get_exponential_shape(theta: np.ndarray, **kwargs) -> np.ndarray:
         return np.exp(theta) * np.heaviside(-1.0 * theta, 1)
 
     @staticmethod
-    def _get_lorentz_shape(theta: np.ndarray, kwargs) -> np.ndarray:
+    def _get_lorentz_shape(theta: np.ndarray, **kwargs) -> np.ndarray:
         return 1 / (np.pi * (1 + theta**2))
 
     @staticmethod
-    def _get_double_exponential_shape(theta: np.ndarray, kwargs) -> np.ndarray:
+    def _get_double_exponential_shape(theta: np.ndarray, **kwargs) -> np.ndarray:
         lam = kwargs["lam"]
         assert (lam > 0.0) & (lam < 1.0)
-        kern = np.zeros(len(theta))
+        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
 
     @staticmethod
-    def _get_gaussian_shape(theta: np.ndarray, kwargs) -> np.ndarray:
+    def _get_gaussian_shape(theta: np.ndarray, **kwargs) -> np.ndarray:
         return 1 / np.sqrt(np.pi) * np.exp(-(theta**2))
 
     @staticmethod
-    def _get_secant_shape(theta: np.ndarray, kwargs) -> np.ndarray:
+    def _get_secant_shape(theta: np.ndarray, **kwargs) -> np.ndarray:
         return 2 / np.pi / (np.exp(theta) + np.exp(-theta))
+
+    __GENERATORS__ = {
+        "exp": _get_exponential_shape,
+        "gauss": _get_gaussian_shape,
+        "2-exp": _get_double_exponential_shape,
+        "lorentz": _get_lorentz_shape,
+        "secant": _get_secant_shape,
+    }

blobmodel/stochasticality.py

@@ -43,11 +43,15 @@ def __init__(
         wy_dist: str = "deg",
         vx_dist: str = "deg",
         vy_dist: str = "normal",
+        spx_dist: str = "deg",
+        spy_dist: str = "deg",
         A_parameter: float = 1.0,
         wx_parameter: float = 1.0,
         wy_parameter: float = 1.0,
         vx_parameter: float = 1.0,
         vy_parameter: float = 1.0,
+        spx_parameter: float = 0.5,
+        spy_parameter: float = 0.5,
     ) -> None:
         """The following distributions are implemented:
 
@@ -64,11 +68,15 @@ def __init__(
         self.wy_dist = wy_dist
         self.vx_dist = vx_dist
         self.vy_dist = vy_dist
+        self.spx_dist = spx_dist
+        self.spy_dist = spy_dist
         self.A_parameter = A_parameter
         self.wx_parameter = wx_parameter
         self.wy_parameter = wy_parameter
         self.vx_parameter = vx_parameter
         self.vy_parameter = vy_parameter
+        self.spx_parameter = spx_parameter
+        self.spy_parameter = spy_parameter
 
     def _draw_random_variables(
         self,
@@ -115,6 +123,11 @@ def sample_blobs(
         wys = self._draw_random_variables(self.wy_dist, self.wy_parameter, num_blobs)
         vxs = self._draw_random_variables(self.vx_dist, self.vx_parameter, num_blobs)
         vys = self._draw_random_variables(self.vy_dist, self.vy_parameter, num_blobs)
+        spxs = self._draw_random_variables(self.spx_dist, self.spx_parameter, num_blobs)
+        spys = self._draw_random_variables(self.spy_dist, self.spy_parameter, num_blobs)
+        # For now, only a lambda parameter is implemented
+        spxs = [{"lam": s} for s in spxs]
+        spys = [{"lam": s} for s in spys]
         posxs = np.zeros(num_blobs)
         posys = np.random.uniform(low=0.0, high=Ly, size=num_blobs)
         t_inits = np.random.uniform(low=0, high=T, size=num_blobs)
@@ -132,6 +145,8 @@ def sample_blobs(
                 pos_y=posys[i],
                 t_init=t_inits[i],
                 t_drain=t_drain,
+                prop_shape_parameters=spxs[i],
+                perp_shape_parameters=spys[i],
             )
             for i in range(num_blobs)
         ]

tests/test_blob.py

@@ -108,6 +108,37 @@ def test_theta_0():
     assert error < 1e-10, "Numerical error too big"
 
 
+def test_kwargs():
+    from unittest.mock import MagicMock
+
+    mock_ps = BlobShapeImpl("2-exp", "2-exp")
+    mock_ps.get_pulse_shape_prop = MagicMock()
+
+    blob_sp = Blob(
+        blob_id=0,
+        blob_shape=mock_ps,
+        amplitude=1,
+        width_prop=1,
+        width_perp=1,
+        v_x=1,
+        v_y=0,
+        pos_x=0,
+        pos_y=0,
+        t_init=0,
+        t_drain=10**100,
+        prop_shape_parameters={"lam": 0.2},
+        perp_shape_parameters={"lam": 0.8},
+    )
+
+    x = 0
+    y = 0
+
+    mesh_x, mesh_y = np.meshgrid(x, y)
+    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()

tests/test_pulse_shape.py

@@ -0,0 +1,28 @@
+import pytest
+from blobmodel import BlobShapeImpl
+import numpy as np
+
+
+def test_gauss_pulse_shape():
+    ps = BlobShapeImpl()
+    x = np.arange(-10, 10, 0.1)
+    values = ps.get_pulse_shape_perp(x)
+    expected_values = 1 / np.sqrt(np.pi) * np.exp(-(x**2))
+    assert np.max(np.abs(values - expected_values)) < 1e-5, "Wrong gaussian shape"
+
+
+def test_throw_unknown_shape():
+    with pytest.raises(NotImplementedError):
+        BlobShapeImpl("LOL")
+
+
+def test_kwargs():
+    lam = 0.5
+    x = np.arange(-10, 10, 0.1)
+    expected_values = np.zeros(len(x))
+    expected_values[x < 0] = np.exp(x[x < 0] / lam)
+    expected_values[x >= 0] = np.exp(-x[x >= 0] / (1 - lam))
+
+    ps = BlobShapeImpl("2-exp", "2-exp")
+    values = ps.get_pulse_shape_perp(x, lam=0.5)
+    assert np.max(np.abs(values - expected_values)) < 1e-5, "Wrong shape"

tests/test_str.py

@@ -1,24 +1,7 @@
-import pytest
 from blobmodel import Model
 from blobmodel.geometry import Geometry
 
 
-def test_blob_shape_exception():
-    with pytest.raises(NotImplementedError):
-        bm = Model(
-            Nx=2,
-            Ny=2,
-            Lx=10,
-            Ly=10,
-            dt=0.5,
-            T=1,
-            periodic_y=False,
-            blob_shape="different_shape",
-            num_blobs=1,
-        )
-        bm.make_realization(speed_up=True, error=0.1)
-
-
 def test_geometry_str():
     geo = Geometry(1, 1, 1, 1, 1, 1, False)
     assert (
@@ -42,6 +25,5 @@ def test_model_str():
     assert str(bm) == "2d Blob Model with num_blobs:1 and t_drain:10"
 
 
-test_blob_shape_exception()
 test_geometry_str()
 test_model_str()