Merge pull request #61 from uit-cosmo/pulse_params

Apply pulse parameters

README.md

@@ -76,8 +76,7 @@ Alternatively, you can specify all blob parameters exactly as you want by writin
             allow periodicity in y-direction 
 
             !!!  this is only a good approximation if Ly is significantly bigger than blobs  !!!
-- `blob_shape`: str, optional,
-            switch between `gauss` and `exp` blob
+- `blob_shape`: AbstractBlobShape or str, optional,
 - `num_blobs`: int, optional
             number of blobs
 - `t_drain`: float, optional,
@@ -109,6 +108,8 @@ Alternatively, you can specify all blob parameters exactly as you want by writin
             `free_parameter` for vx
 - `vy_parameter`: float, optional,
             `free_parameter` for vy
+- `shape_param_x_parameter`: float = 0.5,
+- `shape_param_y_parameter`: float = 0.5
 
 The following distributions are implemented:
 

blobmodel/blobs.py

@@ -15,25 +15,33 @@ def __init__(
         amplitude: float,
         width_prop: float,
         width_perp: float,
-        v_x: float,
-        v_y: float,
+        velocity_x: float,
+        velocity_y: float,
         pos_x: float,
         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
         self.blob_shape = blob_shape
         self.amplitude = amplitude
         self.width_prop = width_prop
         self.width_perp = width_perp
-        self.v_x = v_x
-        self.v_y = v_y
+        self.v_x = velocity_x
+        self.v_y = velocity_y
         self.pos_x = pos_x
         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/model.py

@@ -45,8 +45,8 @@ def __init__(
             Important: only good approximation for Ly >> blob width
         num_blobs:
             number of blobs
-        blob_shape: str, optional
-            see Blob dataclass for available shapes
+        blob_shape: AbstractBlobShape or str, optional
+            see AbstractBlobShape and BlobShapeImpl dataclass for available shapes
         t_drain: float or array of length Nx, optional
             drain time for blobs
         blob_factory: BlobFactory, optional

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,
+        shape_param_x_parameter: float = 0.5,
+        shape_param_y_parameter: float = 0.5,
     ) -> None:
         """The following distributions are implemented:
 
@@ -59,16 +63,20 @@ def __init__(
         deg: degenerate distribution at `free_parameter`
         zeros: array of zeros
         """
-        self.A_dist = A_dist
-        self.wx_dist = wx_dist
-        self.wy_dist = wy_dist
-        self.vx_dist = vx_dist
-        self.vy_dist = vy_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.amplitude_dist = A_dist
+        self.width_x_dist = wx_dist
+        self.width_y_dist = wy_dist
+        self.velocity_x_dist = vx_dist
+        self.velocity_y_dist = vy_dist
+        self.shape_param_x_dist = spx_dist
+        self.shape_param_y_dist = spy_dist
+        self.amplitude_parameter = A_parameter
+        self.width_x_parameter = wx_parameter
+        self.width_y_parameter = wy_parameter
+        self.velocity_x_parameter = vx_parameter
+        self.velocity_y_parameter = vy_parameter
+        self.shape_param_x_parameter = shape_param_x_parameter
+        self.shape_param_y_parameter = shape_param_y_parameter
 
     def _draw_random_variables(
         self,
@@ -109,12 +117,31 @@ def sample_blobs(
         t_drain: float,
     ) -> List[Blob]:
         amps = self._draw_random_variables(
-            dist_type=self.A_dist, free_parameter=self.A_parameter, num_blobs=num_blobs
+            dist_type=self.amplitude_dist,
+            free_parameter=self.amplitude_parameter,
+            num_blobs=num_blobs,
         )
-        wxs = self._draw_random_variables(self.wx_dist, self.wx_parameter, num_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)
+        wxs = self._draw_random_variables(
+            self.width_x_dist, self.width_x_parameter, num_blobs
+        )
+        wys = self._draw_random_variables(
+            self.width_y_dist, self.width_y_parameter, num_blobs
+        )
+        vxs = self._draw_random_variables(
+            self.velocity_x_dist, self.velocity_x_parameter, num_blobs
+        )
+        vys = self._draw_random_variables(
+            self.velocity_y_dist, self.velocity_y_parameter, num_blobs
+        )
+        spxs = self._draw_random_variables(
+            self.shape_param_x_dist, self.shape_param_x_parameter, num_blobs
+        )
+        spys = self._draw_random_variables(
+            self.shape_param_y_dist, self.shape_param_y_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)
@@ -126,12 +153,14 @@ def sample_blobs(
                 amplitude=amps[i],
                 width_prop=wxs[i],
                 width_perp=wys[i],
-                v_x=vxs[i],
-                v_y=vys[i],
+                velocity_x=vxs[i],
+                velocity_y=vys[i],
                 pos_x=posxs[i],
                 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)
         ]
@@ -141,4 +170,4 @@ def sample_blobs(
 
     def is_one_dimensional(self) -> bool:
         # Perpendicular width parameters are irrelevant since perp shape should be ignored by the Bolb class.
-        return self.vy_dist == "zeros"
+        return self.velocity_y_dist == "zeros"

examples/custom_blobfactory.py

@@ -42,8 +42,8 @@ def sample_blobs(
                 amplitude=amp[i],
                 width_prop=width[i],
                 width_perp=width[i],
-                v_x=vx[i],
-                v_y=vy[i],
+                velocity_x=vx[i],
+                velocity_y=vy[i],
                 pos_x=posx[i],
                 pos_y=posy[i],
                 t_init=t_init[i],

examples/single_blob.py

@@ -32,8 +32,8 @@ def sample_blobs(
                 amplitude=amp[i],
                 width_prop=width[i],
                 width_perp=width[i],
-                v_x=vx[i],
-                v_y=vy[i],
+                velocity_x=vx[i],
+                velocity_y=vy[i],
                 pos_x=posx[i],
                 pos_y=posy[i],
                 t_init=t_init[i],

tests/test_blob.py

@@ -7,8 +7,8 @@
     amplitude=1,
     width_prop=1,
     width_perp=1,
-    v_x=1,
-    v_y=5,
+    velocity_x=1,
+    velocity_y=5,
     pos_x=0,
     pos_y=0,
     t_init=0,
@@ -53,8 +53,8 @@ def test_single_point():
         amplitude=1,
         width_prop=1,
         width_perp=1,
-        v_x=1,
-        v_y=1,
+        velocity_x=1,
+        velocity_y=1,
         pos_x=0,
         pos_y=6,
         t_init=0,
@@ -82,8 +82,8 @@ def test_theta_0():
         amplitude=1,
         width_prop=1,
         width_perp=1,
-        v_x=1,
-        v_y=0,
+        velocity_x=1,
+        velocity_y=0,
         pos_x=0,
         pos_y=6,
         t_init=0,
@@ -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,
+        velocity_x=1,
+        velocity_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_labels.py

@@ -34,8 +34,8 @@ def sample_blobs(
                 amplitude=_amp[i],
                 width_prop=_width[i],
                 width_perp=_width[i],
-                v_x=_vx[i],
-                v_y=_vy[i],
+                velocity_x=_vx[i],
+                velocity_y=_vy[i],
                 pos_x=_posx[i],
                 pos_y=_posy[i],
                 t_init=_t_init[i],