diff --git a/README.md b/README.md index cec9353..2540952 100644 --- a/README.md +++ b/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: diff --git a/blobmodel/blobs.py b/blobmodel/blobs.py index d98480c..3b62ea1 100644 --- a/blobmodel/blobs.py +++ b/blobmodel/blobs.py @@ -15,12 +15,14 @@ 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 @@ -28,12 +30,18 @@ def __init__( 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) diff --git a/blobmodel/model.py b/blobmodel/model.py index 61665bf..d80fbdf 100644 --- a/blobmodel/model.py +++ b/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 diff --git a/blobmodel/pulse_shape.py b/blobmodel/pulse_shape.py index d2cd861..f825cc4 100644 --- a/blobmodel/pulse_shape.py +++ b/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, + } diff --git a/blobmodel/stochasticality.py b/blobmodel/stochasticality.py index 63dfada..0e2a80a 100644 --- a/blobmodel/stochasticality.py +++ b/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" diff --git a/examples/custom_blobfactory.py b/examples/custom_blobfactory.py index d75edfa..de457b7 100644 --- a/examples/custom_blobfactory.py +++ b/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], diff --git a/examples/single_blob.py b/examples/single_blob.py index 421393b..e6a5e63 100644 --- a/examples/single_blob.py +++ b/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], diff --git a/tests/test_blob.py b/tests/test_blob.py index 0c580ab..3263baa 100644 --- a/tests/test_blob.py +++ b/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() diff --git a/tests/test_labels.py b/tests/test_labels.py index 3e340c5..90444bc 100644 --- a/tests/test_labels.py +++ b/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], diff --git a/tests/test_pulse_shape.py b/tests/test_pulse_shape.py new file mode 100644 index 0000000..cc75d5c --- /dev/null +++ b/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" diff --git a/tests/test_str.py b/tests/test_str.py index 684abb0..7364877 100644 --- a/tests/test_str.py +++ b/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() diff --git a/tests/test_vx_or_vy=0.py b/tests/test_vx_or_vy=0.py index 12c9d6c..4a189f7 100644 --- a/tests/test_vx_or_vy=0.py +++ b/tests/test_vx_or_vy=0.py @@ -33,8 +33,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], @@ -77,8 +77,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],